Allow subselection of state variables for steady-state simulations

include/amici/model.h

@@ -1481,6 +1481,40 @@ class Model : public AbstractModel, public ModelDimensions {
      */
     virtual std::vector<double> get_trigger_timepoints() const;
 
+    /**
+     * @brief Get steady-state mask as std::vector.
+     *
+     * See `set_steadystate_mask` for details.
+     *
+     * @return Steady-state mask
+     */
+    std::vector<double> get_steadystate_mask() const {
+        return steadystate_mask_.getVector();
+    };
+
+    /**
+     * @brief Get steady-state mask as AmiVector.
+     *
+     * See `set_steadystate_mask` for details.
+     * @return Steady-state mask
+     */
+    AmiVector const& get_steadystate_mask_av() const {
+        return steadystate_mask_;
+    };
+
+    /**
+     * @brief Set steady-state mask.
+     *
+     * The mask is used to exclude certain state variables from the steady-state
+     * convergence check. Positive values indicate that the corresponding state
+     * variable should be included in the convergence check, while non-positive
+     * values indicate that the corresponding state variable should be excluded.
+     * An empty mask is interpreted as including all state variables.
+     *
+     * @param mask Mask of length `nx_solver`.
+     */
+    void set_steadystate_mask(std::vector<double> const& mask);
+
     /**
      * Flag indicating whether for
      * `amici::Solver::sensi_` == `amici::SensitivityOrder::second`
@@ -2087,6 +2121,15 @@ class Model : public AbstractModel, public ModelDimensions {
 
     /** Simulation parameters, initial state, etc. */
     SimulationParameters simulation_parameters_;
+
+    /**
+     * Mask for state variables that should be checked for steady state
+     * during pre-/post-equilibration. Positive values indicate that the
+     * corresponding state variable should be checked for steady state.
+     * Negative values indicate that the corresponding state variable should
+     * be ignored.
+     */
+    AmiVector steadystate_mask_;
 };
 
 bool operator==(Model const& a, Model const& b);

include/amici/serialization.h

@@ -150,6 +150,7 @@ void serialize(Archive& ar, amici::Model& m, unsigned int const /*version*/) {
     ar & m.steadystate_computation_mode_;
     ar & m.steadystate_sensitivity_mode_;
     ar & m.state_independent_events_;
+    ar & m.steadystate_mask_;
 }
 
 /**

include/amici/steadystateproblem.h

@@ -246,14 +246,16 @@ class SteadystateProblem {
      * w_i = 1 / ( rtol * x_i + atol )
      * @param x current state (sx[ip] for sensitivities)
      * @param xdot current rhs (sxdot[ip] for sensitivities)
+     * @param mask mask for state variables to include in WRMS norm.
+     * Positive value: include; non-positive value: exclude; empty: include all.
      * @param atol absolute tolerance
      * @param rtol relative tolerance
      * @param ewt error weight vector
      * @return root-mean-square norm
      */
     realtype getWrmsNorm(
-        AmiVector const& x, AmiVector const& xdot, realtype atol, realtype rtol,
-        AmiVector& ewt
+        AmiVector const& x, AmiVector const& xdot, AmiVector const& mask,
+        realtype atol, realtype rtol, AmiVector& ewt
     ) const;
 
     /**

include/amici/vector.h

@@ -430,8 +430,8 @@ inline span<realtype> make_span(N_Vector nv) {
 }
 
 /**
- * @brief Create span from N_Vector
- * @param nv
+ * @brief Create span from AmiVector
+ * @param av
  *
  */
 inline span<realtype const> make_span(amici::AmiVector const& av) {

python/sdist/amici/swig_wrappers.py

@@ -255,6 +255,7 @@ def writeSolverSettingsToHDF5(
     "SteadyStateSensitivityMode",
     ("t0", "setT0"),
     "Timepoints",
+    "_steadystate_mask",
 ]
 
 

python/tests/test_preequilibration.py

@@ -8,6 +8,10 @@
 from amici.debugging import get_model_for_preeq
 from numpy.testing import assert_allclose, assert_equal
 from test_pysb import get_data
+from amici.testing import (
+    TemporaryDirectoryWinSafe as TemporaryDirectory,
+    skip_on_valgrind,
+)
 
 
 @pytest.fixture
@@ -658,3 +662,62 @@ def test_get_model_for_preeq(preeq_fixture):
         rdata1.sx,
         rdata2.sx,
     )
