Merge pull request #31 from Multiomics-Analytics-Group/feat-batch-friendly-cli

feat: add --output-dir, --skip-plots, and summary.tsv for pipeline integration

Author: marcoreverenna

src/instanexus/consensus.py

@@ -141,10 +141,10 @@ def plot_logo2(pssm_df, output_file):
     plt.close(fig)
 
 
-def run_consensus_generation(align_folder: str, output_folder: str, run_id: str = ""):
+def run_consensus_generation(align_folder: str, output_folder: str, run_id: str = "", skip_plots: bool = False):
     """
     Core logic: Process all .afa files from alignment folder.
-    Generate consensus sequences, heatmaps, and logos.
+    Generate consensus sequences, and optionally heatmaps and logos.
     """
     align_path = Path(align_folder)
     output_path = Path(output_folder)
@@ -155,12 +155,13 @@ def run_consensus_generation(align_folder: str, output_folder: str, run_id: str
         raise FileNotFoundError(f"Alignment folder not found: {align_path}")
 
     consensus_fasta_dir = output_path / "consensus_fasta"
-    heatmap_dir = output_path / "heatmap"
-    logo_dir = output_path / "logo"
-
     consensus_fasta_dir.mkdir(exist_ok=True)
-    heatmap_dir.mkdir(exist_ok=True)
-    logo_dir.mkdir(exist_ok=True)
+
+    if not skip_plots:
+        heatmap_dir = output_path / "heatmap"
+        logo_dir = output_path / "logo"
+        heatmap_dir.mkdir(exist_ok=True)
+        logo_dir.mkdir(exist_ok=True)
 
     alignment_files = [f for f in sorted(os.listdir(align_path)) if f.endswith(".afa")]
 
@@ -190,11 +191,12 @@ def run_consensus_generation(align_folder: str, output_folder: str, run_id: str
         consensus_fasta_path = consensus_fasta_dir / f"{base_filename}_consensus.fasta"
         Bio.SeqIO.write([consensus_record], consensus_fasta_path, "fasta")
 
-        heatmap_path = heatmap_dir / f"{base_filename}_heatmap.svg"
-        plot_heatmap2(pssm_df, heatmap_path)
+        if not skip_plots:
+            heatmap_path = heatmap_dir / f"{base_filename}_heatmap.svg"
+            plot_heatmap2(pssm_df, heatmap_path)
 
-        logo_path = logo_dir / f"{base_filename}_logo.svg"
-        plot_logo2(pssm_df, logo_path)
+            logo_path = logo_dir / f"{base_filename}_logo.svg"
+            plot_logo2(pssm_df, logo_path)
 
     logger.info("All consensus tasks completed.")
     return consensus_fasta_dir
@@ -248,7 +250,7 @@ def generate_consensus_stats(consensus_base_folder):
     logger.info(f"Consensus statistics saved to: {stats_path}")
 
 
-def main(input_alignment_folder: str, output_consensus_folder: str, run_id: str = ""):
+def main(input_alignment_folder: str, output_consensus_folder: str, run_id: str = "", skip_plots: bool = False):
     """
     Main function to run the consensus generation script.
     """
@@ -259,13 +261,16 @@ def main(input_alignment_folder: str, output_consensus_folder: str, run_id: str
 
     logger.info(f"Alignment Folder (Input): {align_folder_in}")
     logger.info(f"Consensus Folder (Output): {consensus_folder_out}")
+    if skip_plots:
+        logger.info("Skipping heatmap and logo generation (--skip-plots)")
 
-    # --- Step 1: Generate consensus, heatmaps, and logos ---
+    # --- Step 1: Generate consensus, and optionally heatmaps and logos ---
     logger.info("Running consensus generation from alignment files...")
     consensus_fasta_dir = run_consensus_generation(
         align_folder=str(align_folder_in),
         output_folder=str(consensus_folder_out),
         run_id=run_id,
+        skip_plots=skip_plots,
     )
 
     # --- Step 2: Generate statistics on the consensus files ---
@@ -300,13 +305,19 @@ def cli():
         default="",
         help="Optional ID to display in the progress bar.",
     )
+    parser.add_argument(
+        "--skip-plots",
+        action="store_true",
+        help="Skip generating heatmap and logo SVG plots.",
+    )
 
     args = parser.parse_args()
 
     main(
         input_alignment_folder=args.input_folder,
         output_consensus_folder=args.output_folder,
         run_id=args.run_id,
+        skip_plots=args.skip_plots,
     )
 
 

src/instanexus/main.py

@@ -143,6 +143,19 @@ def cli():
         action="store_true",
         help="Enables iterative refinement (Overlap Graph) to merge assembled contigs.",
     )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Explicit output directory. When set, overrides the auto-generated path (folder-outputs/run_name/params). "
+        "Useful for pipeline/batch execution where deterministic output paths are required.",
+    )
+    parser.add_argument(
+        "--skip-plots",
+        action="store_true",
+        help="Skip generating heatmap and logo plots in the consensus step. "
+        "Useful for headless/batch execution where visualizations are not needed.",
+    )
 
     args = parser.parse_args()
 
