Correct external process algebraic rows jacobian

include/micm/solver/external_model_dispatcher.hpp

@@ -7,18 +7,29 @@
 
 #include <array>
 #include <memory>
+#include <set>
 #include <tuple>
 #include <utility>
+#include <vector>
 
 namespace micm
 {
   /// @brief Wraps an inner rates policy and a shared tuple of concrete external models.
   ///
   /// Solve-time methods first delegate to the inner (built-in) rates policy, then dispatch
   /// directly on each external model that satisfies `HasProcesses`.
-  template<class InnerRates, class... ExternalModels>
+  ///
+  /// The built-in rates policy skips algebraic rows, but external models do not know which rows
+  /// are algebraic. After the external models run, the bundle sets their algebraic-row
+  /// contributions to zero, so that these rows hold only the constraint terms.
+  template<class InnerRates, class DenseMatrixPolicy, class SparseMatrixPolicy, class... ExternalModels>
   class RatesBundle
   {
+    template<class U>
+    using Vector = typename SparseMatrixPolicy::template VectorType<U>;
+
+    static constexpr bool HAS_EXTERNAL_PROCESSES = (HasProcesses<ExternalModels> || ...);
+
    public:
     using ModelsTuple = std::tuple<ExternalModels...>;
 
@@ -44,22 +55,81 @@ namespace micm
       return inner_;
     }
 
-    template<class State, class DenseMatrixPolicy>
+    /// @brief Records the algebraic rows that external process models must not contribute to.
+    /// @param algebraic_variable_ids Ids of all algebraic variables (built-in and external constraints)
+    /// @param algebraic_jacobian_flat_ids Flat ids of the external process Jacobian elements in algebraic rows
+    void SetAlgebraicRows(const std::set<Index>& algebraic_variable_ids, const std::set<Index>& algebraic_jacobian_flat_ids)
+    {
+      algebraic_variable_ids_ = std::vector<Index>(algebraic_variable_ids.begin(), algebraic_variable_ids.end());
+      algebraic_jacobian_flat_ids_ =
+          std::vector<Index>(algebraic_jacobian_flat_ids.begin(), algebraic_jacobian_flat_ids.end());
+      algebraic_variable_ids_.CopyToDevice();
+      algebraic_jacobian_flat_ids_.CopyToDevice();
+    }
+
+    template<class State>
     void AddForcingTerms(const State& state, const DenseMatrixPolicy& Y, DenseMatrixPolicy& forcing) const
     {
       inner_.AddForcingTerms(state, Y, forcing);
       InvokeProcesses([&](const auto& m) { m.AddForcingTerms(state.custom_rate_parameters_, Y, forcing); });
+      if constexpr (HAS_EXTERNAL_PROCESSES)
+      {
+        ClearAlgebraicForcing(forcing);
+      }
     }
 
-    template<class State, class DenseMatrixPolicy, class SparseMatrixPolicy>
+    template<class State>
     void SubtractJacobianTerms(const State& state, const DenseMatrixPolicy& Y, SparseMatrixPolicy& jacobian) const
     {
       inner_.SubtractJacobianTerms(state, Y, jacobian);
       InvokeProcesses([&](const auto& m) { m.SubtractJacobianTerms(state.custom_rate_parameters_, Y, jacobian); });
+      if constexpr (HAS_EXTERNAL_PROCESSES)
+      {
+        ClearAlgebraicJacobian(jacobian);
+      }
+    }
+
+    // These two methods are public because CUDA does not allow an extended __host__ __device__
+    // lambda in a private or protected member function.
+
+    /// @brief Sets the forcing of each algebraic variable to zero. The constraints set these rows later.
+    void ClearAlgebraicForcing(DenseMatrixPolicy& forcing) const
+    {
+      if (algebraic_variable_ids_.size() == 0)
+      {
+        return;
+      }
+      const auto rows = std::as_const(algebraic_variable_ids_).GetView();
+      DenseMatrixPolicy::Function(
+          MICM_LAMBDA(const typename DenseMatrixPolicy::ViewType& forcing_view) {
+            for (Index i = 0; i < rows.size(); ++i)
+            {
+              forcing_view.ForEachRow([](Real& f) { f = 0; }, forcing_view.GetColumnView(rows[i]));
+            }
+          },
+          forcing)(forcing);
+    }
+
+    /// @brief Sets the external process Jacobian elements in algebraic rows to zero.
+    void ClearAlgebraicJacobian(SparseMatrixPolicy& jacobian) const
+    {
+      if (algebraic_jacobian_flat_ids_.size() == 0)
+      {
+        return;
+      }
+      const auto flat_ids = std::as_const(algebraic_jacobian_flat_ids_).GetView();
+      SparseMatrixPolicy::Function(
+          MICM_LAMBDA(const typename SparseMatrixPolicy::ViewType& jacobian_view) {
+            for (Index i = 0; i < flat_ids.size(); ++i)
+            {
+              jacobian_view.ForEachBlock([](Real& j) { j = 0; }, jacobian_view.GetBlockView(flat_ids[i]));
+            }
+          },
+          jacobian)(jacobian);
     }
 
     /// @brief Called before each solve to refresh temperature-/pressure-dependent parameters.
-    template<class ConditionsVector, class DenseMatrixPolicy>
+    template<class ConditionsVector>
     void UpdateStateParameters(const ConditionsVector& conditions, DenseMatrixPolicy& state_parameters) const
     {
       InvokeProcesses([&](const auto& m) { m.UpdateStateParameters(conditions, state_parameters); });
@@ -106,6 +176,8 @@ namespace micm
 
     InnerRates inner_{};
     std::shared_ptr<ModelsTuple> models_;
+    Vector<Index> algebraic_variable_ids_{};
+    Vector<Index> algebraic_jacobian_flat_ids_{};
   };
 
   /// @brief Wraps an inner constraint set and a shared tuple of concrete external models.

include/micm/solver/solver_builder.inl

@@ -282,7 +282,7 @@ namespace micm
     }
 
     using InnerConstraintSet = ConstraintSet<DenseMatrixPolicy, SparseMatrixPolicy>;
