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.
2026-09-12 07:44:39 +02:00
parent e846d35adb
commit 020b391636
11 changed files with 401 additions and 40 deletions
@@ -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.