Event handling during pre-equilibration with FSA before ASA simulation (#2881)

Implement event handling during pre-equilibration if steady-state sensitivities are computed with FSA, and ASA is used afterwards. (No ASA preeq + ASA main sim yet.)

* Store discontinuities during pre-equilibration (store `PeriodResult` instead of only `SimulationState`)
* Make sure to use the same `EventHandlingSimulator` for initial events and later events, with the correct sensitivity method

Storing the results is not so pretty yet. To be cleaned up once all methods work.

Related to #2775.

Author: dweindl

include/amici/backwardproblem.h

@@ -335,8 +335,8 @@ class BackwardProblem {
     /** current time */
     realtype t_;
 
-    /** array containing the time-points of discontinuities*/
-    std::vector<Discontinuity> discs_;
+    /** The discontinuities encountered during the main simulation. */
+    std::vector<Discontinuity> discs_main_;
 
     /**
      * state derivative of data likelihood

include/amici/forwardproblem.h

@@ -383,9 +383,11 @@ class SteadystateProblem {
      * @param ws Workspace for forward simulation
      * @param solver Solver instance
      * @param model Model instance
+     * @param is_preeq Whether this is a pre-equilibration (`true`) or
+     * post-equilibration problem (`false`).
      */
     explicit SteadystateProblem(
-        FwdSimWorkspace* ws, Solver const& solver, Model& model
+        FwdSimWorkspace* ws, Solver const& solver, Model& model, bool is_preeq
     );
 
     /**
@@ -399,29 +401,30 @@ class SteadystateProblem {
      * @param it Index of the current output time point.
      * @param t0 Initial time for the steady state simulation.
      */
-    void
-    workSteadyStateProblem(Solver& solver, int it, realtype t0);
+    void workSteadyStateProblem(Solver& solver, int it, realtype t0);
 
     /**
      * @brief Return the stored SimulationState.
      * @return stored SimulationState
      */
     [[nodiscard]] SimulationState const& getFinalSimulationState() const {
-        return final_state_;
+        return period_result_.final_state_;
     }
 
     /**
      * @brief Return state at steady state
      * @return x
      */
-    [[nodiscard]] AmiVector const& getState() const { return final_state_.x; }
+    [[nodiscard]] AmiVector const& getState() const {
+        return period_result_.final_state_.x;
+    }
 
     /**
      * @brief Return state sensitivity at steady state
      * @return sx
      */
     [[nodiscard]] AmiVectorArray const& getStateSensitivity() const {
-        return final_state_.sx;
+        return period_result_.final_state_.sx;
     }
 
     /**
@@ -444,7 +447,9 @@ class SteadystateProblem {
      * @brief Get model time at which steady state was found through simulation.
      * @return Time at which steady state was found (model time units).
      */
-    [[nodiscard]] realtype getSteadyStateTime() const { return final_state_.t; }
+    [[nodiscard]] realtype getSteadyStateTime() const {
+        return period_result_.final_state_.t;
+    }
 
     /**
      * @brief Get the weighted root mean square of the residuals.
@@ -546,15 +551,20 @@ class SteadystateProblem {
      */
     void updateRightHandSide();
 
+    /** Whether this is a pre- or post-equilibration problem */
+    bool is_preeq_;
+
     /** Workspace for forward simulation */
     FwdSimWorkspace* ws_;
+
     /** WRMS computer for x */
     WRMSComputer wrms_computer_x_;
+
     /** weighted root-mean-square error */
     realtype wrms_{NAN};
 
-    /** The simulation state at the end of the forward problem. */
-    SimulationState final_state_;
+    /** Results for this period. */
+    PeriodResult period_result_;
 
     /** stores diagnostic information about employed number of steps */
     std::vector<int> numsteps_{std::vector<int>(3, 0)};
@@ -605,6 +615,9 @@ class SteadystateProblem {
 
     /** The model to equilibrate */
     Model* model_{nullptr};
+
+    /** Simulator for event handling */
+    std::optional<EventHandlingSimulator> simulator_;
 };
 
 /**

python/tests/test_preequilibration.py

@@ -13,6 +13,7 @@
     skip_on_valgrind,
 )
 from amici.gradient_check import check_derivatives
+from amici import SensitivityMethod, SensitivityOrder
 
 pytestmark = pytest.mark.filterwarnings(
     # https://github.com/AMICI-dev/AMICI/issues/18
@@ -760,15 +761,18 @@ def test_preequilibration_events(tempdir):
         bolus2 = 1
         bolus3 = 1
         bolus4 = 1
-        at some_time >= 0, t0 = false: target1 = target1 + bolus1
-        at time >= 0, t0 = false: target2 = target2 + bolus2
+        # E1 & E2 will both trigger during pre-equilibration and main
+        #  simulation (Heaviside is reset after pre-equilibration)
+        E1: at some_time >= 0, t0 = false: target1 = target1 + bolus1
+        E2: at time >= 0, t0 = false: target2 = target2 + bolus2
         # requires early time point
         # https://github.com/AMICI-dev/AMICI/issues/2804
         trigger_time2 = 1e-3
-        # will trigger only during preequilibration (some_time is not reset)
-        at some_time >= trigger_time2: target3 = target3 + bolus3
+        # E3 will trigger only during preequilibration
+        # (some_time is not reset and trigger initial value is `true`)
+        E3: at some_time >= trigger_time2: target3 = target3 + bolus3
         # will trigger during preequilibration and main simulation
-        at time >= trigger_time2: target4 = target4 + bolus4
+        E4: at time >= trigger_time2: target4 = target4 + bolus4
     end
     """
     module_name = "test_preequilibration_events"
@@ -798,9 +802,15 @@ def test_preequilibration_events(tempdir):
 
     # Integration-only preequilibration should handle all events
     amici_model = model_module.getModel()
+    amici_model.setSteadyStateSensitivityMode(
+        amici.SteadyStateSensitivityMode.integrationOnly
+    )
+    amici_model.setSteadyStateComputationMode(
+        amici.SteadyStateSensitivityMode.integrationOnly
+    )
     rdata = amici.runAmiciSimulation(amici_model, amici_solver, edata)
     assert rdata.status == amici.AMICI_SUCCESS
-    assert rdata.preeq_t > 1e-3
+    assert rdata.preeq_t > 1e-3  # verifies that integration was done
     assert rdata.x_ss[target1_idx] == 1
     assert rdata.x_ss[target2_idx] == 1
     assert rdata.x_ss[target3_idx] == 1
@@ -814,17 +824,20 @@ def test_preequilibration_events(tempdir):
     edata.fixedParametersPreequilibration = [1.0]
     edata.fixedParameters = [0.0]
 
-    amici_solver.setSensitivityMethod(amici.SensitivityMethod.forward)
-    amici_solver.setSensitivityOrder(amici.SensitivityOrder.first)
+    amici_solver.setSensitivityOrder(SensitivityOrder.first)
 
-    for sensi_meth in [
-        amici.SensitivityMethod.forward,
-        # amici.SensitivityMethod.adjoint,
-    ]:
+    for sensi_meth, sensi_meth_preeq in (
+        (SensitivityMethod.forward, SensitivityMethod.forward),
+        (SensitivityMethod.adjoint, SensitivityMethod.forward),
+        # TODO https://github.com/AMICI-dev/AMICI/issues/2775
+        #  (SensitivityMethod.adjoint, SensitivityMethod.adjoint),
+    ):
         amici_solver.setSensitivityMethod(sensi_meth)
+        amici_solver.setSensitivityMethodPreequilibration(sensi_meth_preeq)
 
         # amici_model.requireSensitivitiesForAllParameters()
-        # FIXME: sensitivities w.r.t. trigger time are off
+        # FIXME: finite differences w.r.t. trigger time are off
+        #   need different epsilon for trigger time
         amici_model.setParameterList(
             [
                 i

src/backwardproblem.cpp

@@ -15,7 +15,7 @@ BackwardProblem::BackwardProblem(ForwardProblem& fwd)
     , solver_(fwd.solver)
     , edata_(fwd.edata)
     , t_(fwd.getTime())
-    , discs_(fwd.getDiscontinuities())
+    , discs_main_(fwd.getDiscontinuities())
     , dJydx_(fwd.getAdjointUpdates(*model_, *edata_))
     , dJzdx_(fwd.getDJzdx())
     , preeq_problem_(fwd.getPreequilibrationProblem())
@@ -45,8 +45,8 @@ void BackwardProblem::workBackwardProblem() {
 
     // initialize state vectors, depending on postequilibration
     model_->initializeB(ws_.xB_, ws_.dxB_, ws_.xQB_, it < model_->nt() - 1);
-    ws_.discs_ = discs_;
-    ws_.nroots_ = compute_nroots(discs_, model_->ne, model_->nMaxEvent());
+    ws_.discs_ = discs_main_;
+    ws_.nroots_ = compute_nroots(discs_main_, model_->ne, model_->nMaxEvent());
     simulator_.run(
         t_, model_->t0(), it, model_->getTimepoints(), &dJydx_, &dJzdx_
     );