Enable presimulation with adjoints (#2786)

Enable presimulation with adjoint sensitivities.

* Use previously introduced `EventHandlingBwdSimulator` for processing events also during presimulation
* Make `EventHandlingBwdSimulator::run` work if there are no datapoints
* Adapt setup/reinitialization logic

Author: dweindl

include/amici/backwardproblem.h

@@ -197,6 +197,9 @@ class BackwardProblem {
     /** The postequilibration steadystate problem from the forward problem. */
     SteadystateProblem* posteq_problem_;
 
+    /** Presimulation results */
+    PeriodResult presim_result;
+
     BwdSimWorkspace ws_;
 
     EventHandlingBwdSimulator simulator_;

include/amici/forwardproblem.h

@@ -502,6 +502,14 @@ class ForwardProblem {
         return nullptr;
     }
 
+    /**
+     * @brief Get the presimulation results.
+     * @return Presimulation results.
+     */
+    PeriodResult const& get_presimulation_result() const {
+        return pre_simulator_.result;
+    }
+
     /** pointer to model instance */
     Model* model;
 

python/tests/test_preequilibration.py

@@ -299,44 +299,6 @@ def test_parameter_in_expdata(preeq_fixture):
         )
 
 
-def test_raise_presimulation_with_adjoints(preeq_fixture):
-    """Test simulation failures with adjoin+presimulation"""
-
-    (
-        model,
-        solver,
-        edata,
-        edata_preeq,
-        edata_presim,
-        edata_sim,
-        pscales,
-        plists,
-    ) = preeq_fixture
-
-    # preequilibration and presimulation with adjoints:
-    # this needs to fail unless we remove presimulation
-    solver.setSensitivityMethod(amici.SensitivityMethod.adjoint)
-
-    rdata = amici.runAmiciSimulation(model, solver, edata)
-    assert rdata["status"] == amici.AMICI_ERROR
-
-    # add postequilibration
-    y = edata.getObservedData()
-    stdy = edata.getObservedDataStdDev()
-    ts = np.hstack([*edata.getTimepoints(), np.inf])
-    edata.setTimepoints(ts)
-    edata.setObservedData(np.hstack([y, y[0]]))
-    edata.setObservedDataStdDev(np.hstack([stdy, stdy[0]]))
-
-    # remove presimulation
-    edata.t_presim = 0
-    edata.fixedParametersPresimulation = ()
-
-    # no presim any more, this should work
-    rdata = amici.runAmiciSimulation(model, solver, edata)
-    assert rdata["status"] == amici.AMICI_SUCCESS
-
-
 def test_equilibration_methods_with_adjoints(preeq_fixture):
     """Test different combinations of equilibration and simulation
     sensitivity methods"""

python/tests/test_sbml_import.py

@@ -251,6 +251,9 @@ def test_presimulation(sbml_example_presimulation_module):
     check_derivatives(model, solver, edata, epsilon=1e-4)
 
 
+@pytest.mark.filterwarnings(
+    "ignore:Adjoint sensitivity analysis for models with discontinuous "
+)
 def test_presimulation_events(tempdir):
     """Test that events are handled during presimulation."""
 
@@ -306,7 +309,8 @@ def test_presimulation_events(tempdir):
 
     for sensi_method in (
         amici.SensitivityMethod.forward,
-        # TODO ASA
+        # FIXME: test with adjoints. currently there is some CVodeF issue
+        #  that fails forward simulation with adjoint sensitivities
         # amici.SensitivityMethod.adjoint,
     ):
         solver.setSensitivityMethod(sensi_method)
@@ -336,6 +340,87 @@ def test_presimulation_events(tempdir):
         )
 
 
