Skip to main content

pmcore/estimation/nonparametric/
psi.rs

1use anyhow::bail;
2use anyhow::Result;
3use faer::Mat;
4use ndarray::Array2;
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}
18
19impl Psi {
20    pub fn new() -> Self {
21        Psi { matrix: Mat::new() }
22    }
23
24    pub fn matrix(&self) -> &Mat<f64> {
25        &self.matrix
26    }
27
28    pub fn nspp(&self) -> usize {
29        self.matrix.nrows()
30    }
31
32    pub fn nsub(&self) -> usize {
33        self.matrix.ncols()
34    }
35
36    pub fn to_ndarray(&self) -> Array2<f64> {
37        let m = &self.matrix;
38        Array2::from_shape_fn((m.nrows(), m.ncols()), |(i, j)| m[(i, j)])
39    }
40
41    pub(crate) fn filter_column_indices(&mut self, indices: &[usize]) {
42        let matrix = self.matrix.to_owned();
43
44        let new = Mat::from_fn(matrix.nrows(), indices.len(), |r, c| {
45            *matrix.get(r, indices[c])
46        });
47
48        self.matrix = new;
49    }
50
51    pub fn write(&self, path: &str) {
52        let mut writer = csv::Writer::from_path(path).unwrap();
53        for row in self.matrix.row_iter() {
54            writer
55                .write_record(row.iter().map(|x| x.to_string()))
56                .unwrap();
57        }
58    }
59
60    pub fn to_csv<W: std::io::Write>(&self, writer: W) -> Result<()> {
61        let mut csv_writer = csv::Writer::from_writer(writer);
62
63        for i in 0..self.matrix.nrows() {
64            let row: Vec<f64> = (0..self.matrix.ncols())
65                .map(|j| *self.matrix.get(i, j))
66                .collect();
67            csv_writer.serialize(row)?;
68        }
69
70        csv_writer.flush()?;
71        Ok(())
72    }
73
74    pub fn from_csv<R: std::io::Read>(reader: R) -> Result<Self> {
75        let mut csv_reader = csv::Reader::from_reader(reader);
76        let mut rows: Vec<Vec<f64>> = Vec::new();
77
78        for result in csv_reader.deserialize() {
79            let row: Vec<f64> = result?;
80            rows.push(row);
81        }
82
83        if rows.is_empty() {
84            bail!("CSV file is empty");
85        }
86
87        let nrows = rows.len();
88        let ncols = rows[0].len();
89
90        for (i, row) in rows.iter().enumerate() {
91            if row.len() != ncols {
92                bail!("Row {} has {} columns, expected {}", i, row.len(), ncols);
93            }
94        }
95
96        let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
97
98        Ok(Psi { matrix: mat })
99    }
100}
101
102impl Default for Psi {
103    fn default() -> Self {
104        Psi::new()
105    }
106}
107
108impl From<Array2<f64>> for Psi {
109    fn from(array: Array2<f64>) -> Self {
110        let matrix = Mat::from_fn(array.nrows(), array.ncols(), |i, j| array[(i, j)]);
111        Psi { matrix }
112    }
113}
114
115impl From<Mat<f64>> for Psi {
116    fn from(matrix: Mat<f64>) -> Self {
117        Psi { matrix }
118    }
119}
120
121impl From<&Array2<f64>> for Psi {
122    fn from(array: &Array2<f64>) -> Self {
123        let matrix = Mat::from_fn(array.nrows(), array.ncols(), |i, j| array[(i, j)]);
124        Psi { matrix }
125    }
126}
127
128impl Serialize for Psi {
129    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
130    where
131        S: serde::Serializer,
132    {
133        use serde::ser::SerializeSeq;
134
135        let mut seq = serializer.serialize_seq(Some(self.matrix.nrows()))?;
136
137        for i in 0..self.matrix.nrows() {
138            let row: Vec<f64> = (0..self.matrix.ncols())
139                .map(|j| *self.matrix.get(i, j))
140                .collect();
141            seq.serialize_element(&row)?;
142        }
143
144        seq.end()
145    }
146}
147
148impl<'de> Deserialize<'de> for Psi {
149    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
150    where
151        D: serde::Deserializer<'de>,
152    {
153        use serde::de::{SeqAccess, Visitor};
154        use std::fmt;
155
156        struct PsiVisitor;
157
158        impl<'de> Visitor<'de> for PsiVisitor {
159            type Value = Psi;
160
161            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
162                formatter.write_str("a sequence of rows (vectors of f64)")
163            }
164
165            fn visit_seq<A>(self, mut seq: A) -> std::result::Result<Self::Value, A::Error>
166            where
167                A: SeqAccess<'de>,
168            {
169                let mut rows: Vec<Vec<f64>> = Vec::new();
170
171                while let Some(row) = seq.next_element::<Vec<f64>>()? {
172                    rows.push(row);
173                }
174
175                if rows.is_empty() {
176                    return Err(serde::de::Error::custom("Empty matrix not allowed"));
177                }
178
179                let nrows = rows.len();
180                let ncols = rows[0].len();
181
182                for (i, row) in rows.iter().enumerate() {
183                    if row.len() != ncols {
184                        return Err(serde::de::Error::custom(format!(
185                            "Row {} has {} columns, expected {}",
186                            i,
187                            row.len(),
188                            ncols
189                        )));
190                    }
191                }
192
193                let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
194
195                Ok(Psi { matrix: mat })
196            }
197        }
198
199        deserializer.deserialize_seq(PsiVisitor)
200    }
201}
202
203pub(crate) fn calculate_psi(
204    equation: &impl Equation,
205    subjects: &Data,
206    theta: &Theta,
207    error_models: &AssayErrorModels,
208    progress: bool,
209) -> Result<Psi> {
210    let tm = theta.matrix();
211    let theta_ndarray = Array2::from_shape_fn((tm.nrows(), tm.ncols()), |(i, j)| tm[(i, j)]);
212    let log_psi =
213        log_likelihood_matrix(equation, subjects, &theta_ndarray, error_models, progress)?;
214    let psi_ndarray = log_psi.mapv(f64::exp);
215
216    Ok(Psi::from(psi_ndarray))
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use ndarray::Array2;
223
224    #[test]
225    fn test_from_array2() {
226        let array = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
227
228        let psi = Psi::from(array.clone());
229
230        assert_eq!(psi.nspp(), 2);
231        assert_eq!(psi.nsub(), 3);
232
233        let m = psi.matrix();
234        for i in 0..2 {
235            for j in 0..3 {
236                assert_eq!(m[(i, j)], array[[i, j]]);
237            }
238        }
239    }
240
241    #[test]
242    fn test_from_array2_ref() {
243        let array =
244            Array2::from_shape_vec((3, 2), vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0]).unwrap();
245
246        let psi = Psi::from(&array);
247
248        assert_eq!(psi.nspp(), 3);
249        assert_eq!(psi.nsub(), 2);
250
251        let m = psi.matrix();
252        for i in 0..3 {
253            for j in 0..2 {
254                assert_eq!(m[(i, j)], array[[i, j]]);
255            }
256        }
257    }
258
259    #[test]
260    fn test_nspp() {
261        let array =
262            Array2::from_shape_vec((4, 2), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();
263        let psi = Psi::from(array);
264
265        assert_eq!(psi.nspp(), 4);
266    }
267
268    #[test]
269    fn test_nspp_empty() {
270        let psi = Psi::new();
271        assert_eq!(psi.nspp(), 0);
272    }
273
274    #[test]
275    fn test_nspp_single_row() {
276        let array = Array2::from_shape_vec((1, 3), vec![1.0, 2.0, 3.0]).unwrap();
277        let psi = Psi::from(array);
278
279        assert_eq!(psi.nspp(), 1);
280    }
281
282    #[test]
283    fn test_nsub() {
284        let array = Array2::from_shape_vec(
285            (2, 5),
286            vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
287        )
288        .unwrap();
289        let psi = Psi::from(array);
290
291        assert_eq!(psi.nsub(), 5);
292    }
293
294    #[test]
295    fn test_nsub_empty() {
296        let psi = Psi::new();
297        assert_eq!(psi.nsub(), 0);
298    }
299}