pmcore/estimation/nonparametric/
result.rs

1use std::path::Path;
2
3use pharmsol::Equation;
4use serde::Serialize;
5
6use crate::algorithms::Status;
7use crate::estimation::nonparametric::{CycleLog, NPPredictions, Posterior, Psi, Theta, Weights};
8
9use pharmsol::{AssayErrorModels, Data};
10
11/// Contains the results of a nonparametric estimation, including the final parameter
12#[derive(Debug)]
13pub struct NonParametricResult<E: Equation> {
14    equation: E,
15    data: Data,
16    error_models: AssayErrorModels,
17    prior: Theta,
18    theta: Theta,
19    psi: Psi,
20    weights: Weights,
21    objf: f64,
22    cycles: usize,
23    status: Status,
24    cyclelog: CycleLog,
25}
26
27impl<E: Equation> NonParametricResult<E> {
28    #[allow(clippy::too_many_arguments)]
29    pub(crate) fn new(
30        equation: E,
31        data: Data,
32        error_models: AssayErrorModels,
33        prior: Theta,
34        theta: Theta,
35        psi: Psi,
36        weights: Weights,
37        objf: f64,
38        cycles: usize,
39        status: Status,
40        cyclelog: CycleLog,
41    ) -> anyhow::Result<Self> {
42        Ok(Self {
43            equation,
44            data,
45            error_models,
46            prior,
47            theta,
48            psi,
49            weights,
50            objf,
51            cycles,
52            status,
53            cyclelog,
54        })
55    }
56
57    pub fn cycles(&self) -> usize {
58        self.cycles
59    }
60
61    pub fn objf(&self) -> f64 {
62        self.objf
63    }
64
65    pub fn converged(&self) -> bool {
66        self.status.converged()
67    }
68
69    pub fn get_theta(&self) -> &Theta {
70        &self.theta
71    }
72
73    /// The prior distribution ([`Theta`]) that seeded the algorithm.
74    ///
75    /// This is the initial set of support points, as opposed to the optimized
76    /// solution returned by [`get_theta`](Self::get_theta).
77    pub fn prior(&self) -> &Theta {
78        &self.prior
79    }
80
81    pub fn data(&self) -> &Data {
82        &self.data
83    }
84
85    pub fn equation(&self) -> &E {
86        &self.equation
87    }
88
89    pub fn cycle_log(&self) -> &CycleLog {
90        &self.cyclelog
91    }
92
93    /// Chain this result into a new fit with a different algorithm.
94    ///
95    /// Uses the optimized theta as the prior for the new run. Keeps the same
96    /// equation, data, and error models. Consumes `self`.
97    ///
98    /// # Example
99    ///
100    /// ```ignore
101    /// let result = problem
102    ///     .fit_with(NpagConfig::new().max_cycles(100))?
103    ///     .chain(NpodConfig::new().max_cycles(50))?;
104    /// ```
105    pub fn chain<A>(self, algorithm: A) -> anyhow::Result<NonParametricResult<E>>
106    where
107        A: crate::algorithms::Algorithm<
108            E,
109            crate::estimation::NonParametric,
110            Output = NonParametricResult<E>,
111        >,
112        E: crate::model::EquationMetadataSource,
113    {
114        crate::estimation::EstimationProblem::nonparametric(
115            self.equation,
116            self.data,
117            self.theta,
118            self.error_models,
119        )?
120        .fit_with(algorithm)
121    }
122
123    pub fn psi(&self) -> &Psi {
124        &self.psi
125    }
126
127    pub fn weights(&self) -> &Weights {
128        &self.weights
129    }
130
131    pub fn error_models(&self) -> &AssayErrorModels {
132        &self.error_models
133    }
134
135    /// Compute the posterior probabilities on demand from [`Psi`] and the
136    /// [`Weights`]. This is a cheap matrix operation and is intentionally not
137    /// cached on the result.
138    pub fn posterior(&self) -> anyhow::Result<Posterior> {
139        Posterior::calculate(&self.psi, &self.weights)
140    }
141
142    /// Compute predictions on demand. Nothing is cached on the result; callers
143    /// that need the predictions repeatedly should hold on to the returned
144    /// value themselves.
145    pub fn predictions(&self, idelta: f64, tad: f64) -> anyhow::Result<NPPredictions> {
146        let posterior = self.posterior()?;
147        self.predictions_with(&posterior, idelta, tad)
148    }
149
150    /// Like [`predictions`](Self::predictions), but reuses an already-computed
151    /// [`Posterior`] instead of recomputing it.
152    fn predictions_with(
153        &self,
154        posterior: &Posterior,
155        idelta: f64,
156        tad: f64,
157    ) -> anyhow::Result<NPPredictions> {
158        NPPredictions::calculate(
159            &self.equation,
160            &self.data,
161            &self.theta,
162            &self.weights,
163            posterior,
164            idelta,
165            tad,
166        )
167    }
168
169    pub fn write_theta(&self, path: &Path) -> anyhow::Result<()> {
170        use anyhow::bail;
171        use csv::WriterBuilder;
172
173        tracing::debug!("Writing population parameter distribution...");
174
175        let w: Vec<f64> = self.weights.to_vec();
176        if w.len() != self.theta.matrix().nrows() {
177            bail!(
178                "Number of weights ({}) and number of support points ({}) do not match.",
179                w.len(),
180                self.theta.matrix().nrows()
181            );
182        }
183
184        crate::estimation::nonparametric::create_parent_dir(path)?;
185        let file = std::fs::File::create(path)?;
186        let mut writer = WriterBuilder::new().has_headers(true).from_writer(file);
187
188        let mut theta_header = self.theta.param_names().clone();
189        theta_header.push("prob".to_string());
190        writer.write_record(&theta_header)?;
191
192        for (theta_row, &w_val) in self.theta.matrix().row_iter().zip(w.iter()) {
193            let mut row: Vec<String> = theta_row.iter().map(|&val| val.to_string()).collect();
194            row.push(w_val.to_string());
195            writer.write_record(&row)?;
196        }
197        writer.flush()?;
198        Ok(())
199    }
200
201    pub fn write_posterior(&self, path: &Path) -> anyhow::Result<()> {
202        let posterior = self.posterior()?;
203        self.write_posterior_with(path, &posterior)
204    }
205
206    fn write_posterior_with(&self, path: &Path, posterior: &Posterior) -> anyhow::Result<()> {
207        use csv::WriterBuilder;
208
209        tracing::debug!("Writing posterior parameter probabilities...");
210
211        crate::estimation::nonparametric::create_parent_dir(path)?;
212        let file = std::fs::File::create(path)?;
213        let mut writer = WriterBuilder::new().has_headers(true).from_writer(file);
214
215        writer.write_field("id")?;
216        writer.write_field("point")?;
217        self.theta.param_names().iter().for_each(|name| {
218            writer.write_field(name).unwrap();
219        });
220        writer.write_field("prob")?;
221        writer.write_record(None::<&[u8]>)?;
222
223        let subjects = self.data.subjects();
224        posterior
225            .matrix()
226            .row_iter()
227            .enumerate()
228            .for_each(|(i, row)| {
229                let subject = subjects.get(i).unwrap();
230                let id = subject.id();
231
232                row.iter().enumerate().for_each(|(spp, prob)| {
233                    writer.write_field(id.clone()).unwrap();
234                    writer.write_field(spp.to_string()).unwrap();
235
236                    self.theta.matrix().row(spp).iter().for_each(|val| {
237                        writer.write_field(val.to_string()).unwrap();
238                    });
239
240                    writer.write_field(prob.to_string()).unwrap();
241                    writer.write_record(None::<&[u8]>).unwrap();
242                });
243            });
244
245        writer.flush()?;
246        Ok(())
247    }
248
249    pub fn write_covariates(&self, path: &Path) -> anyhow::Result<()> {
250        use csv::WriterBuilder;
251        use pharmsol::Event;
252
253        tracing::debug!("Writing covariates...");
254        crate::estimation::nonparametric::create_parent_dir(path)?;
255        let file = std::fs::File::create(path)?;
256        let mut writer = WriterBuilder::new().has_headers(true).from_writer(file);
257
258        let mut covariate_names = std::collections::HashSet::new();
259        for subject in self.data.subjects() {
260            for occasion in subject.occasions() {
261                let covmap = occasion.covariates().covariates();
262                for cov_name in covmap.keys() {
263                    covariate_names.insert(cov_name.clone());
264                }
265            }
266        }
267        let mut covariate_names: Vec<String> = covariate_names.into_iter().collect();
268        covariate_names.sort();
269
270        let mut headers = vec!["id", "time", "block"];
271        headers.extend(covariate_names.iter().map(|s| s.as_str()));
272        writer.write_record(&headers)?;
273
274        for subject in self.data.subjects() {
275            for occasion in subject.occasions() {
276                let covmap = occasion.covariates().covariates();
277
278                for event in occasion.iter() {
279                    let time = match event {
280                        Event::Bolus(bolus) => bolus.time(),
281                        Event::Infusion(infusion) => infusion.time(),
282                        Event::Observation(observation) => observation.time(),
283                    };
284
285                    let mut row: Vec<String> = Vec::new();
286                    row.push(subject.id().clone());
287                    row.push(time.to_string());
288                    row.push(occasion.index().to_string());
289
290                    for cov_name in &covariate_names {
291                        if let Some(cov) = covmap.get(cov_name) {
292                            if let Ok(value) = cov.interpolate(time) {
293                                row.push(value.to_string());
294                            } else {
295                                row.push(String::new());
296                            }
297                        } else {
298                            row.push(String::new());
299                        }
300                    }
301
302                    writer.write_record(&row)?;
303                }
304            }
305        }
306
307        writer.flush()?;
308        Ok(())
309    }
310
311    /// Write the cycle log to a CSV file readable by Pmetrics.
312    pub fn write_cycles(&self, path: &Path) -> anyhow::Result<()> {
313        self.cyclelog.write(path)
314    }
315
316    /// Compute and write the population and posterior predictions to a CSV file
317    /// readable by Pmetrics.
318    ///
319    /// `idelta` is the interval used to densify the prediction grid and `tad` is
320    /// the additional time after the last event to simulate.
321    pub fn write_predictions(&self, path: &Path, idelta: f64, tad: f64) -> anyhow::Result<()> {
322        let predictions = self.predictions(idelta, tad)?;
323        predictions.write(path)
324    }
325
326    /// Serialize the complete result to a single JSON file.
327    ///
328    /// This includes the data, error models, prior, optimized support points,
329    /// likelihoods, weights, objective function, status, cycle log, posterior
330    /// probabilities, and predictions. `idelta` and `tad` control the density of
331    /// the embedded predictions (see [`predictions`](Self::predictions)).
332    pub fn write_json(&self, path: &Path, idelta: f64, tad: f64) -> anyhow::Result<()> {
333        let posterior = self.posterior()?;
334        let predictions = self.predictions_with(&posterior, idelta, tad)?;
335        self.write_json_with(path, &posterior, &predictions)
336    }
337
338    fn write_json_with(
339        &self,
340        path: &Path,
341        posterior: &Posterior,
342        predictions: &NPPredictions,
343    ) -> anyhow::Result<()> {
344        tracing::debug!("Writing result as JSON...");
345
346        crate::estimation::nonparametric::create_parent_dir(path)?;
347
348        let view = NonParametricResultJson {
349            data: &self.data,
350            error_models: &self.error_models,
351            prior: &self.prior,
352            theta: &self.theta,
353            psi: &self.psi,
354            weights: &self.weights,
355            objf: self.objf,
356            cycles: self.cycles,
357            status: &self.status,
358            cyclelog: &self.cyclelog,
359            posterior,
360            predictions,
361        };
362
363        let file = std::fs::File::create(path)?;
364        serde_json::to_writer_pretty(file, &view)?;
365        Ok(())
366    }
367
368    /// Write the full set of result artifacts to `directory`.
369    ///
370    /// The directory is created if it does not exist, and the following files
371    /// are produced (matching the names expected by Pmetrics):
372    /// - `theta.csv` — optimized support points with weights
373    /// - `posterior.csv` — per-subject posterior probabilities
374    /// - `pred.csv` — population and posterior predictions
375    /// - `covs.csv` — interpolated covariates
376    /// - `cycles.csv` — cycle log
377    /// - `result.json` — the complete result as JSON
378    ///
379    /// The posterior and predictions are computed once and shared across the CSV
380    /// and JSON outputs. `idelta` and `tad` control the density of the
381    /// predictions (see [`predictions`](Self::predictions)).
382    pub fn write_outputs(
383        &self,
384        directory: impl AsRef<Path>,
385        idelta: f64,
386        tad: f64,
387    ) -> anyhow::Result<()> {
388        let dir = directory.as_ref();
389        std::fs::create_dir_all(dir)?;
390
391        // Compute the expensive artifacts a single time and reuse them.
392        let posterior = self.posterior()?;
393        let predictions = self.predictions_with(&posterior, idelta, tad)?;
394
395        self.write_theta(&dir.join("theta.csv"))?;
396        self.write_posterior_with(&dir.join("posterior.csv"), &posterior)?;
397        predictions.write(&dir.join("pred.csv"))?;
398        self.write_covariates(&dir.join("covs.csv"))?;
399        self.write_cycles(&dir.join("cycles.csv"))?;
400        self.write_json_with(&dir.join("result.json"), &posterior, &predictions)?;
401
402        tracing::info!("Results written to {}", dir.display());
403        Ok(())
404    }
405}
406
407/// A borrowed, serializable view over [`NonParametricResult`] used to emit a
408/// single JSON document. The equation is intentionally omitted because it is
409/// generic and not necessarily serializable.
410#[derive(Serialize)]
411struct NonParametricResultJson<'a> {
412    data: &'a Data,
413    error_models: &'a AssayErrorModels,
414    prior: &'a Theta,
415    theta: &'a Theta,
416    psi: &'a Psi,
417    weights: &'a Weights,
418    objf: f64,
419    cycles: usize,
420    status: &'a Status,
421    cyclelog: &'a CycleLog,
422    posterior: &'a Posterior,
423    predictions: &'a NPPredictions,
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429    use crate::algorithms::nonparametric::{NpagConfig, NpodConfig};
430    use crate::estimation::EstimationProblem;
431    use crate::model::ParameterSpace;
432    use pharmsol::equation::metadata;
433    use pharmsol::prelude::data::{AssayErrorModel, AssayErrorModels};
434    use pharmsol::{ErrorPoly, SubjectBuilderExt};
435
436    fn minimal_ode() -> pharmsol::ODE {
437        pharmsol::equation::ODE::new(
438            |x, p, _t, dx, b, _rateiv, _cov| {
439                let ke = p[0];
440                dx[0] = -ke * x[0] + b[0];
441            },
442            |_p, _t, _cov| pharmsol::lag! {},
443            |_p, _t, _cov| pharmsol::fa! {},
444            |_p, _t, _cov, _x| {},
445            |x, p, _t, _cov, y| {
446                let v = p[1];
447                y[0] = x[0] / v;
448            },
449        )
450        .with_nstates(1)
451        .with_ndrugs(1)
452        .with_nout(1)
453        .with_metadata(
454            metadata::new("chain_test")
455                .parameters(["ke", "v"])
456                .states(["central"])
457                .outputs(["0"])
458                .route(pharmsol::equation::Route::bolus("0").to_state("central")),
459        )
460        .expect("metadata should validate")
461    }
462
463    fn minimal_data() -> pharmsol::Data {
464        let subject = pharmsol::Subject::builder("1")
465            .bolus(0.0, 100.0, 0)
466            .observation(1.0, 10.0, 0)
467            .observation(2.0, 8.0, 0)
468            .build();
469        pharmsol::Data::new(vec![subject])
470    }
471
472    #[test]
473    fn chain_npag_to_npod_preserves_support_points() {
474        let ode = minimal_ode();
475        let data = minimal_data();
476        let params = ParameterSpace::bounded()
477            .add("ke", 0.001, 3.0)
478            .add("v", 25.0, 250.0);
479        let prior = Theta::sobol_default(&params).unwrap();
480        let err = AssayErrorModels::new()
481            .add(
482                "0",
483                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
484            )
485            .unwrap();
486
487        let r1 = EstimationProblem::nonparametric(ode, data, prior, err)
488            .unwrap()
489            .fit_with(NpagConfig::new().max_cycles(2))
490            .unwrap();
491
492        let n_spp = r1.get_theta().nspp();
493        assert!(n_spp > 0, "NPAG should produce support points");
494
495        // Chain into NPOD — should complete without error
496        let r2 = r1.chain(NpodConfig::new().max_cycles(1)).unwrap();
497
498        assert!(r2.get_theta().nspp() > 0);
499        assert_eq!(r2.data().subjects().len(), 1);
500        assert_eq!(r2.cycles(), 1);
501    }
502
503    #[test]
504    fn chain_npag_to_npag_maintains_or_improves_objf() {
505        let ode = minimal_ode();
506        let data = minimal_data();
507        let params = ParameterSpace::bounded()
508            .add("ke", 0.001, 3.0)
509            .add("v", 25.0, 250.0);
510        let prior = Theta::sobol_with_seed(&params, 5, 42).unwrap();
511        let err = AssayErrorModels::new()
512            .add(
513                "0",
514                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
515            )
516            .unwrap();
517
518        let r1 = EstimationProblem::nonparametric(ode, data, prior, err)
519            .unwrap()
520            .fit_with(NpagConfig::new().max_cycles(5))
521            .unwrap();
522
523        let objf1 = r1.objf();
524
525        // Chain into another NPAG run
526        let r2 = r1.chain(NpagConfig::new().max_cycles(3)).unwrap();
527
528        // Second run should not regress significantly
529        assert!(
530            r2.objf() <= objf1 + 0.5,
531            "OBJF regressed: {} -> {}",
532            objf1,
533            r2.objf()
534        );
535    }
536
537    #[test]
538    fn chain_npag_to_npod_with_unconverged_result() {
539        let ode = minimal_ode();
540        let data = minimal_data();
541        let params = ParameterSpace::bounded()
542            .add("ke", 0.001, 3.0)
543            .add("v", 25.0, 250.0);
544        let prior = Theta::sobol_with_seed(&params, 5, 42).unwrap();
545        let err = AssayErrorModels::new()
546            .add(
547                "0",
548                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
549            )
550            .unwrap();
551
552        // Run only 1 cycle — will be unconverged
553        let r1 = EstimationProblem::nonparametric(ode, data, prior, err)
554            .unwrap()
555            .fit_with(NpagConfig::new().max_cycles(1))
556            .unwrap();
557
558        // Chaining from unconverged should still work
559        let r2 = r1.chain(NpodConfig::new().max_cycles(1));
560        assert!(
561            r2.is_ok(),
562            "Chaining from unconverged result should succeed"
563        );
564    }
565}
pmcore/estimation/nonparametric/result.rs

pmcore/estimation/nonparametric/
result.rs
