Refactor the constraint set to avoid copying for concentration

include/micm/CPU.hpp

@@ -22,19 +22,19 @@ namespace micm
   using StandardState = State<DenseMatrixStandard, SparseMatrixStandard>;
 
   using RosenbrockVectorType = typename RosenbrockSolverParameters::
-      template SolverType<ProcessSet<DenseMatrixVector, SparseMatrixVector>, LinearSolver<SparseMatrixVector, LuDecomposition>>;
+      template SolverType<ProcessSet<DenseMatrixVector, SparseMatrixVector>, LinearSolver<SparseMatrixVector, LuDecomposition>, ConstraintSet<DenseMatrixVector, SparseMatrixVector>>;
   using Rosenbrock = Solver<RosenbrockVectorType, State<DenseMatrixVector, SparseMatrixVector>>;
 
   using RosenbrockStandardType = typename RosenbrockSolverParameters::
-      template SolverType<ProcessSet<DenseMatrixStandard, SparseMatrixStandard>, LinearSolver<SparseMatrixStandard, LuDecomposition>>;
+      template SolverType<ProcessSet<DenseMatrixStandard, SparseMatrixStandard>, LinearSolver<SparseMatrixStandard, LuDecomposition>, ConstraintSet<DenseMatrixStandard, SparseMatrixStandard>>;
   using RosenbrockStandard = Solver<RosenbrockStandardType, State<DenseMatrixStandard, SparseMatrixStandard>>;
 
   using BackwardEulerVectorType = typename BackwardEulerSolverParameters::
-      template SolverType<ProcessSet<DenseMatrixVector, SparseMatrixVector>, LinearSolver<SparseMatrixVector, LuDecomposition>>;
+      template SolverType<ProcessSet<DenseMatrixVector, SparseMatrixVector>, LinearSolver<SparseMatrixVector, LuDecomposition>, ConstraintSet<DenseMatrixVector, SparseMatrixVector>>;
   using BackwardEuler = Solver<BackwardEulerVectorType, State<DenseMatrixVector, SparseMatrixVector>>;
 
   using BackwardEulerStandardType = typename BackwardEulerSolverParameters::
-      template SolverType<ProcessSet<DenseMatrixStandard, SparseMatrixStandard>, LinearSolver<SparseMatrixStandard, LuDecomposition>>;
+      template SolverType<ProcessSet<DenseMatrixStandard, SparseMatrixStandard>, LinearSolver<SparseMatrixStandard, LuDecomposition>, ConstraintSet<DenseMatrixStandard, SparseMatrixStandard>>;
   using BackwardEulerStandard = Solver<BackwardEulerStandardType, State<DenseMatrixStandard, SparseMatrixStandard>>;
 
   using RosenbrockThreeStageBuilder = CpuSolverBuilder<RosenbrockSolverParameters, DenseMatrixVector, SparseMatrixVector>;

include/micm/Constraint.hpp

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 #pragma once
 
+#include <micm/constraint/types/equilibrium_constraint.hpp>
+#include <micm/constraint/types/linear_constraint.hpp>
 #include <micm/constraint/constraint.hpp>
-#include <micm/constraint/constraint_set.hpp>
-#include <micm/constraint/equilibrium_constraint.hpp>
+#include <micm/constraint/constraint_set.hpp>

include/micm/GPU.hpp

@@ -20,6 +20,7 @@
 #include <micm/cuda/util/cuda_param.hpp>
 #include <micm/cuda/util/cuda_sparse_matrix.hpp>
 #include <micm/cuda/util/cuda_util.cuh>
+#include <micm/constraint/constraint_set.hpp>
 #include <micm/solver/solver_builder.hpp>
 
 namespace micm
@@ -30,7 +31,7 @@ namespace micm
   using GpuState = CudaState<CudaDenseMatrixVector, CudaSparseMatrixVector, CudaLuDecompositionMozartInPlace>;
 
   using CudaRosenbrockVectorType = typename CudaRosenbrockSolverParameters::
-      template SolverType<CudaProcessSet<CudaDenseMatrixVector, CudaSparseMatrixVector>, CudaLinearSolverInPlace<CudaSparseMatrixVector>>;
+      template SolverType<CudaProcessSet<CudaDenseMatrixVector, CudaSparseMatrixVector>, CudaLinearSolverInPlace<CudaSparseMatrixVector>, ConstraintSet<CudaDenseMatrixVector, CudaSparseMatrixVector>>;
   using CudaRosenbrock = Solver<CudaRosenbrockVectorType, GpuState>;
 
   using GpuRosenbrockThreeStageBuilder = CudaSolverBuilderInPlace<CudaRosenbrockSolverParameters>;

include/micm/constraint/constraint.hpp

@@ -2,8 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 #pragma once
 
-#include <micm/constraint/equilibrium_constraint.hpp>
-#include <micm/constraint/linear_constraint.hpp>
+#include <micm/constraint/types/equilibrium_constraint.hpp>
+#include <micm/constraint/types/linear_constraint.hpp>
+#include <micm/constraint/constraint_info.hpp>
 
 #include <cstddef>
 #include <string>
