modified files for enabling statistics

Author: marcoreverenna

src/instanexus/assembly.py

@@ -884,6 +884,8 @@ def __init__(
         alpha_len: float = 1.0,
         alpha_cov: float = 1.0,
         alpha_min: float = 0.2,
+        reference_protein: str = None,
+        stats_output_folder: str = None,
     ):
         if mode not in ["greedy", "dbg", "dbg_weighted", "dbgX", "fusion", "multimodal_dbg", "hybrid_dbg"]:
             raise ValueError(
@@ -902,12 +904,37 @@ def __init__(
         self.alpha_len = alpha_len
         self.alpha_cov = alpha_cov
         self.alpha_min = alpha_min
+        self.reference_protein = reference_protein
+        self.stats_output_folder = stats_output_folder
+
+    def _compute_intermediate_stats(self, contigs, label):
+        """Internal wrapper for statistics."""
+        if self.reference_protein and self.stats_output_folder:
+            logger.info(f"Computing intermediate statistics for {label}...")
+            try:
+                mapped = viz.process_protein_contigs_scaffold(
+                    assembled_contigs=contigs,
+                    target_protein=self.reference_protein,
+                    max_mismatches=self.max_mismatches,
+                    min_identity=self.min_identity
+                )
+                df_mapped = viz.create_dataframe_from_mapped_sequences(data=mapped)
+                if not df_mapped.empty:
+                    helpers.compute_assembly_statistics(
+                        df=df_mapped,
+                        sequence_type=label,
+                        output_folder=self.stats_output_folder,
+                        reference=self.reference_protein,
+                    )
+            except Exception as e:
+                logger.warning(f"Could not compute intermediate stats: {e}")
 
     def assemble_greedy(self, sequences):
         logger.info(f"[Assembler] Running Greedy assembly (min_overlap={self.min_overlap})")
         contigs = assemble_contigs_greedy(sequences, self.min_overlap)
         contigs = list(set(contigs))
         contigs = sorted(contigs, key=len, reverse=True)
+        self._compute_intermediate_stats(contigs, label="contig")
 
         scaffolds = scaffold_iterative_greedy(contigs, self.min_overlap, self.size_threshold)
 
@@ -949,6 +976,8 @@ def assemble_dbg_weighted(self, sequences: List[str]) -> List[str]:
 
         logger.info(f"DBG produced {len(contigs)} initial contigs.")
 
+        self._compute_intermediate_stats(contigs, label="contig")
+
         # 2. OVERLAP GRAPH REFINEMENT (The new logic)
         # Only runs if refine_rounds is > 0
         if self.refine_rounds > 0:
@@ -1265,21 +1294,29 @@ def main(
     refine_rounds: int = 0,
 ):
     """Main function for standalone assembly."""
+    output_path = Path(output_scaffolds_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    stats_folder = output_path.parent / "statistics"
 
     protein_norm = None  # None means no reference mode
+
     if reference:
         logger.info("Reference mode enabled. Loading reference protein...")
         if not metadata_json_path:
-            raise ValueError("metadata_json_path is required when reference mode is enabled.")
-
+            raise ValueError("metadata_json_path is required.")
         try:
-            run_stem = Path(input_csv_path).stem  # extract run name from input file
-            run_name = run_stem.replace("_cleaned", "")
-
+            path_obj = Path(input_csv_path)
+            
+            if path_obj.name == "cleaned.csv":
+                run_name = path_obj.parent.parent.name
+            else:
+                run_name = path_obj.stem.replace("_cleaned", "")
+    
             meta = helpers.get_sample_metadata(run=run_name, chain=chain, json_path=metadata_json_path)
             protein = meta["protein"]
             protein_norm = preprocessing.normalize_sequence(protein)
             logger.info("Reference protein loaded and normalized successfully.")
+            stats_folder.mkdir(parents=True, exist_ok=True)
 
         except Exception as e:
             logger.error(f"Failed to get reference protein: {e}")
@@ -1306,6 +1343,8 @@ def main(
         min_identity=min_identity,
         max_mismatches=max_mismatches,
         refine_rounds=refine_rounds,
+        reference_protein=protein_norm,
+        stats_output_folder=str(stats_folder) if protein_norm else None
     )
 
     scaffolds = assembler.run(sequences=sequences, df_full=df)

src/instanexus/main.py

@@ -22,9 +22,10 @@
 
 import argparse
 import logging
+import pandas as pd
 from pathlib import Path
 
-from . import alignment, assembly, clustering, consensus, preprocessing
+from . import alignment, assembly, clustering, consensus, preprocessing, helpers, visualization as viz
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
@@ -164,6 +165,9 @@ def run_pipeline(args):
     # 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:
@@ -172,11 +176,11 @@ def run_pipeline(args):
     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}")
+    # 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}"])
+    # 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'
@@ -186,14 +190,15 @@ def run_pipeline(args):
     scaffolds_folder = experiment_folder / "scaffolds"
     scaffolds_fasta_path = scaffolds_folder / "scaffolds.fasta"
 
+    statistics_folder = scaffolds_folder / "statistics"
+
     clustering_folder = scaffolds_folder / "clustering"  # Clustering output
     cluster_fasta_folder = clustering_folder / "cluster_fasta"  # Input for alignment
 
     alignment_folder = scaffolds_folder / "alignment"
 
     consensus_folder = scaffolds_folder / "consensus"
 
-    # ID for logs (optional)
     run_id_str = f"[{run_name} @ {run_folder_name}]"
 
     logger.info(f"Starting pipeline for run: {run_id_str}")
@@ -214,6 +219,40 @@ def run_pipeline(args):
     except Exception as e:
         logger.error(f"Preprocessing failed: {e}")
         return
+    
+    if args.reference:
+        logger.info("--- [Statistics] Computing INPUT Statistics ---")
+        try:
+            meta = helpers.get_sample_metadata(run=run_name, chain=args.chain, json_path=args.metadata_json_path)
+            protein_seq = meta.get("protein", "")
+            
+            if protein_seq:
+                protein_norm = preprocessing.normalize_sequence(protein_seq)
+                statistics_folder.mkdir(parents=True, exist_ok=True)
+
+                if cleaned_csv_path.exists():
+                    df_input = pd.read_csv(cleaned_csv_path)
+                    if "cleaned_preds" in df_input.columns:
+                        input_seqs = df_input["cleaned_preds"].dropna().unique().tolist()
+                        
+                        mapped_contigs = viz.process_protein_contigs_scaffold(
+                            assembled_contigs=input_seqs,
+                            target_protein=protein_norm,
+                            max_mismatches=args.max_mismatches,
+                            min_identity=args.min_identity
+                        )
+                        df_contigs_mapped = viz.create_dataframe_from_mapped_sequences(data=mapped_contigs)
+                        
+                        if not df_contigs_mapped.empty:
+                            helpers.compute_assembly_statistics(
+                                df=df_contigs_mapped,
+                                sequence_type="peptide", 
+                                output_folder=str(statistics_folder),
+                                reference=protein_norm,
+                            )
+                            logger.info(f"[Statistics] Saved: peptide_stats.json")
+        except Exception as e:
+            logger.error(f"[Statistics] Failed to compute peptide stats: {e}")
 
     try:
         logger.info("--- [Step 2/5] Running Assembly ---")