gamma and invariant

Author: holmrenser

.github/workflows/release.yml

@@ -167,7 +167,7 @@ jobs:
           workspaces: wasm -> wasm/target
       - uses: actions/setup-node@v4
         with:
-          node-version: "22"
+          node-version: "24"
           registry-url: https://registry.npmjs.org
       - run: cp README.md wasm/README.md
       - run: npm ci
@@ -176,5 +176,3 @@ jobs:
         working-directory: wasm
       - run: npm publish --access public --provenance
         working-directory: wasm
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}

CHANGELOG.md

@@ -9,6 +9,16 @@ single version, bumped together via `make bump-{patch,minor,major}`.
 ## [Unreleased]
 
 ### Added
+- Among-site rate-variation corrections for the distance models: a gamma
+  rate-heterogeneity shape parameter (`gamma_shape`, Jin & Nei 1990) and a
+  proportion of invariant sites (`p_invar`). For distance methods both are exact
+  closed-form corrections to each model's `−ln(...)` step, so no discrete rate
+  categories are needed; they apply to every correctable model (`JukesCantor`,
+  `Kimura2P`, `TajimaNei`, `Tamura`, `Poisson`, `KimuraProtein`) and compose as
+  `+I+Γ`. The gamma correction converges back to the uncorrected model as the
+  shape parameter grows; `PDiff` (raw p-distance) is unaffected. Exposed across
+  the CLI (`-g/--gamma-shape`, `-i/--p-invar`), Python, and WASM bindings, with
+  `NJError::InvalidGammaShape` / `NJError::InvalidPInvar` for out-of-range values.
 - Three substitution models: `TajimaNei` (Tajima-Nei 1984, DNA — corrects
   Jukes-Cantor for unequal base frequencies), `Tamura` (Tamura 1992, DNA —
   Kimura two-parameter with a GC-content correction), and `KimuraProtein`

CLAUDE.md

@@ -91,6 +91,8 @@ FASTA input / SequenceObject list
 
 Models implement `ModelCalculation<A: AlphabetEncoding>`. DNA-only: `JukesCantor`, `Kimura2P`, `TajimaNei`, `Tamura`. Protein-only: `Poisson`, `KimuraProtein`. `PDiff` works for both alphabets. Model–alphabet compatibility is enforced at runtime inside `nj()` (the `dispatch_run!` macro in `lib.rs`).
 
+`distance()` takes a `RateHet { gamma_shape: Option<f64>, p_invar: f64 }` value (threaded through `into_dist`/`from_msa`/`pairwise_distance_with`) that applies among-site rate variation to each model's `−c·ln(arg)` term: gamma rate heterogeneity becomes `c·α·(arg^(−1/α) − 1)` via the shared `corrected_term` helper, and invariant sites divide observed proportions by `(1 − p_invar)` then scale the result by `(1 − p_invar)`. `RateHet::NONE` reproduces the classic formulas exactly; `PDiff` ignores it. The `gamma_shape`/`p_invar` config fields are validated in `validate_rate_params` (`NJError::InvalidGammaShape` / `InvalidPInvar`).
+
 ### Bootstrap support
 
 Bootstrap runs N replicate NJ trees on column-resampled MSAs, tallies clade membership using `BitVec`-keyed `HashMap` counters (`count_clades`), then maps support values back onto the main tree's internal nodes (`add_bootstrap_to_tree`).

nj/benches/nj_bench.rs

@@ -47,6 +47,8 @@ fn dist_config(msa: Vec<SequenceObject>) -> DistConfig {
         substitution_model: SubstitutionModel::JukesCantor,
         alphabet: None,
         num_threads: None,
+        gamma_shape: None,
+        p_invar: None,
     }
 }
 
@@ -59,6 +61,8 @@ fn nj_config(msa: Vec<SequenceObject>, n_bootstrap_samples: usize) -> NJConfig {
         num_threads: None,
         return_distance_matrix: false,
         return_average_distance: false,
+        gamma_shape: None,
+        p_invar: None,
     }
 }
 

nj/src/config.rs

@@ -64,6 +64,16 @@ pub struct DistConfig {
     /// When `None` (the default), Rayon uses all available hardware threads.
     #[serde(default)]
     pub num_threads: Option<usize>,
+    /// Gamma rate-heterogeneity shape parameter `α` (must be `> 0`). When `None`
+    /// (the default), substitution rates are uniform across sites. Has no effect
+    /// on the `PDiff` model. See [`crate::models::RateHet`].
+    #[serde(default)]
+    pub gamma_shape: Option<f64>,
+    /// Proportion of invariant sites in `[0, 1)`. When `None` (the default), no
+    /// invariant-sites correction is applied. Has no effect on the `PDiff` model.
+    /// See [`crate::models::RateHet`].
+    #[serde(default)]
+    pub p_invar: Option<f64>,
 }
 
 /// Full configuration for a single Neighbor-Joining run.
@@ -105,6 +115,16 @@ pub struct NJConfig {
     /// Defaults to `false`.
     #[serde(default)]
     pub return_average_distance: bool,
+    /// Gamma rate-heterogeneity shape parameter `α` (must be `> 0`). When `None`
+    /// (the default), substitution rates are uniform across sites. Has no effect
+    /// on the `PDiff` model. See [`crate::models::RateHet`].
+    #[serde(default)]
+    pub gamma_shape: Option<f64>,
+    /// Proportion of invariant sites in `[0, 1)`. When `None` (the default), no
+    /// invariant-sites correction is applied. Has no effect on the `PDiff` model.
+    /// See [`crate::models::RateHet`].
+    #[serde(default)]
+    pub p_invar: Option<f64>,
 }
 
 /// Combined result returned by [`crate::nj`].

