Rosenbrock DAE Constraint Count

include/micm/solver/backward_euler.hpp

@@ -43,7 +43,8 @@ namespace micm
         LinearSolverPolicy&& linear_solver,
         RatesPolicy&& rates,
         auto& jacobian,
-        const size_t number_of_species)
+        const size_t number_of_species,
+	const size_t number_of_constraints)
         : linear_solver_(std::move(linear_solver)),
           rates_(std::move(rates))
     {

include/micm/solver/rosenbrock.hpp

@@ -66,7 +66,8 @@ namespace micm
         LinearSolverPolicy&& linear_solver,
         RatesPolicy&& rates,
         auto& jacobian,
-        const size_t number_of_species)
+        const size_t number_of_species,
+	const size_t number_of_constraints)
         : linear_solver_(std::move(linear_solver)),
           rates_(std::move(rates))
     {
@@ -133,12 +134,13 @@ namespace micm
     /// @param jacobian Jacobian matrix
     ///
     /// Note: This constructor is not intended to be used directly. Instead, use the SolverBuilder to create a solver
-    RosenbrockSolver(LinearSolverPolicy&& linear_solver, RatesPolicy&& rates, auto& jacobian, const size_t number_of_species)
+    RosenbrockSolver(LinearSolverPolicy&& linear_solver, RatesPolicy&& rates, auto& jacobian, const size_t number_of_species, const size_t number_of_constraints)
         : AbstractRosenbrockSolver<RatesPolicy, LinearSolverPolicy, RosenbrockSolver<RatesPolicy, LinearSolverPolicy>>(
               std::move(linear_solver),
               std::move(rates),
               jacobian,
-              number_of_species)
+              number_of_species,
+	      number_of_constraints)
     {
     }
     RosenbrockSolver(const RosenbrockSolver&) = delete;

include/micm/solver/rosenbrock.inl

@@ -211,11 +211,13 @@ namespace micm
       const double& alpha) const
     requires(!VectorizableSparse<SparseMatrixPolicy>)
   {
+    // Only add alpha to rows flagged in the identity mask (currently all state variables)
     for (std::size_t i_block = 0; i_block < state.jacobian_.NumberOfBlocks(); ++i_block)
     {
       auto jacobian_vector = std::next(state.jacobian_.AsVector().begin(), i_block * state.jacobian_.FlatBlockSize());
+      std::size_t diag_row = 0;
       for (const auto& i_elem : state.jacobian_diagonal_elements_)
-        jacobian_vector[i_elem] += alpha;
+        jacobian_vector[i_elem] += alpha * state.upper_left_identity_diagonal_[diag_row++];
     }
   }
 
@@ -227,12 +229,17 @@ namespace micm
     requires(VectorizableSparse<SparseMatrixPolicy>)
   {
     constexpr std::size_t n_cells = SparseMatrixPolicy::GroupVectorSize();
+    // Only add alpha to rows flagged in the identity mask (currently all state variables)
     for (std::size_t i_group = 0; i_group < state.jacobian_.NumberOfGroups(state.jacobian_.NumberOfBlocks()); ++i_group)
     {
       auto jacobian_vector = std::next(state.jacobian_.AsVector().begin(), i_group * state.jacobian_.GroupSize());
+      std::size_t diag_row = 0;
       for (const auto& i_elem : state.jacobian_diagonal_elements_)
+      {
         for (std::size_t i_cell = 0; i_cell < n_cells; ++i_cell)
-          jacobian_vector[i_elem + i_cell] += alpha;
+          jacobian_vector[i_elem + i_cell] += alpha * state.upper_left_identity_diagonal_[diag_row];
+        ++diag_row;
+      }
     }
   }
 

include/micm/solver/solver.hpp

@@ -105,6 +105,11 @@ namespace micm
       return state_parameters_.number_of_species_;
     }
 
+    std::size_t GetNumberOfConstraints() const
+    {
+      return state_parameters_.number_of_constraints_;
+    }
+
     std::size_t GetNumberOfReactions() const
     {
       return state_parameters_.number_of_rate_constants_;

include/micm/solver/solver_builder.hpp

@@ -50,6 +50,8 @@ namespace micm
     SolverParametersPolicy options_;
     System system_;
     std::vector<Process> reactions_;
+    std::size_t constraint_count_ = 0;
+    std::vector<std::string> constraint_names_{};
     bool ignore_unused_species_ = true;
     bool reorder_state_ = true;
     bool valid_system_ = false;
@@ -75,6 +77,16 @@ namespace micm
     /// @return Updated SolverBuilder
     SolverBuilder& SetReactions(const std::vector<Process>& reactions);
 
+    /// @brief Set the number of algebraic constraints (appended after state variables)
+    /// @param number_of_constraints Constraint count
+    /// @return Updated SolverBuilder
+    SolverBuilder& SetConstraintCount(std::size_t number_of_constraints);
+
+    /// @brief Set constraint names (appended after state variables)
+    /// @param names Constraint variable names
+    /// @return Updated SolverBuilder
+    SolverBuilder& SetConstraintNames(const std::vector<std::string>& names);
+
     /// @brief Set whether to ignore unused species
     /// @param ignore_unused_species True if unused species should be ignored
     /// @return Updated SolverBuilder
@@ -152,4 +164,4 @@ namespace micm
 
 }  // namespace micm
 
