Refactor model/solver handling in SteadyStateBackwardProblem (#2869)

dweindl · web-flow · commit 18b27f32bfbc · 2025-07-03T14:38:30.000Z
Store solver and model, use preequilibration solver, simplify parameter lists and initialization.
diff --git a/include/amici/steadystateproblem.h b/include/amici/steadystateproblem.h
@@ -308,6 +308,12 @@ class SteadystateProblem {
      */
     [[nodiscard]] bool checkSteadyStateSuccess() const;
 
+    /**
+     * @brief Get the preequilibration solver.
+     * @return The preequilibration solver.
+     */
+    [[nodiscard]] Solver const* get_solver() const { return solver_; }
+
   private:
     /**
      * @brief Handle the computation of the steady state.
@@ -467,15 +473,10 @@ class SteadyStateBackwardProblem {
      * Integrates over the adjoint state backward in time by solving a linear
      * system of equations, which gives the analytical solution.
      *
-     * @param solver The solver instance
-     * @param model The model instance
      * @param xB0 Initial adjoint state vector.
-     * @param is_preeq Flag indicating whether this is a preequilibration.
      * @param t0 Initial time for the steady state simulation.
      */
-    void
-    run(Solver const& solver, Model& model, AmiVector const& xB0, bool is_preeq,
-        realtype t0);
+    void run(AmiVector const& xB0, realtype t0);
 
     /**
      * @brief Return the quadratures from pre- or postequilibration
@@ -521,37 +522,27 @@ class SteadyStateBackwardProblem {
      * @brief Launch backward simulation if Newton solver or linear system solve
      * fail or are disabled.
      * @param solver Solver instance.
-     * @param model Model instance.
      */
-    void run_simulation(Solver const& solver, Model& model);
+    void run_simulation(Solver const& solver);
 
     /**
      * @brief Compute quadratures in adjoint mode
-     * @param solver Solver instance.
-     * @param model Model instance.
      * @param t0 Initial time for the steady state simulation.
      */
-    void compute_steady_state_quadrature(
-        Solver const& solver, Model& model, realtype t0
-    );
+    void compute_steady_state_quadrature(realtype t0);
 
     /**
      * @brief Compute the quadrature in steady state backward mode by
      * solving the linear system defined by the backward Jacobian.
-     * @param model Model instance.
      */
-    void compute_quadrature_by_lin_solve(Model& model);
+    void compute_quadrature_by_lin_solve();
 
     /**
      * @brief Computes the quadrature in steady state backward mode by
      * numerical integration of xB forward in time.
-     * @param solver Solver instance.
-     * @param model Model instance.
      * @param t0 Initial time for the steady state simulation.
      */
-    void compute_quadrature_by_simulation(
-        Solver const& solver, Model& model, realtype t0
-    );
+    void compute_quadrature_by_simulation(realtype t0);
 
     /** CPU time for solving the backward problem (milliseconds) */
     double cpu_timeB_{0.0};
@@ -582,6 +573,9 @@ class SteadyStateBackwardProblem {
      * checks during simulation.
      */
     bool newton_step_conv_{false};
+
+    Model* model_{nullptr};
+    Solver const* solver_{nullptr};
 };
 
 } // namespace amici
diff --git a/src/backwardproblem.cpp b/src/backwardproblem.cpp
@@ -65,15 +65,31 @@ void BackwardProblem::workBackwardProblem() {
     }
 
     // handle pre-equilibration
-    if (preeq_problem_) {
+    if (preeq_problem_
+        && preeq_problem_->get_solver()->getSensitivityMethodPreequilibration()
+               == SensitivityMethod::adjoint) {
+        auto preeq_solver = preeq_problem_->get_solver();
+
         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)
+            throw NewtonFailure(
+                AMICI_NOT_IMPLEMENTED,
+                "Adjoint preequilibration with reinitialization of "
+                "non-constant states is not yet implemented. Stopping."
+            );
+
+        // only preequilibrations needs a reInit, postequilibration does not
+        preeq_solver->updateAndReinitStatesAndSensitivities(model_);
+
         preeq_problem_bwd_.emplace(*solver_, *model_, final_state);
-        preeq_problem_bwd_->run(*solver_, *model_, ws_.xB_, true, t0);
+        preeq_problem_bwd_->run(ws_.xB_, t0);
     }
 }
 
