Implement SNP distance models with gamma correction and PHYLIP output

Renames the CLI flag from --metric to --distance and introduces eight closed-form SNP distance models with optional Jin-Nei gamma correction. Integrates the ndarray crate for matrix operations and adds relaxed PHYLIP output formatting. Updates architecture and theory documentation to cover the new sparse matrix variants, algorithmic fixes, and distance metric implementations.
This commit is contained in:
Eric Coissac
2026-08-28 21:54:19 +02:00
parent 0b40d2d0da
commit 4f34a646c5
17 changed files with 2199 additions and 73 deletions
+1
View File
@@ -1670,6 +1670,7 @@ version = "1.2.2"
dependencies = [
"clap",
"csv",
"ndarray",
"obifastwrite",
"obikalgorithm",
"obikdump",
+1
View File
@@ -29,6 +29,7 @@ obifastwrite = { path = "../obifastwrite" }
obiskbuilder = { path = "../obiskbuilder" }
clap = { version = "4", features = ["derive"] }
csv = "1"
ndarray = "0.17"
serde = { version = "1", features = ["derive"] }
serde_yaml = "0.9"
tracing = "0.1.44"
+104 -31
View File
@@ -1,10 +1,16 @@
use std::path::PathBuf;
use clap::Args;
use obikphylo::DistanceMetric;
use obikphylo::{DistanceMetric, SnpDistanceKind};
/// `--distance` value — either one of `obikphylo::DistanceMetric`'s
/// whole-index metrics (routed to `IndexCache::distance`) or one of
/// `obikphylo::SnpDistanceKind`'s `snp-*` corrections (routed to
/// `SiblingExt::snp_distance`, the sibling-annex pipeline) — two genuinely
/// different code paths behind one CLI vocabulary, see
/// `DevDocMD/theory/evolutionary_distances.md`, "`--distance` unification".
#[derive(clap::ValueEnum, Clone, Copy, Debug)]
pub enum MetricArg {
pub enum DistanceArg {
Jaccard,
Mash,
Hamming,
@@ -17,35 +23,69 @@ pub enum MetricArg {
Hellinger,
#[value(name = "hellinger-euclidean")]
HellingerEuclidean,
#[value(name = "snp-raw")]
SnpRaw,
#[value(name = "snp-jc")]
SnpJc,
#[value(name = "snp-k2p")]
SnpK2p,
#[value(name = "snp-k81")]
SnpK81,
#[value(name = "snp-f81")]
SnpF81,
#[value(name = "snp-t92")]
SnpT92,
#[value(name = "snp-tn93")]
SnpTn93,
#[value(name = "snp-tv")]
SnpTv,
}
impl From<MetricArg> for DistanceMetric {
fn from(m: MetricArg) -> Self {
match m {
MetricArg::Jaccard => DistanceMetric::Jaccard,
MetricArg::Mash => DistanceMetric::Mash,
MetricArg::Hamming => DistanceMetric::Hamming,
MetricArg::BrayCurtis => DistanceMetric::BrayCurtis,
MetricArg::RelfreqBrayCurtis => DistanceMetric::RelfreqBrayCurtis,
MetricArg::Euclidean => DistanceMetric::Euclidean,
MetricArg::RelfreqEuclidean => DistanceMetric::RelfreqEuclidean,
MetricArg::Hellinger => DistanceMetric::Hellinger,
MetricArg::HellingerEuclidean => DistanceMetric::HellingerEuclidean,
}
impl DistanceArg {
/// `Some` for the whole-index metrics, `None` for `snp-*` values.
pub fn as_classic(self) -> Option<DistanceMetric> {
Some(match self {
DistanceArg::Jaccard => DistanceMetric::Jaccard,
DistanceArg::Mash => DistanceMetric::Mash,
DistanceArg::Hamming => DistanceMetric::Hamming,
DistanceArg::BrayCurtis => DistanceMetric::BrayCurtis,
DistanceArg::RelfreqBrayCurtis => DistanceMetric::RelfreqBrayCurtis,
DistanceArg::Euclidean => DistanceMetric::Euclidean,
DistanceArg::RelfreqEuclidean => DistanceMetric::RelfreqEuclidean,
DistanceArg::Hellinger => DistanceMetric::Hellinger,
DistanceArg::HellingerEuclidean => DistanceMetric::HellingerEuclidean,
_ => return None,
})
}
/// `Some` for the `snp-*` values, `None` for the whole-index metrics.
pub fn as_snp(self) -> Option<SnpDistanceKind> {
Some(match self {
DistanceArg::SnpRaw => SnpDistanceKind::Raw,
DistanceArg::SnpJc => SnpDistanceKind::Jc,
DistanceArg::SnpK2p => SnpDistanceKind::K2p,
DistanceArg::SnpK81 => SnpDistanceKind::K81,
DistanceArg::SnpF81 => SnpDistanceKind::F81,
DistanceArg::SnpT92 => SnpDistanceKind::T92,
DistanceArg::SnpTn93 => SnpDistanceKind::Tn93,
DistanceArg::SnpTv => SnpDistanceKind::Tv,
_ => return None,
})
}
}
/// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric
/// path (`--metric`/NJ/UPGMA), annex construction (`--sibling-annex`),
/// annex diagnostics (`--sibling-stats`, `--sibling-hist`), entropy
/// reporting (`--shannon`), SNP pseudo-alignment sampling
/// (`--pseudo-alignment`, `--subsample`, `--free-loss`, `--no-ambiguity`,
/// `--entropy`/`--entropy-sd`), Sankoff cost-matrix calibration
/// (`--sankoff`, `--sankoff-ratio-ceiling`) and its TNT/PhyG/IQ-TREE exports
/// (`--tnt`, `--phyg`, `--iqtree`/`--iqtree-min-freq`,
/// `--sankoff-cost-scale`) — everything else sibling-annex-based (raw SNP
/// distance, family overlap, ...) stays in `obikmer` until the rest of
/// `obikphylo::siblings` is reconnected (see the project memory on this).
/// Partial transfer of `obikmer`'s `phylo` command: the whole-index
/// `--distance` path (classic metrics + `snp-*` corrections/NJ/UPGMA),
/// annex construction (`--sibling-annex`), annex diagnostics
/// (`--sibling-stats`, `--sibling-hist`), entropy reporting (`--shannon`),
/// SNP pseudo-alignment sampling (`--pseudo-alignment`, `--subsample`,
/// `--free-loss`, `--no-ambiguity`, `--entropy`/`--entropy-sd`), Sankoff
/// cost-matrix calibration (`--sankoff`, `--sankoff-ratio-ceiling`) and its
/// TNT/PhyG/IQ-TREE exports (`--tnt`, `--phyg`, `--iqtree`/
/// `--iqtree-min-freq`, `--sankoff-cost-scale`) — everything else
/// sibling-annex-based (family overlap, ...) stays in `obikmer` until the
/// rest of `obikphylo::siblings` is reconnected (see the project memory on
/// this).
#[derive(Args)]
pub struct PhyloArgs {
/// Index directory
@@ -97,9 +137,13 @@ pub struct PhyloArgs {
pub pseudo_alignment: bool,
/// Target number of variable sites to sample index-wide for
/// `--pseudo-alignment` — a target, not a guarantee (proportional
/// per-layer sampling; see `obikphylo::siblings::SiblingExt::snp_pseudo_alignment`'s
/// own docs).
/// `--pseudo-alignment`/`--sankoff` (mandatory for both) and for a
/// `snp-*` `--distance` value (optional there: omitted means exhaustive
/// — every non-monomorphic minorant of the whole index, not an
/// approximation, see `obikphylo::siblings::SiblingExt::snp_distance`'s
/// own docs) — a target, not a guarantee when given (proportional
/// per-layer sampling; see
/// `obikphylo::siblings::SiblingExt::snp_pseudo_alignment`'s own docs).
#[arg(long)]
pub subsample: Option<usize>,
@@ -205,14 +249,43 @@ pub struct PhyloArgs {
#[arg(long, default_value = "100")]
pub sankoff_cost_scale: f64,
/// Distance metric to compute
/// Distance to compute — either a whole-index metric (`jaccard`,
/// `mash`, `hamming`, `bray-curtis`, ...) or a `snp-*` correction over
/// the central-position SNP substitution spectrum (`snp-raw`, `snp-jc`,
/// `snp-k2p`, `snp-k81`, `snp-f81`, `snp-t92`, `snp-tn93`, `snp-tv`) —
/// the latter route to a different computation entirely
/// (`SiblingExt::snp_distance`, requires `--sibling-annex` first; see
/// `DevDocMD/theory/evolutionary_distances.md`, "`--distance`
/// unification" for the full catalog and why LogDet/Tajima-Nei/F84/
/// HKY85 aren't offered yet).
#[arg(long, value_enum, default_value = "jaccard")]
pub metric: MetricArg,
pub distance: DistanceArg,
/// Rate-heterogeneity correction (Jin-Nei gamma shape parameter `α`)
/// for `snp-*` `--distance` values that support it
/// (`obikphylo::SnpDistanceKind::supports_gamma`: every one except
/// `snp-raw`/`snp-tv`, which have nothing to correct/are deliberately
/// uncorrected). Has no effect on the whole-index metrics. Rejected at
/// runtime if given alongside an unsupported `--distance` value.
#[arg(long, value_name = "ALPHA")]
pub gamma_shape: Option<f64>,
/// Minimum count to consider a kmer present when computing Jaccard on count indexes
#[arg(long, default_value = "1")]
pub presence_threshold: u32,
/// Write the primary distance matrix as plain CSV instead of the
/// default relaxed-PHYLIP format (`n` on the first line, then one
/// `label<TAB>value...` row per genome — no 10-character label
/// truncation, unlike strict PHYLIP, not yet offered here). PHYLIP is
/// the default because it's what external NJ tools (PHYLIP `neighbor`,
/// FastME, T-REX, SplitsTree) actually read; CSV stays available for
/// scripting/inspection. Only affects the primary distance matrix —
/// `--shared-kmers` keeps its own CSV-only format regardless of this
/// flag.
#[arg(long)]
pub csv: bool,
/// Also output the shared-kmer count matrix (CSV)
#[arg(long)]
pub shared_kmers: bool,
+74 -29
View File
@@ -1,6 +1,7 @@
mod args;
mod iqtree;
mod phyg;
mod phylip;
mod sankoff;
mod tnt;
@@ -19,6 +20,7 @@ use tracing::info;
use iqtree::write_iqtree;
use phyg::write_sankoff_phyg;
use phylip::write_phylip_relaxed;
use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params};
use tnt::write_sankoff_tnt;
@@ -271,57 +273,100 @@ pub fn run(args: PhyloArgs) {
}
}
info!("computing {:?} distances for {} genome(s)", args.metric, n);
// ── Distance computation: classic whole-index metric vs. `snp-*` ───────────
// Two genuinely different code paths behind one `--distance` value — see
// `args::DistanceArg`'s own docs.
let (matrix, shared_kmers) = match args.distance.as_classic() {
Some(metric) => {
info!("computing {metric:?} distances for {n} genome(s)");
let need_shared = args.shared_kmers || args.nj || args.upgma;
let t = Stage::start("distance");
let result = cache
.distance(metric, need_shared, args.presence_threshold)
.unwrap_or_else(|e| {
eprintln!("error computing distances: {e}");
std::process::exit(1);
});
rep.push(t.stop());
(result.matrix, result.shared_kmers)
}
None => {
if args.shared_kmers {
eprintln!("error: --shared-kmers has no meaning for a snp-* --distance value");
std::process::exit(1);
}
let kind = args.distance.as_snp().expect("DistanceArg is always classic or snp");
info!(
"computing {kind:?} SNP distance for {n} genome(s){}",
match args.subsample {
Some(n) => format!(" (subsampled, target {n} site(s))"),
None => " (exhaustive)".into(),
}
);
let t = Stage::start("snp_distance");
let matrix = cache
.snp_distance(
kind,
args.subsample,
args.free_loss,
args.no_ambiguity,
&exclude_mask,
entropy_bias,
args.gamma_shape,
)
.unwrap_or_else(|e| {
eprintln!("error computing SNP distance: {e}");
std::process::exit(1);
});
rep.push(t.stop());
(matrix, None)
}
};
let need_shared = args.shared_kmers || args.nj || args.upgma;
let t = Stage::start("distance");
let result = cache
.distance(args.metric.into(), need_shared, args.presence_threshold)
.unwrap_or_else(|e| {
eprintln!("error computing distances: {e}");
std::process::exit(1);
});
rep.push(t.stop());
// Rows/columns kept in every matrix CSV below — `cache.distance(...)`
// above is computed over every genome regardless; only the writers skip
// Rows/columns kept in every matrix output below — the computation
// above runs over every genome regardless; only the writers skip
// excluded ones.
let kept: Vec<usize> = (0..n).filter(|&i| !exclude_mask[i]).collect();
// ── Distance matrix → CSV ─────────────────────────────────────────────────
let write_dist_csv = |w: &mut dyn Write| {
write!(w, "genome").unwrap();
for &j in &kept { write!(w, ",{}", labels[j]).unwrap(); }
writeln!(w).unwrap();
for &i in &kept {
write!(w, "{}", labels[i]).unwrap();
for &j in &kept {
write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap();
}
// ── Distance matrix → relaxed PHYLIP (default) or CSV (`--csv`) ────────────
let write_dist = |w: &mut dyn Write| {
if args.csv {
write!(w, "genome").unwrap();
for &j in &kept { write!(w, ",{}", labels[j]).unwrap(); }
writeln!(w).unwrap();
for &i in &kept {
write!(w, "{}", labels[i]).unwrap();
for &j in &kept {
write!(w, ",{:.6}", matrix[[i, j]]).unwrap();
}
writeln!(w).unwrap();
}
} else {
write_phylip_relaxed(w, &labels, &kept, &matrix);
}
};
match &args.output {
Some(prefix) => {
let path = format!("{}_dist.csv", prefix.display());
let suffix = if args.csv { "_dist.csv" } else { "_dist.phy" };
let path = format!("{}{suffix}", prefix.display());
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
write_dist_csv(&mut f);
write_dist(&mut f);
info!("distance matrix → {path}");
}
None => {
let stdout = io::stdout();
let mut out = BufWriter::new(stdout.lock());
write_dist_csv(&mut out);
write_dist(&mut out);
}
}
// ── Shared-kmer matrix → CSV ──────────────────────────────────────────────
if args.shared_kmers {
if let Some(shared) = &result.shared_kmers {
if let Some(shared) = &shared_kmers {
let path = args.output.as_ref()
.map(|p| format!("{}_shared.csv", p.display()))
.unwrap_or_else(|| "shared.csv".into());
@@ -343,7 +388,7 @@ pub fn run(args: PhyloArgs) {
// ── NJ tree ────────────────────────────────────────────────────────────────
if args.nj {
let tree = neighbor_joining(&result.matrix, &labels).unwrap_or_else(|e| {
let tree = neighbor_joining(&matrix, &labels).unwrap_or_else(|e| {
eprintln!("error computing NJ tree: {e}");
std::process::exit(1);
});
@@ -360,7 +405,7 @@ pub fn run(args: PhyloArgs) {
// ── UPGMA tree ───────────────────────────────────────────────────────────────
if args.upgma {
let newick = upgma(&result.matrix, &labels).to_newick();
let newick = upgma(&matrix, &labels).to_newick();
let path = args.output.as_ref()
.map(|p| format!("{}_upgma.nwk", p.display()))
.unwrap_or_else(|| "upgma.nwk".into());
+22
View File
@@ -0,0 +1,22 @@
//! Relaxed-PHYLIP distance-matrix writer — the default output format for
//! `--distance` (see `args::PhyloArgs::csv`'s own docs for why): `n` on the
//! first line, then one `label<TAB>value...` row per genome. "Relaxed"
//! (unlike strict PHYLIP) means no 10-character label truncation/padding —
//! just whitespace-separated fields, which every external NJ tool this
//! targets (PHYLIP `neighbor`, FastME, T-REX, SplitsTree) already reads,
//! and genome labels here routinely exceed strict PHYLIP's 10 characters.
use std::io::Write;
use ndarray::Array2;
pub(super) fn write_phylip_relaxed(w: &mut dyn Write, labels: &[String], indices: &[usize], matrix: &Array2<f64>) {
writeln!(w, "{}", indices.len()).unwrap();
for &i in indices {
write!(w, "{}", labels[i]).unwrap();
for &j in indices {
write!(w, "\t{:.6}", matrix[[i, j]]).unwrap();
}
writeln!(w).unwrap();
}
}
+1
View File
@@ -19,4 +19,5 @@ mod tree;
pub mod siblings;
pub use distance::{DistanceMetric, DistanceOutput, Metrics};
pub use siblings::SnpDistanceKind;
pub use tree::{Tree, neighbor_joining, upgma};
@@ -22,6 +22,7 @@ mod masking;
mod minorant_selection;
mod pairwise;
mod sankoff;
mod snp_distance;
mod stats;
mod subsample;
@@ -33,6 +34,7 @@ pub use cardcomp::{
};
pub use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput};
pub use sankoff::SankoffBundle;
pub use snp_distance::SnpDistanceKind;
pub use stats::SiblingAnnexStats;
pub use subsample::{EntropyBias, SurvivingFamily};
@@ -41,6 +43,7 @@ pub(crate) use annex::build_layer_sibling_annex;
pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4};
pub(crate) use family_scan::{Selection, scan_layer_families};
pub(crate) use sankoff::sankoff_bundle;
pub(crate) use snp_distance::snp_distance;
pub(crate) use stats::{sibling_annex_stats, sibling_family_size_histogram};
/// Whether every layer number in `cache` fits in a `FamilyMask` field
@@ -175,6 +175,66 @@ impl PairwiseTally {
}
CardinalityTally { counts }
}
/// Categorised substitution counts for pair `(i, j)`, plus its eligible
/// ("shared") locus count — the sufficient statistic every closed-form
/// `algorithms::snp_distance` correction is built from. Base order
/// `0=A,1=C,2=G,3=T` (same convention as `FamilyMask`/`STATE_SYMBOL`).
/// `ts1`/`ts2` split transitions by purine (A↔G) vs pyrimidine (C↔T)
/// pair — required for TN93, collapse to a single `ts1+ts2` for
/// K80/K81/T92, which don't distinguish them. `tv1`/`tv2` split
/// transversions the way Kimura's 3-parameter model does (A↔C/G↔T vs
/// A↔T/C↔G) — collapse to `tv1+tv2` for every model that doesn't need
/// the distinction (K80, F81, T92, TN93).
pub(crate) fn categories(&self, i: usize, j: usize) -> PairCategories {
let stats = self.pair(i, j);
PairCategories {
ts1: stats.subst[0][2],
ts2: stats.subst[1][3],
tv1: stats.subst[0][1] + stats.subst[2][3],
tv2: stats.subst[0][3] + stats.subst[1][2],
shared: stats.same_all.iter().sum(),
}
}
/// Pooled base composition for pair `(i, j)`, estimated from both
/// genomes' observed calls at their `shared`/differing eligible loci —
/// `2 * same_all[a]` (each agreement locus contributes `a` to both
/// genomes) plus `Σ_b subst[a][b]` (each differing locus contributes
/// `a` to whichever of the two genomes carried it — `subst` is
/// symmetric by construction, so summing one row already counts each
/// such event exactly once, see `reduce_pairwise`'s own docs), over
/// `2 * n_eligible_loci` total base observations. Feeds F81/T92/TN93,
/// which all correct for base-composition bias using exactly this kind
/// of pooled empirical frequency (Nei & Kumar's standard estimator for
/// a pairwise comparison, not a whole-index average).
pub(crate) fn base_freq(&self, i: usize, j: usize) -> [f64; 4] {
let stats = self.pair(i, j);
let mut counts = [0u64; 4];
for a in 0..4 {
counts[a] = 2 * stats.same_all[a] + (0..4).map(|b| stats.subst[a][b]).sum::<u64>();
}
let total: u64 = counts.iter().sum();
if total == 0 {
return [0.25; 4];
}
counts.map(|c| c as f64 / total as f64)
}
}
/// See [`PairwiseTally::categories`]'s own docs.
pub(crate) struct PairCategories {
pub ts1: u64,
pub ts2: u64,
pub tv1: u64,
pub tv2: u64,
pub shared: u64,
}
impl PairCategories {
pub(crate) fn n_eligible(&self) -> u64 {
self.ts1 + self.ts2 + self.tv1 + self.tv2 + self.shared
}
}
/// Raw p-distance restricted to loci that are single-copy in **both**
@@ -0,0 +1,398 @@
//! `snp-*` `--distance` values — closed-form (method-of-moments) pairwise
//! corrections over the central-position SNP substitution spectrum, all
//! derived from one shared [`PairwiseTally`] (see its own module docs: a
//! single `O(n²)` structure, built once, from which every correction below
//! is cheap post-processing, never a second index scan). See
//! `DevDocMD/theory/evolutionary_distances.md`, "`--distance` unification"
//! for the design discussion and the literature this implements
//! ([`ape`](https://github.com/emmanuelparadis/ape)'s `dist.dna` C source,
//! `src/dist_dna.c`, verified formula-by-formula against the upstream
//! implementation rather than re-derived from memory).
//!
//! Deliberately **not** included here: LogDet/paralinear (needs the true
//! *directional* per-pair base co-occurrence matrix — `PairwiseTally`
//! only keeps the symmetrised substitution counts `BasePairTally` itself
//! wants, which loses exactly the compositional-asymmetry information
//! LogDet exists to detect), Tajima-Nei (needs each genome's *own* base
//! composition, not the pair-pooled estimate `base_freq` below provides),
//! F84 and the ML-fit HKY85 tree usage (no formula independently verified
//! against a primary source at implementation time). Adding any of these
//! later is a new function in this module plus, for LogDet/Tajima-Nei, a
//! new field on `PairStats`/a per-genome accumulator — not an architecture
//! change.
use ndarray::Array2;
use obikidxcache::index_cache::IndexCache;
use obikindex::OKIResult;
use super::pairwise::PairwiseTally;
use super::sibling_family_size_histogram;
use super::subsample::{EntropyBias, sample_index};
/// One `snp-*` `--distance` value. `pub`: part of
/// [`crate::siblings::extensions::SiblingExt::snp_distance`]'s public
/// signature.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SnpDistanceKind {
/// Uncorrected p-distance (`snp / (snp + shared)`).
Raw,
/// Jukes-Cantor (JC69): corrects for multiple substitutions per site,
/// equal rates, equal base frequencies.
Jc,
/// Kimura 2-parameter (K80): + transition/transversion rate bias.
K2p,
/// Kimura 3-parameter (K81/K3ST): splits transversions into the two
/// categories A↔C/G↔T and A↔T/C↔G, each its own rate.
K81,
/// Felsenstein 81 (F81): JC69 + unequal (pair-pooled empirical) base
/// frequencies, no ts/tv distinction.
F81,
/// Tamura 3-parameter (T92): K80 + GC-content bias.
T92,
/// Tamura-Nei (TN93): unequal base frequencies + separate purine
/// (A↔G) / pyrimidine (C↔T) transition rates + a transversion rate —
/// the richest closed-form correction implemented here.
Tn93,
/// Transversions-only p-distance — diagnostic/deep-divergence variant,
/// deliberately uncorrected (dropping transitions, which saturate
/// first, is itself the correction).
Tv,
}
impl SnpDistanceKind {
/// Whether `--gamma-shape` applies to this correction — every
/// rate-based model here except the two with no "-ln(x)" term to
/// reinterpret as a gamma mixture (`Raw`, nothing to correct; `Tv`,
/// deliberately uncorrected).
pub fn supports_gamma(self) -> bool {
!matches!(self, SnpDistanceKind::Raw | SnpDistanceKind::Tv)
}
}
/// `-ln(x)`, gamma-mixture corrected when `alpha` is given: replaces the
/// single-rate `-ln(x)` term with the standard Jin-Nei (1990) gamma
/// substitution `alpha * (x^(-1/alpha) - 1)` — the mechanical term-wise
/// rewrite every gamma-corrected distance in the literature (JC+Γ, K80+Γ,
/// F81+Γ, ...) applies to each log term of its base formula; verified here
/// against `ape`'s own JC69/K80/F81 gamma branches (`dist_dna.c`), then
/// generalised the same way to every other log term this module uses
/// (K81/T92/TN93 each being multi-term sums of exactly this shape).
fn corrected_log(x: f64, alpha: Option<f64>) -> f64 {
match alpha {
Some(alpha) => alpha * (x.powf(-1.0 / alpha) - 1.0),
None => -x.ln(),
}
}
pub(crate) fn raw(tally: &PairwiseTally, i: usize, j: usize) -> f64 {
let c = tally.categories(i, j);
let l = c.n_eligible();
if l == 0 {
return f64::NAN;
}
(c.ts1 + c.ts2 + c.tv1 + c.tv2) as f64 / l as f64
}
pub(crate) fn tv_only(tally: &PairwiseTally, i: usize, j: usize) -> f64 {
let c = tally.categories(i, j);
let l = c.n_eligible();
if l == 0 {
return f64::NAN;
}
(c.tv1 + c.tv2) as f64 / l as f64
}
pub(crate) fn jc(tally: &PairwiseTally, i: usize, j: usize, alpha: Option<f64>) -> f64 {
let p = raw(tally, i, j);
0.75 * corrected_log(1.0 - 4.0 * p / 3.0, alpha)
}
pub(crate) fn k2p(tally: &PairwiseTally, i: usize, j: usize, alpha: Option<f64>) -> f64 {
let c = tally.categories(i, j);
let l = c.n_eligible();
if l == 0 {
return f64::NAN;
}
let p = (c.ts1 + c.ts2) as f64 / l as f64;
let q = (c.tv1 + c.tv2) as f64 / l as f64;
let a1 = 1.0 - 2.0 * p - q;
let a2 = 1.0 - 2.0 * q;
0.5 * corrected_log(a1, alpha) + 0.25 * corrected_log(a2, alpha)
}
pub(crate) fn k81(tally: &PairwiseTally, i: usize, j: usize, alpha: Option<f64>) -> f64 {
let c = tally.categories(i, j);
let l = c.n_eligible();
if l == 0 {
return f64::NAN;
}
let p = (c.ts1 + c.ts2) as f64 / l as f64;
let q = c.tv1 as f64 / l as f64;
let r = c.tv2 as f64 / l as f64;
let a1 = 1.0 - 2.0 * p - 2.0 * q;
let a2 = 1.0 - 2.0 * p - 2.0 * r;
let a3 = 1.0 - 2.0 * q - 2.0 * r;
0.25 * (corrected_log(a1, alpha) + corrected_log(a2, alpha) + corrected_log(a3, alpha))
}
pub(crate) fn f81(tally: &PairwiseTally, i: usize, j: usize, alpha: Option<f64>) -> f64 {
let p = raw(tally, i, j);
let freq = tally.base_freq(i, j);
let e = 1.0 - freq.iter().map(|f| f * f).sum::<f64>();
match alpha {
Some(a) => e * a * ((1.0 - p / e).powf(-1.0 / a) - 1.0),
None => -e * (1.0 - p / e).ln(),
}
}
pub(crate) fn t92(tally: &PairwiseTally, i: usize, j: usize, alpha: Option<f64>) -> f64 {
let c = tally.categories(i, j);
let l = c.n_eligible();
if l == 0 {
return f64::NAN;
}
let p = (c.ts1 + c.ts2) as f64 / l as f64;
let q = (c.tv1 + c.tv2) as f64 / l as f64;
let freq = tally.base_freq(i, j);
let gc = freq[1] + freq[2]; // C + G
let wg = 2.0 * gc * (1.0 - gc);
let a1 = 1.0 - p / wg - q;
let a2 = 1.0 - 2.0 * q;
wg * corrected_log(a1, alpha) + 0.5 * (1.0 - wg) * corrected_log(a2, alpha)
}
pub(crate) fn tn93(tally: &PairwiseTally, i: usize, j: usize, alpha: Option<f64>) -> f64 {
let c = tally.categories(i, j);
let l = c.n_eligible();
if l == 0 {
return f64::NAN;
}
let freq = tally.base_freq(i, j);
let g_r = freq[0] + freq[2]; // A + G (purines)
let g_y = freq[1] + freq[3]; // C + T (pyrimidines)
let k1 = 2.0 * freq[0] * freq[2] / g_r;
let k2 = 2.0 * freq[1] * freq[3] / g_y;
let k3 = 2.0 * (g_r * g_y - freq[0] * freq[2] * g_y / g_r - freq[1] * freq[3] * g_r / g_y);
let p1 = c.ts1 as f64 / l as f64; // A<->G
let p2 = c.ts2 as f64 / l as f64; // C<->T
let q = (c.tv1 + c.tv2) as f64 / l as f64;
let w1 = 1.0 - p1 / k1 - q / (2.0 * g_r);
let w2 = 1.0 - p2 / k2 - q / (2.0 * g_y);
let w3 = 1.0 - q / (2.0 * g_r * g_y);
k1 * corrected_log(w1, alpha) + k2 * corrected_log(w2, alpha) + k3 * corrected_log(w3, alpha)
}
fn formula(kind: SnpDistanceKind) -> fn(&PairwiseTally, usize, usize, Option<f64>) -> f64 {
match kind {
SnpDistanceKind::Raw => |t, i, j, _| raw(t, i, j),
SnpDistanceKind::Tv => |t, i, j, _| tv_only(t, i, j),
SnpDistanceKind::Jc => jc,
SnpDistanceKind::K2p => k2p,
SnpDistanceKind::K81 => k81,
SnpDistanceKind::F81 => f81,
SnpDistanceKind::T92 => t92,
SnpDistanceKind::Tn93 => tn93,
}
}
/// Build (or reuse) a [`PairwiseTally`] and reduce it to one `n×n` distance
/// matrix under `kind`. `n`: `Some(target)` samples proportionally per
/// layer exactly like `--sankoff`/`--pseudo-alignment` (see
/// `algorithms::subsample::sample_index`'s own docs); `None` means
/// exhaustive — every non-monomorphic minorant of the whole index, achieved
/// by handing `sample_index` a quota equal to the index-wide total (from
/// [`sibling_family_size_histogram`], annex-bits-only, no cross-partition
/// resolution), which makes every per-layer quota equal to that layer's own
/// full eligible count, i.e. Bernoulli `p = 1` everywhere — no separate
/// exhaustive-only driver needed. Entropy-biased sampling
/// (`entropy_bias`) only makes sense for a genuine subsample, so it's
/// ignored (forced to `None`) when `n` is `None`, regardless of what the
/// caller passes.
///
/// `gamma_shape`: `--gamma-shape`, `None` disables the correction. Rejected
/// with [`obikindex::OKIError::InvalidInput`] if given alongside a `kind`
/// that doesn't support it ([`SnpDistanceKind::supports_gamma`]) — checked
/// here rather than left to silently no-op.
pub(crate) fn snp_distance(
cache: &IndexCache,
kind: SnpDistanceKind,
n: Option<usize>,
free_loss: bool,
no_ambiguity: bool,
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
gamma_shape: Option<f64>,
) -> OKIResult<Array2<f64>> {
if gamma_shape.is_some() && !kind.supports_gamma() {
return Err(obikindex::OKIError::InvalidInput(
"--gamma-shape has no effect on this --distance value".into(),
));
}
let n_genomes = cache.meta().genomes().len();
let mut tally = PairwiseTally::new(n_genomes);
let (target, entropy_bias) = match n {
Some(target) => (target, entropy_bias),
None => {
let counts = sibling_family_size_histogram(cache)?;
let total_eligible = (counts[1] + counts[2] + counts[3]) as usize;
(total_eligible, None)
}
};
if target > 0 {
sample_index(
cache,
target,
free_loss,
no_ambiguity,
excluded,
entropy_bias,
|_partition, _layer, survivors| {
super::pairwise::reduce_pairwise(&survivors, &mut tally);
},
)?;
}
let f = formula(kind);
Ok(Array2::from_shape_fn((n_genomes, n_genomes), |(i, j)| {
if i == j { 0.0 } else { f(&tally, i, j, gamma_shape) }
}))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::siblings::FamilyMask;
use crate::siblings::algorithms::SurvivingFamily;
use super::super::pairwise::reduce_pairwise;
// Base bit positions, matching `PairwiseTally::categories`' convention
// (0=A, 1=C, 2=G, 3=T).
const A: u8 = 1 << 0;
const C: u8 = 1 << 1;
const G: u8 = 1 << 2;
const T: u8 = 1 << 3;
/// Two-genome synthetic tally: one variable family per `(genome0,
/// genome1)` pair in `pairs`, plus `n_shared` fully-agreeing families to
/// pad the eligible-locus denominator — cycled over all 4 bases (not a
/// single one) so the pooled base composition stays non-degenerate
/// (`F81`/`T92`/`TN93`'s denominators — `E`/`wg` — are exactly `0` for a
/// single-base composition, an edge case irrelevant to real data, not
/// something these tests need to exercise).
fn tally_2genomes(pairs: &[(u8, u8)], n_shared: usize) -> PairwiseTally {
let mut families = Vec::new();
for &(a, b) in pairs {
families.push(SurvivingFamily {
family_idx: families.len(),
mask: FamilyMask::EMPTY.with(a.trailing_zeros() as u8).with(b.trailing_zeros() as u8),
genome_mask: vec![a, b],
});
}
const BASES: [u8; 4] = [A, C, G, T];
for i in 0..n_shared {
let b = BASES[i % 4];
let other = BASES[(i + 1) % 4];
families.push(SurvivingFamily {
family_idx: families.len(),
mask: FamilyMask::EMPTY.with(b.trailing_zeros() as u8).with(other.trailing_zeros() as u8),
genome_mask: vec![b, b],
});
}
let mut tally = PairwiseTally::new(2);
reduce_pairwise(&families, &mut tally);
tally
}
#[test]
fn categories_classify_each_substitution_type() {
// A<->G (ts1), C<->T (ts2), A<->C (tv1), A<->T (tv2), + 1 shared.
let tally = tally_2genomes(&[(A, G), (C, T), (A, C), (A, T)], 1);
let c = tally.categories(0, 1);
assert_eq!((c.ts1, c.ts2, c.tv1, c.tv2, c.shared), (1, 1, 1, 1, 1));
assert_eq!(c.n_eligible(), 5);
}
#[test]
fn identical_genomes_give_zero_distance_under_every_correction() {
// No substitutions at all — every formula's "-ln(1)" term is 0.
let tally = tally_2genomes(&[], 10);
for kind in [
SnpDistanceKind::Raw,
SnpDistanceKind::Jc,
SnpDistanceKind::K2p,
SnpDistanceKind::K81,
SnpDistanceKind::F81,
SnpDistanceKind::T92,
SnpDistanceKind::Tn93,
SnpDistanceKind::Tv,
] {
let d = formula(kind)(&tally, 0, 1, None);
assert!(d.abs() < 1e-9, "{kind:?} gave {d}, expected ~0");
}
}
#[test]
fn raw_matches_hand_computed_ratio() {
// 1 substitution (A<->G) out of 5 eligible loci (1 subst + 4 shared).
let tally = tally_2genomes(&[(A, G)], 4);
assert!((raw(&tally, 0, 1) - 0.2).abs() < 1e-12);
}
#[test]
fn jc_matches_hand_computed_correction() {
// p = 1/5 = 0.2 -> d = -0.75 * ln(1 - 4*0.2/3) = -0.75 * ln(1 - 4/15).
let tally = tally_2genomes(&[(A, G)], 4);
let expected = -0.75 * (1.0 - 4.0f64 * 0.2 / 3.0).ln();
assert!((jc(&tally, 0, 1, None) - expected).abs() < 1e-12);
}
#[test]
fn jc_uncorrected_undershoots_raw_p_distance() {
// JC always corrects *upward* from the raw p-distance (multiple-hit
// correction only adds distance, never removes it) for any p in its
// domain.
let tally = tally_2genomes(&[(A, G), (C, T)], 20);
let p = raw(&tally, 0, 1);
let d = jc(&tally, 0, 1, None);
assert!(d > p, "JC-corrected {d} should exceed raw p-distance {p}");
}
#[test]
fn gamma_correction_converges_to_uncorrected_as_alpha_grows() {
// The Jin-Nei substitution alpha*(x^(-1/alpha) - 1) -> -ln(x) as
// alpha -> infinity — a basic sanity check on `corrected_log`
// independent of any hand-checked literature value.
let tally = tally_2genomes(&[(A, G), (A, C)], 20);
let uncorrected = jc(&tally, 0, 1, None);
let large_alpha = jc(&tally, 0, 1, Some(1.0e6));
assert!(
(uncorrected - large_alpha).abs() < 1e-3,
"uncorrected {uncorrected} vs large-alpha gamma {large_alpha}"
);
}
#[test]
fn gamma_shape_rejected_for_unsupported_kind() {
assert!(!SnpDistanceKind::Raw.supports_gamma());
assert!(!SnpDistanceKind::Tv.supports_gamma());
assert!(SnpDistanceKind::Jc.supports_gamma());
}
#[test]
fn base_freq_sums_to_one_and_matches_pooled_counts() {
// 2 A/G substitutions (genome0=A, genome1=G each time) + 1 shared
// A/A family -> pooled bases across both genomes: A appears 4 times
// (1 per substitution's genome0 side, 2 from the shared locus's
// both genomes), G appears 2 times (1 per substitution's genome1
// side), total 6 (2 genomes * 3 loci).
let tally = tally_2genomes(&[(A, G), (A, G)], 1);
let freq = tally.base_freq(0, 1);
assert!((freq.iter().sum::<f64>() - 1.0).abs() < 1e-12);
assert!((freq[0] - 4.0 / 6.0).abs() < 1e-12); // A
assert!((freq[2] - 2.0 / 6.0).abs() < 1e-12); // G
assert!((freq[1] - 0.0).abs() < 1e-12); // C
}
}
@@ -9,15 +9,16 @@
use std::io::{BufWriter, Write};
use std::path::Path;
use ndarray::Array2;
use obikidxcache::index_cache::IndexCache;
use obikindex::{OKIError, OKIResult};
use obisys::progress_bar;
use crate::siblings::algorithms::{
EntropyBias, SankoffBundle, Selection, SiblingAnnexStats, SnpAlignment,
EntropyBias, SankoffBundle, Selection, SiblingAnnexStats, SnpAlignment, SnpDistanceKind,
build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode,
sankoff_bundle, scan_layer_families, sibling_annex_stats, sibling_family_size_histogram,
snp_pseudo_alignment,
snp_distance, snp_pseudo_alignment,
};
use crate::siblings::extensions::SiblingBuilder;
use crate::siblings::ENTROPY_ANNEX_FILE_NAME;
@@ -121,6 +122,38 @@ pub trait SiblingExt {
entropy_bias: Option<EntropyBias>,
ratio_ceiling: f64,
) -> OKIResult<SankoffBundle>;
/// A `snp-*` `--distance` matrix — one of the closed-form corrections
/// in [`SnpDistanceKind`], over the central-position SNP substitution
/// spectrum (requires an already-built sibling annex, run
/// [`build_sibling_annex`](Self::build_sibling_annex) first).
///
/// `n`: `Some(target)` samples proportionally per layer exactly like
/// [`sankoff_bundle`](Self::sankoff_bundle)/
/// [`snp_pseudo_alignment`](Self::snp_pseudo_alignment) (`--subsample`);
/// `None` means exhaustive — every non-monomorphic minorant of the
/// whole index, not an approximation (see
/// `algorithms::snp_distance::snp_distance`'s own docs for how this
/// reuses the same sampling machinery at `p = 1` rather than a second
/// driver). `free_loss`/`no_ambiguity`/`excluded`/`entropy_bias` — same
/// meaning as `snp_pseudo_alignment`'s own (the last forced to `None`
/// when `n` is `None`: entropy-biasing only makes sense for a genuine
/// subsample).
///
/// `gamma_shape`: `--gamma-shape`, the Jin-Nei rate-heterogeneity
/// correction — `None` disables it, `Some(alpha)` is rejected with
/// [`OKIError::InvalidInput`] for a `kind` that doesn't support it
/// ([`SnpDistanceKind::supports_gamma`]).
fn snp_distance(
&self,
kind: SnpDistanceKind,
n: Option<usize>,
free_loss: bool,
no_ambiguity: bool,
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
gamma_shape: Option<f64>,
) -> OKIResult<Array2<f64>>;
}
impl SiblingExt for IndexCache {
@@ -265,4 +298,17 @@ impl SiblingExt for IndexCache {
) -> OKIResult<SankoffBundle> {
sankoff_bundle(self, n, free_loss, no_ambiguity, excluded, entropy_bias, ratio_ceiling)
}
fn snp_distance(
&self,
kind: SnpDistanceKind,
n: Option<usize>,
free_loss: bool,
no_ambiguity: bool,
excluded: &[bool],
entropy_bias: Option<EntropyBias>,
gamma_shape: Option<f64>,
) -> OKIResult<Array2<f64>> {
snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape)
}
}
+2 -2
View File
@@ -35,8 +35,8 @@ pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
pub use algorithms::{
BasePairTally, CardinalityTally, EntropyBias, RawSnpDistanceOutput, SankoffBundle,
SiblingAnnexStats, SnpAlignment, cardinality_transition_probs, composition_transition_probs,
pairwise_cost_matrix,
SiblingAnnexStats, SnpAlignment, SnpDistanceKind, cardinality_transition_probs,
composition_transition_probs, pairwise_cost_matrix,
};
pub use extensions::SiblingExt;