fix: reduce memory usage of em step

- Use two passes to build matrix, then calculate coverage after EM.
This avoid storing reduced representations of all alignments in memory.
- Consolidate BAM parsing into common interface.

Author: igboyes

pyproject.toml

@@ -18,7 +18,7 @@ authors = [
 
 [dependency-groups]
 dev = [
-    "maturin>=1.9.2",
+    "maturin[patchelf]>=1.9.2",
     "pysam>=0.22.0",
     "pytest>=8.4.1,<9.0.0",
     "pytest-asyncio<=1.1.0,<2.0.0",

python/workflow_pathoscope/rust.pyi

@@ -24,9 +24,8 @@ class PathoscopeResults:
     coverage: dict[str, list[int]]
 
 def run_expectation_maximization(
-    alignment_path: str,
+    bam_path: str,
     p_score_cutoff: float,
-    ref_lengths: dict[str, int],
 ) -> PathoscopeResults:
     """Run Pathoscope expectation maximization algorithm using Rust on SAM/BAM files."""
 
@@ -45,6 +44,5 @@ def find_candidate_otus_from_bytes(
 def calculate_coverage_from_em_results(
     alignment_path: str,
     p_score_cutoff: float,
-    ref_lengths: dict[str, int],
 ) -> dict[str, list[int]]:
     """Calculate coverage directly from EM results and alignment data."""

python/workflow_pathoscope/utils.py

@@ -140,20 +140,17 @@ def write_report(
 
 
 def run_pathoscope(
-    alignment_path: Path,
+    bam_path: Path,
     p_score_cutoff: float,
-    ref_lengths: dict[str, int],
 ):
     """Run Pathoscope on an alignment file.
 
     Returns PathoscopeResults containing EM results and coverage data.
 
     :param alignment_path: The path to the SAM or BAM file.
     :param p_score_cutoff: The minimum allowed ``p_score`` for an alignment.
-    :param ref_lengths: Dictionary mapping reference IDs to their lengths.
     """
     return run_expectation_maximization(
-        str(alignment_path),
+        str(bam_path),
         p_score_cutoff,
-        ref_lengths,
     )

rustfmt.toml

@@ -0,0 +1 @@
+max_width = 88

src/candidates.rs

@@ -1,15 +1,15 @@
-use std::collections::HashSet;
 use log::info;
-use pyo3::prelude::*;
 use pyo3::exceptions::PyIOError;
+use pyo3::prelude::*;
+use std::collections::HashSet;
 
 const AS_TAG_PREFIX: &str = "AS:i:";
 
 /// Extract AS:i alignment score from SAM optional fields
-/// 
+///
 /// # Arguments
 /// * `fields` - SAM fields starting from the optional fields (field 11+)
-/// 
+///
 /// # Returns
 /// Option containing the AS:i score as f64, None if not found or invalid
 fn extract_as_score(fields: &[&str]) -> Option<f64> {
@@ -24,14 +24,14 @@ fn extract_as_score(fields: &[&str]) -> Option<f64> {
 }
 
 /// Parse a single SAM line and extract candidate OTU information
-/// 
+///
 /// This function processes one SAM line and determines if the read meets the score cutoff.
 /// Used for testing and by the streaming functions.
 ///  
 /// # Arguments
 /// * `line` - A SAM format line as string
 /// * `p_score_cutoff` - Minimum score threshold (AS:i score + read length)
-/// 
+///
 /// # Returns
 /// Option containing the reference name if the read meets the cutoff, None otherwise
 pub fn parse_sam_line(line: &str, p_score_cutoff: f64) -> Option<String> {
@@ -42,14 +42,14 @@ pub fn parse_sam_line(line: &str, p_score_cutoff: f64) -> Option<String> {
 
     // Parse SAM line - tab-separated format
     let fields: Vec<&str> = line.split('\t').collect();
-    
+
     // SAM format requires at least 11 fields
     if fields.len() < 11 {
         return None;
     }
 
     // Extract key fields:
-    // 1: FLAG 
+    // 1: FLAG
     // 2: RNAME (reference name)
     // 9: SEQ (read sequence)
     let flag: u16 = fields[1].parse().unwrap_or(4); // Default to unmapped if parse fails
@@ -75,19 +75,18 @@ pub fn parse_sam_line(line: &str, p_score_cutoff: f64) -> Option<String> {
     None
 }
 
-
 /// Extract candidate OTU reference IDs by running bowtie2 directly with streaming
-/// 
+///
 /// This function spawns a bowtie2 process directly from Rust and streams its output
 /// to avoid memory issues with large SAM files. It processes SAM lines as they arrive
 /// and returns only the unique reference IDs that meet the score cutoff.
-/// 
+///
 /// # Arguments
 /// * `bowtie_index_path` - Path to the bowtie2 index
 /// * `read_paths` - List of paths to the input read files
 /// * `proc` - Number of processor threads for bowtie2
 /// * `p_score_cutoff` - Minimum score threshold (AS:i score + read length)
-/// 
+///
 /// # Returns
 /// Set of reference IDs that have reads meeting the score cutoff
 pub fn find_candidate_otus_with_bowtie2(
@@ -97,52 +96,57 @@ pub fn find_candidate_otus_with_bowtie2(
     proc: i32,
     p_score_cutoff: f64,
 ) -> PyResult<HashSet<String>> {
-    use std::process::{Command, Stdio};
     use std::io::{BufRead, BufReader};
-    
-    info!("running bowtie2: index={}, reads={:?}, cutoff={}", 
-          bowtie_index_path, read_paths, p_score_cutoff);
+    use std::process::{Command, Stdio};
+
+    info!(
+        "running bowtie2: index={}, reads={:?}, cutoff={}",
+        bowtie_index_path, read_paths, p_score_cutoff
+    );
     py.allow_threads(|| {
         let mut cmd = Command::new("bowtie2");
-        cmd.arg("-p").arg(proc.to_string())
-           .arg("--local")
-           .arg("--no-unal")
-           .arg("--score-min").arg("L,20,1.0")
-           .arg("-N").arg("0")
-           .arg("-L").arg("15")
-           .arg("-x").arg(bowtie_index_path)
-           .arg("-U").arg(read_paths.join(","))
-           .stdout(Stdio::piped())
-           .stderr(Stdio::piped());
-           
+        cmd.arg("-p")
+            .arg(proc.to_string())
+            .arg("--local")
+            .arg("--no-unal")
+            .arg("--score-min")
+            .arg("L,20,1.0")
+            .arg("-N")
+            .arg("0")
+            .arg("-L")
+            .arg("15")
+            .arg("-x")
+            .arg(bowtie_index_path)
+            .arg("-U")
+            .arg(read_paths.join(","))
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped());
+
         info!("spawning bowtie2 process");
-        let mut child = cmd.spawn()
-            .map_err(|e| PyErr::new::<PyIOError, _>(format!("Failed to spawn bowtie2: {}", e)))?;
-        
+        let mut child = cmd.spawn()?;
+
         let stdout = child.stdout.take().unwrap();
         let reader = BufReader::new(stdout);
-        
+
         let mut candidate_otus = HashSet::new();
         let mut line_count = 0u64;
         let mut passing_count = 0u64;
-        
+
         for line_result in reader.lines() {
-            let line = line_result
-                .map_err(|e| PyErr::new::<PyIOError, _>(format!("Error reading bowtie2 output: {}", e)))?;
-            
+            let line = line_result?;
+
             line_count += 1;
-            
+
             // Use the extracted SAM parsing function
             if let Some(ref_name) = parse_sam_line(&line, p_score_cutoff) {
                 candidate_otus.insert(ref_name);
                 passing_count += 1;
             }
         }
-        
+
         // Wait for bowtie2 to finish and check exit status
-        let status = child.wait()
-            .map_err(|e| PyErr::new::<PyIOError, _>(format!("Error waiting for bowtie2: {}", e)))?;
-            
+        let status = child.wait()?;
+
         if !status.success() {
             // Read stderr for error details
             let stderr_output = if let Some(mut stderr) = child.stderr.take() {
@@ -152,22 +156,25 @@ pub fn find_candidate_otus_with_bowtie2(
             } else {
                 "Unknown error".to_string()
             };
-            
+
             return Err(PyErr::new::<PyIOError, _>(format!(
-                "bowtie2 failed with exit code {:?}: {}", 
-                status.code(), 
+                "bowtie2 failed with exit code {:?}: {}",
+                status.code(),
                 stderr_output
             )));
         }
-        
-        info!("processed {} sam lines, {} passed cutoff, found {} unique otus", 
-              line_count, passing_count, candidate_otus.len());
-        
+
+        info!(
+            "processed {} sam lines, {} passed cutoff, found {} unique otus",
+            line_count,
+            passing_count,
+            candidate_otus.len()
+        );
+
         Ok(candidate_otus)
     })
 }
 
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -176,7 +183,7 @@ mod tests {
     fn test_parse_sam_line_basic() {
         let line = "read1\t0\tref1\t100\t255\t50M\t*\t0\t0\tAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\t*\tAS:i:45";
         let result = parse_sam_line(line, 0.01);
-        
+
         // AS:i:45 + seq_len(50) = 95.0, should pass cutoff of 0.01
         assert_eq!(result, Some("ref1".to_string()));
     }
@@ -185,7 +192,7 @@ mod tests {
     fn test_parse_sam_line_below_cutoff() {
         let line = "read1\t0\tref1\t100\t255\t50M\t*\t0\t0\tAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\t*\tAS:i:45";
         let result = parse_sam_line(line, 100.0);
-        
+
         // AS:i:45 + seq_len(50) = 95.0, should not pass cutoff of 100.0
         assert_eq!(result, None);
     }
@@ -194,7 +201,7 @@ mod tests {
     fn test_parse_sam_line_unmapped() {
         let line = "read1\t4\t*\t0\t0\t*\t*\t0\t0\tAAAAA\t*";
         let result = parse_sam_line(line, 0.01);
-        
+
         // Unmapped read (flag & 4 != 0), should return None
         assert_eq!(result, None);
     }
@@ -203,7 +210,7 @@ mod tests {
     fn test_parse_sam_line_no_as_score() {
         let line = "read1\t0\tref1\t100\t255\t50M\t*\t0\t0\tAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA\t*";
         let result = parse_sam_line(line, 0.01);
-        
+
         // No AS:i score, should return None
         assert_eq!(result, None);
     }
@@ -212,10 +219,8 @@ mod tests {
     fn test_parse_sam_line_header() {
         let line = "@HD\tVN:1.0\tSO:unsorted";
         let result = parse_sam_line(line, 0.01);
-        
+
         // Header line, should return None
         assert_eq!(result, None);
     }
-
-
-}
+}