feat: add raw SNP distance calculation and CLI flag

Exposes RawSnpDistanceOutput and implements KmerIndex::raw_snp_distance() to compute pairwise single-copy locus counts under a paralogy-aware rule. The implementation leverages ndarray for parallel matrix aggregation, producing raw p-distance matrices for sanity-checking. A --raw-snp-distance CLI flag is added to export results as CSV, mapping zero-eligible pairs to NA.

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::SiblingAnnexStats;
+pub use siblings::{RawSnpDistanceOutput, SiblingAnnexStats};

src/obikindex/src/siblings.rs

@@ -31,6 +31,7 @@ use std::path::Path;
 use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 
+use ndarray::Array2;
 use rayon::prelude::*;
 
 use obicompactvec::{
@@ -609,6 +610,177 @@ impl KmerIndex {
     }
 }
 
+/// Raw p-distance restricted to loci that are single-copy in **both**
+/// genomes of a pair — the "stringent / paralogy-aware" locus eligibility
+/// rule (`docmd/theory/evolutionary_distances.md`, "Locus eligibility"),
+/// without the JC/K2P/LogDet finalisation math: `snp[i,j]` + `shared[i,j]`
+/// is the number of eligible loci, `snp[i,j] / (snp[i,j] + shared[i,j])` is
+/// `p_hat`. A quick, self-contained way to sanity-check the estimator
+/// against a real index before the full `SnpTally` design is built.
+///
+/// A locus (family, tallied once at its minorant) is eligible for pair
+/// `(i, j)` iff genome `i` carries exactly one of the family's observed
+/// forms **and** genome `j` carries exactly one (possibly a different one)
+/// — presence-only: a genome carrying the same form twice (a same-allele
+/// duplicate) is indistinguishable from carrying it once when only a
+/// presence matrix is available, so such cases are not excluded here even
+/// when a count index exists. See "Locus eligibility", stringent rule, for
+/// why this matters and how a count index would close the gap — left as a
+/// follow-up, not applied here.
+pub struct RawSnpDistanceOutput {
+    /// n×n count of eligible loci where the two genomes' single forms differ.
+    pub snp: Array2<u64>,
+    /// n×n count of eligible loci where the two genomes' single forms agree.
+    pub shared: Array2<u64>,
+}
+
+impl KmerIndex {
+    /// Compute [`RawSnpDistanceOutput`] from an already-built sibling annex
+    /// (run [`build_sibling_annex`](Self::build_sibling_annex) first).
+    pub fn raw_snp_distance(&self) -> OKIResult<RawSnpDistanceOutput> {
+        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("raw_snp_distance", layer_dirs.len() as u64, "layers");
+        let partials: Vec<(Array2<u64>, Array2<u64>)> = layer_dirs
+            .par_iter()
+            .map(|layer_dir| -> OKIResult<(Array2<u64>, Array2<u64>)> {
+                let mut snp = Array2::<u64>::zeros((n_genomes, n_genomes));
+                let mut shared = Array2::<u64>::zeros((n_genomes, n_genomes));
+
+                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);
+
+                // Per family: which single form (if exactly one) each genome
+                // carries — `None` once a second form is seen (ambiguous,
+                // not single-copy, ineligible for either side of a pair).
+                let mut single_form: Vec<Option<u8>> = Vec::with_capacity(n_cols);
+                let mut ambiguous: Vec<bool> = Vec::with_capacity(n_cols);
+
+                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
+                    }
+
+                    single_form.clear();
+                    single_form.resize(n_cols, None);
+                    ambiguous.clear();
+                    ambiguous.resize(n_cols, false);
+
+                    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 {
+                                continue;
+                            }
+                            if single_form[g].is_some() {
+                                ambiguous[g] = true;
+                            } else {
+                                single_form[g] = Some(base);
+                            }
+                        }
+                    }
+
+                    for i in 0..n_cols {
+                        if ambiguous[i] {
+                            continue;
+                        }
+                        let Some(bi) = single_form[i] else { continue };
+                        for j in (i + 1)..n_cols {
+                            if ambiguous[j] {
+                                continue;
+                            }
+                            let Some(bj) = single_form[j] else { continue };
+                            if bi == bj {
+                                shared[[i, j]] += 1;
+                                shared[[j, i]] += 1;
+                            } else {
+                                snp[[i, j]] += 1;
+                                snp[[j, i]] += 1;
+                            }
+                        }
+                    }
+                }
+
+                pb.inc(1);
+                Ok((snp, shared))
+            })
+            .collect::<OKIResult<Vec<_>>>()?;
+        pb.finish_and_clear();
+
+        let mut snp = Array2::<u64>::zeros((n_genomes, n_genomes));
+        let mut shared = Array2::<u64>::zeros((n_genomes, n_genomes));
+        for (s, sh) in partials {
+            snp += &s;
+            shared += &sh;
+        }
+        Ok(RawSnpDistanceOutput { snp, shared })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::io::Write;

