diff --git a/include/micm/solver/external_model_dispatcher.hpp b/include/micm/solver/external_model_dispatcher.hpp index 6d0a2962e..7703cf3a4 100644 --- a/include/micm/solver/external_model_dispatcher.hpp +++ b/include/micm/solver/external_model_dispatcher.hpp @@ -7,8 +7,10 @@ #include #include +#include #include #include +#include namespace micm { @@ -16,9 +18,18 @@ namespace micm /// /// Solve-time methods first delegate to the inner (built-in) rates policy, then dispatch /// directly on each external model that satisfies `HasProcesses`. - template + /// + /// 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 RatesBundle { + template + using Vector = typename SparseMatrixPolicy::template VectorType; + + static constexpr bool HAS_EXTERNAL_PROCESSES = (HasProcesses || ...); + public: using ModelsTuple = std::tuple; @@ -44,22 +55,81 @@ namespace micm return inner_; } - template + /// @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& algebraic_variable_ids, const std::set& algebraic_jacobian_flat_ids) + { + algebraic_variable_ids_ = std::vector(algebraic_variable_ids.begin(), algebraic_variable_ids.end()); + algebraic_jacobian_flat_ids_ = + std::vector(algebraic_jacobian_flat_ids.begin(), algebraic_jacobian_flat_ids.end()); + algebraic_variable_ids_.CopyToDevice(); + algebraic_jacobian_flat_ids_.CopyToDevice(); + } + + template 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 + template 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 + template 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 models_; + Vector algebraic_variable_ids_{}; + Vector algebraic_jacobian_flat_ids_{}; }; /// @brief Wraps an inner constraint set and a shared tuple of concrete external models. diff --git a/include/micm/solver/solver_builder.inl b/include/micm/solver/solver_builder.inl index 91fc3ecfe..8baec1d40 100644 --- a/include/micm/solver/solver_builder.inl +++ b/include/micm/solver/solver_builder.inl @@ -282,7 +282,7 @@ namespace micm } using InnerConstraintSet = ConstraintSet; - using RatesBundleType = RatesBundle; + using RatesBundleType = RatesBundle; using ConstraintBundleType = ConstraintBundle; using SolverPolicy = typename SolverParametersPolicy::template SolverType; @@ -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 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, diff --git a/test/integration/kokkos/test_kokkos_external_model_constraints.cpp b/test/integration/kokkos/test_kokkos_external_model_constraints.cpp index 5723538ab..95c156694 100644 --- a/test/integration/kokkos/test_kokkos_external_model_constraints.cpp +++ b/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(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 ? 1.0e-12 : 1.0e-6; + + // Forcing + std::remove_cvref_t 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) { diff --git a/test/integration/test_external_model_constraints.cpp b/test/integration/test_external_model_constraints.cpp index 5f4c178ff..44efce518 100644 --- a/test/integration/test_external_model_constraints.cpp +++ b/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(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 ? 1.0e-12 : 1.0e-6; + + // Forcing + std::remove_cvref_t 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) {