-#include "solver_builder.inl"
+#include "solver_builder.inl"

include/micm/solver/solver_builder.inl

@@ -63,6 +63,65 @@ namespace micm
     return *this;
   }
 
+  template<
+      class SolverParametersPolicy,
+      class DenseMatrixPolicy,
+      class SparseMatrixPolicy,
+      class RatesPolicy,
+      class LuDecompositionPolicy,
+      class LinearSolverPolicy,
+      class StatePolicy>
+  inline SolverBuilder<
+      SolverParametersPolicy,
+      DenseMatrixPolicy,
+      SparseMatrixPolicy,
+      RatesPolicy,
+      LuDecompositionPolicy,
+      LinearSolverPolicy,
+      StatePolicy>&
+  SolverBuilder<
+      SolverParametersPolicy,
+      DenseMatrixPolicy,
+      SparseMatrixPolicy,
+      RatesPolicy,
+      LuDecompositionPolicy,
+      LinearSolverPolicy,
+      StatePolicy>::SetConstraintCount(std::size_t number_of_constraints)
+  {
+    constraint_count_ = number_of_constraints;
+    return *this;
+  }
+
+  template<
+      class SolverParametersPolicy,
+      class DenseMatrixPolicy,
+      class SparseMatrixPolicy,
+      class RatesPolicy,
+      class LuDecompositionPolicy,
+      class LinearSolverPolicy,
+      class StatePolicy>
+  inline SolverBuilder<
+      SolverParametersPolicy,
+      DenseMatrixPolicy,
+      SparseMatrixPolicy,
+      RatesPolicy,
+      LuDecompositionPolicy,
+      LinearSolverPolicy,
+      StatePolicy>&
+  SolverBuilder<
+      SolverParametersPolicy,
+      DenseMatrixPolicy,
+      SparseMatrixPolicy,
+      RatesPolicy,
+      LuDecompositionPolicy,
+      LinearSolverPolicy,
+      StatePolicy>::SetConstraintNames(const std::vector<std::string>& names)
+  {
+    constraint_names_ = names;
+    constraint_count_ = names.size();
+    return *this;
+  }
+
   template<
       class SolverParametersPolicy,
       class DenseMatrixPolicy,
@@ -314,6 +373,7 @@ namespace micm
     auto species_map = this->GetSpeciesMap();
     auto labels = this->GetCustomParameterLabels();
     std::size_t number_of_species = this->system_.StateSize();
+    std::size_t number_of_constraints = constraint_count_;
     if (number_of_species == 0)
     {
       throw std::system_error(
@@ -325,7 +385,7 @@ namespace micm
     RatesPolicy rates(this->reactions_, species_map);
     auto nonzero_elements = rates.NonZeroJacobianElements();
     // The actual number of grid cells is not needed to construct the various solver objects
-    auto jacobian = BuildJacobian<SparseMatrixPolicy>(nonzero_elements, 1, number_of_species, true);
+    auto jacobian = BuildJacobian<SparseMatrixPolicy>(nonzero_elements, 1, number_of_species + number_of_constraints, true);
 
     LinearSolverPolicy linear_solver(jacobian, 0);
     if constexpr (LuDecompositionInPlaceConcept<LuDecompositionPolicy, SparseMatrixPolicy>)
@@ -338,8 +398,19 @@ namespace micm
     std::vector<std::string> variable_names{ number_of_species };
     for (auto& species_pair : species_map)
       variable_names[species_pair.second] = species_pair.first;
+    if (number_of_constraints > 0)
+    {
+      std::vector<std::string> names = constraint_names_;
+      if (names.size() < number_of_constraints)
+      {
+        for (std::size_t i = names.size(); i < number_of_constraints; ++i)
+          names.push_back("constraint_" + std::to_string(i));
+      }
+      variable_names.insert(variable_names.end(), names.begin(), names.begin() + number_of_constraints);
+    }
 
     StateParameters state_parameters = { .number_of_species_ = number_of_species,
+                                         .number_of_constraints_ = number_of_constraints,
                                          .number_of_rate_constants_ = this->reactions_.size(),
                                          .variable_names_ = variable_names,
                                          .custom_rate_parameter_labels_ = labels,
@@ -348,11 +419,12 @@ namespace micm
     this->SetAbsoluteTolerances(state_parameters.absolute_tolerance_, species_map);
 
     return Solver<SolverPolicy, StatePolicy>(
-        SolverPolicy(std::move(linear_solver), std::move(rates), jacobian, number_of_species),
+        SolverPolicy(std::move(linear_solver), std::move(rates), jacobian,
+            number_of_species, number_of_constraints),
         state_parameters,
         options,
         this->reactions_,
         this->system_);
   }
 
-}  // namespace micm
+}  // namespace micm

