From f5e508ed339b753f7d1ce253b87f13dc223ddd65 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Mon, 10 Aug 2026 22:12:12 +0200 Subject: [PATCH] feat: add multi-genome SNP pseudo-alignment and CLI export Introduces a `SnpAlignment` struct and helper methods to construct per-genome SNP pseudo-alignments from sibling k-mer data, filtering monomorphic families and encoding bases as IUPAC ambiguity codes. Exposes the type at the crate root for simplified imports. Adds a `--snp` CLI flag to compute and export these alignments as an IUPAC-coded FASTA file. Updates theory documentation to propose a multi-genome framing approach for joint phylogenetic inference, resolving pairwise correspondence ambiguities through positional homology and partial coverage thresholds. Bumps crate version to 1.1.40. --- docmd/theory/evolutionary_distances.md | 57 +++++++++ src/Cargo.lock | 2 +- src/obikindex/src/lib.rs | 2 +- src/obikindex/src/siblings.rs | 169 +++++++++++++++++++++++++ src/obikmer/Cargo.toml | 2 +- src/obikmer/src/cmd/distance.rs | 64 ++++++++-- 6 files changed, 282 insertions(+), 14 deletions(-) diff --git a/docmd/theory/evolutionary_distances.md b/docmd/theory/evolutionary_distances.md index 2c21a3c..ba63278 100644 --- a/docmd/theory/evolutionary_distances.md +++ b/docmd/theory/evolutionary_distances.md @@ -146,6 +146,63 @@ A once multiplicity > 1 on either side. Any pairing rule invents a correspondence the data cannot support. Multiplicity > 1 is treated as non-identifiable, not as a puzzle to solve with a heuristic. +## Multi-genome framing: family as pseudo-alignment column + +**Idea.** Instead of resolving locus eligibility and correspondence one +genome pair at a time, treat a family as a column of a pseudo multiple +alignment across *all* genomes simultaneously: for each family, each genome +has either a net single-copy state (`A`/`C`/`G`/`T`, when the genome carries +exactly one of the 4 forms) or "missing" (`?`, multi-copy or absent). Flank +conservation (the `2m` bases fixed by construction) supplies positional +homology for free — the same role a real MSA would play, without alignment +software, gap penalties, or progressive-alignment approximations. Stacking +one such column per family, genomes as rows, produces a genuine SNP +pseudo-alignment matrix, not just a bag of pairwise distances. + +**Precedent.** This is the same principle behind reference-free +k-mer-based phylogenomics tools — SKA (Split K-mer Analysis, Harris 2018) and +kSNP: split the k-mer around a variable center, use flank identity to call +homologous columns across arbitrarily many genomes with no reference and no +MSA step, then feed the resulting pseudo-alignment to standard phylogenetic +tools. Landing on the same design independently is a good sign, not a +coincidence. + +**Resolves the pairwise-correspondence problem, properly.** The "Rejected: +parsimony-based multiset pairing" case above failed because, with only two +genomes' cardinalities to look at, there is no external constraint to justify +picking one correspondence between leftover alleles over another — `min(a,b)` +is a lower bound dressed up as a point estimate (see the follow-up discussion +on Felsenstein-style parsimony inconsistency: minimum-event explanations are +systematically biased low whenever homoplasy/multiplicity is real, not +noise-cancelling). With `N` genomes and many families jointly, the same +question can be answered the way real phylogenetics answers it: ancestral +state reconstruction / ML mapping over a tree estimated from the whole +column set. The tree supplies the missing constraint that two isolated +columns cannot — this is the principled way out, not a heuristic replacement +for one. + +**Relation to what's already implemented.** `KmerIndex::raw_snp_distance` +already computes, internally, per family, exactly this row — `single_form: +Vec>`, one entry per genome, `None` where ambiguous/absent — +before immediately collapsing it into pairwise `snp[i,j]`/`shared[i,j]` +tallies. The pivot this section proposes is small at the implementation +level: stop collapsing early, and surface the per-family row as a first-class +artifact (a `families x genomes` matrix). Pairwise raw p-distance becomes one +projection of that matrix (what's computed today), not the primary object; +downstream, the matrix itself could feed real phylogenetic tools (parsimony/ +ML, e.g. RAxML/IQ-TREE-style) instead of only NJ/UPGMA on a homemade +pairwise-distance matrix. + +**Caveat: column completeness shrinks with `N`.** The probability that a +family's flanks stay intact simultaneously across all `N` genomes decays with +`N` (same ascertainment-bias mechanism as Bias 1 above, compounded over more +genomes) — fully-resolved columns (no `?` anywhere) become rare as more +genomes are added. Same missing-data situation any real multi-species +alignment faces, and phylogenetic tools already handle it well; the practical +implication is that columns should be allowed partial coverage (>=2 resolved +genomes, not unanimous) rather than requiring every genome to be net +single-copy at that locus. + ## Heterozygosity, ploidy, and consensus-assembly inputs A within-genome multiplicity signal (more than one of the 4 central forms diff --git a/src/Cargo.lock b/src/Cargo.lock index aa92eee..25955a3 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1715,7 +1715,7 @@ dependencies = [ [[package]] name = "obikmer" -version = "1.1.39" +version = "1.1.40" dependencies = [ "clap", "csv", diff --git a/src/obikindex/src/lib.rs b/src/obikindex/src/lib.rs index 4b09bc3..48daf54 100644 --- a/src/obikindex/src/lib.rs +++ b/src/obikindex/src/lib.rs @@ -19,4 +19,4 @@ pub use merge::MergeMode; pub use meta::{validate_label, GenomeInfo, IndexConfig, IndexMeta, META_FILENAME}; pub use state::{IndexState, SENTINEL_COUNTED, SENTINEL_INDEXED, SENTINEL_SCATTERED}; pub use stats::IndexBitsPerKmer; -pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats}; +pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; diff --git a/src/obikindex/src/siblings.rs b/src/obikindex/src/siblings.rs index 6ab50a8..da51270 100644 --- a/src/obikindex/src/siblings.rs +++ b/src/obikindex/src/siblings.rs @@ -781,6 +781,175 @@ impl KmerIndex { } } +/// IUPAC ambiguity code for a per-genome family presence mask (bit `b` set +/// iff the genome carries the member whose own central base is `b`): +/// single bit -> the plain base; 2 or 3 bits -> the matching IUPAC +/// ambiguity code (preserves partial information instead of collapsing to +/// `N`, the same convention used for diploid heterozygous VCF/FASTA sites); +/// all 4 bits -> `N`; no bits (genome carries none of the family's observed +/// members) -> `-` (no data at this locus for this genome). +fn iupac_code(mask: u8) -> u8 { + match mask & 0b1111 { + 0b0000 => b'-', + 0b0001 => b'A', + 0b0010 => b'C', + 0b0100 => b'G', + 0b1000 => b'T', + 0b0101 => b'R', // A/G + 0b1010 => b'Y', // C/T + 0b0110 => b'S', // C/G + 0b1001 => b'W', // A/T + 0b1100 => b'K', // G/T + 0b0011 => b'M', // A/C + 0b1110 => b'B', // C/G/T + 0b1101 => b'D', // A/G/T + 0b1011 => b'H', // A/C/T + 0b0111 => b'V', // A/C/G + 0b1111 => b'N', + _ => unreachable!("masked to 4 bits"), + } +} + +/// A SNP-only pseudo-alignment: one row (byte sequence, IUPAC-coded) per +/// genome, one column per variable family (`family_size() >= 2` — monomorphic +/// families carry no signal and are skipped, unlike `raw_snp_distance`'s +/// tally which does count them as `shared`). Column order is the same, +/// deterministic sweep order as the annex build (partition, then layer, then +/// slot) — arbitrary but stable and identical across genomes, which is all a +/// pseudo-alignment needs (there is no natural genomic coordinate to sort by +/// once flanks are dropped). See `docmd/theory/evolutionary_distances.md`, +/// "Multi-genome framing: family as pseudo-alignment column". +pub struct SnpAlignment { + /// `sequences[g]` = genome `g`'s IUPAC-coded row, same length for every + /// genome (`sequences.len()` columns). + pub sequences: Vec>, +} + +impl KmerIndex { + /// Build the SNP-only pseudo-alignment from an already-built sibling + /// annex (run [`build_sibling_annex`](Self::build_sibling_annex) first). + pub fn snp_pseudo_alignment(&self) -> OKIResult { + let n_parts = self.n_partitions(); + let n_genomes = self.meta.genomes.len(); + let with_counts = self.meta.config.with_counts; + let k = self.kmer_size(); + let n_bits = n_parts.trailing_zeros() as usize; + + let partition = KmerPartition::open_with_config( + &self.root_path, + self.kmer_size(), + self.minimizer_size(), + n_bits, + ) + .map_err(OKIError::Partition)?; + let cache = PartitionCache::build(&partition, n_parts, with_counts)?; + + let mut layer_dirs = Vec::new(); + for part in 0..n_parts { + let index_dir = self.partition().part_dir(part).join(INDEX_SUBDIR); + if !index_dir.exists() { + continue; + } + let meta = PartitionMeta::load(&index_dir).map_err(olm_to_ok)?; + for l in 0..meta.n_layers { + let layer_dir = index_dir.join(format!("layer_{l}")); + let annex_path = layer_dir.join(ANNEX_FILE_NAME); + if !annex_path.exists() { + return Err(OKIError::InvalidInput(format!( + "no sibling annex at {} — run build_sibling_annex first", + annex_path.display() + ))); + } + layer_dirs.push(layer_dir); + } + } + + let pb = progress_bar("snp_pseudo_alignment", layer_dirs.len() as u64, "layers"); + // `Vec>` per layer, one entry (column) per variable family; + // `par_iter().map(...).collect()` on this indexed source preserves + // input order, so concatenating the results below in order gives a + // single deterministic column order across the whole index. + let partials: Vec>> = layer_dirs + .par_iter() + .map(|layer_dir| -> OKIResult>> { + let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); + let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; + let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; + let mphf = MphfLayer::open(layer_dir, &meta.mode).map_err(olm_to_ok)?; + + let mut slot_kmer: Vec> = vec![None; annex.len()]; + let reader = UnitigFileReader::open_sequential(&layer_dir.join("unitigs.bin")) + .map_err(OKIError::Partition)?; + for (kmer, _, _) in reader.iter_indexed_canonical_kmers() { + if let Some(slot) = mphf.find(kmer) { + slot_kmer[slot] = Some(kmer); + } + } + + let use_counts = with_counts && layer_dir.join("counts").exists(); + let mat = if use_counts { + Mat::Count(PersistentCompactIntMatrix::open(layer_dir)?) + } else { + Mat::Presence(PersistentBitMatrix::open(layer_dir)?) + }; + let n_cols = mat.n_cols().min(n_genomes); + + let mut columns: Vec> = Vec::new(); + let mut genome_mask: Vec = Vec::with_capacity(n_genomes); + + for slot in 0..annex.len() { + let Some(mask) = annex.get(slot) else { continue }; + let Some(kmer) = slot_kmer[slot] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant + } + if mask.family_size() < 2 { + continue; // monomorphic family — no signal, skip + } + + genome_mask.clear(); + genome_mask.resize(n_genomes, 0); + + for other in kmer.central_canonical_neighbors() { + let base = central_base(other, k); + if !mask.has(base) { + continue; + } + let presence: Option> = if other == kmer { + Some((0..n_cols).map(|g| mat.carries(g, slot)).collect()) + } else { + let dest = partition_of(other, n_parts); + cache.find_presence(dest, other, n_genomes) + }; + let Some(presence) = presence else { continue }; + for (g, &present) in presence.iter().enumerate() { + if present { + genome_mask[g] |= 1 << base; + } + } + } + + columns.push(genome_mask.iter().map(|&m| iupac_code(m)).collect()); + } + + pb.inc(1); + Ok(columns) + }) + .collect::>>()?; + pb.finish_and_clear(); + + let mut sequences: Vec> = vec![Vec::new(); n_genomes]; + for layer_columns in partials { + for column in layer_columns { + for (g, &code) in column.iter().enumerate() { + sequences[g].push(code); + } + } + } + Ok(SnpAlignment { sequences }) + } +} + #[cfg(test)] mod tests { use std::io::Write; diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 9d00ffd..8c0e42b 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "obikmer" -version = "1.1.39" +version = "1.1.40" edition = "2024" [[bin]] diff --git a/src/obikmer/src/cmd/distance.rs b/src/obikmer/src/cmd/distance.rs index 9e02552..8852478 100644 --- a/src/obikmer/src/cmd/distance.rs +++ b/src/obikmer/src/cmd/distance.rs @@ -3,7 +3,8 @@ use std::path::PathBuf; use clap::Args; use kodama::{Method, linkage}; -use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats}; +use obifastwrite::{JsonVal, write_record}; +use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats, SnpAlignment}; use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick}; use tracing::info; @@ -85,9 +86,17 @@ pub struct DistanceArgs { #[arg(long)] pub raw_snp_distance: bool, + /// Write a SNP-only pseudo-alignment (FASTA, IUPAC-coded) from an + /// already-built sibling annex — one row per genome, one column per + /// variable family (monomorphic families skipped), no flanking + /// sequence. See `docmd/theory/evolutionary_distances.md`, + /// "Multi-genome framing: family as pseudo-alignment column". + #[arg(long)] + pub snp: bool, + /// Output prefix: _dist.csv, _shared.csv, - /// _siblings.csv, _rawsnp.csv, _nj.nwk, - /// _upgma.nwk. + /// _siblings.csv, _rawsnp.csv, _snp.fasta, + /// _nj.nwk, _upgma.nwk. /// If omitted, the distance matrix is written to stdout. #[arg(short, long)] pub output: Option, @@ -128,15 +137,22 @@ pub fn run(args: DistanceArgs) { }); write_raw_snp_distance_csv(&result, &labels, &args.output); } + if args.snp { + let alignment = idx.snp_pseudo_alignment().unwrap_or_else(|e| { + eprintln!("error computing SNP pseudo-alignment: {e}"); + std::process::exit(1); + }); + write_snp_fasta(&alignment, &labels, &args.output); + } - // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance` are their own - // operation, not a modifier on top of a distance-metric computation — a - // metric was never requested by asking for any of them, so there is - // nothing for the rest of this function to compute. Not a historical - // accident to keep: stop here rather than always also running a Jaccard - // (or whichever `--metric` defaults to) pass and printing an unrequested - // matrix. - if args.sibling_annex || args.sibling_stats || args.raw_snp_distance { + // `--sibling-annex`/`--sibling-stats`/`--raw-snp-distance`/`--snp` are + // their own operation, not a modifier on top of a distance-metric + // computation — a metric was never requested by asking for any of them, + // so there is nothing for the rest of this function to compute. Not a + // historical accident to keep: stop here rather than always also + // running a Jaccard (or whichever `--metric` defaults to) pass and + // printing an unrequested matrix. + if args.sibling_annex || args.sibling_stats || args.raw_snp_distance || args.snp { return; } @@ -322,6 +338,32 @@ fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], info!("raw single-copy SNP distance matrix → {path}"); } +// ── SNP-only pseudo-alignment → FASTA ─────────────────────────────────────── +// +// One record per genome, IUPAC-coded, no flanking sequence — see +// `SnpAlignment` / `KmerIndex::snp_pseudo_alignment`. Uses the project's +// existing FASTA writer (`obifastwrite::write_record`) rather than +// hand-rolling one. + +fn write_snp_fasta(alignment: &SnpAlignment, labels: &[String], output: &Option) { + let path = output.as_ref() + .map(|p| format!("{}_snp.fasta", p.display())) + .unwrap_or_else(|| "snp.fasta".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0); + for (label, seq) in labels.iter().zip(alignment.sequences.iter()) { + write_record(seq, label, &[("n_sites", JsonVal::Num(n_sites as u64))], &mut f).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + } + info!("SNP pseudo-alignment → {path} ({n_sites} site{})", + if n_sites == 1 { "" } else { "s" }); +} + // ── UPGMA Newick from kodama dendrogram ─────────────────────────────────────── fn upgma_to_newick(dendro: &kodama::Dendrogram, names: &[String]) -> String {