feat: Update NPResult to contain the posterior and predictions

Cargo.toml

@@ -29,7 +29,7 @@ tracing-subscriber = { version = "0.3.19", features = [
 ] }
 faer = "0.23.1"
 faer-ext = { version = "0.7.1", features = ["nalgebra", "ndarray"] }
-pharmsol = "=0.20.0"
+pharmsol = "=0.21.0"
 rand = "0.9.0"
 anyhow = "1.0.100"
 rayon = "1.10.0"

examples/bimodal_ke/main.rs

@@ -48,7 +48,7 @@ fn main() -> Result<()> {
     settings.initialize_logs()?;
     let data = data::read_pmetrics("examples/bimodal_ke/bimodal_ke.csv")?;
     let mut algorithm = dispatch_algorithm(settings, eq, data)?;
-    let result = algorithm.fit()?;
+    let mut result = algorithm.fit()?;
     result.write_outputs()?;
 
     Ok(())

examples/drusano/main.rs

@@ -137,7 +137,7 @@ fn main() -> Result<()> {
     algorithm.initialize().unwrap();
     algorithm.fit().unwrap();
     // while !algorithm.next_cycle().unwrap() {}
-    let result = algorithm.into_npresult();
+    let mut result = algorithm.into_npresult()?;
     result.write_outputs().unwrap();
     Ok(())
 }

examples/iov/main.rs

@@ -52,7 +52,7 @@ fn main() -> Result<()> {
     let data = data::read_pmetrics("examples/iov/test.csv").unwrap();
     let mut algorithm = dispatch_algorithm(settings, sde, data).unwrap();
     algorithm.initialize().unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     result.write_outputs().unwrap();
 
     Ok(())

examples/meta/main.rs

@@ -64,6 +64,6 @@ fn main() {
     let mut algorithm = dispatch_algorithm(settings, eq, data).unwrap();
     // let result = algorithm.fit().unwrap();
     algorithm.initialize().unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     result.write_outputs().unwrap();
 }

examples/neely/main.rs

@@ -96,6 +96,6 @@ fn main() {
     settings.initialize_logs().unwrap();
     let data = data::read_pmetrics("examples/neely/data.csv").unwrap();
     let mut algorithm = dispatch_algorithm(settings, ode, data).unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     result.write_outputs().unwrap();
 }

examples/new_iov/main.rs

@@ -55,6 +55,6 @@ fn main() {
     let data = data::read_pmetrics("examples/new_iov/data.csv").unwrap();
     let mut algorithm = dispatch_algorithm(settings, sde, data).unwrap();
     algorithm.initialize().unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     result.write_outputs().unwrap();
 }

examples/theophylline/main.rs

@@ -53,6 +53,6 @@ fn main() {
     let mut algorithm = dispatch_algorithm(settings, eq, data).unwrap();
     // let result = algorithm.fit().unwrap();
     algorithm.initialize().unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     result.write_outputs().unwrap();
 }

examples/two_eq_lag/main.rs

@@ -89,7 +89,7 @@ fn main() {
     settings.initialize_logs().unwrap();
     let data = data::read_pmetrics("examples/two_eq_lag/two_eq_lag.csv").unwrap();
     let mut algorithm = dispatch_algorithm(settings, eq, data).unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     // algorithm.initialize().unwrap();
     // while !algorithm.next_cycle().unwrap() {}
     // let result = algorithm.into_npresult();

examples/vanco_sde/main.rs

@@ -78,6 +78,6 @@ fn main() {
 
     let mut algorithm = dispatch_algorithm(settings, sde, data).unwrap();
     algorithm.initialize().unwrap();
-    let result = algorithm.fit().unwrap();
+    let mut result = algorithm.fit().unwrap();
     result.write_outputs().unwrap();
 }

src/algorithms/mod.rs

@@ -304,11 +304,11 @@ pub trait Algorithms<E: Equation + Send + 'static>: Sync + Send + 'static {
                 Status::Stop(_) => break,
             }
         }
-        Ok(self.into_npresult())
+        Ok(self.into_npresult()?)
     }
 
     #[allow(clippy::wrong_self_convention)]
-    fn into_npresult(&self) -> NPResult<E>;
+    fn into_npresult(&self) -> Result<NPResult<E>>;
 }
 
 pub fn dispatch_algorithm<E: Equation + Send + 'static>(

src/algorithms/npag.rs

@@ -77,7 +77,7 @@ impl<E: Equation + Send + 'static> Algorithms<E> for NPAG<E> {
     fn equation(&self) -> &E {
         &self.equation
     }
