Refactor SteadystateProblem initialization (AMICI-dev#2844)

Remove `SteadystateProblem::initializeForwardProblem` and move initialization to `ForwardProblem`.
This simplifies things since there we don't need to guess whether we are performing pre- or post-equilibration.

Author: dweindl

include/amici/steadystateproblem.h

@@ -232,6 +232,7 @@ class SteadystateProblem {
      *
      * Tries to determine the steady state of the ODE system and computes
      * steady state sensitivities if requested.
+     * Expects that solver, model, and ws_ are already initialized.
      *
      * @param solver The solver instance
      * @param model The model instance
@@ -469,17 +470,6 @@ class SteadystateProblem {
      */
     void runSteadystateSimulationBwd(Solver const& solver, Model& model);
 
-    /**
-     * @brief Initialize forward computation
-     * @param it Index of the current output time point.
-     * @param solver pointer to the solver object
-     * @param model pointer to the model object
-     * @param t0 Initial time for the steady state simulation.
-     */
-    void initializeForwardProblem(
-        int it, Solver const& solver, Model& model, realtype t0
-    );
-
     /**
      * @brief Update member variables to indicate that state_.x has been
      * updated and xdot_, delta_, etc. need to be recomputed.

src/forwardproblem.cpp

@@ -175,6 +175,18 @@ void ForwardProblem::handlePreequilibration() {
 
     preeq_problem_.emplace(&ws_, *solver, *model);
     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.
+    auto roots_found = std::vector<int>(model->ne, 0);
+    model->initialize(
+        t0, ws_.x, ws_.dx, ws_.sx, ws_.sdx,
+        solver->getSensitivityOrder() >= SensitivityOrder::first, roots_found
+    );
+    solver->setup(t0, model, ws_.x, ws_.dx, ws_.sx, ws_.sdx);
+
     preeq_problem_->workSteadyStateProblem(*solver, *model, -1, t0);
 
     ws_.x = preeq_problem_->getState();
@@ -270,6 +282,10 @@ void ForwardProblem::handlePostequilibration() {
         posteq_problem_.emplace(&ws_, *solver, *model);
         auto it = getCurrentTimeIteration();
         auto t0 = it < 1 ? model->t0() : model->getTimepoint(it - 1);
+
+        // The solver was run before, extract current state from solver.
+        solver->writeSolution(ws_.t, ws_.x, ws_.dx, ws_.sx);
+        Expects(t0 == ws_.t);
         posteq_problem_->workSteadyStateProblem(*solver, *model, it, t0);
     }
 }

src/steadystateproblem.cpp

@@ -184,10 +184,11 @@ void SteadystateProblem::workSteadyStateProblem(
         );
     }
 
-    initializeForwardProblem(it, solver, model, t0);
-
     // Compute steady state, track computation time
     CpuTimer cpu_timer;
+    ws_->t = t0;
+    flagUpdatedState();
+    newton_solver_.reinitialize();
     findSteadyState(solver, model, it, t0);
 
     // Check whether state sensitivities still need to be computed.
@@ -433,29 +434,6 @@ SteadyStateStatus SteadystateProblem::findSteadyStateBySimulation(
     }
 }
 
-void SteadystateProblem::initializeForwardProblem(
-    int const it, Solver const& solver, Model& model, realtype const t0
-) {
-    newton_solver_.reinitialize();
-
-    // Process solver handling for pre- or postequilibration.
-    if (it == -1) {
-        // The solver was not run before, set up everything.
-        auto roots_found = std::vector<int>(model.ne, 0);
-        model.initialize(
-            t0, ws_->x, ws_->dx, ws_->sx, ws_->sdx,
-            solver.getSensitivityOrder() >= SensitivityOrder::first, roots_found
-        );
-        solver.setup(t0, &model, ws_->x, ws_->dx, ws_->sx, ws_->sdx);
-    } else {
-        // The solver was run before, extract current state from solver.
-        solver.writeSolution(ws_->t, ws_->x, ws_->dx, ws_->sx);
-    }
-
-    ws_->t = t0;
-    flagUpdatedState();
-}
-
 void SteadystateProblem::computeSteadyStateQuadrature(
     Solver const& solver, Model& model, realtype t0
 ) {