Implement gamma shape correction for SNP distance calculation
Introduces support for rate heterogeneity via a Poisson-Gamma mixture model. This includes new command-line options (`--gamma-shape`, `--gamma-shape auto`) to enable automatic estimation of the shape parameter $\alpha$ based on substitution counts across partitions. The correction is applied to the distance metric, with logic to disable the correction if variance checks fail.
This commit is contained in:
Eric Coissac
@@ -2249,9 +2249,85 @@ F84/TN93 and is included below. `snp-` prefix on every CLI value.
|
||||
**`+Γ` rate-heterogeneity modifier, applicable to `snp-jc`, `snp-k2p`,
|
||||
`snp-k81`, `snp-t92`, `snp-f84`, `snp-hky85`, `snp-tn93`** (not `snp-raw`,
|
||||
nothing to correct; not `snp-logdet`, no standard gamma formulation) — same
|
||||
formula as the base correction, weighted by a shape parameter `α` supplied
|
||||
by the user (`--gamma-shape <alpha>`), not estimated by ML. A modifier on
|
||||
existing values, not a separate enum arm per distance.
|
||||
formula as the base correction, weighted by a shape parameter `α` either
|
||||
supplied by the user (`--gamma-shape <alpha>`) or estimated from the data
|
||||
(`--gamma-shape auto`/`estimate`, method-of-moments — not ML; see
|
||||
"Automatic α estimation" below). A modifier on existing values, not a
|
||||
separate enum arm per distance.
|
||||
|
||||
### Automatic α estimation (`--gamma-shape auto`)
|
||||
|
||||
**Correction (verified against the primary source, 2026-09-11):** Jin &
|
||||
Nei (1990) itself (*"Limitations of the Evolutionary Parsimony Method of
|
||||
Phylogenetic Analysis"*, Mol. Biol. Evol. 7(2):82–102 — the paper this
|
||||
whole `+Γ` correction is cited from, confirmed algebraically to match this
|
||||
codebase's `corrected_log`/`k2p` exactly against their eq. A4, general, and
|
||||
A8, the `a = 1` case) contains **no** data-driven α-estimation procedure.
|
||||
Their own recommendation (p. 98) is a fixed default: *"we suggest that the
|
||||
gamma distance with a = 1 [eq. A8] be used. However, one may choose a
|
||||
different gamma distance, estimating a from data. Wilson et al. (1989)
|
||||
recently used a distance with a = 1/2 for restriction-site data of
|
||||
mitochondrial DNA in hominoids."* — i.e. Jin & Nei explicitly punt
|
||||
data-driven estimation to a *different* paper (Wilson et al. 1989), not
|
||||
read/verified here. The estimator below is therefore **not** "Jin & Nei's
|
||||
method" under any framing — that attribution (present in an earlier
|
||||
revision of this section) was wrong, not just under-cited.
|
||||
|
||||
**Implemented** (`PartitionDispersion`, `obikphylo/src/siblings/algorithms/pairwise.rs`)
|
||||
as an independent method-of-moments estimator, unrelated to any specific
|
||||
published procedure: pools substitution counts by **partition** rather
|
||||
than by genome pair, during the same `reduce_pairwise` pass that builds
|
||||
`PairwiseTally` (no second scan).
|
||||
|
||||
For partition `i`: `n_i` = substitutions pooled over every genome pair,
|
||||
`L_i` = eligible loci pooled over every genome pair, `R_i = n_i / L_i`.
|
||||
Modeling among-site rate heterogeneity the same way as the `+Γ` correction
|
||||
itself (a `Gamma(α, α)`-distributed, mean-1, multiplicative rate shared by
|
||||
every locus in a partition — the classical Poisson-Gamma/negative-binomial
|
||||
mixture, the general identity behind gamma-rate-heterogeneity corrections,
|
||||
also behind Uzzell & Corbin's (1971) original observation that substitution
|
||||
counts across sites/regions are over-dispersed relative to Poisson):
|
||||
|
||||
\[
|
||||
\mathbb{E}[R_i] = \mu \qquad \mathrm{Var}[R_i] = \frac{\mu}{L_i} + \frac{\mu^2}{\alpha}
|
||||
\]
|
||||
|
||||
Weighting each partition's squared deviation by its own `L_i` cancels the
|
||||
Poisson term before attributing what's left to `α`:
|
||||
|
||||
\[
|
||||
\hat\mu = \frac{\sum_i n_i}{\sum_i L_i} \qquad
|
||||
V = \frac{\sum_i L_i (R_i-\hat\mu)^2}{\sum_i L_i} \qquad
|
||||
\hat\alpha = \frac{\hat\mu^2}{V - \hat\mu/\bar L}
|
||||
\]
|
||||
|
||||
where `\bar L` is the mean partition size. Returns "no estimate" (falls
|
||||
back to the uncorrected formula, warns) when fewer than 2 partitions have
|
||||
data, `\hat\mu \le 0`, or `V` doesn't exceed the Poisson floor
|
||||
`\hat\mu/\bar L` — no detectable over-dispersion, `α` would be unbounded.
|
||||
|
||||
**Caveat, stated explicitly rather than left implicit**: unlike every
|
||||
closed-form correction in `snp_distance.rs` (each verified line-by-line
|
||||
against `ape`'s `dist_dna.c`, and now also against Jin & Nei 1990 directly
|
||||
for the base `+Γ` formula), this estimator is derived from first
|
||||
principles (the general Poisson-Gamma/negative-binomial identity) with no
|
||||
primary-source procedure behind it at all — not Jin & Nei's (confirmed
|
||||
above), and Wilson et al. (1989), the paper they point to instead, hasn't
|
||||
been read/verified either. Mathematically self-consistent (re-derived and
|
||||
checked, not guessed), but a from-scratch method, not a literature
|
||||
implementation. If `--gamma-shape` needs a value with a literature
|
||||
pedigree rather than an estimated one, Jin & Nei's own stated default,
|
||||
`α = 1` (`--gamma-shape 1`), is the better-supported choice today.
|
||||
|
||||
Deliberately **not** gated by `--sankoff-ratio-ceiling` the way
|
||||
`base_pair_tally` is (same precedent as `cardinality_tally` — see its own
|
||||
doc comment): that filter excludes individual saturated *pairs* from a
|
||||
composition estimate computed once at the very end, from the complete
|
||||
`PairwiseTally`; the partition axis needed here only exists transiently,
|
||||
one partition at a time, while `PairwiseTally` is still being built — long
|
||||
before any pair's final SNP ratio (and thus its ratio_ceiling eligibility)
|
||||
is known. `--exclude-genome` isn't applied either, matching
|
||||
`reduce_pairwise`'s own raw per-pair fold.
|
||||
|
||||
**Implemented now: `snp-raw`, `snp-jc`, `snp-k2p`, `snp-k81`, `snp-f81`,
|
||||
`snp-t92`, `snp-tn93`, `snp-tv`, all with `+Γ` except `raw`/`tv`** — see
|
||||
|
||||
@@ -17,7 +17,7 @@ obikmer phylo INDEX [OPTIONS]
|
||||
| Option | Default | Description |
|
||||
|---|---|---|
|
||||
| `--distance` | `jaccard` | See the two tables below for the full list of accepted values |
|
||||
| `--gamma-shape ALPHA` | none | Rate-heterogeneity correction, for `snp-*` values that support it (see below). No effect on the other values; rejected if given together with a value that doesn't support it |
|
||||
| `--gamma-shape ALPHA\|auto` | none | Rate-heterogeneity correction, for `snp-*` values that support it (see below). Either a fixed $\alpha$ or `auto`/`estimate` to fit it from the data (see "Automatic $\alpha$ estimation" below). No effect on the other values; rejected if given together with a value that doesn't support it |
|
||||
| `--presence-threshold` | `1` | Minimum count for a kmer to be considered present, for `jaccard`/`mash` on a count index |
|
||||
| `--csv` | off | Write the matrix as plain CSV instead of the default relaxed-PHYLIP format |
|
||||
| `--shared-kmers` | off | Also write the shared-kmer count matrix. Only valid with a whole-index metric, not a `snp-*` value |
|
||||
@@ -121,6 +121,28 @@ $$d = Q$$
|
||||
|
||||
`--gamma-shape ALPHA` applies to every value above except `snp-raw` and `snp-tv`: each $-\ln(x)$ term in the formulas above is replaced by $\alpha\left(x^{-1/\alpha}-1\right)$ (the same weight, same $x$).
|
||||
|
||||
### Automatic $\alpha$ estimation (`--gamma-shape auto`)
|
||||
|
||||
`--gamma-shape auto` (or the equivalent `--gamma-shape estimate`) fits $\alpha$ from the index itself instead of requiring a user-supplied value, using a method-of-moments estimator computed once, from the same sampling pass that builds the pairwise substitution tally — no extra scan of the index.
|
||||
|
||||
The estimator pools substitution counts by **partition** rather than by genome pair: for partition $i$, let $n_i$ be the total number of substitutions observed across every genome pair, and $L_i$ the total number of eligible loci across every genome pair, in that partition. Define the partition's observed substitution rate:
|
||||
|
||||
$$R_i = \frac{n_i}{L_i}$$
|
||||
|
||||
Under a single shared substitution rate with no among-site heterogeneity, each $R_i$ would vary only by Poisson sampling noise. Rate heterogeneity is modeled, as elsewhere in this correction, by a $\mathrm{Gamma}(\alpha,\alpha)$-distributed multiplicative rate (mean 1) shared by every locus in a partition — the classical Poisson–Gamma (negative-binomial) mixture. Under that model:
|
||||
|
||||
$$\mathbb{E}[R_i] = \mu \qquad \mathrm{Var}[R_i] = \frac{\mu}{L_i} + \frac{\mu^2}{\alpha}$$
|
||||
|
||||
where $\mu$ is the pooled substitution rate across every partition. Weighting each partition's squared deviation by its own $L_i$ removes the first (Poisson) term before attributing what's left to genuine rate heterogeneity:
|
||||
|
||||
$$\hat\mu = \frac{\sum_i n_i}{\sum_i L_i} \qquad V = \frac{\sum_i L_i\,(R_i-\hat\mu)^2}{\sum_i L_i} \qquad \bar L = \frac{\sum_i L_i}{\text{number of partitions}}$$
|
||||
|
||||
$$\hat\alpha = \frac{\hat\mu^2}{V - \hat\mu/\bar L}$$
|
||||
|
||||
If the measured variance $V$ doesn't exceed the Poisson floor $\hat\mu/\bar L$ (no detectable over-dispersion across partitions — the data are consistent with a single shared rate), $\alpha$ is left undefined: the correction is silently disabled for that run rather than applying a fabricated value, and a warning is logged. When an estimate is produced, it's logged at the `info` level before the distance matrix is computed.
|
||||
|
||||
Note: this is a method-of-moments estimator derived from the standard Poisson–Gamma relationship between substitution counts and gamma-distributed rate variation, applied per-partition — it is not part of Jin & Nei's (1990) original publication, which only defines the `+Γ` distance formula itself and, absent an estimate, recommends the fixed default $\alpha = 1$ (`--gamma-shape 1`) rather than proposing a way to estimate it from data. `alpha < 1` indicates strong among-site rate heterogeneity (many near-invariant loci, a few fast ones); `alpha` growing large makes the correction converge to the uncorrected formula.
|
||||
|
||||
### Output
|
||||
|
||||
Without `-o`, the matrix goes to stdout in relaxed-PHYLIP format (`n` on the first line, then one `label<TAB>value...` row per genome). With `--csv`, the format is instead a header row `genome,<label1>,<label2>,...` followed by one `<label>,<value1>,<value2>,...` row per genome, 6 decimals. Both formats are symmetric with a zero diagonal, except where noted below.
|
||||
|
||||
@@ -1,7 +1,36 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
use clap::Args;
|
||||
use obikphylo::{DistanceMetric, SnpDistanceKind};
|
||||
use obikphylo::{DistanceMetric, GammaShape, SnpDistanceKind};
|
||||
|
||||
/// `--gamma-shape` CLI value: a fixed `alpha`, or `auto`/`estimate` to fit
|
||||
/// it from the data (see [`GammaShape::Auto`]). Kept as its own tiny type
|
||||
/// (rather than parsing straight to `GammaShape`) so the CLI layer stays
|
||||
/// free of `GammaShape::Disabled`, which the `Option<GammaShapeArg>` around
|
||||
/// this already expresses on its own (`None` = flag not given at all).
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum GammaShapeArg {
|
||||
Fixed(f64),
|
||||
Auto,
|
||||
}
|
||||
|
||||
impl From<GammaShapeArg> for GammaShape {
|
||||
fn from(arg: GammaShapeArg) -> Self {
|
||||
match arg {
|
||||
GammaShapeArg::Fixed(a) => GammaShape::Fixed(a),
|
||||
GammaShapeArg::Auto => GammaShape::Auto,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_gamma_shape(s: &str) -> Result<GammaShapeArg, String> {
|
||||
if s.eq_ignore_ascii_case("auto") || s.eq_ignore_ascii_case("estimate") {
|
||||
return Ok(GammaShapeArg::Auto);
|
||||
}
|
||||
s.parse::<f64>()
|
||||
.map(GammaShapeArg::Fixed)
|
||||
.map_err(|_| format!("invalid --gamma-shape value {s:?}: expected a positive number or \"auto\"/\"estimate\""))
|
||||
}
|
||||
|
||||
/// `--distance` value — either one of `obikphylo::DistanceMetric`'s
|
||||
/// whole-index metrics (routed to `IndexCache::distance`) or one of
|
||||
@@ -292,8 +321,11 @@ pub struct PhyloArgs {
|
||||
/// `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>,
|
||||
/// Either a positive number (fixed `alpha`) or `auto`/`estimate` to fit
|
||||
/// `alpha` from the data itself — method-of-moments over the
|
||||
/// over-dispersion of substitution rate across partitions.
|
||||
#[arg(long, value_name = "ALPHA|auto", value_parser = parse_gamma_shape)]
|
||||
pub gamma_shape: Option<GammaShapeArg>,
|
||||
|
||||
/// Minimum count to consider a kmer present when computing Jaccard on count indexes
|
||||
#[arg(long, default_value = "1")]
|
||||
|
||||
@@ -11,8 +11,8 @@ use std::sync::Arc;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::KmerIndex;
|
||||
use obikphylo::siblings::{
|
||||
EntropyBias, SiblingExt, cardinality_transition_probs, composition_transition_probs,
|
||||
pairwise_cost_matrix,
|
||||
EntropyBias, GammaShape, SiblingExt, cardinality_transition_probs,
|
||||
composition_transition_probs, pairwise_cost_matrix,
|
||||
};
|
||||
use obikphylo::{Metrics, neighbor_joining, upgma};
|
||||
use obisys::{Reporter, Stage};
|
||||
@@ -370,7 +370,7 @@ pub fn run(args: PhyloArgs) {
|
||||
args.no_ambiguity,
|
||||
&snp_exclude_mask,
|
||||
entropy_bias,
|
||||
args.gamma_shape,
|
||||
args.gamma_shape.map_or(GammaShape::Disabled, Into::into),
|
||||
)
|
||||
.unwrap_or_else(|e| {
|
||||
eprintln!("error computing SNP distance: {e}");
|
||||
|
||||
@@ -19,5 +19,5 @@ mod tree;
|
||||
pub mod siblings;
|
||||
|
||||
pub use distance::{DistanceMetric, DistanceOutput, Metrics};
|
||||
pub use siblings::SnpDistanceKind;
|
||||
pub use siblings::{GammaShape, SnpDistanceKind};
|
||||
pub use tree::{Tree, neighbor_joining, upgma};
|
||||
|
||||
@@ -34,7 +34,7 @@ pub use cardcomp::{
|
||||
};
|
||||
pub use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput};
|
||||
pub use sankoff::SankoffBundle;
|
||||
pub use snp_distance::SnpDistanceKind;
|
||||
pub use snp_distance::{GammaShape, SnpDistanceKind};
|
||||
pub use stats::SiblingAnnexStats;
|
||||
pub use subsample::{EntropyBias, SurvivingFamily};
|
||||
|
||||
|
||||
@@ -273,14 +273,137 @@ pub struct CardinalityTally {
|
||||
pub counts: [[u64; 5]; 5],
|
||||
}
|
||||
|
||||
/// One layer's worth of [`SurvivingFamily`] folded into `tally` — the
|
||||
/// "reduce" half of the map/reduce [`super::subsample::sample_index`]
|
||||
/// drives, self-contained (no sampling/masking logic of its own) so it can
|
||||
/// run alongside any other reduce over the same batch (see
|
||||
/// `algorithms::alignment::reduce_alignment`, called from the same
|
||||
/// `on_layer` closure in `algorithms::sankoff::sankoff_bundle`).
|
||||
pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut PairwiseTally) {
|
||||
/// Per-partition, pooled-over-every-genome-pair `(substitutions, eligible
|
||||
/// loci)` counts — feeds [`Self::estimate_alpha`], the method-of-moments
|
||||
/// `--gamma-shape auto` estimator (an independent derivation — see that
|
||||
/// method's own docs for why this is *not* Jin & Nei's (1990) method,
|
||||
/// verified against their paper directly, despite using the same `+Γ`
|
||||
/// distance formula). Pooling across pairs (rather than keeping one series
|
||||
/// per pair) treats every included genome pair as one more observation of
|
||||
/// the same underlying rate-heterogeneity signal — comparable in spirit to
|
||||
/// combining information across several genomic regions/genes for a fixed
|
||||
/// set of taxa, generalised here to every pair at once.
|
||||
///
|
||||
/// Deliberately **not** gated by `--sankoff-ratio-ceiling` the way
|
||||
/// [`PairwiseTally::base_pair_tally`] is (see
|
||||
/// [`PairwiseTally::cardinality_tally`]'s own docs for the same
|
||||
/// precedent): that filter excludes individual saturated *pairs* from a
|
||||
/// composition estimate computed once at the very end, from the complete
|
||||
/// `PairwiseTally`; the partition axis needed here only exists transiently,
|
||||
/// one partition at a time, while `PairwiseTally` is still being built —
|
||||
/// long before any pair's final SNP ratio (and thus its ratio_ceiling
|
||||
/// eligibility) is known. Genome exclusion (`--exclude-genome`) isn't
|
||||
/// applied here either, matching `reduce_pairwise`'s own raw per-pair fold
|
||||
/// (`PairStats` accumulates over every genome pair unconditionally;
|
||||
/// exclusion is a filter applied only in derived, post-hoc views).
|
||||
#[derive(Default, Clone)]
|
||||
pub(crate) struct PartitionDispersion {
|
||||
/// One entry per partition id seen so far (grows lazily — partitions
|
||||
/// are visited in increasing order in practice, but nothing here
|
||||
/// assumes that).
|
||||
counts: Vec<(u64, u64)>,
|
||||
}
|
||||
|
||||
impl PartitionDispersion {
|
||||
/// Add one partition/layer's worth of pooled counts (from one
|
||||
/// [`reduce_pairwise`] call) to that partition's running total —
|
||||
/// additive so a partition with several layers accumulates correctly
|
||||
/// across multiple calls.
|
||||
fn record(&mut self, partition: usize, subst: u64, eligible: u64) {
|
||||
if partition >= self.counts.len() {
|
||||
self.counts.resize(partition + 1, (0, 0));
|
||||
}
|
||||
let entry = &mut self.counts[partition];
|
||||
entry.0 += subst;
|
||||
entry.1 += eligible;
|
||||
}
|
||||
|
||||
/// Method-of-moments estimate of the gamma shape parameter `alpha`
|
||||
/// used by the `+Γ` distance correction, from the over-dispersion of
|
||||
/// per-partition substitution *rates* relative to what a single shared
|
||||
/// rate would produce under pure Poisson sampling noise — the general
|
||||
/// negative-binomial (Poisson-Gamma mixture) identity, the same kind of
|
||||
/// reasoning behind Uzzell & Corbin's (1971) original observation that
|
||||
/// substitution counts across sites/regions are over-dispersed
|
||||
/// relative to Poisson.
|
||||
///
|
||||
/// **Not Jin & Nei's (1990) method** — checked directly against their
|
||||
/// paper (*"Limitations of the Evolutionary Parsimony Method of
|
||||
/// Phylogenetic Analysis"*, Mol. Biol. Evol. 7(2):82-102, the source of
|
||||
/// the `+Γ` distance formula this alpha feeds into, confirmed to match
|
||||
/// `corrected_log`/`k2p` exactly against their eq. A4/A8): that paper
|
||||
/// contains no data-driven alpha estimator. Their own recommendation
|
||||
/// (p. 98) is the fixed default `alpha = 1`; for estimating alpha from
|
||||
/// data they point to a different paper (Wilson et al. 1989, not
|
||||
/// read/verified here) rather than proposing their own procedure. This
|
||||
/// estimator is therefore a from-scratch derivation (Poisson count
|
||||
/// within a partition, with a Gamma(alpha, alpha)-distributed (mean 1)
|
||||
/// multiplicative rate shared by every locus in that partition), not a
|
||||
/// literature implementation — unlike every closed-form correction in
|
||||
/// `snp_distance.rs`, each verified line-by-line against `ape`'s
|
||||
/// `dist_dna.c`.
|
||||
///
|
||||
/// Partitions have unequal eligible-locus counts (`l_i`), so a plain
|
||||
/// unweighted variance across raw per-partition rates would conflate
|
||||
/// genuine rate heterogeneity with the extra Poisson noise smaller
|
||||
/// partitions carry. Weighting each partition's squared deviation by
|
||||
/// its own `l_i` cancels that: for `R_i = subst_i / l_i`, `E[R_i] =
|
||||
/// mu`, `Var[R_i] = mu/l_i + mu²/alpha`, so `E[Σ l_i (R_i-mu)² / Σ l_i]
|
||||
/// = mu/mean(l) + mu²/alpha` — the Poisson floor
|
||||
/// (`mu/mean(l)`) is subtracted below before solving for `alpha`.
|
||||
///
|
||||
/// Returns `None` when there's nothing to estimate from: fewer than 2
|
||||
/// partitions with data, a pooled mean rate of exactly 0, or a
|
||||
/// measured variance at or below the Poisson floor (no detectable
|
||||
/// over-dispersion — `alpha` would be unbounded, not just large).
|
||||
pub(crate) fn estimate_alpha(&self) -> Option<f64> {
|
||||
let rows: Vec<(f64, f64)> = self
|
||||
.counts
|
||||
.iter()
|
||||
.filter(|&&(_, eligible)| eligible > 0)
|
||||
.map(|&(subst, eligible)| (subst as f64 / eligible as f64, eligible as f64))
|
||||
.collect();
|
||||
if rows.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let weight_total: f64 = rows.iter().map(|&(_, w)| w).sum();
|
||||
let mean: f64 = rows.iter().map(|&(rate, w)| rate * w).sum::<f64>() / weight_total;
|
||||
if mean <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let weighted_variance: f64 =
|
||||
rows.iter().map(|&(rate, w)| w * (rate - mean).powi(2)).sum::<f64>() / weight_total;
|
||||
let mean_eligible = weight_total / rows.len() as f64;
|
||||
let poisson_floor = mean / mean_eligible;
|
||||
let excess = weighted_variance - poisson_floor;
|
||||
if excess <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(mean * mean / excess)
|
||||
}
|
||||
}
|
||||
|
||||
/// One layer's worth of [`SurvivingFamily`] folded into `tally` (and
|
||||
/// `dispersion`, pooled under `partition` — see [`PartitionDispersion`]'s
|
||||
/// own docs for why the partition axis must be captured here rather than
|
||||
/// after the fact) — the "reduce" half of the map/reduce
|
||||
/// [`super::subsample::sample_index`] drives, self-contained (no
|
||||
/// sampling/masking logic of its own) so it can run alongside any other
|
||||
/// reduce over the same batch (see `algorithms::alignment::reduce_alignment`,
|
||||
/// called from the same `on_layer` closure in
|
||||
/// `algorithms::sankoff::sankoff_bundle`).
|
||||
pub(crate) fn reduce_pairwise(
|
||||
survivors: &[SurvivingFamily],
|
||||
partition: usize,
|
||||
tally: &mut PairwiseTally,
|
||||
dispersion: &mut PartitionDispersion,
|
||||
) {
|
||||
let n_genomes = tally.n_genomes;
|
||||
let mut partition_subst: u64 = 0;
|
||||
let mut partition_eligible: u64 = 0;
|
||||
for family in survivors {
|
||||
let variable = family.mask.family_size() >= 2;
|
||||
let genome_mask = &family.genome_mask;
|
||||
@@ -296,7 +419,9 @@ pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut Pairwis
|
||||
let stats = tally.pair_mut(i, j);
|
||||
|
||||
if let (Some(bi), Some(bj)) = (bi, single_form(j)) {
|
||||
partition_eligible += 1;
|
||||
if bi != bj {
|
||||
partition_subst += 1;
|
||||
stats.subst[bi as usize][bj as usize] += 1;
|
||||
stats.subst[bj as usize][bi as usize] += 1;
|
||||
} else {
|
||||
@@ -317,4 +442,59 @@ pub(crate) fn reduce_pairwise(survivors: &[SurvivingFamily], tally: &mut Pairwis
|
||||
}
|
||||
}
|
||||
}
|
||||
dispersion.record(partition, partition_subst, partition_eligible);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod dispersion_tests {
|
||||
use super::PartitionDispersion;
|
||||
|
||||
#[test]
|
||||
fn fewer_than_two_partitions_gives_no_estimate() {
|
||||
let mut d = PartitionDispersion::default();
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
d.record(0, 5, 100);
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_mean_rate_gives_no_estimate() {
|
||||
let mut d = PartitionDispersion::default();
|
||||
d.record(0, 0, 100);
|
||||
d.record(1, 0, 100);
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identical_per_partition_rate_gives_no_estimate() {
|
||||
// Every partition sees exactly the same rate — variance across
|
||||
// partitions is *below* the Poisson floor (it's exactly 0), so
|
||||
// there's no over-dispersion to attribute to rate heterogeneity.
|
||||
let mut d = PartitionDispersion::default();
|
||||
for _ in 0..5 {
|
||||
d.record(d.counts.len(), 10, 100);
|
||||
}
|
||||
assert!(d.estimate_alpha().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clearly_over_dispersed_rates_give_a_finite_positive_alpha() {
|
||||
// Two partitions with no substitutions, two with twice the pooled
|
||||
// mean rate — far more spread than Poisson noise alone would
|
||||
// produce at these counts, so a finite alpha should come out.
|
||||
let mut d = PartitionDispersion::default();
|
||||
d.record(0, 0, 100);
|
||||
d.record(1, 20, 100);
|
||||
d.record(2, 0, 100);
|
||||
d.record(3, 20, 100);
|
||||
|
||||
// Hand-computed from the weighted-moments formula in
|
||||
// `estimate_alpha`'s own docs: mean = 0.1, weighted variance =
|
||||
// 0.01, Poisson floor = mean/mean_eligible = 0.001, alpha =
|
||||
// mean² / (variance - floor) = 0.01 / 0.009.
|
||||
let expected = 0.01 / 0.009;
|
||||
let alpha = d.estimate_alpha().expect("clear over-dispersion should yield an estimate");
|
||||
assert!(alpha > 0.0, "alpha should be positive, got {alpha}");
|
||||
assert!((alpha - expected).abs() < 1e-9, "got {alpha}, expected {expected}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,7 +14,10 @@ use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
|
||||
use super::alignment::{SnpAlignment, reduce_alignment};
|
||||
use super::pairwise::{BasePairTally, CardinalityTally, PairwiseTally, RawSnpDistanceOutput, reduce_pairwise};
|
||||
use super::pairwise::{
|
||||
BasePairTally, CardinalityTally, PairwiseTally, PartitionDispersion, RawSnpDistanceOutput,
|
||||
reduce_pairwise,
|
||||
};
|
||||
use super::subsample::{EntropyBias, sample_index};
|
||||
|
||||
/// Every output the `--sankoff`/`--tnt`/`--phyg`/`--iqtree` pipeline needs,
|
||||
@@ -48,6 +51,10 @@ pub(crate) fn sankoff_bundle(
|
||||
.collect();
|
||||
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); genome_indices.len()];
|
||||
let mut tally = PairwiseTally::new(n_genomes);
|
||||
// Unused past this call: the Sankoff/TNT/PhyG/IQ-TREE pipeline has no
|
||||
// `--gamma-shape`, but `reduce_pairwise` folds it in the same pass as
|
||||
// `tally` regardless — cheaper than a second scan just to skip it.
|
||||
let mut dispersion = PartitionDispersion::default();
|
||||
|
||||
sample_index(
|
||||
cache,
|
||||
@@ -56,9 +63,9 @@ pub(crate) fn sankoff_bundle(
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, survivors| {
|
||||
|partition, _layer, survivors| {
|
||||
reduce_alignment(&survivors, &genome_indices, free_loss, no_ambiguity, &mut sequences);
|
||||
reduce_pairwise(&survivors, &mut tally);
|
||||
reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
},
|
||||
)?;
|
||||
|
||||
|
||||
@@ -25,7 +25,7 @@ use ndarray::Array2;
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::OKIResult;
|
||||
|
||||
use super::pairwise::PairwiseTally;
|
||||
use super::pairwise::{PairwiseTally, PartitionDispersion};
|
||||
use super::sibling_family_size_histogram;
|
||||
use super::subsample::{EntropyBias, sample_index};
|
||||
|
||||
@@ -69,6 +69,19 @@ impl SnpDistanceKind {
|
||||
}
|
||||
}
|
||||
|
||||
/// `--gamma-shape` value. `Disabled` (no flag given) leaves every
|
||||
/// `corrected_log` call at its plain `-ln(x)` term; `Fixed` threads the
|
||||
/// given `alpha` through unchanged; `Auto` estimates `alpha` once, up
|
||||
/// front, from [`PartitionDispersion::estimate_alpha`] over the same
|
||||
/// sampling pass that builds the [`PairwiseTally`] — see that method's own
|
||||
/// docs for the estimator itself.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum GammaShape {
|
||||
Disabled,
|
||||
Fixed(f64),
|
||||
Auto,
|
||||
}
|
||||
|
||||
/// `-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
|
||||
@@ -209,10 +222,15 @@ fn formula(kind: SnpDistanceKind) -> fn(&PairwiseTally, usize, usize, Option<f64
|
||||
/// 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.
|
||||
/// `gamma_shape`: `--gamma-shape`, [`GammaShape::Disabled`] leaves the
|
||||
/// correction off. Rejected with [`obikindex::OKIError::InvalidInput`] if
|
||||
/// anything but `Disabled` is given alongside a `kind` that doesn't support
|
||||
/// it ([`SnpDistanceKind::supports_gamma`]) — checked here rather than left
|
||||
/// to silently no-op. [`GammaShape::Auto`] estimates `alpha` once, from the
|
||||
/// same sampling pass, via [`PartitionDispersion::estimate_alpha`]; if that
|
||||
/// estimation finds no measurable signal (see its own docs for when), the
|
||||
/// correction is silently left off rather than applied with a fabricated
|
||||
/// value — logged either way.
|
||||
pub(crate) fn snp_distance(
|
||||
cache: &IndexCache,
|
||||
kind: SnpDistanceKind,
|
||||
@@ -221,9 +239,9 @@ pub(crate) fn snp_distance(
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
gamma_shape: Option<f64>,
|
||||
gamma_shape: GammaShape,
|
||||
) -> OKIResult<Array2<f64>> {
|
||||
if gamma_shape.is_some() && !kind.supports_gamma() {
|
||||
if gamma_shape != GammaShape::Disabled && !kind.supports_gamma() {
|
||||
return Err(obikindex::OKIError::InvalidInput(
|
||||
"--gamma-shape has no effect on this --distance value".into(),
|
||||
));
|
||||
@@ -231,6 +249,7 @@ pub(crate) fn snp_distance(
|
||||
|
||||
let n_genomes = cache.meta().genomes().len();
|
||||
let mut tally = PairwiseTally::new(n_genomes);
|
||||
let mut dispersion = PartitionDispersion::default();
|
||||
|
||||
let (target, entropy_bias) = match n {
|
||||
Some(target) => (target, entropy_bias),
|
||||
@@ -249,15 +268,36 @@ pub(crate) fn snp_distance(
|
||||
no_ambiguity,
|
||||
excluded,
|
||||
entropy_bias,
|
||||
|_partition, _layer, survivors| {
|
||||
super::pairwise::reduce_pairwise(&survivors, &mut tally);
|
||||
|partition, _layer, survivors| {
|
||||
super::pairwise::reduce_pairwise(&survivors, partition, &mut tally, &mut dispersion);
|
||||
},
|
||||
)?;
|
||||
}
|
||||
|
||||
let alpha = match gamma_shape {
|
||||
GammaShape::Disabled => None,
|
||||
GammaShape::Fixed(a) => Some(a),
|
||||
GammaShape::Auto => match dispersion.estimate_alpha() {
|
||||
Some(a) => {
|
||||
tracing::info!(
|
||||
"--gamma-shape auto: estimated alpha = {a:.4} (method-of-moments, \
|
||||
over-dispersion of substitution rate across partitions)"
|
||||
);
|
||||
Some(a)
|
||||
}
|
||||
None => {
|
||||
tracing::warn!(
|
||||
"--gamma-shape auto: no measurable among-partition rate over-dispersion \
|
||||
detected — proceeding without the gamma correction"
|
||||
);
|
||||
None
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
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) }
|
||||
if i == j { 0.0 } else { f(&tally, i, j, alpha) }
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -302,7 +342,8 @@ mod tests {
|
||||
});
|
||||
}
|
||||
let mut tally = PairwiseTally::new(2);
|
||||
reduce_pairwise(&families, &mut tally);
|
||||
let mut dispersion = PartitionDispersion::default();
|
||||
reduce_pairwise(&families, 0, &mut tally, &mut dispersion);
|
||||
tally
|
||||
}
|
||||
|
||||
|
||||
@@ -15,8 +15,8 @@ use obikindex::{OKIError, OKIResult};
|
||||
use obisys::progress_bar;
|
||||
|
||||
use crate::siblings::algorithms::{
|
||||
EntropyBias, SankoffBundle, Selection, SiblingAnnexStats, SnpAlignment, SnpDistanceKind,
|
||||
build_layer_sibling_annex, family_entropy, family_entropy_4, is_fast_mode,
|
||||
EntropyBias, GammaShape, 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_distance, snp_pseudo_alignment,
|
||||
};
|
||||
@@ -141,9 +141,12 @@ pub trait SiblingExt {
|
||||
/// subsample).
|
||||
///
|
||||
/// `gamma_shape`: `--gamma-shape`, the Jin-Nei rate-heterogeneity
|
||||
/// correction — `None` disables it, `Some(alpha)` is rejected with
|
||||
/// correction — [`GammaShape::Disabled`] turns it off,
|
||||
/// [`GammaShape::Fixed`]/[`GammaShape::Auto`] are rejected with
|
||||
/// [`OKIError::InvalidInput`] for a `kind` that doesn't support it
|
||||
/// ([`SnpDistanceKind::supports_gamma`]).
|
||||
/// ([`SnpDistanceKind::supports_gamma`]). `Auto` estimates `alpha` from
|
||||
/// the data itself (method-of-moments over per-partition substitution
|
||||
/// rate dispersion) rather than requiring a user-supplied value.
|
||||
fn snp_distance(
|
||||
&self,
|
||||
kind: SnpDistanceKind,
|
||||
@@ -152,7 +155,7 @@ pub trait SiblingExt {
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
gamma_shape: Option<f64>,
|
||||
gamma_shape: GammaShape,
|
||||
) -> OKIResult<Array2<f64>>;
|
||||
|
||||
/// The Family Overlap annex — number of shared *variable* families per
|
||||
@@ -331,7 +334,7 @@ impl SiblingExt for IndexCache {
|
||||
no_ambiguity: bool,
|
||||
excluded: &[bool],
|
||||
entropy_bias: Option<EntropyBias>,
|
||||
gamma_shape: Option<f64>,
|
||||
gamma_shape: GammaShape,
|
||||
) -> OKIResult<Array2<f64>> {
|
||||
snp_distance(self, kind, n, free_loss, no_ambiguity, excluded, entropy_bias, gamma_shape)
|
||||
}
|
||||
|
||||
@@ -37,7 +37,7 @@ pub use iter::{
|
||||
pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
|
||||
|
||||
pub use algorithms::{
|
||||
BasePairTally, CardinalityTally, EntropyBias, RawSnpDistanceOutput, SankoffBundle,
|
||||
BasePairTally, CardinalityTally, EntropyBias, GammaShape, RawSnpDistanceOutput, SankoffBundle,
|
||||
SiblingAnnexStats, SnpAlignment, SnpDistanceKind, cardinality_transition_probs,
|
||||
composition_transition_probs, pairwise_cost_matrix,
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user