-    using RatesBundleType = RatesBundle<RatesPolicy, ExternalModels...>;
+    using RatesBundleType = RatesBundle<RatesPolicy, DenseMatrixPolicy, SparseMatrixPolicy, ExternalModels...>;
     using ConstraintBundleType = ConstraintBundle<InnerConstraintSet, ExternalModels...>;
     using SolverPolicy =
         typename SolverParametersPolicy::template SolverType<RatesBundleType, LinearSolverPolicy, ConstraintBundleType>;
@@ -628,6 +628,24 @@ namespace micm
         *shared_models);
 
     RatesBundleType rates_bundle(std::move(rates), shared_models);
+
+    // The rates bundle sets the external process terms in algebraic rows to zero at solve time,
+    // so that these rows hold only the constraint terms.
+    if (!algebraic_variable_ids.empty())
+    {
+      std::set<Index> algebraic_jacobian_flat_ids;
+      for (const auto& ps : external_process_sets_)
+      {
+        for (const auto& [row, column] : ps.non_zero_jacobian_elements_func_(species_map))
+        {
+          if (algebraic_variable_ids.count(row) > 0)
+          {
+            algebraic_jacobian_flat_ids.insert(jacobian.VectorIndex(0, row, column));
+          }
+        }
+      }
+      rates_bundle.SetAlgebraicRows(algebraic_variable_ids, algebraic_jacobian_flat_ids);
+    }
     ConstraintBundleType constraint_bundle(std::move(constraint_set), shared_models, constraint_active_mask);
 
     StateParameters state_parameters = { .number_of_species_ = number_of_species,

test/integration/kokkos/test_kokkos_external_model_constraints.cpp

@@ -526,6 +526,64 @@ TEST(ExternalModelConstraints, DAESolveEnforcesConservation)
   }
 }
 
+/// @brief Algebraic rows hold only the constraint terms, not the external process terms.
+///
+/// StubAerosolWithConstraints has the process A_GAS -> A_AQ at rate k and the constraint
+/// [A_GAS] + [A_AQ] - total = 0 in the A_AQ row. The process also writes into the A_AQ row.
+/// After the rates and the constraints run in the Rosenbrock sequence, the A_AQ row must
+/// match the constraint only, and the A_GAS row must keep the process terms.
+TEST(ExternalModelConstraints, AlgebraicRowHoldsOnlyConstraintTerms)
+{
+  auto A_GAS = micm::Species("A_GAS");
+  micm::Phase gas_phase{ "gas", { A_GAS } };
+
+  const micm::Real k = 10.0;
+  const micm::Real total = 1.0;
+  StubAerosolWithConstraints aerosol(k, total);
+
+  auto options = micm::RosenbrockSolverParameters::FourStageDifferentialAlgebraicRosenbrockParameters();
+  auto solver = micm::KokkosSolverBuilder<micm::RosenbrockSolverParameters>(options)
+                    .SetSystem(micm::System(gas_phase))
+                    .SetReactions({})
+                    .SetReorderState(false)
+                    .AddExternalModel(aerosol)
+                    .Build();
+
+  auto state = solver.GetState(1);
+  const auto i_gas = state.variable_map_.at("A_GAS");
+  const auto i_aq = state.variable_map_.at("AEROSOL.A_AQ");
+  const micm::Real gas = 0.7;
+  const micm::Real aq = 0.2;
+  state.variables_[0][i_gas] = gas;
+  state.variables_[0][i_aq] = aq;
+  state.variables_.CopyToDevice();
+  state.custom_rate_parameters_.CopyToDevice();
+
+  const auto& rates = solver.solver_.rates_;
+  const auto& constraints = solver.solver_.constraints_;
+  constexpr micm::Real tol = std::is_same_v<micm::Real, double> ? 1.0e-12 : 1.0e-6;
+
+  // Forcing
+  std::remove_cvref_t<decltype(state.variables_)> forcing(1, state.state_size_, 0.0);
+  forcing.Fill(0);
+  rates.AddForcingTerms(state, state.variables_, forcing);
+  constraints.AddForcingTerms(state.variables_, state.custom_rate_parameters_, forcing);
+  forcing.CopyToHost();
+
+  EXPECT_NEAR(forcing[0][i_gas], -k * gas, tol);
+  EXPECT_NEAR(forcing[0][i_aq], gas + aq - total, tol);
+
+  // Jacobian (the solver stores -J)
+  state.jacobian_.Fill(0);
+  rates.SubtractJacobianTerms(state, state.variables_, state.jacobian_);
+  constraints.SubtractJacobianTerms(state.variables_, state.custom_rate_parameters_, state.jacobian_);
+  state.jacobian_.CopyToHost();
+
+  EXPECT_NEAR(state.jacobian_[0][i_gas][i_gas], k, tol);
+  EXPECT_NEAR(state.jacobian_[0][i_aq][i_gas], -1.0, tol) << "External process term is still in the algebraic row";
+  EXPECT_NEAR(state.jacobian_[0][i_aq][i_aq], -1.0, tol);
+}
+
 /// @brief Verify that external model constraints combine with built-in SetConstraints
 TEST(ExternalModelConstraints, CombinedBuiltInAndExternalConstraints)
 {

test/integration/test_external_model_constraints.cpp

@@ -524,6 +524,64 @@ TEST(ExternalModelConstraints, DAESolveEnforcesConservation)
   }
 }
 
