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;