solow_stats/
proportion.rs1use crate::weightstats::Alternative;
9use solow_distributions::special::betaincinv;
10use solow_distributions::{norm_cdf, norm_isf, norm_sf};
11
12pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum ConfintMethod {
58 Normal,
60 AgrestiCoull,
62 Wilson,
64 Beta,
66 Jeffreys,
68}
69
70fn 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
81fn beta_isf(p: f64, a: f64, b: f64) -> f64 {
83 beta_ppf(1.0 - p, a, b)
84}
85
86pub 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 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}