Skip to main content

solow_stats/
proportion.rs

1//! Tests and confidence intervals for binomial proportions.
2//!
3//! Provides the one- and two-sample proportion z-test
4//! ([`proportions_ztest`]) and several proportion confidence-interval methods
5//! ([`proportion_confint`]). Mirrors the reference `proportions_ztest` and
6//! `proportion_confint`.
7
8use crate::weightstats::Alternative;
9use solow_distributions::special::betaincinv;
10use solow_distributions::{norm_cdf, norm_isf, norm_sf};
11
12/// One- or two-sample test for proportions based on the normal approximation.
13///
14/// `count` and `nobs` are equal-length slices of successes and trials. For a
15/// single element this is the one-sample test against `value` (the null
16/// proportion). For two elements it tests `prop[0] − prop[1] = value`. The
17/// variance uses the pooled proportion. Returns `(zstat, pvalue)`. Mirrors the
18/// reference `proportions_ztest`.
19pub fn proportions_ztest(
20    count: &[f64],
21    nobs: &[f64],
22    value: f64,
23    alternative: Alternative,
24) -> (f64, f64) {
25    assert_eq!(count.len(), nobs.len(), "count and nobs length mismatch");
26    let k = count.len();
27    assert!(
28        k == 1 || k == 2,
29        "only one- and two-sample tests are supported"
30    );
31
32    let prop: Vec<f64> = count.iter().zip(nobs).map(|(&c, &n)| c / n).collect();
33    let diff = if k == 1 {
34        prop[0] - value
35    } else {
36        prop[0] - prop[1] - value
37    };
38
39    let count_sum: f64 = count.iter().sum();
40    let nobs_sum: f64 = nobs.iter().sum();
41    let p_pooled = count_sum / nobs_sum;
42    let nobs_fact: f64 = nobs.iter().map(|&n| 1.0 / n).sum();
43    let var = p_pooled * (1.0 - p_pooled) * nobs_fact;
44    let std_diff = var.sqrt();
45
46    let zstat = diff / std_diff;
47    let pvalue = match alternative {
48        Alternative::TwoSided => norm_sf(zstat.abs()) * 2.0,
49        Alternative::Larger => norm_sf(zstat),
50        Alternative::Smaller => norm_cdf(zstat),
51    };
52    (zstat, pvalue)
53}
54
55/// Confidence-interval method for a binomial proportion.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum ConfintMethod {
58    /// Asymptotic normal (Wald) interval (clipped to `[0, 1]`).
59    Normal,
60    /// Agresti–Coull interval (clipped to `[0, 1]`).
61    AgrestiCoull,
62    /// Wilson score interval.
63    Wilson,
64    /// Clopper–Pearson exact (Beta) interval.
65    Beta,
66    /// Jeffreys Bayesian interval.
67    Jeffreys,
68}
69
70/// Inverse-CDF of the Beta(a, b) distribution at probability `p`.
71fn beta_ppf(p: f64, a: f64, b: f64) -> f64 {
72    if p <= 0.0 {
73        return 0.0;
74    }
75    if p >= 1.0 {
76        return 1.0;
77    }
78    betaincinv(a, b, p)
79}
80
81/// Survival-function inverse (isf) of the Beta(a, b) distribution.
82fn beta_isf(p: f64, a: f64, b: f64) -> f64 {
83    beta_ppf(1.0 - p, a, b)
84}
85
86/// `(1 − alpha)` confidence interval for a binomial proportion.
87///
88/// `count` successes in `nobs` trials. Returns `(lower, upper)`. Mirrors the
89/// reference `proportion_confint` for the supported `method`s.
90pub fn proportion_confint(count: f64, nobs: f64, alpha: f64, method: ConfintMethod) -> (f64, f64) {
91    let q = count / nobs;
92    let alpha_2 = 0.5 * alpha;
93    let crit = norm_isf(alpha / 2.0);
94
95    let (mut lo, mut hi) = match method {
96        ConfintMethod::Normal => {
97            let std = (q * (1.0 - q) / nobs).sqrt();
98            let dist = crit * std;
99            (q - dist, q + dist)
100        }
101        ConfintMethod::AgrestiCoull => {
102            let nobs_c = nobs + crit * crit;
103            let q_c = (count + crit * crit / 2.0) / nobs_c;
104            let std_c = (q_c * (1.0 - q_c) / nobs_c).sqrt();
105            let dist = crit * std_c;
106            (q_c - dist, q_c + dist)
107        }
108        ConfintMethod::Wilson => {
109            let crit2 = crit * crit;
110            let denom = 1.0 + crit2 / nobs;
111            let center = (q + crit2 / (2.0 * nobs)) / denom;
112            let mut dist = crit * (q * (1.0 - q) / nobs + crit2 / (4.0 * nobs * nobs)).sqrt();
113            dist /= denom;
114            (center - dist, center + dist)
115        }
116        ConfintMethod::Beta => {
117            let mut ci_low = beta_ppf(alpha_2, count, nobs - count + 1.0);
118            let mut ci_upp = beta_isf(alpha_2, count + 1.0, nobs - count);
119            if q == 0.0 {
120                ci_low = 0.0;
121            }
122            if q == 1.0 {
123                ci_upp = 1.0;
124            }
125            (ci_low, ci_upp)
126        }
127        ConfintMethod::Jeffreys => {
128            // beta.interval(1 - alpha, count + .5, nobs - count + .5)
129            let a = count + 0.5;
130            let b = nobs - count + 0.5;
131            (beta_ppf(alpha_2, a, b), beta_ppf(1.0 - alpha_2, a, b))
132        }
133    };
134
135    if matches!(method, ConfintMethod::Normal | ConfintMethod::AgrestiCoull) {
136        lo = lo.clamp(0.0, 1.0);
137        hi = hi.clamp(0.0, 1.0);
138    }
139    (lo, hi)
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn one_sample_ztest_sign() {
148        let (z, p) = proportions_ztest(&[45.0], &[100.0], 0.5, Alternative::TwoSided);
149        assert!(z < 0.0);
150        assert!((0.0..=1.0).contains(&p));
151    }
152
153    #[test]
154    fn normal_confint_brackets_estimate() {
155        let (lo, hi) = proportion_confint(45.0, 100.0, 0.05, ConfintMethod::Normal);
156        assert!(lo < 0.45 && hi > 0.45);
157    }
158}