+
+
+@skip_on_valgrind
+def test_partial_eq():
+    """Check that partial equilibration is possible."""
+    from amici.antimony_import import antimony2amici
+
+    ant_str = """
+    model test_partial_eq
+        explodes = 1
+        explodes' = explodes
+        A = 1
+        B = 0
+        R: A -> B; k*A - k*B
+        k = 1
+    end
+    """
+    module_name = "test_partial_eq"
+    with TemporaryDirectory(prefix=module_name) as outdir:
+        antimony2amici(
+            ant_str,
+            model_name=module_name,
+            output_dir=outdir,
+        )
+        model_module = amici.import_model_module(
+            module_name=module_name, module_path=outdir
+        )
+        amici_model = model_module.getModel()
+        amici_model.setTimepoints([np.inf])
+        amici_solver = amici_model.getSolver()
+        amici_solver.setRelativeToleranceSteadyState(1e-12)
+
+        # equilibration of `explodes` will fail
+        rdata = amici.runAmiciSimulation(amici_model, amici_solver)
+        assert rdata.status == amici.AMICI_ERROR
+        assert rdata.messages[0].identifier == "EQUILIBRATION_FAILURE"
+
+        # excluding `explodes` should enable equilibration
+        amici_model.set_steadystate_mask(
+            [
+                0 if state_id == "explodes" else 1
+                for state_id in amici_model.getStateIdsSolver()
+            ]
+        )
+        rdata = amici.runAmiciSimulation(amici_model, amici_solver)
+        assert rdata.status == amici.AMICI_SUCCESS
+        assert_allclose(
+            rdata.by_id("A"),
+            0.5,
+            atol=amici_solver.getAbsoluteToleranceSteadyState(),
+            rtol=amici_solver.getRelativeToleranceSteadyState(),
+        )
+        assert_allclose(
+            rdata.by_id("B"),
+            0.5,
+            atol=amici_solver.getAbsoluteToleranceSteadyState(),
+            rtol=amici_solver.getRelativeToleranceSteadyState(),
+        )
+        assert rdata.t_last < 100

python/tests/test_swig_interface.py

@@ -39,7 +39,7 @@ def test_copy_constructors(pysb_example_presimulation_module):
             val = get_val(obj, attr)
 
             try:
-                modval = get_mod_val(val, attr)
+                modval = get_mod_val(val, attr, obj)
             except ValueError:
                 # happens for everything that is not bool or scalar
                 continue
@@ -79,6 +79,10 @@ def test_copy_constructors(pysb_example_presimulation_module):
         tuple([1.0] + [0.0] * 35),
         tuple([0.1] * 36),
     ],
+    "_steadystate_mask": [
+        (),
+        tuple([0] * 3),
+    ],
     "MinimumSigmaResiduals": [
         50.0,
         60.0,
@@ -361,19 +365,21 @@ def get_val(obj, attr):
         return getattr(obj, attr)
 
 
-def get_mod_val(val, attr):
+def get_mod_val(val, attr, obj):
     if attr == "getReturnDataReportingMode":
         return amici.RDataReporting.likelihood
     elif attr == "getParameterList":
-        return tuple(get_mod_val(val[0], "") for _ in val)
+        return tuple(get_mod_val(val[0], "", obj) for _ in val)
     elif attr == "getStateIsNonNegative":
         raise ValueError("Cannot modify value")
+    elif attr == "get_steadystate_mask":
+        return [0 for _ in range(obj.nx_solver)]
     elif isinstance(val, bool):
         return not val
     elif isinstance(val, numbers.Number):
         return val + 1
     elif isinstance(val, tuple):
-        return tuple(get_mod_val(v, attr) for v in val)
+        return tuple(get_mod_val(v, attr, obj) for v in val)
 
     raise ValueError("Cannot modify value")
 

src/hdf5.cpp

@@ -1244,6 +1244,12 @@
             model.setInitialStates(x0);
     }
 
+    if (locationExists(file, datasetPath + "/steadystate_mask")) {
+        auto mask = getDoubleDataset1D(file, datasetPath + "/steadystate_mask");
+        if (!mask.empty())
+            model.set_steadystate_mask(mask);
+    }
+
     if (locationExists(file, datasetPath + "/sx0")) {
         hsize_t length0 = 0;
         hsize_t length1 = 0;

src/model.cpp

@@ -287,7 +287,9 @@
            && (a.nmaxevent_ == b.nmaxevent_)
            && (a.state_is_non_negative_ == b.state_is_non_negative_)
            && (a.sigma_res_ == b.sigma_res_) && (a.min_sigma_ == b.min_sigma_)
-           && a.state_ == b.state_;
+           && (a.state_ == b.state_)
+           && (a.steadystate_mask_.getVector()
+               == b.steadystate_mask_.getVector());
 }
 
 bool operator==(ModelDimensions const& a, ModelDimensions const& b) {
@@ -3135,6 +3137,23 @@
     return trigger_timepoints;
 }
 