src/obikmer/src/cmd/distance.rs

@@ -3,7 +3,7 @@ use std::path::PathBuf;
 
 use clap::Args;
 use kodama::{Method, linkage};
-use obikindex::{DistanceMetric, KmerIndex, SiblingAnnexStats};
+use obikindex::{DistanceMetric, KmerIndex, RawSnpDistanceOutput, SiblingAnnexStats};
 use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver, to_newick};
 use tracing::info;
 
@@ -77,8 +77,17 @@ pub struct DistanceArgs {
     #[arg(long)]
     pub sibling_stats: bool,
 
+    /// Compute the raw p-distance restricted to loci that are single-copy
+    /// in both genomes of each pair (an already-built sibling annex is
+    /// required — run with `--sibling-annex` first, in this invocation or
+    /// an earlier one). A quick way to test the central-position SNP
+    /// estimator against a real index; not the full `SnpTally` design.
+    #[arg(long)]
+    pub raw_snp_distance: bool,
+
     /// Output prefix: <prefix>_dist.csv, <prefix>_shared.csv,
-    /// <prefix>_siblings.csv, <prefix>_nj.nwk, <prefix>_upgma.nwk.
+    /// <prefix>_siblings.csv, <prefix>_rawsnp.csv, <prefix>_nj.nwk,
+    /// <prefix>_upgma.nwk.
     /// If omitted, the distance matrix is written to stdout.
     #[arg(short, long)]
     pub output: Option<PathBuf>,
@@ -112,15 +121,22 @@ pub fn run(args: DistanceArgs) {
         });
         write_sibling_stats_csv(&stats, &labels, &args.output);
     }
+    if args.raw_snp_distance {
+        let result = idx.raw_snp_distance().unwrap_or_else(|e| {
+            eprintln!("error computing raw SNP distance: {e}");
+            std::process::exit(1);
+        });
+        write_raw_snp_distance_csv(&result, &labels, &args.output);
+    }
 
-    // `--sibling-annex`/`--sibling-stats` are their own operation, not a
-    // modifier on top of a distance-metric computation — a metric was
-    // never requested by asking for either 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
+    // `--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 {
+    if args.sibling_annex || args.sibling_stats || args.raw_snp_distance {
         return;
     }
 
@@ -269,6 +285,43 @@ fn write_sibling_stats_csv(stats: &SiblingAnnexStats, labels: &[String], output:
         if total == 1 { "y" } else { "ies" });
 }
 
+// ── Raw single-copy SNP distance → CSV ──────────────────────────────────────
+//
+// p_hat[i,j] = snp[i,j] / (snp[i,j] + shared[i,j]) over loci single-copy in
+// both i and j — see `RawSnpDistanceOutput` / `KmerIndex::raw_snp_distance`.
+// A single file: the distance matrix, with an eligible-loci count alongside
+// each value so a 0/0 pair (no eligible locus at all) is distinguishable
+// from a genuinely identical pair.
+
+fn write_raw_snp_distance_csv(result: &RawSnpDistanceOutput, labels: &[String], output: &Option<PathBuf>) {
+    let path = output.as_ref()
+        .map(|p| format!("{}_rawsnp.csv", p.display()))
+        .unwrap_or_else(|| "rawsnp.csv".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 = labels.len();
+    write!(f, "genome").unwrap();
+    for g in labels { write!(f, ",{g}").unwrap(); }
+    writeln!(f).unwrap();
+    for (i, g) in labels.iter().enumerate() {
+        write!(f, "{g}").unwrap();
+        for j in 0..n {
+            let snp = result.snp[[i, j]];
+            let shared = result.shared[[i, j]];
+            let eligible = snp + shared;
+            if eligible == 0 {
+                write!(f, ",NA").unwrap();
+            } else {
+                write!(f, ",{:.6}", snp as f64 / eligible as f64).unwrap();
+            }
+        }
+        writeln!(f).unwrap();
+    }
+    info!("raw single-copy SNP distance matrix → {path}");
+}
+
 // ── UPGMA Newick from kodama dendrogram ───────────────────────────────────────
 
 fn upgma_to_newick(dendro: &kodama::Dendrogram<f64>, names: &[String]) -> String {