include/micm/solver/state.hpp

@@ -27,6 +27,7 @@ namespace micm
   struct StateParameters
   {
     std::size_t number_of_species_{ 0 };
+    std::size_t number_of_constraints_{ 0 };
     std::size_t number_of_rate_constants_{ 0 };
     std::vector<std::string> variable_names_{};
     std::vector<std::string> custom_rate_parameter_labels_{};
@@ -58,6 +59,8 @@ namespace micm
     DenseMatrixPolicy rate_constants_;
     /// @brief Atmospheric conditions, varies in time
     std::vector<Conditions> conditions_;
+    /// @brief The block matrix with an upper left identity, zeros elsewhere
+    std::vector<double> upper_left_identity_diagonal_;
     /// @brief The jacobian structure, varies for each solve
     SparseMatrixPolicy jacobian_;
     std::vector<std::size_t> jacobian_diagonal_elements_;
@@ -68,6 +71,7 @@ namespace micm
     LMatrixPolicy lower_matrix_;
     UMatrixPolicy upper_matrix_;
     std::size_t state_size_;
+    std::size_t constraint_size_;
     std::unique_ptr<TemporaryVariables> temporary_variables_;
     double relative_tolerance_;
     std::vector<double> absolute_tolerance_;
@@ -89,6 +93,7 @@ namespace micm
       custom_rate_parameters_ = other.custom_rate_parameters_;
       rate_constants_ = other.rate_constants_;
       conditions_ = other.conditions_;
+      upper_left_identity_diagonal_ = other.upper_left_identity_diagonal_;
       jacobian_ = other.jacobian_;
       jacobian_diagonal_elements_ = other.jacobian_diagonal_elements_;
       variable_map_ = other.variable_map_;
@@ -97,6 +102,7 @@ namespace micm
       lower_matrix_ = other.lower_matrix_;
       upper_matrix_ = other.upper_matrix_;
       state_size_ = other.state_size_;
+      constraint_size_ = other.constraint_size_;
       number_of_grid_cells_ = other.number_of_grid_cells_;
       temporary_variables_ = std::make_unique<TemporaryVariables>(*other.temporary_variables_);
       relative_tolerance_ = other.relative_tolerance_;
@@ -114,6 +120,7 @@ namespace micm
         custom_rate_parameters_ = other.custom_rate_parameters_;
         rate_constants_ = other.rate_constants_;
         conditions_ = other.conditions_;
+        upper_left_identity_diagonal_ = other.upper_left_identity_diagonal_;
         jacobian_ = other.jacobian_;
         jacobian_diagonal_elements_ = other.jacobian_diagonal_elements_;
         variable_map_ = other.variable_map_;
@@ -122,6 +129,7 @@ namespace micm
         lower_matrix_ = other.lower_matrix_;
         upper_matrix_ = other.upper_matrix_;
         state_size_ = other.state_size_;
+        constraint_size_ = other.constraint_size_;
         number_of_grid_cells_ = other.number_of_grid_cells_;
         temporary_variables_ = std::make_unique<TemporaryVariables>(*other.temporary_variables_);
         relative_tolerance_ = other.relative_tolerance_;
@@ -137,6 +145,7 @@ namespace micm
           custom_rate_parameters_(std::move(other.custom_rate_parameters_)),
           rate_constants_(std::move(other.rate_constants_)),
           conditions_(std::move(other.conditions_)),
+          upper_left_identity_diagonal_(std::move(other.upper_left_identity_diagonal_)),
           jacobian_(std::move(other.jacobian_)),
           jacobian_diagonal_elements_(std::move(other.jacobian_diagonal_elements_)),
           variable_map_(std::move(other.variable_map_)),
@@ -145,6 +154,7 @@ namespace micm
           lower_matrix_(std::move(other.lower_matrix_)),
           upper_matrix_(std::move(other.upper_matrix_)),
           state_size_(other.state_size_),
+          constraint_size_(other.constraint_size_),
           number_of_grid_cells_(other.number_of_grid_cells_),
           temporary_variables_(std::move(other.temporary_variables_)),
           relative_tolerance_(other.relative_tolerance_),
@@ -163,6 +173,7 @@ namespace micm
         custom_rate_parameters_ = std::move(other.custom_rate_parameters_);
         rate_constants_ = std::move(other.rate_constants_);
         conditions_ = std::move(other.conditions_);
+        upper_left_identity_diagonal_ = std::move(other.upper_left_identity_diagonal_);
         jacobian_ = std::move(other.jacobian_);
         jacobian_diagonal_elements_ = std::move(other.jacobian_diagonal_elements_);
         variable_map_ = std::move(other.variable_map_);
@@ -171,12 +182,14 @@ namespace micm
         lower_matrix_ = std::move(other.lower_matrix_);
         upper_matrix_ = std::move(other.upper_matrix_);
         state_size_ = other.state_size_;
+        constraint_size_ = other.constraint_size_;
         number_of_grid_cells_ = other.number_of_grid_cells_;
         temporary_variables_ = std::move(other.temporary_variables_);
         relative_tolerance_ = other.relative_tolerance_;
         absolute_tolerance_ = std::move(other.absolute_tolerance_);
 
         other.state_size_ = 0;
+        other.constraint_size_ = 0;
         other.number_of_grid_cells_ = 0;
       }
       return *this;

