Enable setting constraints on state variables (AMICI-dev#2340)

In addition to the current `Model.setStateIsNonNegative`, this adds the option to set additional
(non)negativity/positivity constraints to be enforced by the solver.

See [CVodeSetConstraints](https://sundials.readthedocs.io/en/latest/cvode/Usage/index.html#c.CVodeSetConstraints) for details.

Related to AMICI-dev#2327.

include/amici/defines.h

@@ -255,6 +255,15 @@ enum class SplineExtrapolation {
     periodic        =  3,
 };
 
+/** Constraints on state variables */
+enum class Constraint {
+    none = 0,
+    non_negative = 1,
+    non_positive = -1,
+    positive = 2,
+    negative = -2,
+};
+
 // clang-format on
 
 } // namespace amici

include/amici/serialization.h

@@ -5,6 +5,7 @@
 #include "amici/rdata.h"
 #include "amici/solver.h"
 #include "amici/solver_cvodes.h"
+#include "amici/vector.h"
 
 #include <chrono>
 
@@ -87,6 +88,7 @@ void serialize(Archive& ar, amici::Solver& s, unsigned int const /*version*/) {
     ar & s.maxtime_;
     ar & s.max_conv_fails_;
     ar & s.max_nonlin_iters_;
+    ar & s.constraints_;
     ar & s.max_step_size_;
 }
 
@@ -278,6 +280,25 @@ void serialize(
     ar & m.ubw;
     ar & m.lbw;
 }
+
+/**
+ * @brief Serialize AmiVector to a boost archive
+ * @param ar archive
+ * @param v AmiVector
+ */
+template <class Archive>
+void serialize(
+    Archive& ar, amici::AmiVector& v, unsigned int const /*version*/
+) {
+    if (Archive::is_loading::value) {
+        std::vector<realtype> tmp;
+        ar & tmp;
+        v = amici::AmiVector(tmp);
+    } else {
+        auto tmp = v.getVector();
+        ar & tmp;
+    }
+}
 #endif
 } // namespace serialization
 } // namespace boost

include/amici/solver.h