nj/src/distance_matrix.rs

@@ -9,7 +9,7 @@
 
 use crate::MSA;
 use crate::alphabet::AlphabetEncoding;
-use crate::models::{ModelCalculation, pairwise_distance};
+use crate::models::{ModelCalculation, RateHet, pairwise_distance_with};
 use crate::nj::NJState;
 use crate::tree::TreeNode;
 use serde::{Deserialize, Serialize};
@@ -86,7 +86,7 @@ impl DistMat {
     ///
     /// The `Send + Sync` bounds are required for the `parallel` feature and are
     /// trivially satisfied by all built-in alphabet and model types.
-    pub fn from_msa<M, A>(msa: &MSA<A>) -> DistMat
+    pub fn from_msa<M, A>(msa: &MSA<A>, rates: RateHet) -> DistMat
     where
         M: ModelCalculation<A> + Send + Sync,
         A: AlphabetEncoding + Sync,
@@ -104,7 +104,8 @@ impl DistMat {
                 .collect::<Vec<_>>()
                 .into_par_iter()
                 .map(|(i, j)| {
-                    let d = pairwise_distance::<M, A>(&msa.sequences[i], &msa.sequences[j]);
+                    let d =
+                        pairwise_distance_with::<M, A>(&msa.sequences[i], &msa.sequences[j], rates);
                     (i, j, d)
                 })
                 .collect();
@@ -118,7 +119,11 @@ impl DistMat {
         {
             for i in 0..n {
                 for j in 0..i {
-                    dist.set(i, j, pairwise_distance::<M, A>(&msa.sequences[i], &msa.sequences[j]));
+                    dist.set(
+                        i,
+                        j,
+                        pairwise_distance_with::<M, A>(&msa.sequences[i], &msa.sequences[j], rates),
+                    );
                 }
             }
         }
@@ -234,7 +239,7 @@ mod tests {
         // seq0 vs seq1 differ at middle position only
         let seqs: Vec<String> = vec!["ACG".into(), "ATG".into(), "A-G".into()];
         let msa = MSA::<DNA>::from_unnamed_sequences(seqs).unwrap();
-        let mat = msa.into_dist::<PDiff>();
+        let mat = msa.into_dist::<PDiff>(RateHet::NONE);
         // names default Seq0, Seq1, Seq2
         assert_eq!(mat.names, vec!["Seq0", "Seq1", "Seq2"]);
         // seq0 vs seq1: one mismatch across three aligned positions -> 1/3
@@ -251,7 +256,7 @@ mod tests {
         // n_comparable=2, diffs=1 (T vs C at pos 1), so 1/2.
         let seqs: Vec<String> = vec!["AT-".into(), "ACG".into()];
         let msa = MSA::<DNA>::from_unnamed_sequences(seqs).unwrap();
-        let mat = msa.into_dist::<PDiff>();
+        let mat = msa.into_dist::<PDiff>(RateHet::NONE);
 
         assert!((mat.get(0, 1) - (1.0 / 2.0)).abs() < 1e-12);
     }
@@ -261,7 +266,7 @@ mod tests {
         // no non-gap mismatches are counted, so the distance remains 0.0
         let seqs = vec![("X".into(), "A--".into()), ("Y".into(), "--A".into())];
         let msa = MSA::<DNA>::from_iter(seqs);
-        let mat = msa.into_dist::<PDiff>();
+        let mat = msa.into_dist::<PDiff>(RateHet::NONE);
         assert_eq!(mat.names, vec!["X", "Y"]);
         assert!((mat.get(0, 1) - 0.0).abs() < 1e-12);
     }

nj/src/error.rs

@@ -25,6 +25,10 @@ pub enum NJError {
         model: SubstitutionModel,
         alphabet: Alphabet,
     },
+    /// The gamma rate-heterogeneity shape parameter is not finite and positive.
+    InvalidGammaShape { value: f64 },
+    /// The proportion of invariant sites is not finite and in `[0, 1)`.
+    InvalidPInvar { value: f64 },
     /// An internal NJ algorithm failure (should be unreachable for valid input).
     AlgorithmFailure(String),
     /// Failed to collect entropy from the OS PRNG (bootstrap only).
@@ -46,6 +50,8 @@ impl NJError {
             NJError::SequenceLengthMismatch { .. } => "SequenceLengthMismatch",
             NJError::DuplicateIdentifier { .. } => "DuplicateIdentifier",
             NJError::IncompatibleModel { .. } => "IncompatibleModel",
+            NJError::InvalidGammaShape { .. } => "InvalidGammaShape",
+            NJError::InvalidPInvar { .. } => "InvalidPInvar",
             NJError::AlgorithmFailure(_) => "AlgorithmFailure",
             NJError::RngError(_) => "RngError",
             NJError::ParseError(_) => "ParseError",
@@ -69,6 +75,14 @@ impl fmt::Display for NJError {
                 f,
                 "Substitution model {model:?} is incompatible with {alphabet:?} alphabet"
             ),
+            NJError::InvalidGammaShape { value } => write!(
+                f,
+                "Gamma shape parameter must be a finite positive number, got {value}"
+            ),
+            NJError::InvalidPInvar { value } => write!(
+                f,
+                "Proportion of invariant sites must be a finite number in [0, 1), got {value}"
+            ),
             NJError::AlgorithmFailure(msg) => write!(f, "NJ algorithm failure: {msg}"),
             NJError::RngError(msg) => write!(f, "RNG error: {msg}"),
             NJError::ParseError(msg) => write!(f, "FASTA parse error: {msg}"),