-    fn into_npresult(&self) -> NPResult<E> {
+    fn into_npresult(&self) -> Result<NPResult<E>> {
         NPResult::new(
             self.equation.clone(),
             self.data.clone(),
@@ -329,17 +329,13 @@ impl<E: Equation + Send + 'static> Algorithms<E> for NPAG<E> {
                 let (lambda_up, objf_up) = match burke(&psi_up) {
                     Ok((lambda, objf)) => (lambda, objf),
                     Err(err) => {
-                        //todo: write out report
-                        return Err(anyhow::anyhow!("Error in IPM during optim: {:?}", err));
+                        bail!("Error in IPM during optim: {:?}", err);
                     }
                 };
                 let (lambda_down, objf_down) = match burke(&psi_down) {
                     Ok((lambda, objf)) => (lambda, objf),
                     Err(err) => {
-                        //todo: write out report
-                        //panic!("Error in IPM: {:?}", err);
-                        return Err(anyhow::anyhow!("Error in IPM during optim: {:?}", err));
-                        //(Array1::zeros(1), f64::NEG_INFINITY)
+                        bail!("Error in IPM during optim: {:?}", err);
                     }
                 };
                 if objf_up > self.objf {

src/algorithms/npod.rs

@@ -73,7 +73,7 @@ impl<E: Equation + Send + 'static> Algorithms<E> for NPOD<E> {
             data,
         }))
     }
-    fn into_npresult(&self) -> NPResult<E> {
+    fn into_npresult(&self) -> Result<NPResult<E>> {
         NPResult::new(
             self.equation.clone(),
             self.data.clone(),
@@ -333,17 +333,13 @@ impl<E: Equation + Send + 'static> Algorithms<E> for NPOD<E> {
                 let (lambda_up, objf_up) = match burke(&psi_up) {
                     Ok((lambda, objf)) => (lambda, objf),
                     Err(err) => {
-                        //todo: write out report
-                        return Err(anyhow::anyhow!("Error in IPM during optim: {:?}", err));
+                        bail!("Error in IPM during optim: {:?}", err);
                     }
                 };
                 let (lambda_down, objf_down) = match burke(&psi_down) {
                     Ok((lambda, objf)) => (lambda, objf),
                     Err(err) => {
-                        //todo: write out report
-                        //panic!("Error in IPM: {:?}", err);
-                        return Err(anyhow::anyhow!("Error in IPM during optim: {:?}", err));
-                        //(Array1::zeros(1), f64::NEG_INFINITY)
+                        bail!("Error in IPM during optim: {:?}", err);
                     }
                 };
                 if objf_up > self.objf {

src/algorithms/postprob.rs

@@ -51,7 +51,7 @@ impl<E: Equation + Send + 'static> Algorithms<E> for POSTPROB<E> {
             cyclelog: CycleLog::new(),
         }))
     }
