pmcore/estimation/nonparametric/
psi.rs

1use anyhow::bail;
2use anyhow::Result;
3use faer::Mat;
4use ndarray::{Array2, Axis};
5use pharmsol::prelude::simulator::log_likelihood_matrix;
6use pharmsol::AssayErrorModels;
7use pharmsol::Data;
8use pharmsol::Equation;
9use serde::{Deserialize, Serialize};
10
11use super::theta::Theta;
12
13/// [Psi] is a structure that holds the likelihood for each subject (row), for each support point (column)
14#[derive(Debug, Clone, PartialEq)]
15pub struct Psi {
16    matrix: Mat<f64>,
17    /// Per-subject log-likelihood offsets removed before exponentiation.
18    /// Row scaling leaves the optimal weights unchanged; these offsets restore
19    /// the original likelihood scale where required by an algorithm.
20    row_log_scales: Vec<f64>,
21}
22
23impl Psi {
24    pub fn new() -> Self {
25        Psi {
26            matrix: Mat::new(),
27            row_log_scales: Vec::new(),
28        }
29    }
30
31    pub fn matrix(&self) -> &Mat<f64> {
32        &self.matrix
33    }
34
35    pub(crate) fn log_scale(&self) -> f64 {
36        self.row_log_scales.iter().sum()
37    }
38
39    pub(crate) fn row_log_scales(&self) -> &[f64] {
40        &self.row_log_scales
41    }
42
43    pub(crate) fn from_log_likelihoods(mut log_likelihoods: Array2<f64>) -> Result<Self> {
44        let mut row_log_scales = Vec::with_capacity(log_likelihoods.nrows());
45
46        for mut row in log_likelihoods.axis_iter_mut(Axis(0)) {
47            let row_max = row.iter().copied().fold(f64::NEG_INFINITY, f64::max);
48            if !row_max.is_finite() {
49                bail!("Each subject must have at least one finite log-likelihood");
50            }
51
52            row_log_scales.push(row_max);
53            row.mapv_inplace(|value| (value - row_max).exp());
54        }
55
56        let matrix = Mat::from_fn(log_likelihoods.nrows(), log_likelihoods.ncols(), |i, j| {
57            log_likelihoods[(i, j)]
58        });
59        Ok(Self {
60            matrix,
61            row_log_scales,
62        })
63    }
64
65    pub fn nspp(&self) -> usize {
66        self.matrix.nrows()
67    }
68
69    pub fn nsub(&self) -> usize {
70        self.matrix.ncols()
71    }
72
73    pub fn to_ndarray(&self) -> Array2<f64> {
74        let m = &self.matrix;
75        Array2::from_shape_fn((m.nrows(), m.ncols()), |(i, j)| m[(i, j)])
76    }
77
78    pub(crate) fn filter_column_indices(&mut self, indices: &[usize]) {
79        let matrix = self.matrix.to_owned();
80
81        let new = Mat::from_fn(matrix.nrows(), indices.len(), |r, c| {
82            *matrix.get(r, indices[c])
83        });
84
85        self.matrix = new;
86    }
87
88    pub fn write(&self, path: &str) {
89        let mut writer = csv::Writer::from_path(path).unwrap();
90        for row in self.matrix.row_iter() {
91            writer
92                .write_record(row.iter().map(|x| x.to_string()))
93                .unwrap();
94        }
95    }
96
97    pub fn to_csv<W: std::io::Write>(&self, writer: W) -> Result<()> {
98        let mut csv_writer = csv::Writer::from_writer(writer);
99
100        for i in 0..self.matrix.nrows() {
101            let row: Vec<f64> = (0..self.matrix.ncols())
102                .map(|j| *self.matrix.get(i, j))
103                .collect();
104            csv_writer.serialize(row)?;
105        }
106
107        csv_writer.flush()?;
108        Ok(())
109    }
110
111    pub fn from_csv<R: std::io::Read>(reader: R) -> Result<Self> {
112        let mut csv_reader = csv::Reader::from_reader(reader);
113        let mut rows: Vec<Vec<f64>> = Vec::new();
114
115        for result in csv_reader.deserialize() {
116            let row: Vec<f64> = result?;
117            rows.push(row);
118        }
119
120        if rows.is_empty() {
121            bail!("CSV file is empty");
122        }
123
124        let nrows = rows.len();
125        let ncols = rows[0].len();
126
127        for (i, row) in rows.iter().enumerate() {
128            if row.len() != ncols {
129                bail!("Row {} has {} columns, expected {}", i, row.len(), ncols);
130            }
131        }
132
133        let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
134
135        Ok(Psi {
136            matrix: mat,
137            row_log_scales: vec![0.0; nrows],
138        })
139    }
140}
141
142impl Default for Psi {
143    fn default() -> Self {
144        Psi::new()
145    }
146}
147
148impl From<Array2<f64>> for Psi {
149    fn from(array: Array2<f64>) -> Self {
150        let nrows = array.nrows();
151        let matrix = Mat::from_fn(nrows, array.ncols(), |i, j| array[(i, j)]);
152        Psi {
153            matrix,
154            row_log_scales: vec![0.0; nrows],
155        }
156    }
157}
158
159impl From<Mat<f64>> for Psi {
160    fn from(matrix: Mat<f64>) -> Self {
161        let nrows = matrix.nrows();
162        Psi {
163            matrix,
164            row_log_scales: vec![0.0; nrows],
165        }
166    }
167}
168
169impl From<&Array2<f64>> for Psi {
170    fn from(array: &Array2<f64>) -> Self {
171        let nrows = array.nrows();
172        let matrix = Mat::from_fn(nrows, array.ncols(), |i, j| array[(i, j)]);
173        Psi {
174            matrix,
175            row_log_scales: vec![0.0; nrows],
176        }
177    }
178}
179
180impl Serialize for Psi {
181    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
182    where
183        S: serde::Serializer,
184    {
185        use serde::ser::SerializeSeq;
186
187        let mut seq = serializer.serialize_seq(Some(self.matrix.nrows()))?;
188
189        for i in 0..self.matrix.nrows() {
190            let row: Vec<f64> = (0..self.matrix.ncols())
191                .map(|j| *self.matrix.get(i, j))
192                .collect();
193            seq.serialize_element(&row)?;
194        }
195
196        seq.end()
197    }
198}
199
200impl<'de> Deserialize<'de> for Psi {
201    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
202    where
203        D: serde::Deserializer<'de>,
204    {
205        use serde::de::{SeqAccess, Visitor};
206        use std::fmt;
207
208        struct PsiVisitor;
209
210        impl<'de> Visitor<'de> for PsiVisitor {
211            type Value = Psi;
212
213            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
214                formatter.write_str("a sequence of rows (vectors of f64)")
215            }
216
217            fn visit_seq<A>(self, mut seq: A) -> std::result::Result<Self::Value, A::Error>
218            where
219                A: SeqAccess<'de>,
220            {
221                let mut rows: Vec<Vec<f64>> = Vec::new();
222
223                while let Some(row) = seq.next_element::<Vec<f64>>()? {
224                    rows.push(row);
225                }
226
227                if rows.is_empty() {
228                    return Err(serde::de::Error::custom("Empty matrix not allowed"));
229                }
230
231                let nrows = rows.len();
232                let ncols = rows[0].len();
233
234                for (i, row) in rows.iter().enumerate() {
235                    if row.len() != ncols {
236                        return Err(serde::de::Error::custom(format!(
237                            "Row {} has {} columns, expected {}",
238                            i,
239                            row.len(),
240                            ncols
241                        )));
242                    }
243                }
244
245                let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
246
247                Ok(Psi {
248                    matrix: mat,
249                    row_log_scales: vec![0.0; nrows],
250                })
251            }
252        }
253
254        deserializer.deserialize_seq(PsiVisitor)
255    }
256}
257
258pub(crate) fn calculate_psi(
259    equation: &impl Equation,
260    subjects: &Data,
261    theta: &Theta,
262    error_models: &AssayErrorModels,
263    progress: bool,
264) -> Result<Psi> {
265    let tm = theta.matrix();
266    let theta_ndarray = Array2::from_shape_fn((tm.nrows(), tm.ncols()), |(i, j)| tm[(i, j)]);
267    let log_psi =
268        log_likelihood_matrix(equation, subjects, &theta_ndarray, error_models, progress)?;
269
270    Psi::from_log_likelihoods(log_psi)
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276    use ndarray::Array2;
277
278    #[test]
279    fn log_likelihood_rows_are_scaled_before_exponentiation() -> Result<()> {
280        let log_likelihoods = Array2::from_shape_vec((2, 2), vec![-1000.0, -1001.0, -2.0, -4.0])?;
281
282        let psi = Psi::from_log_likelihoods(log_likelihoods)?;
283
284        assert_eq!(psi.log_scale(), -1002.0);
285        assert_eq!(psi.row_log_scales(), &[-1000.0, -2.0]);
286        assert_eq!(psi.matrix()[(0, 0)], 1.0);
287        assert_eq!(psi.matrix()[(1, 0)], 1.0);
288        assert!((psi.matrix()[(0, 1)] - (-1.0_f64).exp()).abs() < 1e-12);
289        assert!((psi.matrix()[(1, 1)] - (-2.0_f64).exp()).abs() < 1e-12);
290        Ok(())
291    }
292
293    #[test]
294    fn test_from_array2() {
295        let array = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
296
297        let psi = Psi::from(array.clone());
298
299        assert_eq!(psi.nspp(), 2);
300        assert_eq!(psi.nsub(), 3);
301
302        let m = psi.matrix();
303        for i in 0..2 {
304            for j in 0..3 {
305                assert_eq!(m[(i, j)], array[[i, j]]);
306            }
307        }
308    }
309
310    #[test]
311    fn test_from_array2_ref() {
312        let array =
313            Array2::from_shape_vec((3, 2), vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0]).unwrap();
314
315        let psi = Psi::from(&array);
316
317        assert_eq!(psi.nspp(), 3);
318        assert_eq!(psi.nsub(), 2);
319
320        let m = psi.matrix();
321        for i in 0..3 {
322            for j in 0..2 {
323                assert_eq!(m[(i, j)], array[[i, j]]);
324            }
325        }
326    }
327
328    #[test]
329    fn test_nspp() {
330        let array =
331            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
332        let psi = Psi::from(array);
333
334        assert_eq!(psi.nspp(), 4);
335    }
336
337    #[test]
338    fn test_nspp_empty() {
339        let psi = Psi::new();
340        assert_eq!(psi.nspp(), 0);
341    }
342
343    #[test]
344    fn test_nspp_single_row() {
345        let array = Array2::from_shape_vec((1, 3), vec![1.0, 2.0, 3.0]).unwrap();
346        let psi = Psi::from(array);
347
348        assert_eq!(psi.nspp(), 1);
349    }
350
351    #[test]
352    fn test_nsub() {
353        let array = Array2::from_shape_vec(
354            (2, 5),
355            vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
356        )
357        .unwrap();
358        let psi = Psi::from(array);
359
360        assert_eq!(psi.nsub(), 5);
361    }
362
363    #[test]
364    fn test_nsub_empty() {
365        let psi = Psi::new();
366        assert_eq!(psi.nsub(), 0);
367    }
368}
pmcore/estimation/nonparametric/psi.rs

pmcore/estimation/nonparametric/
psi.rs