@@ -964,6 +964,24 @@ class Solver {
      */
     int getMaxConvFails() const;
 
+    /**
+     * @brief Set constraints on the model state.
+     *
+     * See
+     * https://sundials.readthedocs.io/en/latest/cvode/Usage/index.html#c.CVodeSetConstraints.
+     *
+     * @param constraints
+     */
+    void setConstraints(std::vector<realtype> const& constraints);
+
+    /**
+     * @brief Get constraints on the model state.
+     * @return constraints
+     */
+    std::vector<realtype> getConstraints() const {
+        return constraints_.getVector();
+    }
+
     /**
      * @brief Set the maximum step size
      * @param max_step_size maximum step size. `0.0` means no limit.
@@ -1130,7 +1148,7 @@ class Solver {
     virtual void rootInit(int ne) const = 0;
 
     /**
-     * @brief Initalize non-linear solver for sensitivities
+     * @brief Initialize non-linear solver for sensitivities
      * @param model Model instance
      */
     void initializeNonLinearSolverSens(Model const* model) const;
@@ -1648,6 +1666,11 @@ class Solver {
      */
     void applySensitivityTolerances() const;
 
+    /**
+     * @brief Apply the constraints to the solver.
+     */
+    virtual void apply_constraints() const;
+
     /** pointer to solver memory block */
     mutable std::unique_ptr<void, free_solver_ptr> solver_memory_;
 
@@ -1809,6 +1832,9 @@ class Solver {
     /** flag indicating whether sensInit1 was called */
     mutable bool sens_initialized_{false};
 
+    /** Vector of constraints on the solution */
+    mutable AmiVector constraints_;
+
   private:
     /**
      * @brief applies total number of steps for next solver call

include/amici/solver_cvodes.h

@@ -254,6 +254,8 @@ class CVodeSolver : public Solver {
 
     void apply_max_conv_fails() const override;
 
+    void apply_constraints() const override;
+
     void apply_max_step_size() const override;
 };
 

include/amici/solver_idas.h

@@ -234,6 +234,8 @@ class IDASolver : public Solver {
 
     void apply_max_conv_fails() const override;
 
+    void apply_constraints() const override;
+
     void apply_max_step_size() const override;
 };
 

include/amici/vector.h

@@ -10,6 +10,18 @@
 
 #include <gsl/gsl-lite.hpp>
 
+namespace amici {
+class AmiVector;
+}
+
+// for serialization friend
+namespace boost {
+namespace serialization {
+template <class Archive>
+void serialize(Archive& ar, amici::AmiVector& s, unsigned int version);
+}
+} // namespace boost
+
 namespace amici {
 
 /** Since const N_Vector is not what we want */
@@ -54,7 +66,7 @@ class AmiVector {
      * @brief constructor from gsl::span,
      * @param rvec vector from which the data will be copied
      */
-    explicit AmiVector(gsl::span<realtype> rvec)
+    explicit AmiVector(gsl::span<realtype const> rvec)
         : AmiVector(std::vector<realtype>(rvec.begin(), rvec.end())) {}
 
     /**
@@ -213,6 +225,17 @@ class AmiVector {
      */
     void abs() { N_VAbs(getNVector(), getNVector()); };
 
+    /**
+     * @brief Serialize AmiVector (see boost::serialization::serialize)
+     * @param ar Archive to serialize to
+     * @param s Data to serialize
+     * @param version Version number
+     */
+    template <class Archive>
+    friend void boost::serialization::serialize(
+        Archive& ar, AmiVector& s, unsigned int version
+    );
+
   private:
     /** main data storage */
     std::vector<realtype> vec_;
@@ -405,6 +428,16 @@ namespace gsl {
 inline span<realtype> make_span(N_Vector nv) {
     return span<realtype>(N_VGetArrayPointer(nv), N_VGetLength_Serial(nv));
 }
+
+/**
+ * @brief Create span from N_Vector
+ * @param nv
+ *
+ */
+inline span<realtype const> make_span(amici::AmiVector const& av) {
+    return make_span(av.getVector());
+}
+
 } // namespace gsl
 
 #endif /* AMICI_VECTOR_H */

python/tests/test_hdf5.py

@@ -24,6 +24,8 @@ def _modify_solver_attrs(solver):
         elif attr == "setMaxTime":
             # default value is the maximum, must not add to that
             cval = random.random()
+        elif attr == "setConstraints":
+            cval = [1.0, 1.0]
         elif isinstance(val, int):
             cval = val + 1
         else:

python/tests/test_sbml_import.py

@@ -722,3 +722,53 @@ def test_hardcode_parameters(simple_sbml_model):
             constant_parameters=["p1"],
             hardcode_symbols=["p1"],
         )
+
+
+def test_constraints():
+    """Test non-negativity constraint handling."""
+    from amici.antimony_import import antimony2amici
+    from amici import Constraint
+
+    ant_model = """
+    model test_non_negative_species
+        species A = 10
+        species B = 0
+        # R1: A => B; k1f * sqrt(A)
+        R1: A => B; k1f * max(0, A)
+        k1f = 1e10
+    end
+    """
+    module_name = "test_non_negative_species"
+    with TemporaryDirectory(prefix=module_name) as outdir:
+        antimony2amici(
+            ant_model,
+            model_name=module_name,
+            output_dir=outdir,
+            compute_conservation_laws=False,
+        )
+        model_module = amici.import_model_module(
+            module_name=module_name, module_path=outdir
+        )
+        amici_model = model_module.getModel()
+        amici_model.setTimepoints(np.linspace(0, 100, 200))
+        amici_solver = amici_model.getSolver()
+        rdata = amici.runAmiciSimulation(amici_model, amici_solver)
+        assert rdata.status == amici.AMICI_SUCCESS
+        # should be non-negative in theory, but is expected to become negative
+        #  in practice
+        assert np.any(rdata.x < 0)
+
+        amici_solver.setRelativeTolerance(1e-14)
+        amici_solver.setConstraints(
+            [Constraint.non_negative, Constraint.non_negative]
+        )
+        rdata = amici.runAmiciSimulation(amici_model, amici_solver)
+        assert rdata.status == amici.AMICI_SUCCESS
+        assert np.all(rdata.x >= 0)
+        assert np.all(
+            np.sum(rdata.x, axis=1) - np.sum(rdata.x[0])
+            < max(
+                np.sum(rdata.x[0]) * amici_solver.getRelativeTolerance(),
+                amici_solver.getAbsoluteTolerance(),
+            )
+        )

src/hdf5.cpp

@@ -914,6 +914,10 @@ void writeSolverSettingsToHDF5(
     H5LTset_attribute_int(
         file.getId(), hdf5Location.c_str(), "max_conv_fails", &ibuffer, 1
     );
+
+    createAndWriteDouble1DDataset(
+        file, hdf5Location + "/constraints", solver.getConstraints()
+    );
 }
 
 void readSolverSettingsFromHDF5(
@@ -1143,6 +1147,12 @@ void readSolverSettingsFromHDF5(
             getIntScalarAttribute(file, datasetPath, "max_conv_fails")
         );
     }
+
+    if (locationExists(file, datasetPath + "/constraints")) {
+        solver.setConstraints(
+            getDoubleDataset1D(file, datasetPath + "/constraints")
+        );
+    }
 }
 
 void readSolverSettingsFromHDF5(

src/solver.cpp

@@ -200,6 +200,8 @@ void Solver::setup(
 
     cpu_time_ = 0.0;
     cpu_timeB_ = 0.0;
+
+    apply_constraints();
 }
 
 void Solver::setupB(
@@ -646,6 +648,15 @@ void Solver::applySensitivityTolerances() const {
     }
 }
 
+void Solver::apply_constraints() const {
+    if (constraints_.getLength() != 0
+        && gsl::narrow<int>(constraints_.getLength()) != nx()) {
+        throw std::invalid_argument(
+            "Constraints must have the same size as the state vector."
+        );
+    }
+}
+
 SensitivityMethod Solver::getSensitivityMethod() const { return sensi_meth_; }
 
 SensitivityMethod Solver::getSensitivityMethodPreequilibration() const {
@@ -693,6 +704,21 @@ void Solver::setMaxConvFails(int max_conv_fails) {
 
 int Solver::getMaxConvFails() const { return max_conv_fails_; }
 
+void Solver::setConstraints(std::vector<realtype> const& constraints) {
+    auto any_constraint
+        = std::any_of(constraints.begin(), constraints.end(), [](bool x) {
+              return x != 0.0;
+          });
+
+    if (!any_constraint) {
+        // all-0 must be converted to empty, otherwise sundials will fail
+        constraints_ = AmiVector();
+        return;
+    }
+
+    constraints_ = AmiVector(constraints);
+}
+
 void Solver::setMaxStepSize(realtype max_step_size) {
     if (max_step_size < 0)
         throw AmiException("max_step_size must be non-negative.");

src/solver_cvodes.cpp

@@ -270,6 +270,18 @@ void CVodeSolver::apply_max_conv_fails() const {
         throw CvodeException(status, "CVodeSetMaxConvFails");
 }
 
+void CVodeSolver::apply_constraints() const {
+    Solver::apply_constraints();
+
+    int status = CVodeSetConstraints(
+        solver_memory_.get(),
+        constraints_.getLength() > 0 ? constraints_.getNVector() : nullptr
+    );
+    if (status != CV_SUCCESS) {
+        throw CvodeException(status, "CVodeSetConstraints");
+    }
+}
+
 void CVodeSolver::apply_max_step_size() const {
     int status = CVodeSetMaxStep(solver_memory_.get(), getMaxStepSize());
     if (status != CV_SUCCESS)

src/solver_idas.cpp

@@ -267,6 +267,18 @@ void IDASolver::apply_max_conv_fails() const {
         throw IDAException(status, "IDASetMaxConvFails");
 }
 
+void IDASolver::apply_constraints() const {
+    Solver::apply_constraints();
+
+    int status = IDASetConstraints(
+        solver_memory_.get(),
+        constraints_.getLength() > 0 ? constraints_.getNVector() : nullptr
+    );
+    if (status != IDA_SUCCESS) {
+        throw IDAException(status, "IDASetConstraints");
+    }
+}
+
 void IDASolver::apply_max_step_size() const {
     int status = IDASetMaxStep(solver_memory_.get(), getMaxStepSize());
     if (status != IDA_SUCCESS)

swig/amici.i

@@ -327,6 +327,7 @@ SteadyStateStatus = enum('SteadyStateStatus')
 NewtonDampingFactorMode = enum('NewtonDampingFactorMode')
 FixedParameterContext = enum('FixedParameterContext')
 RDataReporting = enum('RDataReporting')
+Constraint = enum('Constraint')
 %}
 
 %template(SteadyStateStatusVector) std::vector<amici::SteadyStateStatus>;