Refactor steady-state simulation (#2855)

Some changes to make it easier to implement event handling during pre-equilibration later on.

* Move getWrmsState to lambda
* Change `t` in `writeSolution` from pointer to reference (missed one overload previously)
* non-const Solver

Author: dweindl

include/amici/forwardproblem.h

@@ -222,8 +222,7 @@ class EventHandlingSimulator {
      */
     EventHandlingSimulator(
         gsl::not_null<Model*> model, gsl::not_null<Solver*> solver,
-        gsl::not_null<FwdSimWorkspace*> ws,
-        gsl::not_null<std::vector<realtype>*> dJzdx
+        gsl::not_null<FwdSimWorkspace*> ws, std::vector<realtype>* dJzdx
     )
         : model_(model)
         , solver_(solver)

include/amici/solver.h

@@ -651,7 +651,7 @@ class Solver {
      * @param xQ quadrature
      */
     void writeSolution(
-        realtype* t, AmiVector& x, AmiVector& dx, AmiVectorArray& sx,
+        realtype& t, AmiVector& x, AmiVector& dx, AmiVectorArray& sx,
         AmiVector& xQ
     ) const;
 

include/amici/steadystateproblem.h

@@ -239,9 +239,8 @@ class SteadystateProblem {
      * @param it Index of the current output time point.
      * @param t0 Initial time for the steady state simulation.
      */
-    void workSteadyStateProblem(
-        Solver const& solver, Model& model, int it, realtype t0
-    );
+    void
+    workSteadyStateProblem(Solver& solver, Model& model, int it, realtype t0);
 
     /**
      * @brief Compute the gradient via adjoint steady state sensitivities.
@@ -375,8 +374,7 @@ class SteadystateProblem {
      * @param it Index of the current output time point.
      * @param t0 Initial time for the steady state simulation.
      */
-    void
-    findSteadyState(Solver const& solver, Model& model, int it, realtype t0);
+    void findSteadyState(Solver& solver, Model& model, int it, realtype t0);
 
     /**
      * @brief Try to determine the steady state by using Newton's method.
@@ -396,7 +394,7 @@ class SteadystateProblem {
      * successfully, or if it failed.
      */
     SteadyStateStatus findSteadyStateBySimulation(
-        Solver const& solver, Model& model, int it, realtype t0
+        Solver& solver, Model& model, int it, realtype t0
     );
 
     /**
@@ -438,13 +436,6 @@ class SteadystateProblem {
         bool tried_newton_1, bool tried_simulation, bool tried_newton_2
     ) const;
 
-    /**
-     * @brief Checks steady-state convergence for state variables
-     * @param model Model instance
-     * @return weighted root mean squared residuals of the RHS
-     */
-    realtype getWrmsState(Model& model);
-
     /**
      * @brief Checks convergence for state sensitivities
      * @param model Model instance
@@ -460,7 +451,7 @@ class SteadystateProblem {
      * @param model Model instance.
      * simulation.
      */
-    void runSteadystateSimulationFwd(Solver const& solver, Model& model);
+    void runSteadystateSimulationFwd(Solver& solver, Model& model);
 
     /**
      * @brief Launch backward simulation if Newton solver or linear system solve

src/forwardproblem.cpp

@@ -450,14 +450,15 @@ void EventHandlingSimulator::store_event(ExpData const* edata) {
             continue;
         }
 
-        if (edata && solver_->computingASA())
+        if (edata && solver_->computingASA()) {
+            Expects(dJzdx_ != nullptr);
             model_->getAdjointStateEventUpdate(
                 slice(
                     *dJzdx_, ws_->nroots.at(ie), model_->nx_solver * model_->nJ
                 ),
                 ie, ws_->nroots.at(ie), t_, ws_->x, *edata
             );
-
+        }
         ws_->nroots.at(ie)++;
     }
 

src/model.cpp

@@ -7,11 +7,11 @@
 
 #include <algorithm>
 #include <cassert>
-#include <sstream>
 #include <cmath>
 #include <cstring>
 #include <numeric>
 #include <regex>
+#include <sstream>
 #include <utility>
 
 namespace amici {

src/solver.cpp

@@ -1307,15 +1307,15 @@ void Solver::resetMutableMemory(
 }
 
 void Solver::writeSolution(
-    realtype* t, AmiVector& x, AmiVector& dx, AmiVectorArray& sx, AmiVector& xQ
+    realtype& t, AmiVector& x, AmiVector& dx, AmiVectorArray& sx, AmiVector& xQ
 ) const {
-    *t = gett();
+    t = gett();
     if (quad_initialized_)
-        xQ.copy(getQuadrature(*t));
+        xQ.copy(getQuadrature(t));
     if (sens_initialized_)
-        sx.copy(getStateSensitivity(*t));
-    x.copy(getState(*t));
-    dx.copy(getDerivativeState(*t));
+        sx.copy(getStateSensitivity(t));
+    x.copy(getState(t));
+    dx.copy(getDerivativeState(t));
 }
 
 void Solver::writeSolution(

src/steadystateproblem.cpp

@@ -173,7 +173,7 @@ SteadystateProblem::SteadystateProblem(
 }
 
 void SteadystateProblem::workSteadyStateProblem(
-    Solver const& solver, Model& model, int it, realtype t0
+    Solver& solver, Model& model, int it, realtype t0
 ) {
     if (model.ne > 0) {
         solver.logger->log(
@@ -280,7 +280,7 @@ void SteadystateProblem::workSteadyStateBackwardProblem(
 }
 
 void SteadystateProblem::findSteadyState(
-    Solver const& solver, Model& model, int it, realtype t0
+    Solver& solver, Model& model, int it, realtype t0
 ) {
     steady_state_status_.resize(3, SteadyStateStatus::not_run);
     // Turn off Newton's method if 'integrationOnly' approach is chosen for
@@ -365,7 +365,7 @@ void SteadystateProblem::findSteadyStateByNewtonsMethod(
 }
 
 SteadyStateStatus SteadystateProblem::findSteadyStateBySimulation(
-    Solver const& solver, Model& model, int const it, realtype const t0
+    Solver& solver, Model& model, int const it, realtype const t0
 ) {
     try {
         if (it < 0) {
@@ -551,17 +551,6 @@ void SteadystateProblem::getQuadratureBySimulation(
     throw AmiException(errorString.c_str());
 }
 
-realtype SteadystateProblem::getWrmsState(Model& model) {
-    updateRightHandSide(model);
-
-    if (newton_step_conv_) {
-        newtons_method_.compute_step(ws_->xdot, {ws_->t, ws_->x, ws_->dx});
-        return wrms_computer_x_.wrms(newtons_method_.get_delta(), ws_->x);
-    }
-
-    return wrms_computer_x_.wrms(ws_->xdot, ws_->x);
-}
-
 realtype
 SteadystateProblem::getWrmsFSA(Model& model, WRMSComputer& wrms_computer_sx) {
     // Forward sensitivities: Compute weighted error norm for their RHS
@@ -601,7 +590,7 @@ bool SteadystateProblem::checkSteadyStateSuccess() const {
 }
 
 void SteadystateProblem::runSteadystateSimulationFwd(
-    Solver const& solver, Model& model
+    Solver& solver, Model& model
 ) {
     if (model.nx_solver == 0)
         return;
@@ -643,14 +632,25 @@ void SteadystateProblem::runSteadystateSimulationFwd(
         sensi_converged = []() { return true; };
     }
 
+    // Returns the WRMS for the current state
+    auto get_wrms_state = [&]() {
+        updateRightHandSide(model);
+        if (newton_step_conv_) {
+            newtons_method_.compute_step(ws_->xdot, {ws_->t, ws_->x, ws_->dx});
+            return wrms_computer_x_.wrms(newtons_method_.get_delta(), ws_->x);
+        }
+
+        return wrms_computer_x_.wrms(ws_->xdot, ws_->x);
+    };
+
     int& sim_steps = numsteps_.at(1);
     int convergence_check_frequency = newton_step_conv_ ? 25 : 1;
 
     while (true) {
         if (sim_steps % convergence_check_frequency == 0) {
             // Check for convergence (already before simulation, since we might
             // start in steady state)
-            wrms_ = getWrmsState(model);
+            wrms_ = get_wrms_state();
             if (wrms_ < conv_thresh && sensi_converged()) {
                 break;
             }
@@ -756,7 +756,7 @@ void SteadystateProblem::runSteadystateSimulationBwd(
         solver.step(std::max(final_state_.t, 1.0) * 10);
 
         solver.writeSolution(
-            &final_state_.t, xB_, final_state_.dx, final_state_.sx, xQ_
+            final_state_.t, xB_, final_state_.dx, final_state_.sx, xQ_
         );
     }
 }