From 949bb7fdf00e0915a39336baacd55cd8ff68131d Mon Sep 17 00:00:00 2001
From: Markus <66058642+mhovd@users.noreply.github.com>
Date: Thu, 17 Oct 2024 20:08:56 +0200
Subject: [PATCH] Write covariates at all event times

---
 .gitignore             |  3 ++-
 src/routines/output.rs | 53 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 98e91c40..1ee158ab 100644
--- a/.gitignore
+++ b/.gitignore
@@ -27,4 +27,5 @@ settings.json
 /benches/*.json
 log.txt
 op.csv
-*results.txt
\ No newline at end of file
+*results.txt
+covs.csv
diff --git a/src/routines/output.rs b/src/routines/output.rs
index d444a10f..c54c2687 100644
--- a/src/routines/output.rs
+++ b/src/routines/output.rs
@@ -83,6 +83,7 @@ impl<E: Equation> NPResult<E> {
                 .context("Failed to write observed-predicted file")?;
             self.write_pred(idelta, tad)
                 .context("Failed to write predictions")?;
+            self.write_covs().context("Failed to write covariates")?;
             if self.w.len() > 0 {
                 //TODO: find a better way to indicate that the run failed
                 self.write_posterior()
@@ -397,6 +398,58 @@ impl<E: Equation> NPResult<E> {
         );
         Ok(())
     }
+
+    /// Writes the covariates
+    pub fn write_covs(&self) -> Result<()> {
+        tracing::debug!("Writing covariates...");
+        let outputfile = OutputFile::new(&self.settings.output.path, "covs.csv")?;
+        let mut writer = WriterBuilder::new()
+            .has_headers(true)
+            .from_writer(&outputfile.file);
+
+        #[derive(Debug, Serialize)]
+        struct Row {
+            id: String,
+            time: f64,
+            block: usize,
+            cov: String,
+            value: f64,
+        }
+
+        for subject in self.data.get_subjects() {
+            for occasion in subject.occasions() {
+                if let Some(cov) = occasion.get_covariates() {
+                    let covmap = cov.covariates();
+
+                    for event in occasion.get_events(&None, &None, false) {
+                        let time = match event {
+                            Event::Bolus(bolus) => bolus.time(),
+                            Event::Infusion(infusion) => infusion.time(),
+                            Event::Observation(observation) => observation.time(),
+                        };
+
+                        for cov in &covmap {
+                            let value = cov.1.interpolate(time).unwrap();
+                            let row = Row {
+                                id: subject.id().clone(),
+                                time,
+                                block: occasion.index(),
+                                cov: cov.0.clone(),
+                                value,
+                            };
+                            writer.serialize(row)?;
+                        }
+                    }
+                }
+            }
+        }
+        writer.flush()?;
+        tracing::info!(
+            "Covariates written to {:?}",
+            &outputfile.get_relative_path()
+        );
+        Ok(())
+    }
 }
 
 /// An [NPCycle] object contains the summary of a cycle