1extern crate alloc;
10use alloc::vec::Vec;
11
12use crate::analysis::{compute_effect_estimate, compute_effect_estimate_analytical};
13use crate::math::sqrt;
14use crate::result::{EffectEstimate, MeasurementQuality};
15use crate::types::{Matrix9, Vector9};
16
17#[derive(Clone, Debug)]
24pub struct Posterior {
25 pub delta_post: Vector9,
27
28 pub lambda_post: Matrix9,
30
31 pub delta_draws: Vec<Vector9>,
34
35 pub leak_probability: f64,
38
39 pub theta: f64,
41
42 pub n: usize,
44
45 pub lambda_mean: Option<f64>,
49
50 pub lambda_mixing_ok: Option<bool>,
54
55 pub kappa_mean: Option<f64>,
58
59 pub kappa_cv: Option<f64>,
62
63 pub kappa_ess: Option<f64>,
66
67 pub kappa_mixing_ok: Option<bool>,
70}
71
72impl Posterior {
73 #[allow(clippy::too_many_arguments)]
75 pub fn new(
76 delta_post: Vector9,
77 lambda_post: Matrix9,
78 delta_draws: Vec<Vector9>,
79 leak_probability: f64,
80 theta: f64,
81 n: usize,
82 ) -> Self {
83 Self {
84 delta_post,
85 lambda_post,
86 delta_draws,
87 leak_probability,
88 theta,
89 n,
90 lambda_mean: None, lambda_mixing_ok: None, kappa_mean: None, kappa_cv: None, kappa_ess: None, kappa_mixing_ok: None, }
97 }
98
99 #[allow(clippy::too_many_arguments)]
101 pub fn new_with_gibbs(
102 delta_post: Vector9,
103 lambda_post: Matrix9,
104 delta_draws: Vec<Vector9>,
105 leak_probability: f64,
106 theta: f64,
107 n: usize,
108 lambda_mean: f64,
109 lambda_mixing_ok: bool,
110 kappa_mean: f64,
111 kappa_cv: f64,
112 kappa_ess: f64,
113 kappa_mixing_ok: bool,
114 ) -> Self {
115 Self {
116 delta_post,
117 lambda_post,
118 delta_draws,
119 leak_probability,
120 theta,
121 n,
122 lambda_mean: Some(lambda_mean),
123 lambda_mixing_ok: Some(lambda_mixing_ok),
124 kappa_mean: Some(kappa_mean),
125 kappa_cv: Some(kappa_cv),
126 kappa_ess: Some(kappa_ess),
127 kappa_mixing_ok: Some(kappa_mixing_ok),
128 }
129 }
130
131 pub fn max_effect_ns(&self) -> f64 {
133 self.delta_post
134 .iter()
135 .map(|x| x.abs())
136 .fold(0.0_f64, f64::max)
137 }
138
139 pub fn to_effect_estimate(&self) -> EffectEstimate {
143 if !self.delta_draws.is_empty() {
144 compute_effect_estimate(&self.delta_draws, self.theta)
145 } else {
146 compute_effect_estimate_analytical(&self.delta_post, &self.lambda_post, self.theta)
147 }
148 }
149
150 pub fn measurement_quality(&self) -> MeasurementQuality {
155 let max_se = (0..9)
157 .map(|k| sqrt(self.lambda_post[(k, k)].max(1e-12)))
158 .fold(0.0_f64, f64::max);
159 MeasurementQuality::from_mde_ns(max_se * 2.0)
160 }
161
162 pub fn to_summary(&self) -> crate::ffi_summary::PosteriorSummary {
164 let effect = self.to_effect_estimate();
165
166 crate::ffi_summary::PosteriorSummary {
167 max_effect_ns: effect.max_effect_ns,
168 ci_low_ns: effect.credible_interval_ns.0,
169 ci_high_ns: effect.credible_interval_ns.1,
170 leak_probability: self.leak_probability,
171 n: self.n,
172 lambda_mean: self.lambda_mean.unwrap_or(1.0),
173 lambda_mixing_ok: self.lambda_mixing_ok.unwrap_or(true),
174 kappa_mean: self.kappa_mean.unwrap_or(1.0),
175 kappa_cv: self.kappa_cv.unwrap_or(0.0),
176 kappa_ess: self.kappa_ess.unwrap_or(0.0),
177 kappa_mixing_ok: self.kappa_mixing_ok.unwrap_or(true),
178 }
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185
186 #[test]
187 fn test_posterior_accessors() {
188 let delta_post =
189 Vector9::from_row_slice(&[10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0]);
190 let lambda_post = Matrix9::identity();
191
192 let posterior = Posterior::new(
193 delta_post,
194 lambda_post,
195 Vec::new(), 0.75,
197 5.0, 1000,
199 );
200
201 assert_eq!(posterior.leak_probability, 0.75);
202 assert_eq!(posterior.n, 1000);
203 assert!((posterior.max_effect_ns() - 10.0).abs() < 1e-10);
204 }
205
206 #[test]
207 fn test_posterior_clone() {
208 let delta_post = Vector9::from_row_slice(&[5.0; 9]);
209 let lambda_post = Matrix9::identity();
210
211 let posterior = Posterior::new(
212 delta_post,
213 lambda_post,
214 Vec::new(), 0.5,
216 5.0, 500,
218 );
219
220 let cloned = posterior.clone();
221 assert_eq!(cloned.leak_probability, posterior.leak_probability);
222 assert_eq!(cloned.max_effect_ns(), posterior.max_effect_ns());
223 }
224
225 #[test]
226 fn test_effect_estimate_from_draws() {
227 let delta_post = Vector9::from_row_slice(&[10.0; 9]);
228 let lambda_post = Matrix9::identity();
229
230 let delta_draws: Vec<Vector9> = (0..100)
232 .map(|_| {
233 Vector9::from_row_slice(&[10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0])
234 })
235 .collect();
236
237 let posterior = Posterior::new(delta_post, lambda_post, delta_draws, 0.99, 5.0, 1000);
238
239 let effect = posterior.to_effect_estimate();
240 assert!(effect.max_effect_ns > 9.0, "max effect should be around 10");
241 }
242}