Extract WRMS computation and buffers (#2761)

Also fix an issue caused by incorrectly using the state-mask for computing the WRMS for `xQB`, which should only be used for `x`.

Author: dweindl

include/amici/steadystateproblem.h

@@ -15,6 +15,54 @@ class Solver;
 class Model;
 class BackwardProblem;
 
+/**
+ * @brief Computes the weighted root mean square norm.
+ *
+ * This class is used to compute the weighted root mean square of the residuals
+ * and maintains its work space to avoid reallocation.
+ */
+class WRMSComputer {
+  public:
+    /**
+     * @brief Constructor.
+     * @param n The length of the vectors for which to compute the WRMS.
+     * @param sunctx A SUNDIALS context for the NVector.
+     * @param atol Absolute tolerance to compute error weights.
+     * @param rtol Relative tolerance to compute error weights.
+     * @param mask Mask for entries to include in the WRMS norm.
+     * Positive value: include; non-positive value: exclude; empty: include all.
+     */
+    WRMSComputer(
+        int n, SUNContext sunctx, realtype atol, realtype rtol, AmiVector mask
+    )
+        : ewt_(n, sunctx)
+        , rtol_(rtol)
+        , atol_(atol)
+        , mask_(mask) {}
+
+    /**
+     * @brief Compute the weighted root mean square of the residuals.
+     * @param x Vector to compute the WRMS for.
+     * @param x_ref The reference vector from which to compute the error
+     * weights.
+     * @return The WRMS norm.
+     */
+    realtype wrms(AmiVector const& x, AmiVector const& x_ref);
+
+  private:
+    /** Error weights for the residuals. */
+    AmiVector ewt_;
+    /** Relative tolerance to compute error weights. */
+    realtype rtol_;
+    /** Absolute tolerance to compute error weights. */
+    realtype atol_;
+    /**
+     * Mask for entries to include in the WRMS norm.
+     * Positive value: include; non-positive value: exclude; empty: include all.
+     */
+    AmiVector mask_;
+};
+
 /**
  * @brief The SteadystateProblem class solves a steady-state problem using
  * Newton's method and falls back to integration on failure.
@@ -374,10 +422,12 @@ class SteadystateProblem {
     AmiVector delta_;
     /** previous newton step (size: nx_solver). */
     AmiVector delta_old_;
-    /** error weights for solver state, dimension nx_solver */
-    AmiVector ewt_;
-    /** error weights for backward quadratures, dimension nplist() */
-    AmiVector ewtQB_;
+    /** WRMS computer for x */
+    WRMSComputer wrms_computer_x_;
+    /** WRMS computer for xQB */
+    WRMSComputer wrms_computer_xQB_;
+    /** WRMS computer for sx */
+    WRMSComputer wrms_computer_sx_;
     /** old state vector */
     AmiVector x_old_;
     /** time derivative state vector */
@@ -430,19 +480,6 @@ class SteadystateProblem {
      */
     std::vector<SteadyStateStatus> steady_state_status_;
 
-    /** absolute tolerance for convergence check (state)*/
-    realtype atol_{NAN};
-    /** relative tolerance for convergence check (state)*/
-    realtype rtol_{NAN};
-    /** absolute tolerance for convergence check (state sensi)*/
-    realtype atol_sensi_{NAN};
-    /** relative tolerance for convergence check (state sensi)*/
-    realtype rtol_sensi_{NAN};
-    /** absolute tolerance for convergence check (quadratures)*/
-    realtype atol_quad_{NAN};
-    /** relative tolerance for convergence check (quadratures)*/
-    realtype rtol_quad_{NAN};
-
     /** Newton solver */
     NewtonSolver newton_solver_;
 

src/steadystateproblem.cpp

@@ -61,9 +61,6 @@ realtype getWrmsNorm(
     AmiVector const& x, AmiVector const& xdot, AmiVector const& mask,
     realtype atol, realtype rtol, AmiVector& ewt
 ) {
-    // Depending on what convergence we want to check (xdot, sxdot, xQBdot)
-    //  we need to pass ewt[QB], as xdot and xQBdot have different sizes.
-
     // ewt = x
     N_VAbs(const_cast<N_Vector>(x.getNVector()), ewt.getNVector());
     // ewt *= rtol
@@ -85,6 +82,10 @@ realtype getWrmsNorm(
     );
 }
 
+realtype WRMSComputer::wrms(AmiVector const& x, AmiVector const& x_ref) {
+    return getWrmsNorm(x_ref, x, mask_, atol_, rtol_, ewt_);
+}
+
 /**
  * @brief Compute the backward quadratures, which contribute to the
  * gradient (xQB) from the quadrature over the backward state itself (xQ)
@@ -118,8 +119,23 @@ void computeQBfromQ(
 SteadystateProblem::SteadystateProblem(Solver const& solver, Model const& model)
     : delta_(model.nx_solver, solver.getSunContext())
     , delta_old_(model.nx_solver, solver.getSunContext())
-    , ewt_(model.nx_solver, solver.getSunContext())
-    , ewtQB_(model.nplist(), solver.getSunContext())
+    , wrms_computer_x_(
+          model.nx_solver, solver.getSunContext(),
+          solver.getAbsoluteToleranceSteadyState(),
+          solver.getRelativeToleranceSteadyState(),
+          AmiVector(model.get_steadystate_mask(), solver.getSunContext())
+      )
+    , wrms_computer_xQB_(
+          model.nplist(), solver.getSunContext(),
+          solver.getAbsoluteToleranceQuadratures(),
+          solver.getRelativeToleranceQuadratures(), AmiVector()
+      )
+    , wrms_computer_sx_(
+          model.nx_solver, solver.getSunContext(),
+          solver.getAbsoluteToleranceSteadyStateSensi(),
+          solver.getRelativeToleranceSteadyStateSensi(),
+          AmiVector(model.get_steadystate_mask(), solver.getSunContext())
+      )
     , x_old_(model.nx_solver, solver.getSunContext())
     , xdot_(model.nx_solver, solver.getSunContext())
     , sdx_(model.nx_solver, model.nplist(), solver.getSunContext())
@@ -141,12 +157,6 @@ SteadystateProblem::SteadystateProblem(Solver const& solver, Model const& model)
            ),
            .state = model.getModelState()}
       )
-    , atol_(solver.getAbsoluteToleranceSteadyState())
-    , rtol_(solver.getRelativeToleranceSteadyState())
-    , atol_sensi_(solver.getAbsoluteToleranceSteadyStateSensi())
-    , rtol_sensi_(solver.getRelativeToleranceSteadyStateSensi())
-    , atol_quad_(solver.getAbsoluteToleranceQuadratures())
-    , rtol_quad_(solver.getRelativeToleranceQuadratures())
     , newton_solver_(
           NewtonSolver(model, solver.getLinearSolver(), solver.getSunContext())
       )
@@ -607,7 +617,6 @@ bool SteadystateProblem::requires_state_sensitivities(
 
 realtype
 SteadystateProblem::getWrms(Model& model, SensitivityMethod sensi_method) {
-    realtype wrms = INFINITY;
     if (sensi_method == SensitivityMethod::adjoint) {
         if (newton_step_conv_) {
             throw NewtonFailure(
@@ -622,22 +631,18 @@ SteadystateProblem::getWrms(Model& model, SensitivityMethod sensi_method) {
         // to zero at all. So we need xQBdot, hence compute xQB first.
         computeQBfromQ(model, xQ_, xQB_, state_);
         computeQBfromQ(model, xB_, xQBdot_, state_);
-        wrms = getWrmsNorm(
-            xQB_, xQBdot_, steadystate_mask_, atol_quad_, rtol_quad_, ewtQB_
-        );
-    } else {
-        // If we're doing a forward simulation (with or without sensitivities),
-        // get RHS and compute weighted error norm.
-        if (newton_step_conv_)
-            getNewtonStep(model);
-        else
-            updateRightHandSide(model);
-        wrms = getWrmsNorm(
-            state_.x, newton_step_conv_ ? delta_ : xdot_, steadystate_mask_,
-            atol_, rtol_, ewt_
-        );
+        return wrms_computer_xQB_.wrms(xQBdot_, xQB_);
     }
-    return wrms;
+
+    if (newton_step_conv_) {
+        getNewtonStep(model);
+        return wrms_computer_x_.wrms(delta_, state_.x);
+    }
+
+    // If we're doing a forward simulation (with or without sensitivities),
+    // get RHS and compute weighted error norm.
+    updateRightHandSide(model);
+    return wrms_computer_x_.wrms(xdot_, state_.x);
 }
 
 realtype SteadystateProblem::getWrmsFSA(Model& model) {
@@ -655,10 +660,7 @@ realtype SteadystateProblem::getWrmsFSA(Model& model) {
         );
         if (newton_step_conv_)
             newton_solver_.solveLinearSystem(xdot_);
-        wrms = getWrmsNorm(
-            state_.sx[ip], xdot_, steadystate_mask_, atol_sensi_, rtol_sensi_,
-            ewt_
-        );
+        wrms = wrms_computer_sx_.wrms(xdot_, state_.sx[ip]);
         // ideally this function would report the maximum of all wrms over
         // all ip, but for practical purposes we can just report the wrms for
         // the first ip where we know that the convergence threshold is not
@@ -939,4 +941,5 @@ void SteadystateProblem::getNewtonStep(Model& model) {
     newton_solver_.getStep(delta_, model, state_);
     delta_updated_ = true;
 }
+
 } // namespace amici