argenus/
reassemble.rs

1//! Per-locus reassembly (v3 core/flank read split) — opt-in via `--reassemble`.
2//!
3//! Rationale (validated in the low-depth ablation, see project memory
4//! `perlocus-reassembly-and-emit`): reads that align >=70% of their length to the
5//! ARG are the shared, conserved *core* that BRIDGES different genomes' flanking
6//! during assembly and collapses them into chimeras. If we split reads per-mate,
7//! drop pairs where both mates are core, keep the flanking mate of a half-core pair
8//! as a singleton, and assemble the remaining *flanking* reads alone, the flanking
9//! separates by genome. We then stitch the ARG gene back onto each flanking contig
10//! so ARGenus can re-classify it with the full 4-axis (genus/species/context)
11//! machinery — which honestly labels promiscuous plasmid genes rather than forcing
12//! a single (usually wrong) genus.
13//!
14//! This targets STALLED loci only (short flanking on the strict contig): they are
15//! `NA` on every axis today, and reassembly is what converts them into something
16//! classifiable. Plasmid-context loci are skipped by the caller (reassembly cannot
17//! fix a mobile-element genus). All artifacts are written under a per-sample
18//! `reassemble/` dir and the step is resumable (existing `spades/contigs.fasta`
19//! is reused).
20
21use anyhow::{Context, Result};
22use std::fs::{self, File};
23use std::io::{BufWriter, Write};
24use std::path::{Path, PathBuf};
25use std::process::Command;
26
27use rustc_hash::FxHashMap;
28
29use crate::seqio::{FastaReader, FastqFile};
30
31/// Tool paths + tuning for the reassembly step.
32pub struct ReasmConfig {
33    pub minimap2: String,
34    /// Path to `spades.py`.
35    pub spades_py: PathBuf,
36    /// Python interpreter that can run SPAdes (its pixi env python; the system
37    /// python is often too old). Empty => call `spades_py` directly on PATH.
38    pub spades_python: PathBuf,
39    pub threads: usize,
40    /// Fraction of a read covered by the ARG alignment to call it "core" (0.70).
41    pub core_cov: f64,
42    /// Concurrent SPAdes jobs.
43    pub jobs: usize,
44}
45
46/// A stalled locus to reassemble: the ARG gene region plus the whole strict contig
47/// (used as the orientation seed when stitching).
48pub struct StalledLocus {
49    /// Unique key, `arg_name:contig` — also the output name prefix.
50    pub key: String,
51    pub arg_name: String,
52    pub arg_class: String,
53    pub identity: f64,
54    pub coverage: f64,
55    /// ARG gene sequence (substring of the contig).
56    pub gene_seq: String,
57    /// Whole strict contig sequence (orientation seed).
58    pub seed_seq: String,
59}
60
61/// One stitched flanking contig for a locus, with the ARG gene glued back on and
62/// the gene's coordinates on the stitched sequence (for re-classification).
63pub struct StitchedContig {
64    pub name: String,
65    pub seq: String,
66    pub arg_start: usize,
67    pub arg_end: usize,
68}
69
70/// Reassembly output for one original locus (may yield several flanking contigs =
71/// genuine multi-genome presence).
72pub struct StitchedLocus {
73    pub key: String,
74    pub arg_name: String,
75    pub arg_class: String,
76    pub identity: f64,
77    pub coverage: f64,
78    pub contigs: Vec<StitchedContig>,
79}
80
81fn norm_id(h: &str) -> String {
82    let mut t = h.split_whitespace().next().unwrap_or("").to_string();
83    if let Some(first) = t.chars().next() {
84        if first == '@' || first == '>' {
85            t = t[1..].to_string();
86        }
87    }
88    // strip trailing /1 or /2 (possibly repeated: minimap2 can append its own)
89    loop {
90        let b = t.as_bytes();
91        if b.len() > 2 && b[b.len() - 2] == b'/' && (b[b.len() - 1] == b'1' || b[b.len() - 1] == b'2') {
92            t.truncate(b.len() - 2);
93        } else {
94            break;
95        }
96    }
97    t
98}
99
100fn revcomp(s: &str) -> String {
101    s.bytes()
102        .rev()
103        .map(|b| match b {
104            b'A' => 'T', b'a' => 't',
105            b'C' => 'G', b'c' => 'g',
106            b'G' => 'C', b'g' => 'c',
107            b'T' => 'A', b't' => 'a',
108            _ => 'N',
109        })
110        .collect()
111}
112
113/// Run minimap2 (query -> target) writing PAF to `out`. Returns Ok even if the PAF
114/// is empty (no alignments); errors only on process failure.
115fn minimap2_paf(mm2: &str, preset: &str, target: &Path, query: &Path, out: &Path, threads: usize) -> Result<()> {
116    let f = File::create(out).with_context(|| format!("create {:?}", out))?;
117    let status = Command::new(mm2)
118        .args(["-x", preset, "-t", &threads.to_string()])
119        .arg(target)
120        .arg(query)
121        .stdout(std::process::Stdio::from(f))
122        .stderr(std::process::Stdio::null())
123        .status()
124        .with_context(|| format!("run minimap2 ({} {} -> {})", preset, query.display(), target.display()))?;
125    if !status.success() {
126        anyhow::bail!("minimap2 exited with {}", status);
127    }
128    Ok(())
129}
130
131/// Parse a reads-vs-genes PAF into best (locus_idx, query_coverage) per read id.
132/// `locus_idx` is decoded from the target name `L<idx>`.
133fn best_locus_per_read(paf: &Path) -> Result<FxHashMap<String, (usize, f64)>> {
134    let mut best: FxHashMap<String, (usize, f64)> = FxHashMap::default();
135    let content = fs::read_to_string(paf).with_context(|| format!("read {:?}", paf))?;
136    for line in content.lines() {
137        let c: Vec<&str> = line.split('\t').collect();
138        if c.len() < 11 {
139            continue;
140        }
141        let qlen: f64 = c[1].parse().unwrap_or(0.0);
142        let qs: f64 = c[2].parse().unwrap_or(0.0);
143        let qe: f64 = c[3].parse().unwrap_or(0.0);
144        if qlen <= 0.0 {
145            continue;
146        }
147        let cov = (qe - qs) / qlen;
148        let target = c[5];
149        let idx: usize = match target.strip_prefix('L').and_then(|s| s.parse().ok()) {
150            Some(i) => i,
151            None => continue,
152        };
153        let rid = norm_id(c[0]);
154        let e = best.entry(rid).or_insert((idx, 0.0));
155        if cov > e.1 {
156            *e = (idx, cov);
157        }
158    }
159    Ok(best)
160}
161
162/// Split filtered reads into per-locus flanking read sets (core/flank per mate).
163/// Returns, per locus index, (has_pairs, has_singles).
164fn split_reads(
165    loci: &[StalledLocus],
166    filtered_r1: &Path,
167    filtered_r2: &Path,
168    workdir: &Path,
169    cfg: &ReasmConfig,
170) -> Result<Vec<(bool, bool)>> {
171    // 1. gene multi-FASTA (record name L<idx>)
172    let genes_fa = workdir.join("genes.fa");
173    {
174        let mut w = BufWriter::new(File::create(&genes_fa)?);
175        for (i, loc) in loci.iter().enumerate() {
176            writeln!(w, ">L{}\n{}", i, loc.gene_seq)?;
177        }
178    }
179
180    // 2. map R1 and R2 to the genes separately
181    let r1_paf = workdir.join("r1_to_genes.paf");
182    let r2_paf = workdir.join("r2_to_genes.paf");
183    minimap2_paf(&cfg.minimap2, "sr", &genes_fa, filtered_r1, &r1_paf, cfg.threads)?;
184    minimap2_paf(&cfg.minimap2, "sr", &genes_fa, filtered_r2, &r2_paf, cfg.threads)?;
185    let r1_best = best_locus_per_read(&r1_paf)?;
186    let r2_best = best_locus_per_read(&r2_paf)?;
187
188    // 3. stream both mates in lockstep, route each recruited pair to its locus
189    let dirs: Vec<PathBuf> = (0..loci.len()).map(|i| workdir.join(format!("L{}", i))).collect();
190    for d in &dirs {
191        fs::create_dir_all(d)?;
192    }
193    // lazily-opened per-locus writers: (R1, R2, S)
194    let mut w1: Vec<Option<BufWriter<File>>> = (0..loci.len()).map(|_| None).collect();
195    let mut w2: Vec<Option<BufWriter<File>>> = (0..loci.len()).map(|_| None).collect();
196    let mut ws: Vec<Option<BufWriter<File>>> = (0..loci.len()).map(|_| None).collect();
197    let mut npair = vec![0u64; loci.len()];
198    let mut nsing = vec![0u64; loci.len()];
199
200    let mut f1 = FastqFile::open(filtered_r1)?;
201    let mut f2 = FastqFile::open(filtered_r2)?;
202    let core = cfg.core_cov;
203    loop {
204        let a = f1.read_next()?;
205        let b = f2.read_next()?;
206        let (a, b) = match (a, b) {
207            (Some(a), Some(b)) => (a, b),
208            _ => break,
209        };
210        let rid = norm_id(&a.name);
211        let o1 = r1_best.get(&rid).copied();
212        let o2 = r2_best.get(&rid).copied();
213        // choose the locus this pair belongs to
214        let locus = match (o1, o2) {
215            (Some((l1, c1)), Some((l2, c2))) => {
216                if l1 == l2 { Some(l1) } else if c1 >= c2 { Some(l1) } else { Some(l2) }
217            }
218            (Some((l1, _)), None) => Some(l1),
219            (None, Some((l2, _))) => Some(l2),
220            (None, None) => None,
221        };
222        let l = match locus {
223            Some(l) => l,
224            None => continue,
225        };
226        let r1core = o1.map_or(false, |(li, cv)| li == l && cv >= core);
227        let r2core = o2.map_or(false, |(li, cv)| li == l && cv >= core);
228        if r1core && r2core {
229            continue; // both core: shared bridge — drop
230        } else if r1core {
231            // keep flanking mate (R2) as a singleton
232            let w = ws[l].get_or_insert_with(|| BufWriter::new(File::create(dirs[l].join("flank_S.fq")).unwrap()));
233            write_fq(w, &b)?;
234            nsing[l] += 1;
235        } else if r2core {
236            let w = ws[l].get_or_insert_with(|| BufWriter::new(File::create(dirs[l].join("flank_S.fq")).unwrap()));
237            write_fq(w, &a)?;
238            nsing[l] += 1;
239        } else {
240            let x1 = w1[l].get_or_insert_with(|| BufWriter::new(File::create(dirs[l].join("flank_R1.fq")).unwrap()));
241            write_fq(x1, &a)?;
242            let x2 = w2[l].get_or_insert_with(|| BufWriter::new(File::create(dirs[l].join("flank_R2.fq")).unwrap()));
243            write_fq(x2, &b)?;
244            npair[l] += 1;
245        }
246    }
247    // flush
248    drop(w1); drop(w2); drop(ws);
249    Ok((0..loci.len()).map(|i| (npair[i] > 0, nsing[i] > 0)).collect())
250}
251
252fn write_fq(w: &mut BufWriter<File>, r: &crate::seqio::FastqRecord) -> Result<()> {
253    writeln!(w, "@{}\n{}\n+\n{}", r.name, r.seq, r.qual)?;
254    Ok(())
255}
256
257/// Run SPAdes (`--only-assembler`) on a locus's flanking reads. Resumable: if
258/// `spades/contigs.fasta` already exists and is non-empty, it is reused.
259fn run_spades(dir: &Path, has_pairs: bool, has_singles: bool, cfg: &ReasmConfig, per_job_threads: usize) -> Result<Option<PathBuf>> {
260    let out = dir.join("spades");
261    let contigs = out.join("contigs.fasta");
262    if contigs.exists() && fs::metadata(&contigs).map(|m| m.len() > 0).unwrap_or(false) {
263        return Ok(Some(contigs));
264    }
265    if !has_pairs && !has_singles {
266        return Ok(None);
267    }
268    let mut cmd = if cfg.spades_python.as_os_str().is_empty() {
269        let mut c = Command::new(&cfg.spades_py);
270        c.args(["--only-assembler", "-t", &per_job_threads.to_string(), "-m", "8", "-o"]).arg(&out);
271        c
272    } else {
273        let mut c = Command::new(&cfg.spades_python);
274        c.arg(&cfg.spades_py).args(["--only-assembler", "-t", &per_job_threads.to_string(), "-m", "8", "-o"]).arg(&out);
275        c
276    };
277    if has_pairs {
278        cmd.arg("-1").arg(dir.join("flank_R1.fq")).arg("-2").arg(dir.join("flank_R2.fq"));
279    }
280    if has_singles {
281        cmd.arg("-s").arg(dir.join("flank_S.fq"));
282    }
283    let status = cmd.stdout(std::process::Stdio::null()).stderr(std::process::Stdio::null()).status();
284    match status {
285        Ok(s) if s.success() && contigs.exists() => Ok(Some(contigs)),
286        _ => Ok(None), // spades_fail — skip this locus
287    }
288}
289
290/// Orient each flanking contig against the seed and glue the ARG gene onto the end
291/// nearest the gene span, recording the gene's coordinates on the stitched contig.
292fn stitch(loc: &StalledLocus, contigs: &Path, dir: &Path, cfg: &ReasmConfig) -> Result<Vec<StitchedContig>> {
293    // orientation info: contig -> (strand, target_start)
294    let mut side: FxHashMap<String, (char, usize)> = FxHashMap::default();
295    if !loc.seed_seq.is_empty() {
296        let seed_fa = dir.join("seed.fa");
297        {
298            let mut w = BufWriter::new(File::create(&seed_fa)?);
299            writeln!(w, ">seed\n{}", loc.seed_seq)?;
300        }
301        let cpaf = dir.join("contig2seed.paf");
302        minimap2_paf(&cfg.minimap2, "asm10", &seed_fa, contigs, &cpaf, 2)?;
303        let content = fs::read_to_string(&cpaf)?;
304        for line in content.lines() {
305            let c: Vec<&str> = line.split('\t').collect();
306            if c.len() < 9 {
307                continue;
308            }
309            let strand = c[4].chars().next().unwrap_or('+');
310            let ts: usize = c[7].parse().unwrap_or(0);
311            side.entry(c[0].to_string()).or_insert((strand, ts));
312        }
313    }
314    let gene = &loc.gene_seq;
315    let seedlen = loc.seed_seq.len();
316    let mut out = Vec::new();
317    let mut reader = FastaReader::open(contigs)?;
318    let mut n = 0;
319    while let Some(rec) = reader.read_next()? {
320        if rec.seq.len() < 200 {
321            continue;
322        }
323        let cname = rec.name.split_whitespace().next().unwrap_or(&rec.name).to_string();
324        let mut oriented = rec.seq.clone();
325        // default (no seed hit): gene_down -> gene + oriented
326        let mut gene_up = false;
327        if let Some((strand, ts)) = side.get(&cname) {
328            if *strand == '-' {
329                oriented = revcomp(&rec.seq);
330            }
331            gene_up = *ts < seedlen / 2;
332        }
333        n += 1;
334        let (seq, arg_start, arg_end) = if gene_up {
335            let start = oriented.len();
336            (format!("{}{}", oriented, gene), start, start + gene.len())
337        } else {
338            (format!("{}{}", gene, oriented), 0, gene.len())
339        };
340        out.push(StitchedContig {
341            name: format!("{}__c{}", loc.key.replace([':', '|', ' '], "_"), n),
342            seq,
343            arg_start,
344            arg_end,
345        });
346    }
347    Ok(out)
348}
349
350/// Reassemble a batch of stalled loci. Returns one `StitchedLocus` per input locus
351/// that produced at least one stitched flanking contig.
352pub fn reassemble_stalled(
353    loci: &[StalledLocus],
354    filtered_r1: &Path,
355    filtered_r2: &Path,
356    workdir: &Path,
357    cfg: &ReasmConfig,
358) -> Result<Vec<StitchedLocus>> {
359    if loci.is_empty() {
360        return Ok(Vec::new());
361    }
362    fs::create_dir_all(workdir)?;
363    let flags = split_reads(loci, filtered_r1, filtered_r2, workdir, cfg)?;
364
365    let per_job_threads = (cfg.threads / cfg.jobs.max(1)).max(1);
366    let pool = rayon::ThreadPoolBuilder::new().num_threads(cfg.jobs.max(1)).build()?;
367
368    // Assemble + stitch each locus; parallel over loci with a bounded pool.
369    let results: Vec<Option<StitchedLocus>> = pool.install(|| {
370        use rayon::prelude::*;
371        (0..loci.len())
372            .into_par_iter()
373            .map(|i| {
374                let loc = &loci[i];
375                let dir = workdir.join(format!("L{}", i));
376                let (has_pairs, has_singles) = flags[i];
377                let contigs = match run_spades(&dir, has_pairs, has_singles, cfg, per_job_threads) {
378                    Ok(Some(c)) => c,
379                    _ => return None,
380                };
381                let stitched = match stitch(loc, &contigs, &dir, cfg) {
382                    Ok(s) if !s.is_empty() => s,
383                    _ => return None,
384                };
385                Some(StitchedLocus {
386                    key: loc.key.clone(),
387                    arg_name: loc.arg_name.clone(),
388                    arg_class: loc.arg_class.clone(),
389                    identity: loc.identity,
390                    coverage: loc.coverage,
391                    contigs: stitched,
392                })
393            })
394            .collect()
395    });
396
397    Ok(results.into_iter().flatten().collect())
398}
argenus/reassemble.rs

argenus/
reassemble.rs