@@ -42,9 +43,16 @@ namespace micm
       return std::visit([](const auto& c) { return c.name_; }, constraint_);
     }
 
+    /// @brief Returns the species whose state row should be replaced by this algebraic constraint
+    /// @return Algebraic species name
+    const std::string& AlgebraicSpecies() const
+    {
+      return std::visit([](const auto& c) -> const std::string& { return c.AlgebraicSpecies(); }, constraint_);
+    }
+
     /// @brief Get species dependencies
     /// @return Vector of species names this constraint depends on
-    const std::vector<std::string>& GetSpeciesDependencies() const
+    const std::vector<std::string>& SpeciesDependencies() const
     {
       return std::visit([](const auto& c) -> const std::vector<std::string>& { return c.species_dependencies_; }, constraint_);
     }
@@ -56,29 +64,41 @@ namespace micm
       return std::visit([](const auto& c) { return c.species_dependencies_.size(); }, constraint_);
     }
 
-    /// @brief Evaluate the constraint residual G(y)
-    /// @param concentrations Pointer to species concentrations (row of state matrix)
-    /// @param indices Pointer to indices mapping species_dependencies_ to positions in concentrations
-    /// @return Residual value (should be 0 when constraint is satisfied)
-    double Residual(const double* concentrations, const std::size_t* indices) const
+    /// @brief Get a function object to compute the constraint residual
+    ///        This returns a reusable function that can be invoked multiple times
+    /// @param info Constraint information including species indices and row index
+    /// @param state_variable_indices Map from species names to state variable indices
+    /// @return Function object that takes (state_variables, forcing) and computes the residual
+    template<typename DenseMatrixPolicy>
+    auto ResidualFunction(const ConstraintInfo& info, const auto& state_variable_indices) const
     {
-      return std::visit([&](const auto& c) { return c.Residual(concentrations, indices); }, constraint_);
+      return std::visit(
+        [&info, &state_variable_indices](const auto& c) { 
+          return c.template ResidualFunction<DenseMatrixPolicy>(info, state_variable_indices); 
+        }, 
+        constraint_);
     }
 
-    /// @brief Compute partial derivatives dG/d[species] for each dependent species
-    /// @param concentrations Pointer to species concentrations (row of state matrix)
-    /// @param indices Pointer to indices mapping species_dependencies_ to positions in concentrations
-    /// @param jacobian Output buffer for partial derivatives dG/d[species] (same order as species_dependencies_)
-    void Jacobian(const double* concentrations, const std::size_t* indices, double* jacobian) const
+    /// @brief Get a function object to compute the constraint Jacobian
+    ///        This returns a reusable function that can be invoked multiple times
+    /// @param info Constraint information including species indices and Jacobian flat IDs
+    /// @param state_variable_indices Map from species names to state variable indices
+    /// @param jacobian_flat_ids Iterator to the jacobian flat IDs for this constraint
+    /// @param jacobian Sparse matrix to store Jacobian values
+    /// @return Function object that takes (state_variables, jacobian) and computes partials
+    template<typename DenseMatrixPolicy, typename SparseMatrixPolicy>
+    auto JacobianFunction(
+        const ConstraintInfo& info, 
+        const auto& state_variable_indices,
+        auto jacobian_flat_ids,
+        SparseMatrixPolicy& jacobian) const
     {
-      std::visit([&](const auto& c) { c.Jacobian(concentrations, indices, jacobian); }, constraint_);
-    }
-
-    /// @brief Returns the species whose state row should be replaced by this algebraic constraint
-    /// @return Algebraic species name
-    const std::string& GetAlgebraicSpecies() const
-    {
-      return std::visit([](const auto& c) -> const std::string& { return c.AlgebraicSpecies(); }, constraint_);
+      return std::visit(
+        [&info, &state_variable_indices, jacobian_flat_ids, &jacobian](const auto& c) { 
+          return c.template JacobianFunction<DenseMatrixPolicy, SparseMatrixPolicy>(
+              info, state_variable_indices, jacobian_flat_ids, jacobian); 
+        }, 
+        constraint_);
     }
   };
 

include/micm/constraint/constraint_info.hpp

@@ -0,0 +1,20 @@
+// Copyright (C) 2023-2026 University Corporation for Atmospheric Research
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <vector>
+#include <cstddef>
+
+namespace micm
+{
+/// @brief Information for each constraint (built during ConstraintSet construction)
+  struct ConstraintInfo
+  {
+    std::size_t index_;                       // Index in constraints_ vector
+    std::size_t row_index_;                   // Row in the forcing/Jacobian
+    std::size_t number_of_dependencies_;      // Number of species this constraint depends on
+    std::size_t dependency_offset_;           // Starting offset in dependency_ids_
+    std::size_t jacobian_flat_offset_;        // Starting offset in jacobian_flat_ids_
+    std::vector<std::size_t> state_indices_;  // Dependency indices in state_variables_
+  };
+}