-    fn into_npresult(&self) -> NPResult<E> {
+    fn into_npresult(&self) -> Result<NPResult<E>> {
         NPResult::new(
             self.equation.clone(),
             self.data.clone(),

src/bestdose/predictions.rs

@@ -319,7 +319,7 @@ pub fn calculate_final_predictions(
     let concentration_preds = NPPredictions::calculate(
         &problem.eq,
         &Data::new(vec![target_with_optimal.clone()]),
-        problem.theta.clone(),
+        &problem.theta,
         weights,
         &posterior,
         0.0,

src/routines/output/mod.rs

@@ -1,6 +1,7 @@
 use crate::algorithms::{Status, StopReason};
 use crate::prelude::*;
 use crate::routines::output::cycles::CycleLog;
+use crate::routines::output::posterior::Posterior;
 use crate::routines::output::predictions::NPPredictions;
 use crate::routines::settings::Settings;
 use crate::structs::psi::Psi;
@@ -36,15 +37,18 @@ pub struct NPResult<E: Equation> {
     objf: f64,
     cycles: usize,
     status: Status,
-    par_names: Vec<String>,
     settings: Settings,
     cyclelog: CycleLog,
+    predictions: Option<NPPredictions>,
+    posterior: Posterior,
 }
 
 #[allow(clippy::too_many_arguments)]
 impl<E: Equation> NPResult<E> {
     /// Create a new NPResult object
-    pub fn new(
+    ///
+    /// This will also calculate the [Posterior] structure and add it to the NPResult
+    pub(crate) fn new(
         equation: E,
         data: Data,
         theta: Theta,
@@ -55,12 +59,12 @@ impl<E: Equation> NPResult<E> {
         status: Status,
         settings: Settings,
         cyclelog: CycleLog,
-    ) -> Self {
-        // TODO: Add support for fixed and constant parameters
-
-        let par_names = settings.parameters().names();
+    ) -> Result<Self> {
+        // Calculate the posterior probabilities
+        let posterior = posterior(&psi, &w)
+            .context("Failed to calculate posterior during initialization of NPResult")?;
 
-        Self {
+        let result = Self {
             equation,
             data,
             theta,
@@ -69,10 +73,13 @@ impl<E: Equation> NPResult<E> {
             objf,
             cycles,
             status,
-            par_names,
             settings,
             cyclelog,
-        }
+            predictions: None,
+            posterior,
+        };
+
+        Ok(result)
     }
 
     pub fn cycles(&self) -> usize {
@@ -91,6 +98,18 @@ impl<E: Equation> NPResult<E> {
         &self.theta
     }
 
+    pub fn data(&self) -> &Data {
+        &self.data
+    }
+
+    pub fn cycle_log(&self) -> &CycleLog {
+        &self.cyclelog
+    }
+
+    pub fn settings(&self) -> &Settings {
+        &self.settings
+    }
+
     /// Get the [Psi] structure
     pub fn psi(&self) -> &Psi {
         &self.psi
@@ -101,7 +120,24 @@ impl<E: Equation> NPResult<E> {
         &self.w
     }
 
-    pub fn write_outputs(&self) -> Result<()> {
+    /// Calculate and store the [NPPredictions] in the [NPResult]
+    ///
+    /// This will overwrite any existing predictions stored in the result!
+    pub fn calculate_predictions(&mut self, idelta: f64, tad: f64) -> Result<()> {
+        let predictions = NPPredictions::calculate(
+            &self.equation,
+            &self.data,
+            &self.theta,
+            &self.w,
+            &self.posterior,
+            idelta,
+            tad,
+        )?;
+        self.predictions = Some(predictions);
+        Ok(())
+    }
+
+    pub fn write_outputs(&mut self) -> Result<()> {
         if self.settings.output().write {
             tracing::debug!("Writing outputs to {:?}", self.settings.output().path);
             self.settings.write()?;
@@ -288,7 +324,7 @@ impl<E: Equation> NPResult<E> {
             .from_writer(&outputfile.file);
 
         // Create the headers
-        let mut theta_header = self.par_names.clone();
+        let mut theta_header = self.settings.parameters().names();
         theta_header.push("prob".to_string());
         writer.write_record(&theta_header)?;
 
@@ -310,11 +346,9 @@ impl<E: Equation> NPResult<E> {
     pub fn write_posterior(&self) -> Result<()> {
         tracing::debug!("Writing posterior parameter probabilities...");
         let theta = &self.theta;
-        let w = &self.w;
-        let psi = &self.psi;
 
         // Calculate the posterior probabilities
-        let posterior = posterior(psi, w)?;
+        let posterior = self.posterior.clone();
 
         // Create the output folder if it doesn't exist
         let outputfile = match OutputFile::new(&self.settings.output().path, "posterior.csv") {
@@ -372,21 +406,15 @@ impl<E: Equation> NPResult<E> {
     }
 
     /// Writes the predictions
-    pub fn write_predictions(&self, idelta: f64, tad: f64) -> Result<()> {
+    pub fn write_predictions(&mut self, idelta: f64, tad: f64) -> Result<()> {
         tracing::debug!("Writing predictions...");
 
-        let posterior = posterior(&self.psi, &self.w)?;
+        self.calculate_predictions(idelta, tad)?;
 
-        // Calculate the predictions
-        let predictions = NPPredictions::calculate(
-            &self.equation,
-            &self.data,
-            self.theta.clone(),
-            &self.w,
-            &posterior,
-            idelta,
-            tad,
-        )?;
+        let predictions = self
+            .predictions
+            .as_ref()
+            .expect("Predictions should have been calculated, but are of type None.");
 
         // Write (full) predictions to pred.csv
         let outputfile_pred = OutputFile::new(&self.settings.output().path, "pred.csv")?;

src/routines/output/predictions.rs

@@ -7,7 +7,11 @@ use crate::{
     structs::{theta::Theta, weights::Weights},
 };
 
-// Structure for the output
+/// Container for the multiple model estimated predictions
+///
+/// Each row contains the predictions for a single time point for a single subject
+/// It includes the population and posterior mean and median predictions
+/// These are defined by the mean and median of the prediction for each model, weighted by the population or posterior weights
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct NPPredictionRow {
     /// The subject ID
@@ -60,6 +64,10 @@ impl NPPredictionRow {
     pub fn post_median(&self) -> f64 {
         self.post_median
     }
+
+    pub fn censoring(&self) -> Censor {
+        self.cens
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -114,7 +122,7 @@ impl NPPredictions {
     pub fn calculate(
         equation: &impl pharmsol::prelude::simulator::Equation,
         data: &Data,
-        theta: Theta,
+        theta: &Theta,
         w: &Weights,
         posterior: &Posterior,
         idelta: f64,