@@ -160,30 +173,28 @@ def run_pipeline(args):
     logger.info("--- InstaNexus Pipeline started ---")
 
     run_name = Path(args.input_csv).stem
-    base_output_folder = Path(args.folder_outputs) / run_name  # e.g., 'outputs/bsa'
-
-    # Build the experiment folder name based on parameters
-    folder_name_parts = [f"{args.assembly_mode}"]
 
-    if args.chain:
-        folder_name_parts.append(f"{args.chain}")
-
-    if args.fdr is not None:
-        folder_name_parts.append(f"fdr{args.fdr}")
-    elif args.conf is not None:
-        folder_name_parts.append(f"c{args.conf}")
-
-    if "dbg" in args.assembly_mode:
-        folder_name_parts.append(f"ks{args.kmer_size}")
-
-    # folder_name_parts.append(f"mo{args.min_overlap}")
-    # folder_name_parts.append(f"ts{args.size_threshold}")
-
-    # if args.reference:
-    #     folder_name_parts.extend([f"mi{args.min_identity}", f"mm{args.max_mismatches}"])
-
-    run_folder_name = "_".join(folder_name_parts)
-    experiment_folder = base_output_folder / run_folder_name  # e.g., 'outputs/bsa/greedy_c0.9_mo4_ts10'
+    # Determine experiment output folder
+    if args.output_dir:
+        # Explicit output directory — deterministic path for pipeline/batch use
+        experiment_folder = Path(args.output_dir)
+        run_folder_name = experiment_folder.name
+    else:
+        # Auto-generated path from input name + parameters (interactive use)
+        base_output_folder = Path(args.folder_outputs) / run_name
+
+        folder_name_parts = [f"{args.assembly_mode}"]
+        if args.chain:
+            folder_name_parts.append(f"{args.chain}")
+        if args.fdr is not None:
+            folder_name_parts.append(f"fdr{args.fdr}")
+        elif args.conf is not None:
+            folder_name_parts.append(f"c{args.conf}")
+        if "dbg" in args.assembly_mode:
+            folder_name_parts.append(f"ks{args.kmer_size}")
+
+        run_folder_name = "_".join(folder_name_parts)
+        experiment_folder = base_output_folder / run_folder_name
 
     cleaned_csv_path = experiment_folder / "cleaned.csv"
 
@@ -300,12 +311,38 @@ def run_pipeline(args):
         consensus.main(
             input_alignment_folder=str(alignment_folder),
             output_consensus_folder=str(consensus_folder),
-            run_id=run_id_str,  # Pass ID for logs
+            run_id=run_id_str,
+            skip_plots=getattr(args, "skip_plots", False),
         )
     except Exception as e:
         logger.error(f"Consensus failed: {e}")
         return
 
+    # Write a stable summary file at a predictable path for pipeline integration
+    summary_path = experiment_folder / "summary.tsv"
+    try:
+        summary_data = {
+            "run_name": run_name,
+            "assembly_mode": args.assembly_mode,
+            "output_dir": str(experiment_folder),
+            "scaffolds_fasta": str(scaffolds_fasta_path),
+            "consensus_dir": str(consensus_folder),
+        }
+        # Include consensus stats if they exist
+        consensus_stats_path = consensus_folder / "consensus_stats.json"
+        if consensus_stats_path.exists():
+            import json
+
+            with open(consensus_stats_path) as f:
+                stats = json.load(f)
+            summary_data.update(stats)
+
+        summary_df = pd.DataFrame([summary_data])
+        summary_df.to_csv(summary_path, sep="\t", index=False)
+        logger.info(f"Summary written to: {summary_path}")
+    except Exception as e:
+        logger.warning(f"Failed to write summary: {e}")
+
     logger.info("--- InstaNexus Pipeline finished successfully! ---")
     logger.info(f"Final results in: {experiment_folder}")
 

src/instanexus/preprocessing.py

@@ -358,10 +358,10 @@ def main(
     if metadata_json is not None and "experiment_name" in df.columns:
         df["protease"] = df["experiment_name"].apply(lambda name: extract_protease(name, proteases))
 
-    if "preds" in df.columns:
-        df["cleaned_preds"] = df["preds"].apply(remove_modifications)
-    elif "prediction_untokenised" in df.columns:
-        df["cleaned_preds"] = df["prediction_untokenised"].apply(remove_modifications)
+    seq_candidates = ["preds", "prediction_untokenised", "prediction", "Peptide", "sequence"]
+    seq_col = next((c for c in seq_candidates if c in df.columns), None)
+    if seq_col is not None:
+        df["cleaned_preds"] = df[seq_col].apply(remove_modifications)
     else:
         raise ValueError("No suitable column found for peptide sequences.")