+@pytest.mark.filterwarnings(
+    "ignore:Adjoint sensitivity analysis for models with discontinuous "
+)
+def test_presimulation_events_and_sensitivities(tempdir):
+    """Test that presimulation with adjoint sensitivities works
+    and test that events are handled during presimulation."""
+
+    from amici.antimony_import import antimony2amici
+
+    model_name = "test_presim_events2"
+    antimony2amici(
+        """
+    some_time = time
+    some_time' = 1
+    bolus = 1
+
+    k_pre = 3
+    k_main = 2
+    xx = 0
+    xx' = piecewise(k_pre, time < 0, k_main)
+
+    # this will trigger twice, once in presimulation
+    # and once in the main simulation
+    at time >= -1 , t0=false: some_time = some_time + bolus
+    """,
+        model_name=model_name,
+        output_dir=tempdir,
+    )
+
+    model_module = import_model_module(model_name, tempdir)
+
+    model = model_module.get_model()
+    model.setTimepoints([0, 1, 2])
+    edata = amici.ExpData(model)
+    edata.t_presim = 2
+    solver = model.getSolver()
+
+    # generate artificial data
+    rdata = amici.runAmiciSimulation(model, solver, edata)
+    edata_tmp = amici.ExpData(rdata, 1, 0)
+    edata.setTimepoints(np.array(edata_tmp.getTimepoints()) + 0.1)
+    edata.setObservedData(edata_tmp.getObservedData())
+    edata.setObservedDataStdDev(edata_tmp.getObservedDataStdDev())
+
+    solver.setSensitivityOrder(amici.SensitivityOrder.first)
+
+    for sensi_method in (
+        amici.SensitivityMethod.forward,
+        amici.SensitivityMethod.adjoint,
+    ):
+        solver.setSensitivityMethod(sensi_method)
+        rdata = amici.runAmiciSimulation(model, solver, edata)
+
+        assert rdata.status == amici.AMICI_SUCCESS
+        assert_allclose(
+            rdata.by_id("some_time"), np.array([0, 1, 2]) + 2.1, atol=1e-14
+        )
+
+        if sensi_method == amici.SensitivityMethod.forward:
+            model.requireSensitivitiesForAllParameters()
+        else:
+            # FIXME ASA with events:
+            #   https://github.com/AMICI-dev/AMICI/pull/1539
+            model.setParameterList(
+                [
+                    i
+                    for i, p in enumerate(model.getParameterIds())
+                    if p != "bolus"
+                ]
+            )
+
+        check_derivatives(
+            model,
+            solver,
+            edata=edata,
+            atol=1e-6,
+            rtol=1e-6,
+            epsilon=1e-8,
+        )
+
+
 def test_steadystate_simulation(model_steadystate_module):
     model = model_steadystate_module.getModel()
     model.setTimepoints(np.linspace(0, 60, 60))

src/backwardproblem.cpp

@@ -20,6 +20,7 @@ BackwardProblem::BackwardProblem(ForwardProblem& fwd)
     , dJzdx_(fwd.getDJzdx())
     , preeq_problem_(fwd.getPreequilibrationProblem())
     , posteq_problem_(fwd.getPostequilibrationProblem())
+    , presim_result(fwd.get_presimulation_result())
     , ws_(model_, solver_)
     , simulator_(model_, solver_, &ws_) {}
 