include/micm/solver/state.inl

@@ -69,6 +69,7 @@ namespace micm
         custom_rate_parameters_(),
         rate_constants_(),
         conditions_(),
+        upper_left_identity_diagonal_(),
         jacobian_(),
         jacobian_diagonal_elements_(),
         variable_map_(),
@@ -77,6 +78,7 @@ namespace micm
         lower_matrix_(),
         upper_matrix_(),
         state_size_(0),
+        constraint_size_(0),
         number_of_grid_cells_(0),
         temporary_variables_(nullptr),
         relative_tolerance_(1e-06),
@@ -100,11 +102,13 @@ namespace micm
         variable_map_(),
         custom_rate_parameter_map_(),
         variable_names_(parameters.variable_names_),
+        upper_left_identity_diagonal_(),
         jacobian_(),
         jacobian_diagonal_elements_(),
         lower_matrix_(),
         upper_matrix_(),
-        state_size_(parameters.variable_names_.size()),
+        state_size_(parameters.number_of_species_),
+        constraint_size_(parameters.number_of_constraints_),
         number_of_grid_cells_(number_of_grid_cells),
         relative_tolerance_(parameters.relative_tolerance_),
         absolute_tolerance_(parameters.absolute_tolerance_)
@@ -116,17 +120,24 @@ namespace micm
     for (auto& label : parameters.custom_rate_parameter_labels_)
       custom_rate_parameter_map_[label] = index++;
 
+    for (std::size_t i = 0; i < state_size_; i++) {
+      upper_left_identity_diagonal_.push_back(1.0);
+    }
+    for (std::size_t i = 0; i < constraint_size_; i++) {
+      upper_left_identity_diagonal_.push_back(0.0);
+    }
+
     if constexpr (LuDecompositionInPlaceConcept<LuDecompositionPolicy, SparseMatrixPolicy>)
     {
       jacobian_ =
-          BuildJacobian<SparseMatrixPolicy>(parameters.nonzero_jacobian_elements_, number_of_grid_cells, state_size_, true);
+          BuildJacobian<SparseMatrixPolicy>(parameters.nonzero_jacobian_elements_, number_of_grid_cells, state_size_ + constraint_size_, true);
       auto lu = LuDecompositionPolicy::template GetLUMatrix<SparseMatrixPolicy>(jacobian_, 0, false);
       jacobian_ = std::move(lu);
     }
     else
     {
       jacobian_ =
-          BuildJacobian<SparseMatrixPolicy>(parameters.nonzero_jacobian_elements_, number_of_grid_cells, state_size_, false);
+          BuildJacobian<SparseMatrixPolicy>(parameters.nonzero_jacobian_elements_, number_of_grid_cells, state_size_ + constraint_size_, false);
       auto lu = LuDecompositionPolicy::template GetLUMatrices<SparseMatrixPolicy, LMatrixPolicy, UMatrixPolicy>(
           jacobian_, 0, false);
       auto lower_matrix = std::move(lu.first);

test/unit/solver/test_linear_solver_in_place_policy.hpp

@@ -4,6 +4,7 @@
 #include <gtest/gtest.h>
 
 #include <functional>
+#include <iomanip>
 #include <random>
 
 template<typename T, class MatrixPolicy, class SparseMatrixPolicy>