Refactor SimulationState, add SolutionState (#2887)

Often `SimulationState` is passed around where only a subset of the data is needed. Extract that to a separate `SolutionState`. This makes storing intermediate results a bit easier and avoids missing individual components.

Author: dweindl

include/amici/backwardproblem.h

@@ -151,7 +151,7 @@ class SteadyStateBackwardProblem {
      * @param ws Workspace for backward simulation
      */
     SteadyStateBackwardProblem(
-        Solver const& solver, Model& model, SimulationState& final_state,
+        Solver const& solver, Model& model, SolutionState& final_state,
         gsl::not_null<BwdSimWorkspace*> ws
     );
 
@@ -248,7 +248,7 @@ class SteadyStateBackwardProblem {
     AmiVector xQ_;
 
     /** Final state from pre/post-equilibration forward problem */
-    SimulationState& final_state_;
+    SolutionState& final_state_;
 
     /** Newton solver */
     NewtonSolver newton_solver_;

include/amici/forwardproblem.h

@@ -113,12 +113,10 @@ struct FwdSimWorkspace {
     FwdSimWorkspace(
         gsl::not_null<Model*> const& model, gsl::not_null<Solver*> solver
     )
-        : x(model->nx_solver, solver->getSunContext())
+        : sol(NAN, model->nx_solver, model->nplist(), solver->getSunContext())
         , x_old(model->nx_solver, solver->getSunContext())
-        , dx(model->nx_solver, solver->getSunContext())
         , xdot(model->nx_solver, solver->getSunContext())
         , xdot_old(model->nx_solver, solver->getSunContext())
-        , sx(model->nx_solver, model->nplist(), solver->getSunContext())
         , sdx(model->nx_solver, model->nplist(), solver->getSunContext())
         , stau(model->nplist())
         , roots_found(model->ne, 0)
@@ -127,27 +125,18 @@ struct FwdSimWorkspace {
         , rootvals(gsl::narrow<decltype(rootvals)::size_type>(model->ne), 0.0)
 
     {}
-    /** current simulation time */
-    realtype t{NAN};
-
-    /** state vector (dimension: nx_solver) */
-    AmiVector x;
+    /** Current solution state */
+    SolutionState sol;
 
     /** old state vector (dimension: nx_solver) */
     AmiVector x_old;
 
-    /** differential state vector (dimension: nx_solver) */
-    AmiVector dx;
-
     /** time derivative state vector (dimension: nx_solver) */
     AmiVector xdot;
 
     /** old time derivative state vector (dimension: nx_solver) */
     AmiVector xdot_old;
 
-    /** sensitivity state vector array (dimension: nx_cl x nplist, row-major) */
-    AmiVectorArray sx;
-
     /** differential sensitivity state vector array
      * (dimension: nx_cl x nplist, row-major) */
     AmiVectorArray sdx;
@@ -362,7 +351,7 @@ class EventHandlingSimulator {
             }))
             return;
 
-        result.discs.emplace_back(ws_->t);
+        result.discs.emplace_back(ws_->sol.t);
         store_event(edata);
     }
 
@@ -428,15 +417,15 @@ class SteadystateProblem {
      * @return x
      */
     [[nodiscard]] AmiVector const& getState() const {
-        return period_result_.final_state_.x;
+        return period_result_.final_state_.sol.x;
     }
 
     /**
      * @brief Return state sensitivity at steady state
      * @return sx
      */
     [[nodiscard]] AmiVectorArray const& getStateSensitivity() const {
-        return period_result_.final_state_.sx;
+        return period_result_.final_state_.sol.sx;
     }
 
     /**
@@ -460,7 +449,7 @@ class SteadystateProblem {
      * @return Time at which steady state was found (model time units).
      */
     [[nodiscard]] realtype getSteadyStateTime() const {
-        return period_result_.final_state_.t;
+        return period_result_.final_state_.sol.t;
     }
 
     /**
@@ -504,7 +493,6 @@ class SteadystateProblem {
 
     /**
      * @brief Try to determine the steady state by using Newton's method.
-     * @param model Model instance.
      * @param newton_retry Flag indicating whether Newton's method is being
      * relaunched.
      */
@@ -672,18 +660,12 @@ class ForwardProblem {
      */
     std::vector<realtype> getAdjointUpdates(Model& model, ExpData const& edata);
 
-    /**
-     * @brief Accessor for t
-     * @return t
-     */
-    [[nodiscard]] realtype getTime() const { return t_; }
-
     /**
      * @brief Accessor for sx
      * @return sx
      */
     [[nodiscard]] AmiVectorArray const& getStateSensitivity() const {
-        return ws_.sx;
+        return ws_.sol.sx;
     }
 
     /**
@@ -713,7 +695,7 @@ class ForwardProblem {
      * @return time point
      */
     [[nodiscard]] realtype getFinalTime() const {
-        return main_simulator_.result.final_state_.t;
+        return main_simulator_.result.final_state_.sol.t;
     }
 
     /**
@@ -732,7 +714,8 @@ class ForwardProblem {
      */
     [[nodiscard]] SimulationState const&
     getSimulationStateTimepoint(int const it) const {
-        if (model->getTimepoint(it) == main_simulator_.result.initial_state_.t)
+        if (model->getTimepoint(it)
+            == main_simulator_.result.initial_state_.sol.t)
             return getInitialSimulationState();
         auto const map_iter = main_simulator_.result.timepoint_states_.find(
             model->getTimepoint(it)
@@ -877,9 +860,6 @@ class ForwardProblem {
      * (dimension nJ x nx x nMaxEvent, ordering =?) */
     std::vector<realtype> dJzdx_;
 
-    /** current time */
-    realtype t_;
-
     /** flag to indicate whether solver was preeinitialized via preequilibration
      */
     bool preequilibrated_{false};

include/amici/model_state.h

@@ -462,12 +462,11 @@ struct ModelStateDerived {
 };
 
 /**
- * @brief implements an exchange format to store and transfer the state of a
- * simulation at a specific timepoint.
+ * @brief Container for the IVP solution state at a specific timepoint.
  */
-struct SimulationState {
+struct SolutionState {
     /** timepoint */
-    realtype t;
+    realtype t{NAN};
     /**
      * partial state vector, excluding states eliminated from conservation laws
      */
@@ -482,8 +481,35 @@ struct SimulationState {
      * conservation laws
      */
     AmiVectorArray sx;
+
+    SolutionState() = default;
+
+    /**
+     * @brief Constructor.
+     * @param t_ Current timepoint.
+     * @param nx_solver Number of solver state variables.
+     * @param nplist Number of parameter w.r.t. which to compute sensitivities.
+     * @param ctx SUNDIALS context.
+     */
+    SolutionState(
+        realtype t_, long int nx_solver, long int nplist, SUNContext ctx
+    )
+        : t(t_)
+        , x(nx_solver, ctx)
+        , dx(nx_solver, ctx)
+        , sx(nx_solver, nplist, ctx) {}
+};
+
+/**
+ * @brief implements an exchange format to store and transfer the state of a
+ * simulation at a specific timepoint.
+ */
+struct SimulationState {
+    /** Solution state */
+    SolutionState sol;
+
     /** state of the model that was used for simulation */
-    ModelState state;
+    ModelState mod;
 };
 
 } // namespace amici

include/amici/newton_solver.h

@@ -2,6 +2,7 @@
 #define amici_newton_solver_h
 
 #include "amici/defines.h"
+#include "amici/model_state.h"
 #include "amici/sundials_linsol_wrapper.h"
 #include "amici/vector.h"
 
@@ -46,6 +47,15 @@ struct DEStateView {
         : t(t_)
         , x(x_)
         , dx(dx_) {}
+
+    /**
+     * @brief Construct a DEStateView from a SolutionState reference.
+     * @param sol
+     */
+    DEStateView(SolutionState& sol)
+        : t(sol.t)
+        , x(sol.x)
+        , dx(sol.dx) {}
 };
 
 /**

include/amici/rdata.h

@@ -642,26 +642,20 @@ class ReturnData : public ModelDimensions {
     template <class T>
     void
     storeJacobianAndDerivativeInReturnData(T const& problem, Model& model) {
-        auto simulation_state = problem.getFinalSimulationState();
-        model.setModelState(simulation_state.state);
-
+        auto const& simulation_state = problem.getFinalSimulationState();
+        model.setModelState(simulation_state.mod);
+        auto const& sol = simulation_state.sol;
         sundials::Context sunctx;
         AmiVector xdot(nx_solver, sunctx);
         if (!this->xdot.empty() || !this->J.empty())
-            model.fxdot(
-                simulation_state.t, simulation_state.x, simulation_state.dx,
-                xdot
-            );
+            model.fxdot(sol.t, sol.x, sol.dx, xdot);
 
         if (!this->xdot.empty())
             writeSlice(xdot, this->xdot);
 
         if (!this->J.empty()) {
             SUNMatrixWrapper J(nx_solver, nx_solver, sunctx);
-            model.fJ(
-                simulation_state.t, 0.0, simulation_state.x,
-                simulation_state.dx, xdot, J
-            );
+            model.fJ(sol.t, 0.0, sol.x, sol.dx, xdot, J);
             // CVODES uses colmajor, so we need to transform to rowmajor
             for (int ix = 0; ix < model.nx_solver; ix++)
                 for (int jx = 0; jx < model.nx_solver; jx++)
@@ -674,13 +668,11 @@ class ReturnData : public ModelDimensions {
      * @brief Residual function
      * @param it time index
      * @param model model that was used for forward/backward simulation
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol Solution state the timepoint `it`
      * @param edata ExpData instance containing observable data
      */
-    void fres(
-        int it, Model& model, SimulationState const& simulation_state,
-        ExpData const& edata
-    );
+    void
+    fres(int it, Model& model, SolutionState const& sol, ExpData const& edata);
 
     /**
      * @brief Chi-squared function
@@ -693,25 +685,21 @@ class ReturnData : public ModelDimensions {
      * @brief Residual sensitivity function
      * @param it time index
      * @param model model that was used for forward/backward simulation
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol solution state the timepoint `it`
      * @param edata ExpData instance containing observable data
      */
-    void fsres(
-        int it, Model& model, SimulationState const& simulation_state,
-        ExpData const& edata
-    );
+    void
+    fsres(int it, Model& model, SolutionState const& sol, ExpData const& edata);
 
     /**
      * @brief Fisher information matrix function
      * @param it time index
      * @param model model that was used for forward/backward simulation
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol Solution state the timepoint `it`
      * @param edata ExpData instance containing observable data
      */
-    void fFIM(
-        int it, Model& model, SimulationState const& simulation_state,
-        ExpData const& edata
-    );
+    void
+    fFIM(int it, Model& model, SolutionState const& sol, ExpData const& edata);
 
     /**
      * @brief Set likelihood, state variables, outputs and respective
@@ -755,54 +743,49 @@ class ReturnData : public ModelDimensions {
      * the model state was set appropriately
      * @param it timepoint index
      * @param model model that was used in forward solve
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol solution state the timepoint `it`
      * @param edata ExpData instance carrying experimental data
      */
     void getDataOutput(
-        int it, Model& model, SimulationState const& simulation_state,
-        ExpData const* edata
+        int it, Model& model, SolutionState const& sol, ExpData const* edata
     );
 
     /**
      * @brief Extracts data information for forward sensitivity analysis,
      * expects that the model state was set appropriately
      * @param it index of current timepoint
      * @param model model that was used in forward solve
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol Solution state the timepoint `it`
      * @param edata ExpData instance carrying experimental data
      */
     void getDataSensisFSA(
-        int it, Model& model, SimulationState const& simulation_state,
-        ExpData const* edata
+        int it, Model& model, SolutionState const& sol, ExpData const* edata
     );
 
     /**
      * @brief Extracts output information for events, expects that the model
      * state was set appropriately
-     * @param t event timepoint
      * @param rootidx information about which roots fired
      * (1 indicating fired, 0/-1 for not)
      * @param model model that was used in forward solve
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol Solution state the timepoint `it`
      * @param edata ExpData instance carrying experimental data
      */
     void getEventOutput(
-        realtype t, std::vector<int> const& rootidx, Model& model,
-        SimulationState const& simulation_state, ExpData const* edata
+        std::vector<int> const& rootidx, Model& model, SolutionState const& sol,
+        ExpData const* edata
     );
 
     /**
      * @brief Extracts event information for forward sensitivity analysis,
      * expects the model state was set appropriately
      * @param ie index of event type
-     * @param t event timepoint
      * @param model model that was used in forward solve
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol Solution state the timepoint `it`
      * @param edata ExpData instance carrying experimental data
      */
     void getEventSensisFSA(
-        int ie, realtype t, Model& model,
-        SimulationState const& simulation_state, ExpData const* edata
+        int ie, Model& model, SolutionState const& sol, ExpData const* edata
     );
 
     /**
@@ -826,13 +809,13 @@ class ReturnData : public ModelDimensions {
      * (llhS0), if no preequilibration was run or if forward sensitivities were
      * used
      * @param model model that was used for forward/backward simulation
-     * @param simulation_state simulation state the timepoint `it`
+     * @param sol Solution state the timepoint `it`
      * @param llhS0 contribution to likelihood for initial state sensitivities
      * @param xB vector with final adjoint state
      * (excluding conservation laws)
      */
     void handleSx0Forward(
-        Model const& model, SimulationState const& simulation_state,
+        Model const& model, SolutionState const& sol,
         std::vector<realtype>& llhS0, AmiVector const& xB
     ) const;
 };