From 351b5b19520d2f58b6041337429cb0e50325bd3c Mon Sep 17 00:00:00 2001 From: Daniel Weindl Date: Mon, 7 Jul 2025 09:57:11 +0200 Subject: [PATCH] Event-handling during pre-equilibration ASA / ASA Handle events during pre-equilibration when sensitivities for both pre-equilibration and main simulation are computed via adjoint sensitivity analysis. --- include/amici/backwardproblem.h | 14 ++++++ include/amici/forwardproblem.h | 8 ++++ include/amici/rdata.h | 10 ++--- python/tests/test_preequilibration.py | 3 +- src/backwardproblem.cpp | 45 ++++++++++++++++--- src/forwardproblem.cpp | 27 ++++++----- src/rdata.cpp | 64 +++++++++++++++++---------- 7 files changed, 122 insertions(+), 49 deletions(-) diff --git a/include/amici/backwardproblem.h b/include/amici/backwardproblem.h index 401340669f..98201b8910 100644 --- a/include/amici/backwardproblem.h +++ b/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. diff --git a/include/amici/forwardproblem.h b/include/amici/forwardproblem.h index 9389c68918..61acb290f0 100644 --- a/include/amici/forwardproblem.h +++ b/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. diff --git a/include/amici/rdata.h b/include/amici/rdata.h index 0e29fca7c0..9945771dd7 100644 --- a/include/amici/rdata.h +++ b/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& llhS0, AmiVector& xQB + Model const& model, AmiVectorArray const& sx0, AmiVector const& xB, + std::vector& llhS0 ) const; /** diff --git a/python/tests/test_preequilibration.py b/python/tests/test_preequilibration.py index 8871002669..afe3fcba47 100644 --- a/python/tests/test_preequilibration.py +++ b/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) diff --git a/src/backwardproblem.cpp b/src/backwardproblem.cpp index 0984e8d3a1..412d2c8080 100644 --- a/src/backwardproblem.cpp +++ b/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()" + ); + } } } diff --git a/src/forwardproblem.cpp b/src/forwardproblem.cpp index 610150c4e9..9f00707113 100644 --- a/src/forwardproblem.cpp +++ b/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; diff --git a/src/rdata.cpp b/src/rdata.cpp index 4a7f532307..3e09aede03 100644 --- a/src/rdata.cpp +++ b/src/rdata.cpp @@ -488,11 +488,44 @@ void ReturnData::processBackwardProblem( model.setModelState(simulation_state.mod); std::vector 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& llhS0, - AmiVector& xQB + Model const& model, AmiVectorArray const& sx0, AmiVector const& xB, + std::vector& 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); } } }