+void Model::set_steadystate_mask(std::vector<double> const& mask) {
+    if (mask.size() == 0) {
+        if (steadystate_mask_.getLength() != 0) {
+            steadystate_mask_ = AmiVector();
+        }
+        return;
+    }
+
+    if (gsl::narrow<int>(mask.size()) != nx_solver)
+        throw AmiException(
+            "Steadystate mask has wrong size: %d, expected %d",
+            gsl::narrow<int>(mask.size()), nx_solver
+        );
+
+    steadystate_mask_ = AmiVector(mask);
+}
+
 const_N_Vector Model::computeX_pos(const_N_Vector x) {
     if (any_state_non_negative_) {
         for (int ix = 0; ix < derived_state_.x_pos_tmp_.getLength(); ++ix) {

src/steadystateproblem.cpp

@@ -509,8 +509,8 @@ bool SteadystateProblem::getSensitivityFlag(
 }
 
 realtype SteadystateProblem::getWrmsNorm(
-    AmiVector const& x, AmiVector const& xdot, realtype atol, realtype rtol,
-    AmiVector& ewt
+    AmiVector const& x, AmiVector const& xdot, AmiVector const& mask,
+    realtype atol, realtype rtol, AmiVector& ewt
 ) const {
     /* Depending on what convergence we want to check (xdot, sxdot, xQBdot)
        we need to pass ewt[QB], as xdot and xQBdot have different sizes */
@@ -522,7 +522,14 @@ realtype SteadystateProblem::getWrmsNorm(
     N_VAddConst(ewt.getNVector(), atol, ewt.getNVector());
     /* ewt = 1/ewt (ewt = 1/(rtol*x+atol)) */
     N_VInv(ewt.getNVector(), ewt.getNVector());
-    /* wrms = sqrt(sum((xdot/ewt)**2)/n) where n = size of state vector */
+
+    // wrms = sqrt(sum((xdot/ewt)**2)/n) where n = size of state vector
+    if (mask.getLength()) {
+        return N_VWrmsNormMask(
+            const_cast<N_Vector>(xdot.getNVector()), ewt.getNVector(),
+            const_cast<N_Vector>(mask.getNVector())
+        );
+    }
     return N_VWrmsNorm(
         const_cast<N_Vector>(xdot.getNVector()), ewt.getNVector()
     );
@@ -543,7 +550,10 @@ SteadystateProblem::getWrms(Model& model, SensitivityMethod sensi_method) {
                 "Newton type convergence check is not implemented for adjoint "
                 "steady state computations. Stopping."
             );
-        wrms = getWrmsNorm(xQB_, xQBdot_, atol_quad_, rtol_quad_, ewtQB_);
+        wrms = getWrmsNorm(
+            xQB_, xQBdot_, model.get_steadystate_mask_av(), atol_quad_,
+            rtol_quad_, ewtQB_
+        );
     } else {
         /* If we're doing a forward simulation (with or without sensitivities:
            Get RHS and compute weighted error norm */
@@ -552,7 +562,8 @@ SteadystateProblem::getWrms(Model& model, SensitivityMethod sensi_method) {
         else
             updateRightHandSide(model);
         wrms = getWrmsNorm(
-            state_.x, newton_step_conv_ ? delta_ : xdot_, atol_, rtol_, ewt_
+            state_.x, newton_step_conv_ ? delta_ : xdot_,
+            model.get_steadystate_mask_av(), atol_, rtol_, ewt_
         );
     }
     return wrms;
@@ -573,8 +584,10 @@ realtype SteadystateProblem::getWrmsFSA(Model& model) {
         );
         if (newton_step_conv_)
             newton_solver_->solveLinearSystem(xdot_);
-        wrms
-            = getWrmsNorm(state_.sx[ip], xdot_, atol_sensi_, rtol_sensi_, ewt_);
+        wrms = getWrmsNorm(
+            state_.sx[ip], xdot_, model.get_steadystate_mask_av(), atol_sensi_,
+            rtol_sensi_, ewt_
+        );
         /* 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

swig/model.i

@@ -94,6 +94,7 @@ using namespace amici;
 %ignore fdx_rdatadtcl;
 %ignore fdx_rdatadx_solver;
 %ignore fdsigmaydy;
+%ignore get_steadystate_mask_av;
 
 %newobject amici::Model::clone;