From a3f0c377357ee5e7e4d6a43e84dfbe6259c3e7da Mon Sep 17 00:00:00 2001
From: Daniel Weindl <daniel.weindl@uni-bonn.de>
Date: Sun, 24 Nov 2024 11:00:55 +0100
Subject: [PATCH] Add missing ReturnData fields in writeReturnData

Closes #2591.
---
 include/amici/hdf5.h  | 13 +++++++
 include/amici/rdata.h |  4 +--
 src/hdf5.cpp          | 84 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 2 deletions(-)

diff --git a/include/amici/hdf5.h b/include/amici/hdf5.h
index d8acf8881e..58736e2d12 100644
--- a/include/amici/hdf5.h
+++ b/include/amici/hdf5.h
@@ -26,6 +26,7 @@ class ReturnData;
 class ExpData;
 class Model;
 class Solver;
+class LogItem;
 
 namespace hdf5 {
 
@@ -137,6 +138,18 @@ void writeReturnDataDiagnosis(
     std::string const& hdf5Location
 );
 
+
+/**
+ * @brief Write log message to HDF5 file
+ * @param file HDF5 file to write to
+ * @param logItems Log items to write
+ * @param hdf5Location Full dataset path inside the HDF5 file (will be created)
+ */
+void writeLogItemsToHDF5(
+    H5::H5File const& file,
+    std::vector<amici::LogItem> const& logItems,
+    std::string const& hdf5Location);
+
 /**
  * @brief Create the given group and possibly parents.
  * @param file HDF5 file to write to
diff --git a/include/amici/rdata.h b/include/amici/rdata.h
index df25f39923..793be9435a 100644
--- a/include/amici/rdata.h
+++ b/include/amici/rdata.h
@@ -453,7 +453,7 @@ class ReturnData : public ModelDimensions {
 
     /** boolean indicating whether residuals for standard deviations have been
      * added */
-    bool sigma_res;
+    bool sigma_res{false};
 
     /** log messages */
     std::vector<LogItem> messages;
@@ -463,7 +463,7 @@ class ReturnData : public ModelDimensions {
 
   protected:
     /** offset for sigma_residuals */
-    realtype sigma_offset;
+    realtype sigma_offset{0.0};
 
     /** array of number of found roots for a certain event type
      * (shape `ne`) */
diff --git a/src/hdf5.cpp b/src/hdf5.cpp
index 8ee20c0dbe..009c667913 100644
--- a/src/hdf5.cpp
+++ b/src/hdf5.cpp
@@ -13,6 +13,7 @@
 #include <amici/model.h>
 #include <amici/rdata.h>
 #include <amici/solver.h>
+#include <amici/logging.h>
 
 #include <hdf5_hl.h>
 
@@ -412,6 +413,11 @@ void writeReturnData(
             file, hdf5Location + "/y", rdata.y, rdata.nt, rdata.ny
         );
 
+    if (!rdata.w.empty())
+        createAndWriteDouble2DDataset(
+            file, hdf5Location + "/w", rdata.w, rdata.nt, rdata.nw
+        );
+
     if (!rdata.z.empty())
         createAndWriteDouble2DDataset(
             file, hdf5Location + "/z", rdata.z, rdata.nmaxevent, rdata.nz
@@ -480,6 +486,52 @@ void writeReturnData(
             rdata.nplist, rdata.nz
         );
 
+    // TODO currently unused
+    /*
+    if (!rdata.s2rz.empty())
+        createAndWriteDouble4DDataset(
+            file, hdf5Location + "/s2rz", rdata.s2rz, rdata.nmaxevent,
+            rdata.nztrue, rdata.nplist, rdata.nplist
+        );
+    */
+
+    std::vector<int> int_buffer(1);
+
+    int_buffer[0] = gsl::narrow<int>(rdata.newton_maxsteps);
+    H5LTset_attribute_int(
+        file.getId(), hdf5Location.c_str(), "newton_maxsteps",
+        int_buffer.data(), 1
+    );
+
+    if (!rdata.pscale.empty()) {
+        int_buffer.resize(rdata.pscale.size());
+        for (int i = 0; (unsigned)i < rdata.pscale.size(); i++)
+            int_buffer[i] = static_cast<int>(rdata.pscale[i]);
+        createAndWriteInt1DDataset(file, hdf5Location + "/pscale", int_buffer);
+    }
+
+    int_buffer[0] = static_cast<int>(rdata.o2mode);
+    H5LTset_attribute_int(
+        file.getId(), hdf5Location.c_str(), "o2mode", int_buffer.data(), 1
+    );
+
+    int_buffer[0] = static_cast<int>(rdata.sensi);
+    H5LTset_attribute_int(
+        file.getId(), hdf5Location.c_str(), "sensi", int_buffer.data(), 1
+    );
+
+    int_buffer[0] = static_cast<int>(rdata.sensi_meth);
+    H5LTset_attribute_int(
+        file.getId(), hdf5Location.c_str(), "sensi_meth", int_buffer.data(), 1
+    );
+
+    int_buffer[0] = static_cast<int>(rdata.rdata_reporting);
+    H5LTset_attribute_int(
+        file.getId(), hdf5Location.c_str(), "rdrm", int_buffer.data(), 1
+    );
+
+    writeLogItemsToHDF5(file, rdata.messages, hdf5Location + "/messages");
+
     writeReturnDataDiagnosis(rdata, file, hdf5Location + "/diagnosis");
 }
 
@@ -634,6 +686,38 @@ void writeReturnDataDiagnosis(
         createAndWriteDouble2DDataset(
             file, hdf5Location + "/J", rdata.J, rdata.nx, rdata.nx
         );
+
+    if (!rdata.x_ss.empty())
+        createAndWriteDouble1DDataset(file, hdf5Location + "/x_ss", rdata.x_ss);
+
+    if (!rdata.sx_ss.empty())
+        createAndWriteDouble2DDataset(
+            file, hdf5Location + "/sx_ss", rdata.sx_ss, rdata.nplist,
+            rdata.nx_rdata
+        );
+}
+
+void writeLogItemsToHDF5(
+    H5::H5File const& file,
+    std::vector<amici::LogItem> const& logItems,
+    std::string const& hdf5Location) {
+    try {
+        hsize_t dims[1] = { logItems.size() };
+        H5::DataSpace dataspace(1, dims);
+
+        // Create a compound datatype for the LogItem struct.
+        H5::CompType logItemType(sizeof(amici::LogItem));
+        logItemType.insertMember("severity", HOFFSET(amici::LogItem, severity), H5::PredType::NATIVE_INT);
+        logItemType.insertMember("identifier", HOFFSET(amici::LogItem, identifier), H5::StrType(H5::PredType::C_S1, H5T_VARIABLE));
+        logItemType.insertMember("message", HOFFSET(amici::LogItem, message), H5::StrType(H5::PredType::C_S1, H5T_VARIABLE));
+
+        H5::DataSet dataset = file.createDataSet(hdf5Location, logItemType, dataspace);
+
+        dataset.write(logItems.data(), logItemType);
+    } catch (H5::Exception& e) {
+        // Handle exceptions
+        throw AmiException(e.getCDetailMsg());
+    }
 }
 
 void writeReturnData(