Event-handling during pre-equilibration ASA / ASA

include/amici/backwardproblem.h

@@ -307,6 +307,20 @@ class BackwardProblem {
         return xQB_pre_preeq_;
     }
 
+    /**
+     * @brief The final adjoint state vector
+     * @return xB
+     */
+    [[nodiscard]] AmiVector const& getAdjointState() const { return ws_.xB_; }
+
+    /**
+     * @brief The final quadrature state vector.
+     * @return xQB
+     */
+    [[nodiscard]] AmiVector const& getAdjointQuadrature() const {
+        return ws_.xQB_;
+    }
+
     /**
      * @brief Return the postequilibration SteadyStateBwdProblem.
      * @return The postequilibration SteadyStateBackwardProblem, if any.

include/amici/forwardproblem.h

@@ -476,6 +476,14 @@ class SteadystateProblem {
      */
     [[nodiscard]] Solver const* get_solver() const { return solver_; }
 
+    /**
+     * @brief Get the preequilibration result.
+     * @return
+     */
+    [[nodiscard]] PeriodResult const& get_result() const {
+        return period_result_;
+    }
+
   private:
     /**
      * @brief Handle the computation of the steady state.

include/amici/rdata.h

@@ -792,16 +792,14 @@ class ReturnData : public ModelDimensions {
      * @brief Updates contribution to likelihood from quadratures (xQB),
      * if preequilibration was run in adjoint mode
      * @param model model that was used for forward/backward simulation
-     * @param preeq SteadyStateProblem for preequilibration
-     * @param preeq_bwd SteadyStateBackwardProblem for preequilibration
+     * @param sx0 State sensitivities at pre-equilibration steady state
+     * @param xB Adjoint state from pre-equilibration.
      * @param llhS0 contribution to likelihood for initial state sensitivities
      * of preequilibration
-     * @param xQB vector with quadratures from adjoint computation
      */
     void handleSx0Backward(
-        Model const& model, SteadystateProblem const& preeq,
-        SteadyStateBackwardProblem const* preeq_bwd,
-        std::vector<realtype>& llhS0, AmiVector& xQB
+        Model const& model, AmiVectorArray const& sx0, AmiVector const& xB,
+        std::vector<realtype>& llhS0
     ) const;
 
     /**

python/tests/test_preequilibration.py

@@ -829,8 +829,7 @@ def test_preequilibration_events(tempdir):
     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),
+        (SensitivityMethod.adjoint, SensitivityMethod.adjoint),
     ):
         amici_solver.setSensitivityMethod(sensi_meth)
         amici_solver.setSensitivityMethodPreequilibration(sensi_meth_preeq)

src/backwardproblem.cpp

@@ -79,9 +79,6 @@ void BackwardProblem::workBackwardProblem() {
         ConditionContext cc2(
             model_, edata_, FixedParameterContext::preequilibration
         );
-        auto const t0
-            = std::isnan(model_->t0Preeq()) ? model_->t0() : model_->t0Preeq();
-        auto final_state = preeq_problem_->getFinalSimulationState();
 
         // If we need to reinitialize solver states, this won't work yet.
         if (model_->nx_reinit() > 0)
@@ -91,11 +88,45 @@ void BackwardProblem::workBackwardProblem() {
                 "non-constant states is not yet implemented. Stopping."
             );
 
-        // only preequilibrations needs a reInit, postequilibration does not
-        preeq_solver->updateAndReinitStatesAndSensitivities(model_);
+        auto const t0
+            = std::isnan(model_->t0Preeq()) ? model_->t0() : model_->t0Preeq();
+
+        auto const& preeq_result = preeq_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 (preeq_problem_->getSteadyStateStatus()[1]
+                == SteadyStateStatus::not_run
+            || preeq_result.discs.empty()) {
+            preeq_solver->updateAndReinitStatesAndSensitivities(model_);
 
-        preeq_problem_bwd_.emplace(*solver_, *model_, final_state.sol, &ws_);
-        preeq_problem_bwd_->run(t0);
+            auto preeq_final_fwd_state = preeq_result.final_state_;
+            preeq_problem_bwd_.emplace(
+                *solver_, *model_, preeq_final_fwd_state.sol, &ws_
+            );
+            preeq_problem_bwd_->run(t0);
+        } else if (preeq_problem_->getSteadyStateStatus()[1]
+                   != SteadyStateStatus::not_run) {
+            // backward integration of the pre-equilibration problem
+            // (only if preequilibration was done via simulation)
+            ws_.discs_ = preeq_result.discs;
+            ws_.nroots_
+                = compute_nroots(ws_.discs_, model_->ne, model_->nMaxEvent());
+            EventHandlingBwdSimulator preeq_simulator(
+                model_, preeq_solver, &ws_
+            );
+            preeq_simulator.run(
+                preeq_result.final_state_.sol.t,
+                preeq_result.initial_state_.sol.t, -1, {}, &dJydx_, &dJzdx_
+            );
+        } else {
+            Expects(
+                false
+                && "Unhandled preequilibration case in "
+                   "BackwardProblem::workBackwardProblem()"
+            );
+        }
     }
 }
 

src/forwardproblem.cpp

@@ -257,14 +257,12 @@ void ForwardProblem::handlePreequilibration() {
     auto t0 = std::isnan(model->t0Preeq()) ? model->t0() : model->t0Preeq();
 
     // The solver was not run before, set up everything.
-    // TODO: For pre-equilibration in combination with adjoint sensitivities,
-    // we will need to use a separate solver instance because we still need the
-    // forward solver for each period for backward integration.
     model->initialize(
         t0, ws_.sol.x, ws_.sol.dx, ws_.sol.sx, ws_.sdx,
         solver->getSensitivityOrder() >= SensitivityOrder::first,
         ws_.roots_found
     );
+    model->initializeStateSensitivities(t0, ws_.sol.sx, ws_.sol.x);
     solver->setup(t0, model, ws_.sol.x, ws_.sol.dx, ws_.sol.sx, ws_.sdx);
 
     preeq_problem_->workSteadyStateProblem(*solver, -1, t0);
@@ -874,16 +872,17 @@ SteadyStateStatus SteadystateProblem::findSteadyStateBySimulation(
             new_solver->logger = main_solver->logger;
 
             // do we need sensitivities?
-            if (new_solver->getSensitivityMethodPreequilibration()
-                    == SensitivityMethod::forward
-                && new_solver->getSensitivityOrder() >= SensitivityOrder::first
+            if (new_solver->getSensitivityOrder() >= SensitivityOrder::first
                 && (model_->getSteadyStateSensitivityMode()
                         == SteadyStateSensitivityMode::integrationOnly
                     || model_->getSteadyStateSensitivityMode()
                            == SteadyStateSensitivityMode::integrateIfNewtonFails
                 )) {
-                // need forward to compute sx0 for the pre/main simulation
-                new_solver->setSensitivityMethod(SensitivityMethod::forward);
+                // need FSA to compute sx0 for the pre/main simulation,
+                // or enable ASA for backward integration of pre-equibration
+                new_solver->setSensitivityMethod(
+                    new_solver->getSensitivityMethodPreequilibration()
+                );
             } else {
                 new_solver->setSensitivityMethod(SensitivityMethod::none);
                 new_solver->setSensitivityOrder(SensitivityOrder::none);
@@ -934,8 +933,6 @@ SteadyStateStatus SteadystateProblem::findSteadyStateBySimulation(
             );
         return SteadyStateStatus::failed;
     }
-
-    period_result_ = simulator_->result;
 }
 
 [[noreturn]] void SteadystateProblem::handleSteadyStateFailure(
@@ -1068,7 +1065,11 @@ void SteadystateProblem::runSteadystateSimulationFwd() {
     // Create a new simulator that uses the (potentially changed) solver
     // with the right sensitivity settings.
     simulator_.emplace(model_, solver_, ws_, nullptr);
-    simulator_->result.initial_state_ = period_result_.initial_state_;
+    simulator_->result.initial_state_.mod = model_->getModelState();
+    simulator_->result.initial_state_.sol = ws_->sol;
+
+    // store results on (any) exit
+    auto _ = gsl::finally([&]() { period_result_ = simulator_->result; });
 
     simulator_->handle_initial_events(nullptr);
 
@@ -1146,6 +1147,10 @@ void SteadystateProblem::findSteadyStateByNewtonsMethod(bool newton_retry) {
 
         simulator_->handle_initial_events(nullptr);
         period_result_ = simulator_->result;
+        // set time to infinity after processing initial events,
+        //  so we'll get warnings in case of explicit time dependencies
+        //  which are incompatible with Newton's method
+        ws_->sol.t = INFINITY;
     }
 
     int const stage = newton_retry ? 2 : 0;

src/rdata.cpp

@@ -488,11 +488,44 @@ void ReturnData::processBackwardProblem(
     model.setModelState(simulation_state.mod);
 
     std::vector<realtype> llhS0(model.nJ * model.nplist(), 0.0);
+    // Adjoint quadratures before pre-equilibration
     auto xQB = bwd.getAdjointQuadraturePrePreeq();
 
     if (preeq_bwd && preeq_bwd->hasQuadrature()) {
-        handleSx0Backward(model, *preeq, preeq_bwd, llhS0, xQB);
+        // Pre-equilibration with ASA steady-state shortcuts
+
+        // If pre-equilibration is run in adjoint mode, the scalar product of
+        // sx0 with its adjoint counterpart (see handleSx0Forward()) is not
+        // necessary:
+        // the actual simulation is "extended" by the pre-equilibration time.
+        // At initialization (at t=-inf), the adjoint state is in steady state
+        // (= 0) and so is the scalar product.
+        // Instead of the scalar product, the quadratures xQB from
+        // pre-equilibration contribute to the gradient
+        // (see example notebook on equilibration for further documentation).
+
+        // Adjoint quadratures after pre-equilibration
+        auto const& xQB_post_preeq = preeq_bwd->getAdjointQuadrature();
+        for (int ip = 0; ip < model.nplist(); ++ip)
+            xQB[ip] += xQB_post_preeq.at(ip);
+
+        handleSx0Backward(
+            model, preeq->getStateSensitivity(), bwd.getAdjointState(), llhS0
+        );
+    } else if (preeq
+               && preeq->getSteadyStateStatus()[1]
+                      != SteadyStateStatus::not_run) {
+        // Pre-equilibration with ASA backward simulation
+
+        // Adjoint quadratures after pre-equilibration
+        xQB = bwd.getAdjointQuadrature();
+
+        handleSx0Backward(
+            model, preeq->getStateSensitivity(), bwd.getAdjointState(), llhS0
+        );
+
     } else {
+        // No pre-equilibration, or pre-equilibration with FSA
         handleSx0Forward(
             model, simulation_state.sol, llhS0, bwd.getAdjointStatePrePreeq()
         );
@@ -512,32 +545,17 @@ void ReturnData::processBackwardProblem(
 }
 
 void ReturnData::handleSx0Backward(
-    Model const& model, SteadystateProblem const& preeq,
-    SteadyStateBackwardProblem const* preeq_bwd, std::vector<realtype>& llhS0,
-    AmiVector& xQB
+    Model const& model, AmiVectorArray const& sx0, AmiVector const& xB,
+    std::vector<realtype>& llhS0
 ) const {
-    /* If preequilibration is run in adjoint mode, the scalar product of sx0
-       with its adjoint counterpart (see handleSx0Forward()) is not necessary:
-       the actual simulation is "extended" by the preequilibration time.
-       At initialization (at t=-inf), the adjoint state is in steady state (= 0)
-       and so is the scalar product. Instead of the scalar product, the
-       quadratures xQB from preequilibration contribute to the gradient
-       (see example notebook on equilibration for further documentation). */
-    auto const& xQBpreeq = preeq_bwd->getAdjointQuadrature();
-    for (int ip = 0; ip < model.nplist(); ++ip)
-        xQB[ip] += xQBpreeq.at(ip);
-
-    /* We really need references here, as sx0 can be large... */
-    auto const& sx0preeq = preeq.getStateSensitivity();
-    auto const& xBpreeq = preeq_bwd->getAdjointState();
-
-    /* Add the contribution for sx0 from preequilibration. If backward
-     * preequilibration was done by simulation due to a singular Jacobian,
-     * xB is not necessarily 0 and we may get a non-zero contribution here. */
+
+    // Add the contribution for sx0 from preequilibration. If backward
+    // preequilibration was done by simulation due to a singular Jacobian,
+    // xB is not necessarily 0 and we may get a non-zero contribution here.
     for (int ip = 0; ip < model.nplist(); ++ip) {
         llhS0[ip] = 0.0;
         for (int ix = 0; ix < model.nxtrue_solver; ++ix) {
-            llhS0[ip] += xBpreeq.at(ix) * sx0preeq.at(ix, ip);
+            llhS0[ip] += xB.at(ix) * sx0.at(ix, ip);
         }
     }
 }