Handle events during backward simulation of post-equilibration

python/tests/test_events.py

@@ -1166,19 +1166,18 @@ def test_posteq_events_are_handled(tempdir):
     edata = amici.ExpData(rdata, 1, 0, 0)
     for sens_method in (
         SensitivityMethod.forward,
-        # FIXME: sensitivities w.r.t. the bolus parameter are off for ASA (0.0)
-        # SensitivityMethod.adjoint,
+        SensitivityMethod.adjoint,
     ):
         solver.setSensitivityOrder(SensitivityOrder.first)
         solver.setSensitivityMethod(sens_method)
+
         check_derivatives(
             model,
             solver,
             edata=edata,
             atol=1e-12,
             rtol=1e-7,
             epsilon=1e-8,
-            skip_fields=["res"],
         )
 
 

src/backwardproblem.cpp

@@ -43,8 +43,6 @@ void BackwardProblem::workBackwardProblem() {
         --it;
     }
 
-    // initialize state vectors, depending on postequilibration
-    model_->initializeB(ws_.xB_, ws_.dxB_, ws_.xQB_, it < model_->nt() - 1);
     ws_.discs_ = discs_main_;
     ws_.nroots_ = compute_nroots(discs_main_, model_->ne, model_->nMaxEvent());
     simulator_.run(
@@ -135,18 +133,47 @@ void BackwardProblem::handlePostequilibration() {
         return;
     }
 
-    // initialize xB - only process the postequilibration timepoints
+    // initialize xB - only process the post-equilibration timepoints
     for (int it = 0; it < model_->nt(); it++) {
         if (std::isinf(model_->getTimepoint(it))) {
             for (int ix = 0; ix < model_->nxtrue_solver; ix++)
                 ws_.xB_[ix] += dJydx_[ix + it * model_->nx_solver];
         }
     }
 
-    // TODO handle any events
-    auto final_state = posteq_problem_->getFinalSimulationState();
-    posteq_problem_bwd_.emplace(*solver_, *model_, final_state.sol, &ws_);
-    posteq_problem_bwd_->run(model_->t0());
+    auto const& posteq_result = posteq_problem_->get_result();
+
+    // If there were no discontinuities or no simulation was performed,
+    // we can use the steady-state shortcuts.
+    // If not we need to do the regular backward integration.
+    if (posteq_problem_->getSteadyStateStatus()[1] == SteadyStateStatus::not_run
+        || posteq_result.discs.empty()) {
+
+        auto final_state = posteq_problem_->getFinalSimulationState();
+        posteq_problem_bwd_.emplace(*solver_, *model_, final_state.sol, &ws_);
+        posteq_problem_bwd_->run(model_->t0());
+
+        // re-initialize state vectors for main simulation
+        model_->initializeB(ws_.xB_, ws_.dxB_, ws_.xQB_, true);
+    } else if (posteq_problem_->getSteadyStateStatus()[1]
+               != SteadyStateStatus::not_run) {
+        // backward integration of the post-equilibration problem
+        // (only if post-equilibration was done via simulation)
+        ws_.discs_ = posteq_result.discs;
+        ws_.nroots_
+            = compute_nroots(ws_.discs_, model_->ne, model_->nMaxEvent());
+        EventHandlingBwdSimulator posteq_simulator(model_, solver_, &ws_);
+        posteq_simulator.run(
+            posteq_result.final_state_.sol.t,
+            posteq_result.initial_state_.sol.t, -1, {}, &dJydx_, &dJzdx_
+        );
+    } else {
+        Expects(
+            false
+            && "Unhandled post-equilibration case in "
+               "BackwardProblem::handlePostequilibration()"
+        );
+    }
 }
 
 void EventHandlingBwdSimulator::handleEventB(
@@ -364,7 +391,9 @@ AmiVector const& SteadyStateBackwardProblem::getAdjointQuadrature() const {
     return ws_->xQB_;
 }
 
-void SteadyStateBackwardProblem::compute_steady_state_quadrature(realtype const t0) {
+void SteadyStateBackwardProblem::compute_steady_state_quadrature(
+    realtype const t0
+) {
     // This routine computes the quadratures:
     //     xQB = Integral[ xB(x(t), t, p) * dxdot/dp(x(t), t, p) | dt ]
     // As we're in steady state, we have x(t) = x_ss (x_steadystate), hence
@@ -430,7 +459,9 @@ void SteadyStateBackwardProblem::compute_quadrature_by_lin_solve() {
     }
 }
 
-void SteadyStateBackwardProblem::compute_quadrature_by_simulation(realtype const t0) {
+void SteadyStateBackwardProblem::compute_quadrature_by_simulation(
+    realtype const t0
+) {
     // If the Jacobian is singular, the integral over xB must be computed
     // by usual integration over time, but simplifications can be applied:
     // x is not time-dependent, no forward trajectory is needed.
@@ -486,8 +517,8 @@ void SteadyStateBackwardProblem::run_simulation(Solver const& solver) {
 
     int const convergence_check_frequency = newton_step_conv_ ? 25 : 1;
     auto const max_steps = (solver.getMaxStepsBackwardProblem() > 0)
-                         ? solver.getMaxStepsBackwardProblem()
-                         : solver.getMaxSteps() * 100;
+                               ? solver.getMaxStepsBackwardProblem()
+                               : solver.getMaxSteps() * 100;
 
     while (true) {
         if (sim_steps % convergence_check_frequency == 0) {