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.

src/Cargo.lock

@@ -1715,7 +1715,7 @@ dependencies = [
 
 [[package]]
 name = "obikmer"
-version = "1.1.39"
+version = "1.1.40"
 dependencies = [
  "clap",
  "csv",

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};

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<Vec<u8>>,
+}
+
+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<SnpAlignment> {
+        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<Vec<u8>>` 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<Vec<Vec<u8>>> = layer_dirs
+            .par_iter()
+            .map(|layer_dir| -> OKIResult<Vec<Vec<u8>>> {
+                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<Option<CanonicalKmer>> = 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<u8>> = Vec::new();
+                let mut genome_mask: Vec<u8> = 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<Vec<bool>> = 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::<OKIResult<Vec<_>>>()?;
+        pb.finish_and_clear();
+
+        let mut sequences: Vec<Vec<u8>> = 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;

src/obikmer/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "obikmer"
-version = "1.1.39"
+version = "1.1.40"
 edition = "2024"
 
 [[bin]]

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: <prefix>_dist.csv, <prefix>_shared.csv,
-    /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_nj.nwk,
-    /// <prefix>_upgma.nwk.
+    /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_snp.fasta,
+    /// <prefix>_nj.nwk, <prefix>_upgma.nwk.
     /// If omitted, the distance matrix is written to stdout.
     #[arg(short, long)]
     pub output: Option<PathBuf>,
@@ -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<PathBuf>) {
+    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<f64>, names: &[String]) -> String {