+/// @brief Algebraic rows hold only the constraint terms, not the external process terms.
+///
+/// StubAerosolWithConstraints has the process A_GAS -> A_AQ at rate k and the constraint
+/// [A_GAS] + [A_AQ] - total = 0 in the A_AQ row. The process also writes into the A_AQ row.
+/// After the rates and the constraints run in the Rosenbrock sequence, the A_AQ row must
+/// match the constraint only, and the A_GAS row must keep the process terms.
+TEST(ExternalModelConstraints, AlgebraicRowHoldsOnlyConstraintTerms)
+{
+  auto A_GAS = micm::Species("A_GAS");
+  micm::Phase gas_phase{ "gas", { A_GAS } };
+
+  const micm::Real k = 10.0;
+  const micm::Real total = 1.0;
+  StubAerosolWithConstraints aerosol(k, total);
+
+  auto options = micm::RosenbrockSolverParameters::FourStageDifferentialAlgebraicRosenbrockParameters();
+  auto solver = micm::CpuSolverBuilder<micm::RosenbrockSolverParameters>(options)
+                    .SetSystem(micm::System(gas_phase))
+                    .SetReactions({})
+                    .SetReorderState(false)
+                    .AddExternalModel(aerosol)
+                    .Build();
+
+  auto state = solver.GetState(1);
+  const auto i_gas = state.variable_map_.at("A_GAS");
+  const auto i_aq = state.variable_map_.at("AEROSOL.A_AQ");
+  const micm::Real gas = 0.7;
+  const micm::Real aq = 0.2;
+  state.variables_[0][i_gas] = gas;
+  state.variables_[0][i_aq] = aq;
+  state.variables_.CopyToDevice();
+  state.custom_rate_parameters_.CopyToDevice();
+
+  const auto& rates = solver.solver_.rates_;
+  const auto& constraints = solver.solver_.constraints_;
+  constexpr micm::Real tol = std::is_same_v<micm::Real, double> ? 1.0e-12 : 1.0e-6;
+
+  // Forcing
+  std::remove_cvref_t<decltype(state.variables_)> forcing(1, state.state_size_, 0.0);
+  forcing.Fill(0);
+  rates.AddForcingTerms(state, state.variables_, forcing);
+  constraints.AddForcingTerms(state.variables_, state.custom_rate_parameters_, forcing);
+  forcing.CopyToHost();
+
+  EXPECT_NEAR(forcing[0][i_gas], -k * gas, tol);
+  EXPECT_NEAR(forcing[0][i_aq], gas + aq - total, tol);
+
+  // Jacobian (the solver stores -J)
+  state.jacobian_.Fill(0);
+  rates.SubtractJacobianTerms(state, state.variables_, state.jacobian_);
+  constraints.SubtractJacobianTerms(state.variables_, state.custom_rate_parameters_, state.jacobian_);
+  state.jacobian_.CopyToHost();
+
+  EXPECT_NEAR(state.jacobian_[0][i_gas][i_gas], k, tol);
+  EXPECT_NEAR(state.jacobian_[0][i_aq][i_gas], -1.0, tol) << "External process term is still in the algebraic row";
+  EXPECT_NEAR(state.jacobian_[0][i_aq][i_aq], -1.0, tol);
+}
+
 /// @brief Verify that external model constraints combine with built-in SetConstraints
 TEST(ExternalModelConstraints, CombinedBuiltInAndExternalConstraints)
 {