pmcore/estimation/nonparametric/
statistics.rs1use anyhow::{bail, Result};
2use ndarray::{Array, Array1, Array2, Axis};
3
4pub fn median(data: &[f64]) -> f64 {
5 let mut data: Vec<f64> = data.to_vec();
6 data.sort_by(|a, b| a.partial_cmp(b).unwrap());
7
8 let size = data.len();
9 match size {
10 even if even % 2 == 0 => {
11 let fst = data.get(even / 2 - 1).unwrap();
12 let snd = data.get(even / 2).unwrap();
13 (fst + snd) / 2.0
14 }
15 odd => *data.get(odd / 2_usize).unwrap(),
16 }
17}
18
19pub fn weighted_median(data: &[f64], weights: &[f64]) -> f64 {
20 assert_eq!(
21 data.len(),
22 weights.len(),
23 "The length of data and weights must be the same"
24 );
25 assert!(
26 weights.iter().all(|&x| x >= 0.0),
27 "Weights must be non-negative, weights: {:?}",
28 weights
29 );
30
31 let mut weighted_data: Vec<(f64, f64)> = data
32 .iter()
33 .zip(weights.iter())
34 .map(|(&d, &w)| (d, w))
35 .collect();
36
37 weighted_data.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
38
39 let total_weight: f64 = weights.iter().sum();
40 let mut cumulative_sum = 0.0;
41
42 for (i, &(_, weight)) in weighted_data.iter().enumerate() {
43 cumulative_sum += weight;
44
45 if cumulative_sum == total_weight / 2.0 {
46 if i + 1 < weighted_data.len() {
47 return (weighted_data[i].0 + weighted_data[i + 1].0) / 2.0;
48 } else {
49 return weighted_data[i].0;
50 }
51 } else if cumulative_sum > total_weight / 2.0 {
52 return weighted_data[i].0;
53 }
54 }
55
56 unreachable!("The function should have returned a value before reaching this point.");
57}
58
59pub fn population_mean_median(
60 theta: &Array2<f64>,
61 w: &Array1<f64>,
62) -> Result<(Array1<f64>, Array1<f64>)> {
63 let w = if w.is_empty() {
64 tracing::warn!("w.len() == 0, setting all weights to 1/n");
65 Array1::from_elem(theta.nrows(), 1.0 / theta.nrows() as f64)
66 } else {
67 w.clone()
68 };
69
70 if theta.nrows() != w.len() {
71 bail!(
72 "Number of parameters and number of weights do not match. Theta: {}, w: {}",
73 theta.nrows(),
74 w.len()
75 );
76 }
77
78 let mut mean = Array1::zeros(theta.ncols());
79 let mut median = Array1::zeros(theta.ncols());
80
81 for (i, (mn, mdn)) in mean.iter_mut().zip(&mut median).enumerate() {
82 let col = theta.column(i).to_owned() * w.to_owned();
83 *mn = col.sum();
84
85 let ct = theta.column(i);
86 let mut params = vec![];
87 let mut weights = vec![];
88 for (ti, wi) in ct.iter().zip(w.clone()) {
89 params.push(*ti);
90 weights.push(wi);
91 }
92
93 *mdn = weighted_median(¶ms, &weights);
94 }
95
96 Ok((mean, median))
97}
98
99pub fn posterior_mean_median(
100 theta: &Array2<f64>,
101 psi: &Array2<f64>,
102 w: &Array1<f64>,
103) -> Result<(Array2<f64>, Array2<f64>)> {
104 let mut mean = Array2::zeros((0, theta.ncols()));
105 let mut median = Array2::zeros((0, theta.ncols()));
106
107 let w = if w.is_empty() {
108 tracing::warn!("w is empty, setting all weights to 1/n");
109 Array1::from_elem(theta.nrows(), 1.0 / theta.nrows() as f64)
110 } else {
111 w.clone()
112 };
113
114 if theta.nrows() != w.len() || theta.nrows() != psi.ncols() || psi.ncols() != w.len() {
115 bail!("Number of parameters and number of weights do not match, theta.nrows(): {}, w.len(): {}, psi.ncols(): {}", theta.nrows(), w.len(), psi.ncols());
116 }
117
118 let mut psi_norm: Array2<f64> = Array2::zeros((0, psi.ncols()));
119 for (i, row) in psi.axis_iter(Axis(0)).enumerate() {
120 let row_w = row.to_owned() * w.to_owned();
121 let row_sum = row_w.sum();
122 let row_norm = if row_sum == 0.0 {
123 tracing::warn!("Sum of row {} of psi is 0.0, setting that row to 1/n", i);
124 Array1::from_elem(psi.ncols(), 1.0 / psi.ncols() as f64)
125 } else {
126 &row_w / row_sum
127 };
128 psi_norm.push_row(row_norm.view())?;
129 }
130 if psi_norm.iter().any(|&x| x.is_nan()) {
131 dbg!(&psi);
132 bail!("NaN values found in psi_norm");
133 };
134
135 for probs in psi_norm.axis_iter(Axis(0)) {
136 let mut post_mean: Vec<f64> = Vec::new();
137 let mut post_median: Vec<f64> = Vec::new();
138
139 for pars in theta.axis_iter(Axis(1)) {
140 let weighted_par = &probs * &pars;
141 let the_mean = weighted_par.sum();
142 post_mean.push(the_mean);
143
144 let median = weighted_median(&pars.to_vec(), &probs.to_vec());
145 post_median.push(median);
146 }
147
148 mean.push_row(Array::from(post_mean.clone()).view())?;
149 median.push_row(Array::from(post_median.clone()).view())?;
150 }
151
152 Ok((mean, median))
153}
154
155#[cfg(test)]
156mod tests {
157 use super::{median, weighted_median};
158
159 #[test]
160 fn test_median_odd() {
161 let data = vec![1.0, 3.0, 2.0];
162 assert_eq!(median(&data), 2.0);
163 }
164
165 #[test]
166 fn test_median_even() {
167 let data = vec![1.0, 2.0, 3.0, 4.0];
168 assert_eq!(median(&data), 2.5);
169 }
170
171 #[test]
172 fn test_median_single() {
173 let data = vec![42.0];
174 assert_eq!(median(&data), 42.0);
175 }
176
177 #[test]
178 fn test_median_sorted() {
179 let data = vec![5.0, 10.0, 15.0, 20.0, 25.0];
180 assert_eq!(median(&data), 15.0);
181 }
182
183 #[test]
184 fn test_median_unsorted() {
185 let data = vec![10.0, 30.0, 20.0, 50.0, 40.0];
186 assert_eq!(median(&data), 30.0);
187 }
188
189 #[test]
190 fn test_median_with_duplicates() {
191 let data = vec![1.0, 2.0, 2.0, 3.0, 4.0];
192 assert_eq!(median(&data), 2.0);
193 }
194
195 #[test]
196 fn test_weighted_median_simple() {
197 let data = vec![1.0, 2.0, 3.0];
198 let weights = vec![0.2, 0.5, 0.3];
199 assert_eq!(weighted_median(&data, &weights), 2.0);
200 }
201
202 #[test]
203 fn test_weighted_median_even_weights() {
204 let data = vec![1.0, 2.0, 3.0, 4.0];
205 let weights = vec![0.25, 0.25, 0.25, 0.25];
206 assert_eq!(weighted_median(&data, &weights), 2.5);
207 }
208
209 #[test]
210 fn test_weighted_median_single_element() {
211 let data = vec![42.0];
212 let weights = vec![1.0];
213 assert_eq!(weighted_median(&data, &weights), 42.0);
214 }
215
216 #[test]
217 #[should_panic(expected = "The length of data and weights must be the same")]
218 fn test_weighted_median_mismatched_lengths() {
219 let data = vec![1.0, 2.0, 3.0];
220 let weights = vec![0.1, 0.2];
221 weighted_median(&data, &weights);
222 }
223
224 #[test]
225 fn test_weighted_median_all_same_elements() {
226 let data = vec![5.0, 5.0, 5.0, 5.0];
227 let weights = vec![0.1, 0.2, 0.3, 0.4];
228 assert_eq!(weighted_median(&data, &weights), 5.0);
229 }
230
231 #[test]
232 #[should_panic(expected = "Weights must be non-negative")]
233 fn test_weighted_median_negative_weights() {
234 let data = vec![1.0, 2.0, 3.0, 4.0];
235 let weights = vec![0.2, -0.5, 0.5, 0.8];
236 assert_eq!(weighted_median(&data, &weights), 4.0);
237 }
238}