@@ -34,6 +35,8 @@ void BackwardProblem::workBackwardProblem() {
 
     handlePostequilibration();
 
+    // handle main simulation
+
     // If we have posteq, infinity timepoints were already treated
     int it = model_->nt() - 1;
     while (it >= 0 && std::isinf(model_->getTimepoint(it))) {
@@ -53,8 +56,13 @@ void BackwardProblem::workBackwardProblem() {
         ConditionContext cc(
             model_, edata_, FixedParameterContext::presimulation
         );
-        solver_->runB(model_->t0() - edata_->t_presim);
-        solver_->writeSolutionB(&t_, ws_.xB_, ws_.dxB_, ws_.xQB_, ws_.which);
+        ws_.discs_ = presim_result.discs;
+        ws_.nroots_
+            = compute_nroots(ws_.discs_, model_->ne, model_->nMaxEvent());
+        simulator_.run(
+            model_->t0(), model_->t0() - edata_->t_presim, -1, {}, &dJydx_,
+            &dJzdx_
+        );
     }
 
     // handle pre-equilibration
@@ -178,7 +186,8 @@ void EventHandlingBwdSimulator::run(
 
     t_ = t_start;
 
-    if ((it >= 0 || !ws_->discs_.empty()) && timepoints[it] > t_end) {
+    // datapoint at t_start?
+    if (it >= 0 && timepoints[it] == t_start) {
         handleDataPointB(it, dJydx);
         solver_->setupB(
             &ws_->which, timepoints[it], model_, ws_->xB_, ws_->dxB_, ws_->xQB_
@@ -187,34 +196,40 @@ void EventHandlingBwdSimulator::run(
         // as it is not called in handleDataPointB
         solver_->storeDiagnosisB(ws_->which);
         --it;
+    } else {
+        // no data points, only discontinuities, just set up the solver
+        // (e.g., during presimulation)
+        solver_->setupB(
+            &ws_->which, t_start, model_, ws_->xB_, ws_->dxB_, ws_->xQB_
+        );
+    }
 
-        while (it >= 0 || !ws_->discs_.empty()) {
-            // check if next timepoint is a discontinuity or a data-point
-            double tnext = getTnext(it);
-
-            if (tnext < t_) {
-                solver_->runB(tnext);
-                solver_->writeSolutionB(
-                    &t_, ws_->xB_, ws_->dxB_, ws_->xQB_, ws_->which
-                );
-            }
+    while (it >= 0 || !ws_->discs_.empty()) {
+        // check if next timepoint is a discontinuity or a data-point
+        double tnext = getTnext(it);
 
-            // handle discontinuity
-            if (!ws_->discs_.empty() && tnext == ws_->discs_.back().time) {
-                handleEventB(ws_->discs_.back(), dJzdx);
-                ws_->discs_.pop_back();
-            }
+        if (tnext < t_) {
+            solver_->runB(tnext);
+            solver_->writeSolutionB(
+                &t_, ws_->xB_, ws_->dxB_, ws_->xQB_, ws_->which
+            );
+        }
 
-            // handle data-point
-            if (it >= 0 && tnext == timepoints[it]) {
-                handleDataPointB(it, dJydx);
-                it--;
-            }
+        // handle discontinuity
+        if (!ws_->discs_.empty() && tnext == ws_->discs_.back().time) {
+            handleEventB(ws_->discs_.back(), dJzdx);
+            ws_->discs_.pop_back();
+        }
 
-            // reinitialize state
-            solver_->reInitB(ws_->which, t_, ws_->xB_, ws_->dxB_);
-            solver_->quadReInitB(ws_->which, ws_->xQB_);
+        // handle data-point
+        if (it >= 0 && tnext == timepoints[it]) {
+            handleDataPointB(it, dJydx);
+            it--;
         }
+
+        // reinitialize state
+        solver_->reInitB(ws_->which, t_, ws_->xB_, ws_->dxB_);
+        solver_->quadReInitB(ws_->which, ws_->xQB_);
     }
 
     // we still need to integrate from first datapoint to t_start

src/forwardproblem.cpp

@@ -181,39 +181,35 @@ void ForwardProblem::handlePresimulation() {
     if (!uses_presimulation_)
         return;
 
-    if (solver->computingASA()) {
-        throw AmiException(
-            "Presimulation with adjoint sensitivities"
-            " is currently not implemented."
-        );
-    }
+    // Are there dedicated condition preequilibration parameters provided?
+    ConditionContext cond(model, edata, FixedParameterContext::presimulation);
 
-    {
-        // Are there dedicated condition preequilibration parameters provided?
-        ConditionContext cond(
-            model, edata, FixedParameterContext::presimulation
+    // If we need to reinitialize solver states, this won't work yet.
+    if (model->nx_reinit() > 0)
+        throw AmiException(
+            "Adjoint presimulation with reinitialization of "
+            "non-constant states is not yet implemented. Stopping."
         );
 
-        // compute initial time and setup solver for (pre-)simulation
-        t_ = model->t0() - edata->t_presim;
-
-        // if preequilibration was done, model was already initialized
-        if (!preequilibrated_) {
-            model->initialize(
-                t_, ws_.x, ws_.dx, ws_.sx, ws_.sdx,
-                solver->getSensitivityOrder() >= SensitivityOrder::first,
-                ws_.roots_found
-            );
-        } else if (model->ne) {
-            model->initEvents(t_, ws_.x, ws_.dx, ws_.roots_found);
-        }
-        solver->setup(t_, model, ws_.x, ws_.dx, ws_.sx, ws_.sdx);
-        solver->updateAndReinitStatesAndSensitivities(model);
+    // compute initial time and setup solver for (pre-)simulation
+    t_ = model->t0() - edata->t_presim;
 
-        std::vector<realtype> const timepoints{model->t0()};
-        pre_simulator_.run(t_, edata, timepoints);
-        solver->writeSolution(&t_, ws_.x, ws_.dx, ws_.sx, ws_.dx);
+    // if preequilibration was done, model was already initialized
+    if (!preequilibrated_) {
+        model->initialize(
+            t_, ws_.x, ws_.dx, ws_.sx, ws_.sdx,
+            solver->getSensitivityOrder() >= SensitivityOrder::first,
+            ws_.roots_found
+        );
+    } else if (model->ne) {
+        model->initEvents(t_, ws_.x, ws_.dx, ws_.roots_found);
     }
+    solver->setup(t_, model, ws_.x, ws_.dx, ws_.sx, ws_.sdx);
+    solver->updateAndReinitStatesAndSensitivities(model);
+
+    std::vector<realtype> const timepoints{model->t0()};
+    pre_simulator_.run(t_, edata, timepoints);
+    solver->writeSolution(&t_, ws_.x, ws_.dx, ws_.sx, ws_.dx);
 }
 
 void ForwardProblem::handleMainSimulation() {
@@ -236,7 +232,10 @@ void ForwardProblem::handleMainSimulation() {
 
     t_ = model->t0();
 
-    solver->setup(t_, model, ws_.x, ws_.dx, ws_.sx, ws_.sdx);
+    // in case of presimulation, the solver was set up already
+    if (!uses_presimulation_) {
+        solver->setup(t_, model, ws_.x, ws_.dx, ws_.sx, ws_.sdx);
+    }
 
     if (preequilibrated_ || uses_presimulation_) {
         // Reset the time and re-initialize events for the main simulation
@@ -278,10 +277,11 @@ void EventHandlingSimulator::handle_event(
 
     if (!initial_event && t_ == ws_->tlastroot) {
         throw AmiException(
-            "AMICI is stuck in an event, as the initial "
+            "AMICI is stuck in an event at time %g, as the initial "
             "step-size after the event is too small. "
             "To fix this, increase absolute and relative "
-            "tolerances!"
+            "tolerances!",
+            t_
         );
     }
     ws_->tlastroot = t_;