pmcore/estimation/nonparametric/
psi.rs1use 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#[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}