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    ///
378    /// Use [`write_json`](Self::write_json) to explicitly write the complete result.
379    /// `idelta` and `tad` control the density of the predictions (see
380    /// [`predictions`](Self::predictions)).
381    pub fn write_outputs(
382        &self,
383        directory: impl AsRef<Path>,
384        idelta: f64,
385        tad: f64,
386    ) -> anyhow::Result<()> {
387        let dir = directory.as_ref();
388        std::fs::create_dir_all(dir)?;
389
390        // Compute the expensive artifacts a single time and reuse them.
391        let posterior = self.posterior()?;
392        let predictions = self.predictions_with(&posterior, idelta, tad)?;
393
394        self.write_theta(&dir.join("theta.csv"))?;
395        self.write_posterior_with(&dir.join("posterior.csv"), &posterior)?;
396        predictions.write(&dir.join("pred.csv"))?;
397        self.write_covariates(&dir.join("covs.csv"))?;
398        self.write_cycles(&dir.join("cycles.csv"))?;
399
400        tracing::info!("Results written to {}", dir.display());
401        Ok(())
402    }
403}
404
405/// A borrowed, serializable view over [`NonParametricResult`] used to emit a
406/// single JSON document. The equation is intentionally omitted because it is
407/// generic and not necessarily serializable.
408#[derive(Serialize)]
409struct NonParametricResultJson<'a> {
410    data: &'a Data,
411    error_models: &'a AssayErrorModels,
412    prior: &'a Theta,
413    theta: &'a Theta,
414    psi: &'a Psi,
415    weights: &'a Weights,
416    objf: f64,
417    cycles: usize,
418    status: &'a Status,
419    cyclelog: &'a CycleLog,
420    posterior: &'a Posterior,
421    predictions: &'a NPPredictions,
422}
423
424#[cfg(test)]
425mod tests {
426    use super::*;
427    use crate::algorithms::nonparametric::{NpagConfig, NpodConfig};
428    use crate::estimation::EstimationProblem;
429    use crate::model::ParameterSpace;
430    use pharmsol::equation::metadata;
431    use pharmsol::prelude::data::{AssayErrorModel, AssayErrorModels};
432    use pharmsol::{ErrorPoly, SubjectBuilderExt};
433
434    fn minimal_ode() -> pharmsol::ODE {
435        pharmsol::equation::ODE::new(
436            |x, p, _t, dx, b, _rateiv, _cov| {
437                let ke = p[0];
438                dx[0] = -ke * x[0] + b[0];
439            },
440            |_p, _t, _cov| pharmsol::lag! {},
441            |_p, _t, _cov| pharmsol::fa! {},
442            |_p, _t, _cov, _x| {},
443            |x, p, _t, _cov, y| {
444                let v = p[1];
445                y[0] = x[0] / v;
446            },
447        )
448        .with_nstates(1)
449        .with_ndrugs(1)
450        .with_nout(1)
451        .with_metadata(
452            metadata::new("chain_test")
453                .parameters(["ke", "v"])
454                .states(["central"])
455                .outputs(["0"])
456                .route(pharmsol::equation::Route::bolus("0").to_state("central")),
457        )
458        .expect("metadata should validate")
459    }
460
461    fn minimal_data() -> pharmsol::Data {
462        let subject = pharmsol::Subject::builder("1")
463            .bolus(0.0, 100.0, 0)
464            .observation(1.0, 10.0, 0)
465            .observation(2.0, 8.0, 0)
466            .build();
467        pharmsol::Data::new(vec![subject])
468    }
469
470    #[test]
471    fn chain_npag_to_npod_preserves_support_points() {
472        let ode = minimal_ode();
473        let data = minimal_data();
474        let params = ParameterSpace::bounded()
475            .add("ke", 0.001, 3.0)
476            .add("v", 25.0, 250.0);
477        let prior = Theta::sobol_default(&params).unwrap();
478        let err = AssayErrorModels::new()
479            .add(
480                "0",
481                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
482            )
483            .unwrap();
484
485        let r1 = EstimationProblem::nonparametric(ode, data, prior, err)
486            .unwrap()
487            .fit_with(NpagConfig::new().max_cycles(2))
488            .unwrap();
489
490        let n_spp = r1.get_theta().nspp();
491        assert!(n_spp > 0, "NPAG should produce support points");
492
493        // Chain into NPOD — should complete without error
494        let r2 = r1.chain(NpodConfig::new().max_cycles(1)).unwrap();
495
496        assert!(r2.get_theta().nspp() > 0);
497        assert_eq!(r2.data().subjects().len(), 1);
498        assert_eq!(r2.cycles(), 1);
499    }
500
501    #[test]
502    fn chain_npag_to_npag_maintains_or_improves_objf() {
503        let ode = minimal_ode();
504        let data = minimal_data();
505        let params = ParameterSpace::bounded()
506            .add("ke", 0.001, 3.0)
507            .add("v", 25.0, 250.0);
508        let prior = Theta::sobol_with_seed(&params, 5, 42).unwrap();
509        let err = AssayErrorModels::new()
510            .add(
511                "0",
512                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
513            )
514            .unwrap();
515
516        let r1 = EstimationProblem::nonparametric(ode, data, prior, err)
517            .unwrap()
518            .fit_with(NpagConfig::new().max_cycles(5))
519            .unwrap();
520
521        let objf1 = r1.objf();
522
523        // Chain into another NPAG run
524        let r2 = r1.chain(NpagConfig::new().max_cycles(3)).unwrap();
525
526        // Second run should not regress significantly
527        assert!(
528            r2.objf() <= objf1 + 0.5,
529            "OBJF regressed: {} -> {}",
530            objf1,
531            r2.objf()
532        );
533    }
534
535    #[test]
536    fn chain_npag_to_npod_with_unconverged_result() {
537        let ode = minimal_ode();
538        let data = minimal_data();
539        let params = ParameterSpace::bounded()
540            .add("ke", 0.001, 3.0)
541            .add("v", 25.0, 250.0);
542        let prior = Theta::sobol_with_seed(&params, 5, 42).unwrap();
543        let err = AssayErrorModels::new()
544            .add(
545                "0",
546                AssayErrorModel::additive(ErrorPoly::new(0.0, 0.5, 0.0, 0.0), 0.0),
547            )
548            .unwrap();
549
550        // Run only 1 cycle — will be unconverged
551        let r1 = EstimationProblem::nonparametric(ode, data, prior, err)
552            .unwrap()
553            .fit_with(NpagConfig::new().max_cycles(1))
554            .unwrap();
555
556        // Chaining from unconverged should still work
557        let r2 = r1.chain(NpodConfig::new().max_cycles(1));
558        assert!(
559            r2.is_ok(),
560            "Chaining from unconverged result should succeed"
561        );
562    }
563}
pmcore/estimation/nonparametric/result.rs

pmcore/estimation/nonparametric/
result.rs
