pmcore/estimation/nonparametric/
ipm.rs1use crate::estimation::nonparametric::{Psi, Weights};
2use anyhow::bail;
3use faer::linalg::triangular_solve::solve_lower_triangular_in_place;
4use faer::linalg::triangular_solve::solve_upper_triangular_in_place;
5use faer::{Col, Mat, Row};
6use rayon::prelude::*;
7
8pub fn burke(psi: &Psi) -> anyhow::Result<(Weights, f64)> {
10 let mut psi = psi.matrix().to_owned();
11
12 psi.row_iter_mut().try_for_each(|row| {
13 row.iter_mut().try_for_each(|x| {
14 if !x.is_finite() {
15 bail!("Input matrix must have finite entries")
16 } else {
17 *x = x.abs();
18 Ok(())
19 }
20 })
21 })?;
22
23 let (n_sub, n_point) = psi.shape();
24 let ecol: Col<f64> = Col::from_fn(n_point, |_| 1.0);
25 let erow: Row<f64> = Row::from_fn(n_sub, |_| 1.0);
26 let mut plam: Col<f64> = &psi * &ecol;
27 let eps: f64 = 1e-8;
28 let mut sig: f64 = 0.0;
29 let mut lam = ecol.clone();
30 let mut w: Col<f64> = Col::from_fn(plam.nrows(), |i| 1.0 / plam.get(i));
31 let mut ptw: Col<f64> = psi.transpose() * &w;
32
33 let ptw_max = ptw.iter().fold(f64::NEG_INFINITY, |acc, &x| x.max(acc));
34 let shrink = 2.0 * ptw_max;
35 lam *= shrink;
36 plam *= shrink;
37 w /= shrink;
38 ptw /= shrink;
39
40 let mut y: Col<f64> = &ecol - &ptw;
41 let mut r: Col<f64> = Col::from_fn(n_sub, |i| erow.get(i) - w.get(i) * plam.get(i));
42 let mut norm_r: f64 = r.iter().fold(0.0, |max, &val| max.max(val.abs()));
43 let sum_log_plam: f64 = plam.iter().map(|x| x.ln()).sum();
44 let sum_log_w: f64 = w.iter().map(|x| x.ln()).sum();
45 let mut gap: f64 = (sum_log_w + sum_log_plam).abs() / (1.0 + sum_log_plam);
46 let mut mu = lam.transpose() * &y / n_point as f64;
47
48 let mut psi_inner: Mat<f64> = Mat::zeros(psi.nrows(), psi.ncols());
49 let n_threads = faer::get_global_parallelism().degree();
50 let rows = psi.nrows();
51 let mut output: Vec<Mat<f64>> = (0..n_threads).map(|_| Mat::zeros(rows, rows)).collect();
52 let mut h: Mat<f64> = Mat::zeros(rows, rows);
53
54 while mu > eps || norm_r > eps || gap > eps {
55 let smu = sig * mu;
56 let inner = Col::from_fn(lam.nrows(), |i| lam.get(i) / y.get(i));
57 let w_plam = Col::from_fn(plam.nrows(), |i| plam.get(i) / w.get(i));
58
59 if psi.ncols() > n_threads * 128 {
60 psi_inner
61 .par_col_partition_mut(n_threads)
62 .zip(psi.par_col_partition(n_threads))
63 .zip(inner.par_partition(n_threads))
64 .zip(output.par_iter_mut())
65 .for_each(|(((mut psi_inner, psi), inner), output)| {
66 psi_inner
67 .as_mut()
68 .col_iter_mut()
69 .zip(psi.col_iter())
70 .zip(inner.iter())
71 .for_each(|((col, psi_col), inner_val)| {
72 col.iter_mut().zip(psi_col.iter()).for_each(|(x, psi_val)| {
73 *x = psi_val * inner_val;
74 });
75 });
76 faer::linalg::matmul::triangular::matmul(
77 output.as_mut(),
78 faer::linalg::matmul::triangular::BlockStructure::TriangularLower,
79 faer::Accum::Replace,
80 &psi_inner,
81 faer::linalg::matmul::triangular::BlockStructure::Rectangular,
82 psi.transpose(),
83 faer::linalg::matmul::triangular::BlockStructure::Rectangular,
84 1.0,
85 faer::Par::Seq,
86 );
87 });
88
89 let mut first_iter = true;
90 for output in &output {
91 if first_iter {
92 h.copy_from(output);
93 first_iter = false;
94 } else {
95 h += output;
96 }
97 }
98 } else {
99 psi_inner
100 .as_mut()
101 .col_iter_mut()
102 .zip(psi.col_iter())
103 .zip(inner.iter())
104 .for_each(|((col, psi_col), inner_val)| {
105 col.iter_mut().zip(psi_col.iter()).for_each(|(x, psi_val)| {
106 *x = psi_val * inner_val;
107 });
108 });
109 faer::linalg::matmul::triangular::matmul(
110 h.as_mut(),
111 faer::linalg::matmul::triangular::BlockStructure::TriangularLower,
112 faer::Accum::Replace,
113 &psi_inner,
114 faer::linalg::matmul::triangular::BlockStructure::Rectangular,
115 psi.transpose(),
116 faer::linalg::matmul::triangular::BlockStructure::Rectangular,
117 1.0,
118 faer::Par::Seq,
119 );
120 }
121
122 for i in 0..h.nrows() {
123 h[(i, i)] += w_plam[i];
124 }
125
126 let uph = match h.llt(faer::Side::Lower) {
127 Ok(llt) => llt,
128 Err(_) => {
129 bail!("Error during Cholesky decomposition. The matrix might not be positive definite. This is usually due to model misspecification or numerical issues.")
130 }
131 };
132 let uph = uph.L().transpose().to_owned();
133
134 let smuyinv: Col<f64> = Col::from_fn(ecol.nrows(), |i| smu * (ecol[i] / y[i]));
135 let psi_dot_muyinv: Col<f64> = &psi * &smuyinv;
136 let rhsdw: Row<f64> = Row::from_fn(erow.ncols(), |i| erow[i] / w[i] - psi_dot_muyinv[i]);
137 let mut dw = Mat::from_fn(rhsdw.ncols(), 1, |i, _j| *rhsdw.get(i));
138
139 solve_lower_triangular_in_place(uph.transpose().as_ref(), dw.as_mut(), faer::Par::rayon(0));
140 solve_upper_triangular_in_place(uph.as_ref(), dw.as_mut(), faer::Par::rayon(0));
141
142 let dw = dw.col(0);
143 let dy = -(psi.transpose() * dw);
144 let inner_times_dy = Col::from_fn(ecol.nrows(), |i| inner[i] * dy[i]);
145 let dlam: Row<f64> =
146 Row::from_fn(ecol.nrows(), |i| smuyinv[i] - lam[i] - inner_times_dy[i]);
147
148 let ratio_dlam_lam = Row::from_fn(lam.nrows(), |i| dlam[i] / lam[i]);
149 let min_ratio_dlam = ratio_dlam_lam.iter().cloned().fold(f64::INFINITY, f64::min);
150 let mut alfpri: f64 = -1.0 / min_ratio_dlam.min(-0.5);
151 alfpri = (0.99995 * alfpri).min(1.0);
152
153 let ratio_dy_y = Row::from_fn(y.nrows(), |i| dy[i] / y[i]);
154 let min_ratio_dy = ratio_dy_y.iter().cloned().fold(f64::INFINITY, f64::min);
155 let ratio_dw_w = Row::from_fn(dw.nrows(), |i| dw[i] / w[i]);
156 let min_ratio_dw = ratio_dw_w.iter().cloned().fold(f64::INFINITY, f64::min);
157 let mut alfdual = -1.0 / min_ratio_dy.min(-0.5);
158 alfdual = alfdual.min(-1.0 / min_ratio_dw.min(-0.5));
159 alfdual = (0.99995 * alfdual).min(1.0);
160
161 lam += alfpri * dlam.transpose();
162 w += alfdual * dw;
163 y += alfdual * &dy;
164
165 mu = lam.transpose() * &y / n_point as f64;
166 plam = &psi * &lam;
167 r = Col::from_fn(n_sub, |i| erow.get(i) - w.get(i) * plam.get(i));
168 ptw -= alfdual * dy;
169
170 norm_r = r.norm_max();
171 let sum_log_plam: f64 = plam.iter().map(|x| x.ln()).sum();
172 let sum_log_w: f64 = w.iter().map(|x| x.ln()).sum();
173 gap = (sum_log_w + sum_log_plam).abs() / (1.0 + sum_log_plam);
174
175 if mu < eps && norm_r > eps {
176 sig = 1.0;
177 } else {
178 let candidate1 = (1.0 - alfpri).powi(2);
179 let candidate2 = (1.0 - alfdual).powi(2);
180 let candidate3 = (norm_r - mu) / (norm_r + 100.0 * mu);
181 sig = candidate1.max(candidate2).max(candidate3).min(0.3);
182 }
183 }
184
185 lam /= n_sub as f64;
186 let obj = (psi * &lam).iter().map(|x| x.ln()).sum();
187 let lam_sum: f64 = lam.iter().sum();
188 lam = &lam / lam_sum;
189
190 Ok((lam.into(), obj))
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use approx::assert_relative_eq;
197 use faer::Mat;
198
199 #[test]
200 fn test_burke_identity() {
201 let n = 100;
202 let mat = Mat::identity(n, n);
203 let psi = Psi::from(mat);
204 let (lam, _) = burke(&psi).unwrap();
205
206 let expected = 1.0 / n as f64;
207 for i in 0..n {
208 assert_relative_eq!(lam[i], expected, epsilon = 1e-10);
209 }
210 assert_relative_eq!(lam.iter().sum::<f64>(), 1.0, epsilon = 1e-10);
211 }
212
213 #[test]
214 fn test_burke_uniform_square() {
215 let n_sub = 10;
216 let n_point = 10;
217 let mat = Mat::from_fn(n_sub, n_point, |_, _| 1.0);
218 let psi = Psi::from(mat);
219 let (lam, _) = burke(&psi).unwrap();
220
221 assert_relative_eq!(lam.iter().sum::<f64>(), 1.0, epsilon = 1e-10);
222 let expected = 1.0 / n_point as f64;
223 for i in 0..n_point {
224 assert_relative_eq!(lam[i], expected, epsilon = 1e-10);
225 }
226 }
227
228 #[test]
229 fn test_burke_with_non_finite_values() {
230 let n_sub = 10;
231 let n_point = 10;
232 let mat = Mat::from_fn(n_sub, n_point, |i, j| {
233 if i == 0 && j == 0 {
234 f64::NAN
235 } else {
236 1.0
237 }
238 });
239 let psi = Psi::from(mat);
240 assert!(burke(&psi).is_err());
241 }
242}