1use prospicio_core::{Error, Result};
12
13use crate::family::Family;
14
15#[derive(Debug, Clone, PartialEq)]
17pub struct PeriodSummary<K> {
18 pub period: Option<K>,
20 pub n: usize,
22 pub weight: f64,
24 pub actual: f64,
26 pub expected: f64,
28 pub std_dev: f64,
31}
32
33impl<K> PeriodSummary<K> {
34 pub fn ratio(&self) -> f64 {
36 self.actual / self.expected
37 }
38
39 pub fn z(&self) -> f64 {
42 (self.actual - self.expected) / self.std_dev
43 }
44}
45
46#[derive(Debug, Clone, PartialEq)]
48pub struct Monitor<K> {
49 pub periods: Vec<PeriodSummary<K>>,
51 pub total: PeriodSummary<K>,
53 pub trend: f64,
57 pub trend_std_error: f64,
59}
60
61impl<K> Monitor<K> {
62 pub fn trend_z(&self) -> f64 {
64 self.trend / self.trend_std_error
65 }
66}
67
68pub fn actual_vs_expected<K: Ord + Clone>(
91 periods: &[K],
92 y: &[f64],
93 mu: &[f64],
94 weights: Option<&[f64]>,
95 family: Family,
96 dispersion: f64,
97) -> Result<Monitor<K>> {
98 let n = y.len();
99 if periods.len() != n || mu.len() != n || weights.is_some_and(|w| w.len() != n) {
100 return Err(Error::Data(
101 "periods, y, mu and weights need one entry per row".into(),
102 ));
103 }
104 if n == 0 {
105 return Err(Error::Data("no rows to monitor".into()));
106 }
107 if !(dispersion.is_finite() && dispersion > 0.0) {
108 return Err(Error::InvalidParameter {
109 name: "dispersion",
110 value: dispersion,
111 reason: "must be finite and positive",
112 });
113 }
114 let mut keys: Vec<K> = periods.to_vec();
115 keys.sort();
116 keys.dedup();
117 let empty = |period| PeriodSummary {
118 period,
119 n: 0,
120 weight: 0.0,
121 actual: 0.0,
122 expected: 0.0,
123 std_dev: 0.0,
124 };
125 let mut rows: Vec<PeriodSummary<K>> = keys.iter().cloned().map(|k| empty(Some(k))).collect();
126 let mut var = vec![0.0; keys.len()];
127 for i in 0..n {
128 let w = weights.map_or(1.0, |w| w[i]);
129 let t = keys
130 .binary_search(&periods[i])
131 .expect("every period is a key");
132 let r = &mut rows[t];
133 r.n += 1;
134 r.weight += w;
135 r.actual += w * y[i];
136 r.expected += w * mu[i];
137 var[t] += dispersion * w * family.variance(mu[i]);
138 }
139 let mut total = empty(None);
140 for (r, v) in rows.iter_mut().zip(&var) {
141 r.std_dev = v.sqrt();
142 total.n += r.n;
143 total.weight += r.weight;
144 total.actual += r.actual;
145 total.expected += r.expected;
146 }
147 total.std_dev = var.iter().sum::<f64>().sqrt();
148
149 let (mut sw, mut st, mut sr, mut stt, mut str_) = (0.0, 0.0, 0.0, 0.0, 0.0);
152 for (t, row) in rows.iter().enumerate() {
153 let w = (row.expected / row.std_dev).powi(2);
154 let (tf, r) = (t as f64, row.ratio() - 1.0);
155 sw += w;
156 st += w * tf;
157 sr += w * r;
158 stt += w * tf * tf;
159 str_ += w * tf * r;
160 }
161 let sxx = stt - st * st / sw;
162 let (trend, trend_std_error) = if rows.len() >= 3 && sxx > 0.0 {
163 ((str_ - st * sr / sw) / sxx, (1.0 / sxx).sqrt())
164 } else {
165 (f64::NAN, f64::NAN)
166 };
167 Ok(Monitor {
168 periods: rows,
169 total,
170 trend,
171 trend_std_error,
172 })
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 #[test]
180 fn totals_weights_and_trend() {
181 let periods: Vec<&str> = ["2021", "2022", "2023", "2024"]
183 .iter()
184 .flat_map(|p| [*p, *p])
185 .collect();
186 let mu = vec![5.0; 8];
187 let y = [4.0, 6.0, 5.0, 6.0, 6.0, 6.0, 6.0, 7.0];
188 let m = actual_vs_expected(&periods, &y, &mu, None, Family::Poisson, 1.0).unwrap();
189 assert_eq!(m.periods.len(), 4);
190 assert_eq!(m.periods[0].period, Some("2021"));
191 assert_eq!(
192 m.periods.iter().map(|p| p.actual).collect::<Vec<_>>(),
193 [10.0, 11.0, 12.0, 13.0]
194 );
195 assert_eq!(m.total.actual, 46.0);
196 assert!((m.total.std_dev - 40f64.sqrt()).abs() < 1e-12);
197 assert!((m.trend - 0.1).abs() < 1e-12);
199 assert!((m.trend_std_error - (0.1f64 / 5.0).sqrt()).abs() < 1e-12);
201
202 let w = vec![2.0; 8];
204 let g = actual_vs_expected(&periods, &y, &mu, Some(&w), Family::Gamma, 0.5).unwrap();
205 assert_eq!(g.periods[3].expected, 20.0);
206 assert!((g.periods[3].std_dev - (0.5f64 * 2.0 * 2.0 * 25.0).sqrt()).abs() < 1e-12);
207 }
208
209 #[test]
210 fn rejects_bad_input() {
211 let p = [1, 2];
212 assert!(actual_vs_expected(&p, &[1.0], &[1.0, 1.0], None, Family::Poisson, 1.0).is_err());
213 assert!(
214 actual_vs_expected(&p, &[1.0, 1.0], &[1.0, 1.0], None, Family::Poisson, 0.0).is_err()
215 );
216 let two =
217 actual_vs_expected(&p, &[1.0, 1.0], &[1.0, 1.0], None, Family::Poisson, 1.0).unwrap();
218 assert!(two.trend.is_nan());
219 }
220}