pmcore/estimation/nonparametric/
theta.rs

1use std::{fmt::Debug, fs::File, path::Path};
2
3use anyhow::{bail, Context, Result};
4use faer::Mat;
5use serde::{Deserialize, Serialize};
6
7use super::sampling::{self, latin, sobol};
8use super::weights::Weights;
9use crate::model::{BoundedParameter, ParameterSpace};
10
11/// [Theta] is a structure that holds the support points
12/// These represent the joint population parameter distribution
13///
14/// Each row represents a support points, and each column a parameter
15#[derive(Clone, PartialEq)]
16pub struct Theta {
17    matrix: Mat<f64>,
18    parameters: ParameterSpace<BoundedParameter>,
19}
20
21impl Default for Theta {
22    fn default() -> Self {
23        Theta {
24            matrix: Mat::new(),
25            parameters: ParameterSpace::<BoundedParameter>::new(),
26        }
27    }
28}
29
30impl Theta {
31    pub fn new() -> Self {
32        Theta::default()
33    }
34
35    /// Create a new [Theta] from a matrix and [ParameterSpace]
36    ///
37    /// It is important that the number of columns in the matrix matches the number of parameters
38    /// in the [ParameterSpace]
39    ///
40    /// The order of parameters in the [ParameterSpace] should match the order of columns in the matrix
41    pub fn from_parts(
42        matrix: Mat<f64>,
43        parameters: ParameterSpace<BoundedParameter>,
44    ) -> Result<Self> {
45        if matrix.ncols() != parameters.len() {
46            bail!(
47                "Number of columns in matrix ({}) does not match number of parameters ({})",
48                matrix.ncols(),
49                parameters.len()
50            );
51        }
52
53        Ok(Theta { matrix, parameters })
54    }
55
56    /// Get the matrix containing parameter values
57    ///
58    /// The matrix is a 2D array where each row represents a support point, and each column a parameter
59    pub fn matrix(&self) -> &Mat<f64> {
60        &self.matrix
61    }
62
63    /// Get a mutable reference to the matrix
64    pub fn matrix_mut(&mut self) -> &mut Mat<f64> {
65        &mut self.matrix
66    }
67
68    /// Get the [ParameterSpace] associated with this [Theta]
69    pub fn parameters(&self) -> &ParameterSpace<BoundedParameter> {
70        &self.parameters
71    }
72
73    /// Get a mutable reference to the [ParameterSpace]
74    pub fn parameters_mut(&mut self) -> &mut ParameterSpace<BoundedParameter> {
75        &mut self.parameters
76    }
77
78    /// Get the number of support points, equal to the number of rows in the matrix
79    pub fn nspp(&self) -> usize {
80        self.matrix.nrows()
81    }
82
83    /// Get the parameter names
84    pub fn param_names(&self) -> Vec<String> {
85        self.parameters.names()
86    }
87
88    /// Modify the [Theta::matrix] to only include the rows specified by `indices`
89    pub(crate) fn filter_indices(&mut self, indices: &[usize]) {
90        let matrix = self.matrix.to_owned();
91
92        let new = Mat::from_fn(indices.len(), matrix.ncols(), |r, c| {
93            *matrix.get(indices[r], c)
94        });
95
96        self.matrix = new;
97    }
98
99    /// Forcibly add a support point to the matrix
100    pub fn add_point(&mut self, spp: &[f64]) -> Result<()> {
101        if spp.len() != self.matrix.ncols() {
102            bail!(
103                "Support point length ({}) does not match number of parameters ({})",
104                spp.len(),
105                self.matrix.ncols()
106            );
107        }
108
109        self.matrix
110            .resize_with(self.matrix.nrows() + 1, self.matrix.ncols(), |_, i| spp[i]);
111        Ok(())
112    }
113
114    /// Suggest a new support point to add to the matrix
115    /// The point is only added if it is at least `min_dist` away from all existing support points
116    /// and within the limits specified by `limits`
117    pub(crate) fn suggest_point(&mut self, spp: &[f64], min_dist: f64) -> Result<()> {
118        if self.check_point(spp, min_dist) {
119            self.add_point(spp)?;
120        }
121        Ok(())
122    }
123
124    /// Check if a point is at least `min_dist` away from all existing support points
125    pub(crate) fn check_point(&self, spp: &[f64], min_dist: f64) -> bool {
126        if self.matrix.nrows() == 0 {
127            return true;
128        }
129
130        let limits = self.parameters.finite_ranges();
131
132        for row_idx in 0..self.matrix.nrows() {
133            let mut squared_dist = 0.0;
134            for (i, val) in spp.iter().enumerate() {
135                let normalized_diff =
136                    (val - self.matrix.get(row_idx, i)) / (limits[i].1 - limits[i].0);
137                squared_dist += normalized_diff * normalized_diff;
138            }
139            let dist = squared_dist.sqrt();
140            if dist <= min_dist {
141                return false;
142            }
143        }
144        true
145    }
146
147    /// Create a new Theta with an additional parameter column.
148    ///
149    /// All existing rows keep their values. The new column is filled with
150    /// `initial_value` for every row. Returns a new `Theta` — does not
151    /// mutate in place.
152    ///
153    /// # Errors
154    ///
155    /// - `name` already exists in the parameter space
156    /// - `lower` or `upper` are non-finite
157    /// - `lower >= upper`
158    pub fn with_added_parameter(
159        &self,
160        name: &str,
161        lower: f64,
162        upper: f64,
163        initial_value: f64,
164    ) -> Result<Theta> {
165        // Validate uniqueness
166        if self.parameters().iter().any(|p| p.name.as_str() == name) {
167            bail!("parameter '{}' already exists in theta", name);
168        }
169
170        // Validate bounds
171        if !lower.is_finite() || !upper.is_finite() {
172            bail!(
173                "bounds must be finite for parameter '{}': [{}, {}]",
174                name,
175                lower,
176                upper
177            );
178        }
179        if lower >= upper {
180            bail!(
181                "lower bound ({}) must be strictly less than upper bound ({}) for parameter '{}'",
182                lower,
183                upper,
184                name
185            );
186        }
187
188        let (nrows, ncols) = (self.matrix().nrows(), self.matrix().ncols());
189        let new_matrix = faer::Mat::from_fn(nrows, ncols + 1, |r, c| {
190            if c < ncols {
191                self.matrix()[(r, c)]
192            } else {
193                initial_value
194            }
195        });
196
197        let new_params = self.parameters().clone().add(name, lower, upper);
198
199        Theta::from_parts(new_matrix, new_params)
200    }
201
202    /// Write the matrix to a CSV file
203    pub fn write(&self, path: &str) {
204        let mut writer = csv::Writer::from_path(path).unwrap();
205        for row in self.matrix.row_iter() {
206            writer
207                .write_record(row.iter().map(|x| x.to_string()))
208                .unwrap();
209        }
210    }
211
212    /// Write the matrix to a CSV file with weights
213    pub fn write_with_weights(&self, path: &str, weights: &Weights) -> Result<()> {
214        if self.nspp() != weights.len() {
215            bail!(
216                "Number of support points ({}) does not match number of weights ({})",
217                self.nspp(),
218                weights.len()
219            );
220        }
221
222        let mut writer = csv::Writer::from_path(path)?;
223
224        let header: Vec<String> = self
225            .parameters
226            .names()
227            .iter()
228            .cloned()
229            .chain(std::iter::once("prob".to_string()))
230            .collect();
231
232        writer.write_record(header)?;
233
234        for (row_idx, row) in self.matrix.row_iter().enumerate() {
235            let mut record: Vec<String> = row.iter().map(|x| x.to_string()).collect();
236            record.push(weights[row_idx].to_string());
237            writer.write_record(record)?;
238        }
239        Ok(())
240    }
241
242    /// Write the theta matrix to a CSV writer
243    /// Each row represents a support point, each column represents a parameter
244    pub fn to_csv<W: std::io::Write>(&self, writer: W) -> Result<()> {
245        let mut csv_writer = csv::Writer::from_writer(writer);
246
247        for i in 0..self.matrix.nrows() {
248            let row: Vec<f64> = (0..self.matrix.ncols())
249                .map(|j| *self.matrix.get(i, j))
250                .collect();
251            csv_writer.serialize(row)?;
252        }
253
254        csv_writer.flush()?;
255        Ok(())
256    }
257
258    /// Read theta matrix from a CSV reader
259    /// Each row represents a support point, each column represents a parameter
260    /// Note: This only reads the matrix values, not the parameter metadata
261    pub fn from_csv<R: std::io::Read>(reader: R) -> Result<Self> {
262        let mut csv_reader = csv::Reader::from_reader(reader);
263        let mut rows: Vec<Vec<f64>> = Vec::new();
264
265        for result in csv_reader.deserialize() {
266            let row: Vec<f64> = result?;
267            rows.push(row);
268        }
269
270        if rows.is_empty() {
271            bail!("CSV file is empty");
272        }
273
274        let nrows = rows.len();
275        let ncols = rows[0].len();
276
277        for (i, row) in rows.iter().enumerate() {
278            if row.len() != ncols {
279                bail!("Row {} has {} columns, expected {}", i, row.len(), ncols);
280            }
281        }
282
283        let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
284        let parameters = ParameterSpace::<BoundedParameter>::new();
285
286        Theta::from_parts(mat, parameters)
287    }
288
289    /// Generate a starting grid of `points` support points over `parameters`
290    /// using a Sobol sequence and the default seed ([`sampling::DEFAULT_SEED`]).
291    ///
292    /// The returned [Theta] carries `parameters`, so the chosen grid is explicit
293    /// and self-describing.
294    pub fn sobol(parameters: &ParameterSpace<BoundedParameter>, points: usize) -> Result<Self> {
295        Self::sobol_with_seed(parameters, points, sampling::DEFAULT_SEED)
296    }
297
298    /// Generate a starting grid over `parameters` using a Sobol sequence with the
299    /// default number of support points ([`sampling::DEFAULT_POINTS`]) and the
300    /// default seed ([`sampling::DEFAULT_SEED`]).
301    pub fn sobol_default(parameters: &ParameterSpace<BoundedParameter>) -> Result<Self> {
302        Self::sobol(parameters, sampling::DEFAULT_POINTS)
303    }
304
305    /// Like [`Theta::sobol`], with an explicit seed for the quasi-random sequence.
306    pub fn sobol_with_seed(
307        parameters: &ParameterSpace<BoundedParameter>,
308        points: usize,
309        seed: usize,
310    ) -> Result<Self> {
311        validate_bounds(parameters)?;
312        sobol::generate(parameters, points, seed)
313    }
314
315    /// Generate a starting grid of `points` support points over `parameters`
316    /// using Latin Hypercube Sampling and the default seed ([`sampling::DEFAULT_SEED`]).
317    ///
318    /// The returned [Theta] carries `parameters`, so the chosen grid is explicit
319    /// and self-describing.
320    pub fn latin(parameters: &ParameterSpace<BoundedParameter>, points: usize) -> Result<Self> {
321        Self::latin_with_seed(parameters, points, sampling::DEFAULT_SEED)
322    }
323
324    /// Like [`Theta::latin`], with an explicit seed for the quasi-random sequence.
325    pub fn latin_with_seed(
326        parameters: &ParameterSpace<BoundedParameter>,
327        points: usize,
328        seed: usize,
329    ) -> Result<Self> {
330        validate_bounds(parameters)?;
331        latin::generate(parameters, points, seed)
332    }
333
334    pub fn from_file(
335        path: impl AsRef<Path>,
336        parameters: &ParameterSpace<BoundedParameter>,
337    ) -> Result<(Theta, Option<Weights>)> {
338        let path = path.as_ref();
339        tracing::info!("Reading prior from {}", path.display());
340        let file = File::open(path).context(format!(
341            "Unable to open the prior file '{}'",
342            path.display()
343        ))?;
344        let mut reader = csv::ReaderBuilder::new()
345            .has_headers(true)
346            .from_reader(file);
347
348        let mut parameter_names: Vec<String> = reader
349            .headers()?
350            .clone()
351            .into_iter()
352            .map(|s| s.trim().to_owned())
353            .collect();
354
355        let prob_index = parameter_names.iter().position(|name| name == "prob");
356        if let Some(index) = prob_index {
357            parameter_names.remove(index);
358        }
359
360        let random_names: Vec<String> = parameters.names();
361
362        let mut reordered_indices: Vec<usize> = Vec::new();
363        for random_name in &random_names {
364            match parameter_names.iter().position(|name| name == random_name) {
365                Some(index) => {
366                    let adjusted_index = if let Some(prob_idx) = prob_index {
367                        if index >= prob_idx {
368                            index + 1
369                        } else {
370                            index
371                        }
372                    } else {
373                        index
374                    };
375                    reordered_indices.push(adjusted_index);
376                }
377                None => bail!("Parameter {} is not present in the CSV file.", random_name),
378            }
379        }
380
381        if parameter_names.len() > random_names.len() {
382            let extra_parameters: Vec<&String> = parameter_names.iter().collect();
383            bail!(
384                "Found parameters in the prior not present in configuration: {:?}",
385                extra_parameters
386            );
387        }
388
389        let mut theta_values = Vec::new();
390        let mut prob_values = Vec::new();
391
392        for result in reader.records() {
393            let record = result.unwrap();
394            let values: Vec<f64> = reordered_indices
395                .iter()
396                .map(|&i| record[i].parse::<f64>().unwrap())
397                .collect();
398            theta_values.push(values);
399
400            if let Some(prob_idx) = prob_index {
401                let prob_value: f64 = record[prob_idx].parse::<f64>().unwrap();
402                prob_values.push(prob_value);
403            }
404        }
405
406        let n_points = theta_values.len();
407        let n_params = random_names.len();
408        let theta_values: Vec<f64> = theta_values.into_iter().flatten().collect();
409        let theta_matrix: Mat<f64> =
410            Mat::from_fn(n_points, n_params, |i, j| theta_values[i * n_params + j]);
411
412        let theta = Theta::from_parts(theta_matrix, parameters.clone())?;
413        let weights = if !prob_values.is_empty() {
414            Some(Weights::from_vec(prob_values))
415        } else {
416            None
417        };
418
419        Ok((theta, weights))
420    }
421}
422
423impl Debug for Theta {
424    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
425        writeln!(f, "\nTheta contains {} support points\n", self.nspp())?;
426
427        for name in self.parameters.names().iter() {
428            write!(f, "\t{}", name)?;
429        }
430        writeln!(f)?;
431        self.matrix.row_iter().enumerate().for_each(|(index, row)| {
432            write!(f, "{}", index).unwrap();
433            for val in row.iter() {
434                write!(f, "\t{:.2}", val).unwrap();
435            }
436            writeln!(f).unwrap();
437        });
438        Ok(())
439    }
440}
441
442impl Serialize for Theta {
443    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
444    where
445        S: serde::Serializer,
446    {
447        use serde::ser::SerializeStruct;
448
449        let rows: Vec<Vec<f64>> = (0..self.matrix.nrows())
450            .map(|i| {
451                (0..self.matrix.ncols())
452                    .map(|j| *self.matrix.get(i, j))
453                    .collect()
454            })
455            .collect();
456
457        let mut state = serializer.serialize_struct("Theta", 2)?;
458        state.serialize_field("matrix", &rows)?;
459        state.serialize_field("parameters", &self.parameters)?;
460        state.end()
461    }
462}
463
464impl<'de> Deserialize<'de> for Theta {
465    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
466    where
467        D: serde::Deserializer<'de>,
468    {
469        #[derive(Deserialize)]
470        struct ThetaSerde {
471            matrix: Vec<Vec<f64>>,
472            parameters: ParameterSpace<BoundedParameter>,
473        }
474
475        let decoded = ThetaSerde::deserialize(deserializer)?;
476
477        if decoded.matrix.is_empty() {
478            return Ok(Self {
479                matrix: Mat::new(),
480                parameters: decoded.parameters,
481            });
482        }
483
484        let nrows = decoded.matrix.len();
485        let ncols = decoded.matrix[0].len();
486        for (index, row) in decoded.matrix.iter().enumerate() {
487            if row.len() != ncols {
488                return Err(serde::de::Error::custom(format!(
489                    "Row {} has {} columns, expected {}",
490                    index,
491                    row.len(),
492                    ncols
493                )));
494            }
495        }
496
497        let matrix = Mat::from_fn(nrows, ncols, |i, j| decoded.matrix[i][j]);
498        Self::from_parts(matrix, decoded.parameters).map_err(serde::de::Error::custom)
499    }
500}
501
502/// Validates that every parameter has a strictly-ordered, finite bound interval.
503fn validate_bounds(parameters: &ParameterSpace<BoundedParameter>) -> Result<()> {
504    for parameter in parameters.iter() {
505        if parameter.lower >= parameter.upper {
506            bail!(
507                "Parameter '{}' has invalid bounds: [{}, {}]. Lower bound must be less than upper bound.",
508                parameter.name,
509                parameter.lower,
510                parameter.upper
511            );
512        }
513    }
514    Ok(())
515}
516
517#[cfg(test)]
518mod tests {
519    use super::*;
520    use std::fs;
521
522    fn parameters() -> ParameterSpace<BoundedParameter> {
523        ParameterSpace::<BoundedParameter>::new()
524            .add("ke", 0.1, 1.0)
525            .add("v", 5.0, 50.0)
526    }
527
528    fn temp_csv_path() -> String {
529        format!("test_temp_theta_{}.csv", rand::random::<u32>())
530    }
531
532    #[test]
533    fn sobol_generates_expected_shape() {
534        let theta = Theta::sobol_with_seed(&parameters(), 10, 42).unwrap();
535        assert_eq!(theta.nspp(), 10);
536        assert_eq!(theta.matrix().ncols(), 2);
537    }
538
539    #[test]
540    fn latin_generates_expected_shape() {
541        let theta = Theta::latin(&parameters(), 10).unwrap();
542        assert_eq!(theta.nspp(), 10);
543        assert_eq!(theta.matrix().ncols(), 2);
544    }
545
546    #[test]
547    fn sampling_rejects_invalid_bounds() {
548        let bad = ParameterSpace::<BoundedParameter>::new().add("ke", 1.0, 1.0);
549        let err = Theta::sobol(&bad, 10).unwrap_err();
550        assert!(err.to_string().contains("invalid bounds"));
551    }
552
553    #[test]
554    fn from_file_parses_weights_and_reorders_columns() {
555        let path = temp_csv_path();
556        fs::write(&path, "v,ke,prob\n10.0,0.5,0.3\n15.0,0.7,0.7\n").unwrap();
557
558        let (theta, weights) = Theta::from_file(&path, &parameters()).unwrap();
559        let _ = fs::remove_file(&path);
560
561        assert_eq!(theta.nspp(), 2);
562        assert_eq!(theta.matrix()[(0, 0)], 0.5);
563        assert_eq!(theta.matrix()[(0, 1)], 10.0);
564
565        let weights = weights.expect("weights should be parsed from prob column");
566        assert_eq!(weights.len(), 2);
567        assert_eq!(weights[0], 0.3);
568        assert_eq!(weights[1], 0.7);
569    }
570
571    #[test]
572    fn from_file_rejects_extra_parameters() {
573        let path = temp_csv_path();
574        fs::write(&path, "ke,v,extra\n0.5,10.0,1.0\n").unwrap();
575
576        let err = Theta::from_file(&path, &parameters()).unwrap_err();
577        let _ = fs::remove_file(&path);
578
579        assert!(err
580            .to_string()
581            .contains("Found parameters in the prior not present in configuration"));
582    }
583}
pmcore/estimation/nonparametric/theta.rs

pmcore/estimation/nonparametric/
theta.rs
