solow_stats/
rates.rs

1//! Two-sample Poisson rate comparison.
2//!
3//! [`test_poisson_2indep`] tests the equality (ratio or difference) of two
4//! independent Poisson intensity rates, mirroring the reference
5//! `rates.test_poisson_2indep`. The closed-form score / Wald / log / sqrt
6//! statistics referenced to the normal distribution are implemented for both
7//! the ratio and difference comparisons, along with the exact-conditional and
8//! conditional mid-p tests based on the binomial distribution. The simulation
9//! /grid `etest` variants are out of scope.
10
11use crate::weightstats::Alternative;
12use solow_distributions::special::{betainc, lgamma};
13use solow_distributions::{norm_cdf, norm_sf};
14
15/// Comparison target for [`test_poisson_2indep`].
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum Compare {
18    /// Test the ratio `rate1 / rate2` against `value` (default `value = 1`).
19    Ratio,
20    /// Test the difference `rate1 - rate2` against `value` (default `value = 0`).
21    Diff,
22}
23
24/// Test statistic / p-value method for [`test_poisson_2indep`].
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum PoissonMethod {
27    /// Wald test, variance from the observed rates (`ratio` and `diff`).
28    Wald,
29    /// Wald test with a continuity-corrected variance (`diff` only).
30    WaldCcv,
31    /// Score test, variance from the null-constrained estimate (`ratio`, `diff`).
32    Score,
33    /// Wald test on the log-ratio (`ratio` only).
34    WaldLog,
35    /// Score test on the log-ratio (`ratio` only).
36    ScoreLog,
37    /// Variance-stabilising square-root transformation test (`ratio` only).
38    Sqrt,
39    /// Exact conditional test based on the binomial distribution (`ratio` only).
40    ExactCond,
41    /// Mid-p value of the exact conditional test (`ratio` only).
42    CondMidp,
43}
44
45/// Result of [`test_poisson_2indep`].
46#[derive(Debug, Clone, Copy)]
47pub struct PoissonResult {
48    /// Test statistic. `NaN` for the binomial (`ExactCond` / `CondMidp`) tests.
49    pub statistic: f64,
50    /// p-value of the test.
51    pub pvalue: f64,
52    /// Estimated rate of sample 1, `count1 / exposure1`.
53    pub rate1: f64,
54    /// Estimated rate of sample 2, `count2 / exposure2`.
55    pub rate2: f64,
56    /// Observed rate ratio `rate1 / rate2`.
57    pub ratio: f64,
58    /// Observed rate difference `rate1 - rate2`.
59    pub diff: f64,
60}
61
62/// log of the binomial coefficient `C(n, k)`.
63fn ln_binom(n: f64, k: f64) -> f64 {
64    lgamma(n + 1.0) - lgamma(k + 1.0) - lgamma(n - k + 1.0)
65}
66
67/// Binomial probability mass `P(X = k)` for `X ~ Binom(n, p)`.
68fn binom_pmf(k: f64, n: f64, p: f64) -> f64 {
69    if k < 0.0 || k > n {
70        return 0.0;
71    }
72    if p <= 0.0 {
73        return if k == 0.0 { 1.0 } else { 0.0 };
74    }
75    if p >= 1.0 {
76        return if k == n { 1.0 } else { 0.0 };
77    }
78    (ln_binom(n, k) + k * p.ln() + (n - k) * (1.0 - p).ln()).exp()
79}
80
81/// Binomial CDF `P(X <= k)` for `X ~ Binom(n, p)` via the regularised
82/// incomplete beta function: `P(X <= k) = I_{1-p}(n - k, k + 1)`.
83fn binom_cdf(k: f64, n: f64, p: f64) -> f64 {
84    let k = k.floor();
85    if k < 0.0 {
86        return 0.0;
87    }
88    if k >= n {
89        return 1.0;
90    }
91    betainc(n - k, k + 1.0, 1.0 - p)
92}
93
94/// Binomial survival function `P(X > k) = 1 - P(X <= k)`.
95fn binom_sf(k: f64, n: f64, p: f64) -> f64 {
96    1.0 - binom_cdf(k, n, p)
97}
98
99/// Implicit binary search used by the two-sided binomial test: returns the
100/// index `i` in `[lo, hi]` such that `a(i) <= d < a(i+1)`, where `a` is assumed
101/// monotone increasing over the range. Mirrors the reference helper.
102fn binary_search_binom(a: &dyn Fn(f64) -> f64, d: f64, mut lo: f64, mut hi: f64) -> f64 {
103    while lo < hi {
104        let mid = lo + ((hi - lo) / 2.0).floor();
105        let midval = a(mid);
106        if midval < d {
107            lo = mid + 1.0;
108        } else if midval > d {
109            hi = mid - 1.0;
110        } else {
111            return mid;
112        }
113    }
114    if a(lo) <= d {
115        lo
116    } else {
117        lo - 1.0
118    }
119}
120
121/// Two-sided exact binomial p-value using the "minlike" method, reproducing
122/// `scipy.stats.binomtest(k, n, p).pvalue`.
123fn binom_test_two_sided(k: f64, n: f64, p: f64) -> f64 {
124    let d = binom_pmf(k, n, p);
125    let rerr = 1.0 + 1e-7;
126    let pval = if k == p * n {
127        1.0
128    } else if k < p * n {
129        // Search the upper tail (mode .. n) for terms <= d*rerr.
130        let neg_pmf = |x: f64| -binom_pmf(x, n, p);
131        let ix = binary_search_binom(&neg_pmf, -d * rerr, (p * n).ceil(), n);
132        let y = n - ix
133            + if d * rerr == binom_pmf(ix, n, p) {
134                1.0
135            } else {
136                0.0
137            };
138        binom_cdf(k, n, p) + binom_sf(n - y, n, p)
139    } else {
140        // Search the lower tail (0 .. mode) for terms <= d*rerr.
141        let pmf = |x: f64| binom_pmf(x, n, p);
142        let ix = binary_search_binom(&pmf, d * rerr, 0.0, (p * n).floor());
143        let y = ix + 1.0;
144        binom_cdf(y - 1.0, n, p) + binom_sf(k - 1.0, n, p)
145    };
146    pval.min(1.0)
147}
148
149/// Binomial p-value used by the conditional Poisson tests for a given
150/// alternative; `count` successes in `total` trials under success probability
151/// `prop`.
152fn binom_test(count: f64, total: f64, prop: f64, alternative: Alternative) -> f64 {
153    match alternative {
154        Alternative::TwoSided => binom_test_two_sided(count, total, prop),
155        Alternative::Larger => binom_sf(count - 1.0, total, prop),
156        Alternative::Smaller => binom_cdf(count, total, prop),
157    }
158}
159
160/// p-value of a normal (z) test statistic for the given alternative.
161fn z_pvalue(stat: f64, alternative: Alternative) -> f64 {
162    match alternative {
163        Alternative::TwoSided => norm_sf(stat.abs()) * 2.0,
164        Alternative::Larger => norm_sf(stat),
165        Alternative::Smaller => norm_cdf(stat),
166    }
167}
168
169/// Test the equality of two independent Poisson rates.
170///
171/// `count1`/`exposure1` and `count2`/`exposure2` are the event counts and total
172/// exposures of the two samples; `value` is the null ratio (default `1.0` for
173/// [`Compare::Ratio`]) or difference (default `0.0` for [`Compare::Diff`]). The
174/// `method` selects the test statistic (see [`PoissonMethod`]). Mirrors the
175/// reference `test_poisson_2indep`. Panics if a method is not valid for the
176/// chosen comparison.
177#[allow(clippy::too_many_arguments)]
178pub fn test_poisson_2indep(
179    count1: f64,
180    exposure1: f64,
181    count2: f64,
182    exposure2: f64,
183    value: Option<f64>,
184    method: PoissonMethod,
185    compare: Compare,
186    alternative: Alternative,
187) -> PoissonResult {
188    let (y1, n1, y2, n2) = (count1, exposure1, count2, exposure2);
189    let d = n2 / n1;
190    let rate1 = y1 / n1;
191    let rate2 = y2 / n2;
192
193    let (stat, pvalue) = match compare {
194        Compare::Ratio => {
195            let r = value.unwrap_or(1.0);
196            let r_d = r / d; // r1 * n1 / (r2 * n2)
197            match method {
198                PoissonMethod::Score => {
199                    let stat = (y1 - y2 * r_d) / ((y1 + y2) * r_d).sqrt();
200                    (stat, z_pvalue(stat, alternative))
201                }
202                PoissonMethod::Wald => {
203                    let stat = (y1 - y2 * r_d) / (y1 + y2 * r_d * r_d).sqrt();
204                    (stat, z_pvalue(stat, alternative))
205                }
206                PoissonMethod::ScoreLog => {
207                    let stat =
208                        ((y1 / y2).ln() - r_d.ln()) / ((2.0 + 1.0 / r_d + r_d) / (y1 + y2)).sqrt();
209                    (stat, z_pvalue(stat, alternative))
210                }
211                PoissonMethod::WaldLog => {
212                    let stat = ((y1 / y2).ln() - r_d.ln()) / (1.0 / y1 + 1.0 / y2).sqrt();
213                    (stat, z_pvalue(stat, alternative))
214                }
215                PoissonMethod::Sqrt => {
216                    let stat = 2.0 * ((y1 + 3.0 / 8.0).sqrt() - ((y2 + 3.0 / 8.0) * r_d).sqrt())
217                        / (1.0 + r_d).sqrt();
218                    (stat, z_pvalue(stat, alternative))
219                }
220                PoissonMethod::ExactCond => {
221                    let bp = r_d / (1.0 + r_d);
222                    let y_total = y1 + y2;
223                    (f64::NAN, binom_test(y1, y_total, bp, alternative))
224                }
225                PoissonMethod::CondMidp => {
226                    let bp = r_d / (1.0 + r_d);
227                    let y_total = y1 + y2;
228                    let p =
229                        binom_test(y1, y_total, bp, alternative) - 0.5 * binom_pmf(y1, y_total, bp);
230                    (f64::NAN, p)
231                }
232                PoissonMethod::WaldCcv => {
233                    panic!("waldccv is only defined for compare = diff");
234                }
235            }
236        }
237        Compare::Diff => {
238            let v = value.unwrap_or(0.0);
239            match method {
240                PoissonMethod::Wald => {
241                    let stat = (rate1 - rate2 - v) / (rate1 / n1 + rate2 / n2).sqrt();
242                    (stat, z_pvalue(stat, alternative))
243                }
244                PoissonMethod::WaldCcv => {
245                    let stat = (rate1 - rate2 - v)
246                        / ((y1 + 0.5) / (n1 * n1) + (y2 + 0.5) / (n2 * n2)).sqrt();
247                    (stat, z_pvalue(stat, alternative))
248                }
249                PoissonMethod::Score => {
250                    let count_pooled = y1 + y2;
251                    let rate_pooled = count_pooled / (n1 + n2);
252                    let dt = rate_pooled - v;
253                    let r2_cmle = 0.5 * (dt + (dt * dt + 4.0 * v * y2 / (n1 + n2)).sqrt());
254                    let r1_cmle = r2_cmle + v;
255                    let stat = (rate1 - rate2 - v) / (r1_cmle / n1 + r2_cmle / n2).sqrt();
256                    (stat, z_pvalue(stat, alternative))
257                }
258                _ => panic!("method is not valid for compare = diff"),
259            }
260        }
261    };
262
263    PoissonResult {
264        statistic: stat,
265        pvalue,
266        rate1,
267        rate2,
268        ratio: rate1 / rate2,
269        diff: rate1 - rate2,
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn score_ratio_runs() {
279        let r = test_poisson_2indep(
280            60.0,
281            51477.5,
282            30.0,
283            54308.7,
284            None,
285            PoissonMethod::Score,
286            Compare::Ratio,
287            Alternative::TwoSided,
288        );
289        assert!(r.statistic > 0.0);
290        assert!((0.0..=1.0).contains(&r.pvalue));
291    }
292
293    #[test]
294    fn binom_cdf_matches_pmf_sum() {
295        // CDF(3) should equal sum of PMF(0..=3).
296        let (n, p) = (10.0, 0.3);
297        let direct: f64 = (0..=3).map(|k| binom_pmf(k as f64, n, p)).sum();
298        assert!((binom_cdf(3.0, n, p) - direct).abs() < 1e-12);
299    }
300}
solow_stats/rates.rs

solow_stats/
rates.rs