@@ -92,7 +108,7 @@ void BackwardProblem::handlePostequilibration() {
 
     auto final_state = posteq_problem_->getFinalSimulationState();
     posteq_problem_bwd_.emplace(*solver_, *model_, final_state);
-    posteq_problem_bwd_->run(*solver_, *model_, ws_.xB_, false, model_->t0());
+    posteq_problem_bwd_->run(ws_.xB_, model_->t0());
     ws_.xQB_ = posteq_problem_bwd_->getEquilibrationQuadratures();
 }
 
diff --git a/src/steadystateproblem.cpp b/src/steadystateproblem.cpp
@@ -245,32 +245,11 @@ SteadyStateBackwardProblem::SteadyStateBackwardProblem(
     , newton_solver_(
           NewtonSolver(model, solver.getLinearSolver(), solver.getSunContext())
       )
-    , newton_step_conv_(solver.getNewtonStepSteadyStateCheck()) {}
-
-void SteadyStateBackwardProblem::run(
-    Solver const& solver, Model& model, AmiVector const& xB0, bool is_preeq,
-    realtype t0
-) {
-    if (is_preeq) {
-        if (solver.getSensitivityMethodPreequilibration()
-            != SensitivityMethod::adjoint) {
-            // if not adjoint mode, there's nothing to do
-            return;
-        }
-
-        // If we need to reinitialize solver states, this won't work yet.
-        if (model.nx_reinit() > 0)
-            throw NewtonFailure(
-                AMICI_NOT_IMPLEMENTED,
-                "Adjoint preequilibration with reinitialization of "
-                "non-constant states is not yet implemented. Stopping."
-            );
-
-        // only preequilibrations needs a reInit,
-        // postequilibration does not
-        solver.updateAndReinitStatesAndSensitivities(&model);
-    }
+    , newton_step_conv_(solver.getNewtonStepSteadyStateCheck())
+    , model_(&model)
+    , solver_(&solver) {}
 
+void SteadyStateBackwardProblem::run(AmiVector const& xB0, realtype t0) {
     newton_solver_.reinitialize();
     xB_.copy(xB0);
 
@@ -280,7 +259,7 @@ void SteadyStateBackwardProblem::run(
 
     // Compute quadratures, track computation time
     CpuTimer cpu_timer;
-    compute_steady_state_quadrature(solver, model, t0);
+    compute_steady_state_quadrature(t0);
     cpu_timeB_ = cpu_timer.elapsed_milliseconds();
 }
 
@@ -440,30 +419,28 @@ SteadyStateStatus SteadystateProblem::findSteadyStateBySimulation(
     }
 }
 
-void SteadyStateBackwardProblem::compute_steady_state_quadrature(
-    Solver const& solver, Model& model, realtype t0
-) {
+void SteadyStateBackwardProblem::compute_steady_state_quadrature(realtype 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
     //     xQB = Integral[ xB(x_ss, t, p) | dt ] * dxdot/dp(x_ss, t, p)
     // We therefore compute the integral over xB first and then do a
     // matrix-vector multiplication.
 
-    auto const sensitivityMode = model.getSteadyStateSensitivityMode();
+    auto const sensitivityMode = model_->getSteadyStateSensitivityMode();
 
     // Try to compute the analytical solution for quadrature algebraically
     if (sensitivityMode == SteadyStateSensitivityMode::newtonOnly
         || sensitivityMode
                == SteadyStateSensitivityMode::integrateIfNewtonFails)
-        compute_quadrature_by_lin_solve(model);
+        compute_quadrature_by_lin_solve();
 
     // Perform simulation if necessary
     if (sensitivityMode == SteadyStateSensitivityMode::integrationOnly
         || (sensitivityMode
                 == SteadyStateSensitivityMode::integrateIfNewtonFails
             && !hasQuadrature()))
-        compute_quadrature_by_simulation(solver, model, t0);
+        compute_quadrature_by_simulation(t0);
 
     // If the analytic solution and integration did not work, throw
     if (!hasQuadrature())
@@ -474,7 +451,7 @@ void SteadyStateBackwardProblem::compute_steady_state_quadrature(
         );
 }
 
-void SteadyStateBackwardProblem::compute_quadrature_by_lin_solve(Model& model) {
+void SteadyStateBackwardProblem::compute_quadrature_by_lin_solve() {
     // Computes the integral over the adjoint state xB:
     // If the Jacobian has full rank, this has an analytical solution, since
     //   d/dt[ xB(t) ] = JB^T(x(t), p) xB(t) = JB^T(x_ss, p) xB(t)
@@ -491,12 +468,13 @@ void SteadyStateBackwardProblem::compute_quadrature_by_lin_solve(Model& model) {
     try {
         // compute integral over xB and write to xQ
         newton_solver_.prepareLinearSystemB(
-            model, {final_state_.t, final_state_.x, final_state_.dx}
+            *model_, {final_state_.t, final_state_.x, final_state_.dx}
         );
         newton_solver_.solveLinearSystem(xQ_);
         // Compute the quadrature as the inner product xQ * dxdotdp
         computeQBfromQ(
-            model, xQ_, xQB_, {final_state_.t, final_state_.x, final_state_.dx}
+            *model_, xQ_, xQB_,
+            {final_state_.t, final_state_.x, final_state_.dx}
         );
         has_quadrature_ = true;
 
@@ -507,9 +485,7 @@ void SteadyStateBackwardProblem::compute_quadrature_by_lin_solve(Model& model) {
     }
 }
 
-void SteadyStateBackwardProblem::compute_quadrature_by_simulation(
-    Solver const& solver, Model& model, realtype t0
-) {
+void SteadyStateBackwardProblem::compute_quadrature_by_simulation(realtype 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.
@@ -519,18 +495,18 @@ void SteadyStateBackwardProblem::compute_quadrature_by_simulation(
     // xQ was written in getQuadratureByLinSolve() -> set to zero
     xQ_.zero();
 
-    auto sim_solver = std::unique_ptr<Solver>(solver.clone());
-    sim_solver->logger = solver.logger;
+    auto sim_solver = std::unique_ptr<Solver>(solver_->clone());
+    sim_solver->logger = solver_->logger;
     sim_solver->setSensitivityMethod(SensitivityMethod::none);
     sim_solver->setSensitivityOrder(SensitivityOrder::none);
-    sim_solver->setup(t0, &model, xB_, xB_, final_state_.sx, final_state_.sx);
+    sim_solver->setup(t0, model_, xB_, xB_, final_state_.sx, final_state_.sx);
     sim_solver->setupSteadystate(
-        t0, &model, final_state_.x, final_state_.dx, xB_, xB_, xQ_
+        t0, model_, final_state_.x, final_state_.dx, xB_, xB_, xQ_
     );
 
     // perform integration and quadrature
     try {
-        run_simulation(*sim_solver, model);
+        run_simulation(*sim_solver);
         has_quadrature_ = true;
     } catch (NewtonFailure const&) {
         has_quadrature_ = false;
@@ -681,10 +657,8 @@ void SteadystateProblem::runSteadystateSimulationFwd(Model& model) {
         updateSensiSimulation();
 }
 
-void SteadyStateBackwardProblem::run_simulation(
-    Solver const& solver, Model& model
-) {
-    if (model.nx_solver == 0)
+void SteadyStateBackwardProblem::run_simulation(Solver const& solver) {
+    if (model_->nx_solver == 0)
         return;
 
     if (newton_step_conv_) {
@@ -699,13 +673,13 @@ void SteadyStateBackwardProblem::run_simulation(
 
     // WRMS computer for xQB
     WRMSComputer wrms_computer_xQB_(
-        model.nplist(), solver.getSunContext(),
+        model_->nplist(), solver.getSunContext(),
         solver.getAbsoluteToleranceQuadratures(),
         solver.getRelativeToleranceQuadratures(), AmiVector()
     );
 
     // time-derivative of quadrature state vector
-    AmiVector xQBdot(model.nplist(), solver.getSunContext());
+    AmiVector xQBdot(model_->nplist(), solver.getSunContext());
 
     int const convergence_check_frequency = newton_step_conv_ ? 25 : 1;
     auto max_steps = (solver.getMaxStepsBackwardProblem() > 0)
@@ -722,11 +696,11 @@ void SteadyStateBackwardProblem::run_simulation(
             // converge to zero at all. So we need xQBdot, hence compute xQB
             // first.
             computeQBfromQ(
-                model, xQ_, xQB_,
+                *model_, xQ_, xQB_,
                 {final_state_.t, final_state_.x, final_state_.dx}
             );
             computeQBfromQ(
-                model, xB_, xQBdot,
+                *model_, xB_, xQBdot,
                 {final_state_.t, final_state_.x, final_state_.dx}
             );
             auto wrms = wrms_computer_xQB_.wrms(xQBdot, xQB_);
