diff --git a/include/micm/CPU.hpp b/include/micm/CPU.hpp index f57a1f9e8..334ab7aff 100644 --- a/include/micm/CPU.hpp +++ b/include/micm/CPU.hpp @@ -22,19 +22,19 @@ namespace micm using StandardState = State; using RosenbrockVectorType = typename RosenbrockSolverParameters:: - template SolverType, LinearSolver>; + template SolverType, LinearSolver, ConstraintSet>; using Rosenbrock = Solver>; using RosenbrockStandardType = typename RosenbrockSolverParameters:: - template SolverType, LinearSolver>; + template SolverType, LinearSolver, ConstraintSet>; using RosenbrockStandard = Solver>; using BackwardEulerVectorType = typename BackwardEulerSolverParameters:: - template SolverType, LinearSolver>; + template SolverType, LinearSolver, ConstraintSet>; using BackwardEuler = Solver>; using BackwardEulerStandardType = typename BackwardEulerSolverParameters:: - template SolverType, LinearSolver>; + template SolverType, LinearSolver, ConstraintSet>; using BackwardEulerStandard = Solver>; using RosenbrockThreeStageBuilder = CpuSolverBuilder; diff --git a/include/micm/Constraint.hpp b/include/micm/Constraint.hpp index c957f2e1b..dd20b82b5 100644 --- a/include/micm/Constraint.hpp +++ b/include/micm/Constraint.hpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#include +#include #include -#include -#include +#include \ No newline at end of file diff --git a/include/micm/GPU.hpp b/include/micm/GPU.hpp index e0402ed91..6f45ce61e 100644 --- a/include/micm/GPU.hpp +++ b/include/micm/GPU.hpp @@ -20,6 +20,7 @@ #include #include #include +#include #include namespace micm @@ -30,7 +31,7 @@ namespace micm using GpuState = CudaState; using CudaRosenbrockVectorType = typename CudaRosenbrockSolverParameters:: - template SolverType, CudaLinearSolverInPlace>; + template SolverType, CudaLinearSolverInPlace, ConstraintSet>; using CudaRosenbrock = Solver; using GpuRosenbrockThreeStageBuilder = CudaSolverBuilderInPlace; diff --git a/include/micm/constraint/constraint.hpp b/include/micm/constraint/constraint.hpp index baec3a1c4..35065907f 100644 --- a/include/micm/constraint/constraint.hpp +++ b/include/micm/constraint/constraint.hpp @@ -2,8 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once -#include -#include +#include +#include +#include #include #include @@ -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& GetSpeciesDependencies() const + const std::vector& SpeciesDependencies() const { return std::visit([](const auto& c) -> const std::vector& { 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 + 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(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 + 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( + info, state_variable_indices, jacobian_flat_ids, jacobian); + }, + constraint_); } }; diff --git a/include/micm/constraint/constraint_info.hpp b/include/micm/constraint/constraint_info.hpp new file mode 100644 index 000000000..5edae4091 --- /dev/null +++ b/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 +#include + +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 state_indices_; // Dependency indices in state_variables_ + }; +} \ No newline at end of file diff --git a/include/micm/constraint/constraint_set.hpp b/include/micm/constraint/constraint_set.hpp index 3201122d8..faf7ac2e4 100644 --- a/include/micm/constraint/constraint_set.hpp +++ b/include/micm/constraint/constraint_set.hpp @@ -3,11 +3,13 @@ #pragma once #include +#include #include #include #include #include +#include #include #include #include @@ -22,36 +24,30 @@ namespace micm /// ConstraintSet handles the computation of constraint residuals (forcing terms) /// and Jacobian contributions for a set of constraints. It follows the same /// pattern as ProcessSet for integration with the Rosenbrock solver. + template class ConstraintSet { - protected: - /// @brief Information for each constraint - struct ConstraintInfo - { - std::size_t constraint_index_; // Index in constraints_ vector - std::size_t constraint_row_; // 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_ - }; - + private: /// @brief The constraints std::vector constraints_; - /// @brief Flat list of species indices for each constraint's dependencies - std::vector dependency_ids_; - /// @brief Information about each constraint for forcing/Jacobian computation std::vector constraint_info_; + /// @brief Flat list of species indices for each constraint's dependencies + std::vector dependency_ids_; + /// @brief Flat indices into the Jacobian sparse matrix for each constraint's Jacobian entries std::vector jacobian_flat_ids_; /// @brief Species variable ids whose ODE rows are replaced by constraints std::set algebraic_variable_ids_; - /// @brief Maximum number of dependencies across all constraints (for buffer allocation) - std::size_t max_dependencies_{ 0 }; + /// @brief Pre-compiled constraint residual functions (initialized during solver build via SetConstraintFunctions) + std::vector> constraint_forcing_functions_; + + /// @brief Pre-compiled constraint Jacobian functions (initialized during solver build via SetConstraintFunctions) + std::vector> constraint_jacobian_functions_; public: /// @brief Default constructor @@ -63,12 +59,68 @@ namespace micm /// @param variable_map Map from species names to state variable indices ConstraintSet( std::vector&& constraints, - const std::unordered_map& variable_map); + const std::unordered_map& variable_map) + : constraints_(std::move(constraints)) + { + // Build constraint info and dependency indices + std::size_t dependency_offset = 0; + for (std::size_t i = 0; i < constraints_.size(); ++i) + { + const auto& constraint = constraints_[i]; + + ConstraintInfo info; + info.index_ = i; - /// @brief Move constructor + const auto& algebraic_species = constraint.AlgebraicSpecies(); + auto row_it = variable_map.find(algebraic_species); + if (row_it == variable_map.end()) + { + throw MicmException( + MicmSeverity::Error, + MICM_ERROR_CATEGORY_CONSTRAINT, + MICM_CONSTRAINT_ERROR_CODE_UNKNOWN_SPECIES, + "Constraint '" + constraint.GetName() + "' targets unknown algebraic species '" + algebraic_species + "'"); + } + info.row_index_ = row_it->second; + + if (!algebraic_variable_ids_.insert(info.row_index_).second) + { + throw MicmException( + MicmSeverity::Error, + MICM_ERROR_CATEGORY_CONSTRAINT, + MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY, + "Multiple constraints map to the same algebraic species row '" + algebraic_species + "'"); + } + + info.number_of_dependencies_ = constraint.NumberOfDependencies(); + info.dependency_offset_ = dependency_offset; + info.jacobian_flat_offset_ = 0; // Set later in SetJacobianFlatIds + + // Map species dependencies to variable indices + for (const auto& species_name : constraint.SpeciesDependencies()) + { + auto it = variable_map.find(species_name); + if (it == variable_map.end()) + { + throw MicmException( + MicmSeverity::Error, + MICM_ERROR_CATEGORY_CONSTRAINT, + MICM_CONSTRAINT_ERROR_CODE_UNKNOWN_SPECIES, + "Constraint '" + constraint.GetName() + "' depends on unknown species '" + species_name + "'"); + } + dependency_ids_.push_back(it->second); + info.state_indices_.push_back(it->second); // Also store in ConstraintInfo + } + + dependency_offset += info.number_of_dependencies_; + constraint_info_.push_back(info); + } + } + + /// @brief Move constructor - default implementation ConstraintSet(ConstraintSet&& other) noexcept = default; - /// @brief Move assignment + /// @brief Move assignment operator ConstraintSet& operator=(ConstraintSet&& other) noexcept = default; /// @brief Copy constructor @@ -83,10 +135,6 @@ namespace micm return constraints_.size(); } - /// @brief Returns positions of all non-zero Jacobian elements for constraint rows - /// @return Set of (row, column) index pairs - std::set> NonZeroJacobianElements() const; - /// @brief Returns species ids whose rows are algebraic when constraints replace state rows /// @return Set of variable ids for algebraic rows const std::set& AlgebraicVariableIds() const @@ -94,248 +142,96 @@ namespace micm return algebraic_variable_ids_; } - /// @brief Computes and stores flat indices for Jacobian elements - /// @param matrix The sparse Jacobian matrix - template - void SetJacobianFlatIds(const SparseMatrix& matrix); - /// @brief Add constraint residuals to forcing vector (constraint rows) /// For each constraint G_i, writes or adds G_i(x) to forcing[constraint_row] /// @param state_variables Current species concentrations (grid cell, species) /// @param forcing Forcing terms (grid cell, state variable) - constraint rows will be modified - template - void AddForcingTerms(const DenseMatrixPolicy& state_variables, DenseMatrixPolicy& forcing) const; + void AddForcingTerms(const DenseMatrixPolicy& state_variables, DenseMatrixPolicy& forcing) const + { + for (const auto& forcing_fn : constraint_forcing_functions_) + forcing_fn(state_variables, forcing); + } /// @brief Subtract constraint Jacobian terms from Jacobian matrix /// For each constraint G_i, subtracts dG_i/dx_j from jacobian[constraint_row, j] /// (Subtraction matches the convention used by ProcessSet) /// @param state_variables Current species concentrations (grid cell, species) /// @param jacobian Sparse Jacobian matrix (grid cell, row, column) - template - void SubtractJacobianTerms(const DenseMatrixPolicy& state_variables, SparseMatrixPolicy& jacobian) const; - }; - - inline ConstraintSet::ConstraintSet( - std::vector&& constraints, - const std::unordered_map& variable_map) - : constraints_(std::move(constraints)) - { - // Build constraint info and dependency indices - std::size_t dependency_offset = 0; - for (std::size_t i = 0; i < constraints_.size(); ++i) + void SubtractJacobianTerms(const DenseMatrixPolicy& state_variables, SparseMatrixPolicy& jacobian) const { - const auto& constraint = constraints_[i]; - - ConstraintInfo info; - info.constraint_index_ = i; - - const auto& algebraic_species = constraint.GetAlgebraicSpecies(); - auto row_it = variable_map.find(algebraic_species); - if (row_it == variable_map.end()) - { - throw MicmException( - MicmSeverity::Error, - MICM_ERROR_CATEGORY_CONSTRAINT, - MICM_CONSTRAINT_ERROR_CODE_UNKNOWN_SPECIES, - "Constraint '" + constraint.GetName() + "' targets unknown algebraic species '" + algebraic_species + "'"); - } - info.constraint_row_ = row_it->second; - - if (!algebraic_variable_ids_.insert(info.constraint_row_).second) - { - throw MicmException( - MicmSeverity::Error, - MICM_ERROR_CATEGORY_CONSTRAINT, - MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY, - "Multiple constraints map to the same algebraic species row '" + algebraic_species + "'"); - } - - info.number_of_dependencies_ = constraint.NumberOfDependencies(); - info.dependency_offset_ = dependency_offset; - info.jacobian_flat_offset_ = 0; // Set later in SetJacobianFlatIds + for (const auto& jacobian_fn : constraint_jacobian_functions_) + jacobian_fn(state_variables, jacobian); + } - // Track maximum dependencies for buffer allocation - if (info.number_of_dependencies_ > max_dependencies_) - { - max_dependencies_ = info.number_of_dependencies_; - } + /// @brief Returns positions of all non-zero Jacobian elements for constraint rows + /// @return Set of (row, column) index pairs + std::set> NonZeroJacobianElements() const + { + std::set> ids; - // Map species dependencies to variable indices - for (const auto& species_name : constraint.GetSpeciesDependencies()) + auto dep_id = dependency_ids_.begin(); + for (const auto& info : constraint_info_) { - auto it = variable_map.find(species_name); - if (it == variable_map.end()) + // Ensure the diagonal element exists for the constraint row (required by AlphaMinusJacobian and LU decomposition) + ids.insert(std::make_pair(info.row_index_, info.row_index_)); + // Each constraint contributes Jacobian entries at (constraint_row, dependency_column) + for (std::size_t i = 0; i < info.number_of_dependencies_; ++i) { - throw MicmException( - MicmSeverity::Error, - MICM_ERROR_CATEGORY_CONSTRAINT, - MICM_CONSTRAINT_ERROR_CODE_UNKNOWN_SPECIES, - "Constraint '" + constraint.GetName() + "' depends on unknown species '" + species_name + "'"); + ids.insert(std::make_pair(info.row_index_, dep_id[i])); } - dependency_ids_.push_back(it->second); + dep_id += info.number_of_dependencies_; } - dependency_offset += info.number_of_dependencies_; - constraint_info_.push_back(info); + return ids; } - } - - inline std::set> ConstraintSet::NonZeroJacobianElements() const - { - std::set> ids; - - auto dep_id = dependency_ids_.begin(); - for (const auto& info : constraint_info_) - { - // Ensure the diagonal element exists for the constraint row (required by AlphaMinusJacobian and LU decomposition) - ids.insert(std::make_pair(info.constraint_row_, info.constraint_row_)); - // Each constraint contributes Jacobian entries at (constraint_row, dependency_column) - for (std::size_t i = 0; i < info.number_of_dependencies_; ++i) - { - ids.insert(std::make_pair(info.constraint_row_, dep_id[i])); - } - dep_id += info.number_of_dependencies_; - } - - return ids; - } - template - inline void ConstraintSet::SetJacobianFlatIds(const SparseMatrix& matrix) - { - jacobian_flat_ids_.clear(); - - std::size_t flat_offset = 0; - for (auto& info : constraint_info_) + /// @brief Computes and stores flat indices for Jacobian elements + /// @param matrix The sparse Jacobian matrix + template + void SetJacobianFlatIds(const SparseMatrix& matrix) { - info.jacobian_flat_offset_ = flat_offset; + jacobian_flat_ids_.clear(); - // Store flat indices for each dependency of this constraint - const std::size_t* dep_id = dependency_ids_.data() + info.dependency_offset_; - for (std::size_t i = 0; i < info.number_of_dependencies_; ++i) + std::size_t flat_offset = 0; + for (auto& info : constraint_info_) { - jacobian_flat_ids_.push_back(matrix.VectorIndex(0, info.constraint_row_, dep_id[i])); - } - - flat_offset += info.number_of_dependencies_; - } - } + info.jacobian_flat_offset_ = flat_offset; - template - inline void ConstraintSet::AddForcingTerms( - const DenseMatrixPolicy& state_variables, - DenseMatrixPolicy& forcing) const - { - if constexpr (VectorizableDense) - { - // Vectorized dense layouts are not row-contiguous, so build a contiguous row buffer. - std::vector concentrations(state_variables.NumColumns()); - for (std::size_t i_cell = 0; i_cell < state_variables.NumRows(); ++i_cell) - { - for (std::size_t i_var = 0; i_var < state_variables.NumColumns(); ++i_var) + // Store flat indices for each dependency of this constraint + const std::size_t* dep_id = dependency_ids_.data() + info.dependency_offset_; + for (std::size_t i = 0; i < info.number_of_dependencies_; ++i) { - concentrations[i_var] = state_variables[i_cell][i_var]; + jacobian_flat_ids_.push_back(matrix.VectorIndex(0, info.row_index_, dep_id[i])); } - auto cell_forcing = forcing[i_cell]; - for (const auto& info : constraint_info_) - { - const std::size_t* indices = dependency_ids_.data() + info.dependency_offset_; - double residual = constraints_[info.constraint_index_].Residual(concentrations.data(), indices); - cell_forcing[info.constraint_row_] = residual; - } + flat_offset += info.number_of_dependencies_; } } - else - { - // Loop over grid cells - for (std::size_t i_cell = 0; i_cell < state_variables.NumRows(); ++i_cell) - { - auto cell_state = state_variables[i_cell]; - auto cell_forcing = forcing[i_cell]; - - // Get pointer to concentration data for this cell - const double* concentrations = &cell_state[0]; - for (const auto& info : constraint_info_) - { - // Get pointer to indices for this constraint - const std::size_t* indices = dependency_ids_.data() + info.dependency_offset_; - - // Evaluate constraint residual: replaces the algebraic species row - double residual = constraints_[info.constraint_index_].Residual(concentrations, indices); - cell_forcing[info.constraint_row_] = residual; - } - } - } - } - - template - inline void ConstraintSet::SubtractJacobianTerms( - const DenseMatrixPolicy& state_variables, - SparseMatrixPolicy& jacobian) const - { - // Allocate reusable buffer for constraint Jacobian values - std::vector jac_buffer(max_dependencies_); - - [[maybe_unused]] std::vector concentrations; - if constexpr (VectorizableDense) + /// @brief Pre-compiles constraint residual and Jacobian functions for efficient evaluation + /// Creates reusable function objects from each constraint's ResidualFunction and JacobianFunction. + /// Must be called after SetJacobianFlatIds and before solver execution. + /// @param state_variable_indices Map from species names to state variable indices + /// @param jacobian The sparse Jacobian matrix (used for function template instantiation) + void SetConstraintFunctions( + const auto& state_variable_indices, // acts like std::unordered_map + SparseMatrixPolicy& jacobian) { - concentrations.resize(state_variables.NumColumns()); - } - - [[maybe_unused]] auto cell_jacobian = jacobian.AsVector().begin(); - - // Loop over grid cells - for (std::size_t i_cell = 0; i_cell < state_variables.NumRows(); ++i_cell) - { - const double* concentration_ptr; - if constexpr (VectorizableDense) - { - // Vectorized dense layouts are not row-contiguous, so build a contiguous row buffer. - for (std::size_t i_var = 0; i_var < state_variables.NumColumns(); ++i_var) - { - concentrations[i_var] = state_variables[i_cell][i_var]; - } - concentration_ptr = concentrations.data(); - } - else - { - auto cell_state = state_variables[i_cell]; - concentration_ptr = &cell_state[0]; - } - + constraint_forcing_functions_.clear(); + constraint_jacobian_functions_.clear(); for (const auto& info : constraint_info_) { - const std::size_t* indices = dependency_ids_.data() + info.dependency_offset_; - constraints_[info.constraint_index_].Jacobian(concentration_ptr, indices, jac_buffer.data()); - - // In replace mode, the Jacobian row is already zeroed (Fill(0) + ProcessSet skips algebraic rows), - // so -= is safe and correctly accumulates when a species appears in multiple dependency slots. - if constexpr (VectorizableSparse) - { - const std::size_t* dep_id = dependency_ids_.data() + info.dependency_offset_; - for (std::size_t i = 0; i < info.number_of_dependencies_; ++i) - { - jacobian[i_cell][info.constraint_row_][dep_id[i]] -= jac_buffer[i]; - } - } - else - { - const std::size_t* flat_ids = jacobian_flat_ids_.data() + info.jacobian_flat_offset_; - for (std::size_t i = 0; i < info.number_of_dependencies_; ++i) - { - cell_jacobian[flat_ids[i]] -= jac_buffer[i]; - } - } - } - - if constexpr (!VectorizableSparse) - { - // Advance to next grid cell's Jacobian block - cell_jacobian += jacobian.FlatBlockSize(); + constraint_forcing_functions_.push_back( + constraints_[info.index_].template ResidualFunction(info, state_variable_indices)); + + constraint_jacobian_functions_.push_back( + constraints_[info.index_].template JacobianFunction( + info, + state_variable_indices, + jacobian_flat_ids_.begin() + info.jacobian_flat_offset_, + jacobian)); } } - } + }; } // namespace micm \ No newline at end of file diff --git a/include/micm/constraint/equilibrium_constraint.hpp b/include/micm/constraint/equilibrium_constraint.hpp deleted file mode 100644 index 14d36423e..000000000 --- a/include/micm/constraint/equilibrium_constraint.hpp +++ /dev/null @@ -1,239 +0,0 @@ -// Copyright (C) 2023-2026 University Corporation for Atmospheric Research -// SPDX-License-Identifier: Apache-2.0 -#pragma once - -#include -#include - -#include -#include -#include -#include - -namespace micm -{ - - /// @brief Constraint for chemical equilibrium: K_eq = [products]^stoich / [reactants]^stoich - /// For a reversible reaction: aA + bB <-> cC + dD - /// The equilibrium constraint is: G = K_eq * [A]^a * [B]^b - [C]^c * [D]^d = 0 - /// This can also be written in terms of forward/backward rate constants: - /// G = k_f * [A]^a * [B]^b - k_b * [C]^c * [D]^d = 0 - /// where K_eq = k_f / k_b - class EquilibriumConstraint - { - public: - /// @brief Name of the constraint (for identification) - std::string name_; - - /// @brief Names of species this constraint depends on - std::vector species_dependencies_; - - /// @brief Reactant species and their stoichiometric coefficients - std::vector reactants_; - - /// @brief Product species and their stoichiometric coefficients - std::vector products_; - - /// @brief Equilibrium constant K_eq = k_forward / k_backward - double equilibrium_constant_; - - private: - /// @brief Indices into the reactants_ vector for each species dependency - std::vector reactant_dependency_indices_; - - /// @brief Indices into the products_ vector for each species dependency - std::vector product_dependency_indices_; - - public: - /// @brief Default constructor - EquilibriumConstraint() = default; - - /// @brief Construct an equilibrium constraint - /// Validates that equilibrium constraint > 0 - /// Builds species_dependencies_ by concatenating reactants then products - /// Stores index mappings for efficient Jacobian computation - /// @param name Constraint identifier - /// @param reactants Vector of StoichSpecies (species, stoichiometry) for reactants - /// @param products Vector of StoichSpecies (species, stoichiometry) for products - /// @param equilibrium_constant K_eq = [products]/[reactants] at equilibrium - EquilibriumConstraint( - const std::string& name, - const std::vector& reactants, - const std::vector& products, - double equilibrium_constant) - : name_(name), - reactants_(reactants), - products_(products), - equilibrium_constant_(equilibrium_constant) - { - if (reactants_.empty()) - { - throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_EMPTY_REACTANTS, "Equilibrium constraint requires at least one reactant"); - } - if (products_.empty()) - { - throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_EMPTY_PRODUCTS, "Equilibrium constraint requires at least one product"); - } - if (equilibrium_constant_ <= 0) - { - throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_INVALID_EQUILIBRIUM_CONSTANT, "Equilibrium constant must be positive"); - } - for (const auto& r : reactants_) - { - if (r.coefficient_ <= 0) - { - throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY, "Stoichiometric coefficients must be positive"); - } - } - for (const auto& p : products_) - { - if (p.coefficient_ <= 0) - { - throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY, "Stoichiometric coefficients must be positive"); - } - } - - // Build species dependencies list (reactants first, then products) - std::size_t idx = 0; - for (const auto& r : reactants_) - { - species_dependencies_.push_back(r.species_.name_); - reactant_dependency_indices_.push_back(idx++); - } - for (const auto& p : products_) - { - species_dependencies_.push_back(p.species_.name_); - product_dependency_indices_.push_back(idx++); - } - } - - /// @brief Evaluate the equilibrium constraint residual - /// G = K_eq * prod([reactants]^stoich) - prod([products]^stoich) - /// At equilibrium, G = 0 - /// @param concentrations Pointer to species concentrations (row of state matrix) - /// @param indices Pointer to indices mapping species_dependencies_ to concentrations - /// @return Residual value - double Residual(const double* concentrations, const std::size_t* indices) const - { - // Compute product of reactant concentrations raised to stoichiometric powers - // Guard against negative concentrations (transient solver artifacts) to avoid NaN from std::pow - double reactant_product = 1.0; - for (std::size_t i = 0; i < reactants_.size(); ++i) - { - double conc = std::max(0.0, concentrations[indices[reactant_dependency_indices_[i]]]); - reactant_product *= std::pow(conc, reactants_[i].coefficient_); - } - - // Compute product of product concentrations raised to stoichiometric powers - double product_product = 1.0; - for (std::size_t i = 0; i < products_.size(); ++i) - { - double conc = std::max(0.0, concentrations[indices[product_dependency_indices_[i]]]); - product_product *= std::pow(conc, products_[i].coefficient_); - } - - // G = K_eq * [reactants] - [products] - return equilibrium_constant_ * reactant_product - product_product; - } - - /// @brief Compute Jacobian entries dG/d[species] - /// For reactant R with stoichiometry n: - /// dG/d[R] = K_eq * n * [R]^(n-1) * prod([other_reactants]^stoich) - /// For product P with stoichiometry m: - /// dG/d[P] = -m * [P]^(m-1) * prod([other_products]^stoich) - /// @param concentrations Pointer to species concentrations (row of state matrix) - /// @param indices Pointer to indices mapping species_dependencies_ to concentrations - /// @param jacobian Output buffer for partial derivatives (same order as species_dependencies_) - void Jacobian(const double* concentrations, const std::size_t* indices, double* jacobian) const - { - // Initialize jacobian entries to zero - for (std::size_t i = 0; i < species_dependencies_.size(); ++i) - { - jacobian[i] = 0.0; - } - - // Compute full reactant and product terms - // Guard against negative concentrations (consistent with Residual policy) - double reactant_product = 1.0; - for (std::size_t i = 0; i < reactants_.size(); ++i) - { - double conc = std::max(0.0, concentrations[indices[reactant_dependency_indices_[i]]]); - reactant_product *= std::pow(conc, reactants_[i].coefficient_); - } - - double product_product = 1.0; - for (std::size_t i = 0; i < products_.size(); ++i) - { - double conc = std::max(0.0, concentrations[indices[product_dependency_indices_[i]]]); - product_product *= std::pow(conc, products_[i].coefficient_); - } - - // Jacobian for reactants: dG/d[R] = K_eq * n * [R]^(n-1) * prod([others]) - for (std::size_t i = 0; i < reactants_.size(); ++i) - { - double conc = std::max(0.0, concentrations[indices[reactant_dependency_indices_[i]]]); - double stoich = reactants_[i].coefficient_; - - if (conc > 0) - { - // dG/d[R] = K_eq * stoich * reactant_product / [R] - jacobian[reactant_dependency_indices_[i]] = equilibrium_constant_ * stoich * reactant_product / conc; - } - else if (stoich == 1.0) - { - // Special case: if conc = 0 and stoich = 1, derivative is K_eq * prod(others) - double others = equilibrium_constant_; - for (std::size_t j = 0; j < reactants_.size(); ++j) - { - if (j != i) - { - double other_conc = std::max(0.0, concentrations[indices[reactant_dependency_indices_[j]]]); - others *= std::pow(other_conc, reactants_[j].coefficient_); - } - } - jacobian[reactant_dependency_indices_[i]] = others; - } - // else: derivative is 0 when conc <= 0 and stoich > 1 - } - - // Jacobian for products: dG/d[P] = -m * [P]^(m-1) * prod([others]) - for (std::size_t i = 0; i < products_.size(); ++i) - { - double conc = std::max(0.0, concentrations[indices[product_dependency_indices_[i]]]); - double stoich = products_[i].coefficient_; - - if (conc > 0) - { - // dG/d[P] = -stoich * product_product / [P] - jacobian[product_dependency_indices_[i]] = -stoich * product_product / conc; - } - else if (stoich == 1.0) - { - // Special case: if conc = 0 and stoich = 1, derivative is -prod(others) - double others = -1.0; - for (std::size_t j = 0; j < products_.size(); ++j) - { - if (j != i) - { - double other_conc = std::max(0.0, concentrations[indices[product_dependency_indices_[j]]]); - others *= std::pow(other_conc, products_[j].coefficient_); - } - } - jacobian[product_dependency_indices_[i]] = others; - } - // else: derivative is 0 when conc <= 0 and stoich > 1 - } - } - - /// @brief Returns the species whose row should be replaced by this algebraic constraint - /// @return Species name of the primary algebraic variable - /// - /// For equilibrium constraints, we use the first product species as the algebraic row target. - /// This supports common forms such as K_eq * [B] - [C] = 0 where C is algebraic. - const std::string& AlgebraicSpecies() const - { - return products_[0].species_.name_; - } - }; - -} // namespace micm diff --git a/include/micm/constraint/linear_constraint.hpp b/include/micm/constraint/linear_constraint.hpp deleted file mode 100644 index 0ff38540e..000000000 --- a/include/micm/constraint/linear_constraint.hpp +++ /dev/null @@ -1,100 +0,0 @@ -// Copyright (C) 2023-2026 University Corporation for Atmospheric Research -// SPDX-License-Identifier: Apache-2.0 -#pragma once - -#include -#include - -#include -#include -#include - -namespace micm -{ - - /// @brief Constraint for linear relationships: sum(coeff[i] * [species[i]]) = constant - /// For example: A + B + C = 1.0 represents a conservation law - /// The linear constraint is: G = c1*[A] + c2*[B] + c3*[C] - constant = 0 - class LinearConstraint - { - public: - /// @brief Name of the constraint - std::string name_; - - /// @brief Names of species this constraint depends on - std::vector species_dependencies_; - - /// @brief Species and their coefficients in the linear sum - std::vector terms_; - - /// @brief The constant value the linear sum should equal - double constant_; - - public: - /// @brief Default constructor - LinearConstraint() = default; - - /// @brief Construct a linear constraint - /// Validates that terms are non-empty - /// Builds species_dependencies_ from terms - /// @param name Constraint identifier - /// @param terms Vector of StoichSpecies (species, coefficient) in the linear sum - /// @param constant The value that sum(coeff[i] * [species[i]]) should equal - LinearConstraint( - const std::string& name, - const std::vector& terms, - double constant) - : name_(name), - terms_(terms), - constant_(constant) - { - if (terms_.empty()) - { - throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_EMPTY_REACTANTS, ""); - } - for (const auto& term : terms_) - { - species_dependencies_.push_back(term.species_.name_); - } - } - - /// @brief Evaluate the linear constraint residual - /// G = sum(coeff[i] * [species[i]]) - constant - /// When satisfied, G = 0 - /// @param concentrations Pointer to species concentrations (row of state matrix) - /// @param indices Pointer to indices mapping species_dependencies_ to concentrations - /// @return Residual value - double Residual(const double* concentrations, const std::size_t* indices) const - { - double sum = 0.0; - for (std::size_t i = 0; i < terms_.size(); ++i) - { - sum += terms_[i].coefficient_ * concentrations[indices[i]]; - } - return sum - constant_; - } - - /// @brief Compute Jacobian entries dG/d[species] - /// For a linear constraint, the Jacobian is simply the coefficients: - /// dG/d[species[i]] = coeff[i] - /// @param concentrations Pointer to species concentrations (row of state matrix) - /// @param indices Pointer to indices mapping species_dependencies_ to concentrations - /// @param jacobian Output buffer for partial derivatives (same order as species_dependencies_) - void Jacobian(const double* concentrations, const std::size_t* indices, double* jacobian) const - { - for (std::size_t i = 0; i < terms_.size(); ++i) - { - jacobian[i] = terms_[i].coefficient_; - } - } - - /// @brief Returns the species whose row should be replaced by this algebraic constraint - /// For linear constraints, the last species is used in the terms list as the algebraic row target. - /// @return Species name of the primary algebraic variable - const std::string& AlgebraicSpecies() const - { - return terms_.back().species_.name_; - } - }; - -} // namespace micm diff --git a/include/micm/constraint/types/equilibrium_constraint.hpp b/include/micm/constraint/types/equilibrium_constraint.hpp new file mode 100644 index 000000000..6cfe9980f --- /dev/null +++ b/include/micm/constraint/types/equilibrium_constraint.hpp @@ -0,0 +1,415 @@ +// Copyright (C) 2023-2026 University Corporation for Atmospheric Research +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace micm +{ + /// @brief Constraint for chemical equilibrium: K_eq = [products]^stoich / [reactants]^stoich + /// For a reversible reaction: aA + bB <-> cC + dD + /// The equilibrium constraint is: G = K_eq * [A]^a * [B]^b - [C]^c * [D]^d = 0 + /// This can also be written in terms of forward/backward rate constants: + /// G = k_f * [A]^a * [B]^b - k_b * [C]^c * [D]^d = 0 + /// where K_eq = k_f / k_b + class EquilibriumConstraint + { + public: + /// @brief Name of the constraint (for identification) + std::string name_; + + /// @brief Names of species this constraint depends on + std::vector species_dependencies_; + + /// @brief Reactant species and their stoichiometric coefficients + std::vector reactants_; + + /// @brief Product species and their stoichiometric coefficients + std::vector products_; + + /// @brief Equilibrium constant K_eq = k_forward / k_backward + double equilibrium_constant_; + + private: + /// @brief Indices into the reactants_ vector for each species dependency + std::vector reactant_dependency_indices_; + + /// @brief Indices into the products_ vector for each species dependency + std::vector product_dependency_indices_; + + public: + /// @brief Default constructor + EquilibriumConstraint() = default; + + /// @brief Construct an equilibrium constraint + /// Validates that equilibrium constraint > 0 + /// Builds species_dependencies_ by concatenating reactants then products + /// Stores index mappings for efficient Jacobian computation + /// @param name Constraint identifier + /// @param reactants Vector of StoichSpecies (species, stoichiometry) for reactants + /// @param products Vector of StoichSpecies (species, stoichiometry) for products + /// @param equilibrium_constant K_eq = [products]/[reactants] at equilibrium + EquilibriumConstraint( + const std::string& name, + const std::vector& reactants, + const std::vector& products, + double equilibrium_constant) + : name_(name), + reactants_(reactants), + products_(products), + equilibrium_constant_(equilibrium_constant) + { + if (reactants_.empty()) + { + throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_EMPTY_REACTANTS, "Equilibrium constraint requires at least one reactant"); + } + if (products_.empty()) + { + throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_EMPTY_PRODUCTS, "Equilibrium constraint requires at least one product"); + } + if (equilibrium_constant_ <= 0) + { + throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_INVALID_EQUILIBRIUM_CONSTANT, "Equilibrium constant must be positive"); + } + for (const auto& r : reactants_) + { + if (r.coefficient_ <= 0) + { + throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY, "Stoichiometric coefficients must be positive"); + } + } + for (const auto& p : products_) + { + if (p.coefficient_ <= 0) + { + throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY, "Stoichiometric coefficients must be positive"); + } + } + + // Build species dependencies list (reactants first, then products) + std::size_t idx = 0; + for (const auto& r : reactants_) + { + species_dependencies_.push_back(r.species_.name_); + reactant_dependency_indices_.push_back(idx++); + } + for (const auto& p : products_) + { + species_dependencies_.push_back(p.species_.name_); + product_dependency_indices_.push_back(idx++); + } + } + + /// @brief Returns the species whose row should be replaced by this algebraic constraint + /// @return Species name of the primary algebraic variable + /// + /// For equilibrium constraints, we use the first product species as the algebraic row target. + /// This supports common forms such as K_eq * [B] - [C] = 0 where C is algebraic. + const std::string& AlgebraicSpecies() const + { + return products_[0].species_.name_; + } + + /// @brief Create function object to compute constraint residual G = K_eq * [reactants]^stoich - [products]^stoich + /// Called during solver build (SetConstraintFunctions) to pre-compile residual computation + /// @param info Constraint information including row index and species indices + /// @param state_variable_indices Mapping of state variable names to indices + /// @return Function object that takes (state, forcing) and writes G to forcing[constraint_row] + template + std::function + ResidualFunction( + const ConstraintInfo& info, + const auto& state_variable_indices) const + { + DenseMatrixPolicy temp_state_variables{1, state_variable_indices.size(), 0.0}; + + // Copy data to avoid issues when ConstraintSet is moved + double K_eq = this->equilibrium_constant_; + std::vector reactant_stoich; + std::vector reactant_state_idx; + for (std::size_t i = 0; i < this->reactants_.size(); ++i) + { + reactant_stoich.push_back(this->reactants_[i].coefficient_); + reactant_state_idx.push_back(info.state_indices_[reactant_dependency_indices_[i]]); + } + std::vector product_stoich; + std::vector product_state_idx; + for (std::size_t i = 0; i < this->products_.size(); ++i) + { + product_stoich.push_back(this->products_[i].coefficient_); + product_state_idx.push_back(info.state_indices_[product_dependency_indices_[i]]); + } + std::size_t row_idx = info.row_index_; + + return DenseMatrixPolicy::Function( + [K_eq, reactant_stoich, reactant_state_idx, product_stoich, product_state_idx, row_idx](auto&& state, auto&& force) + { + auto reactant_product = force.GetRowVariable(); + auto product_product = force.GetRowVariable(); + + // Initialize reactant_product to K_eq and product_product to 1.0 + state.ForEachRow([K_eq](double& rp, double& pp) + { + rp = K_eq; + pp = 1.0; + }, reactant_product, product_product); + + // Multiply in each reactant concentration raised to its stoichiometry + for (std::size_t i = 0; i < reactant_stoich.size(); ++i) + { + const double stoich = reactant_stoich[i]; + const std::size_t species_idx = reactant_state_idx[i]; + + state.ForEachRow([stoich](const double& conc, double& product) + { + product *= std::pow(std::max(0.0, conc), stoich); + }, state.GetConstColumnView(species_idx), reactant_product); + } + + // Multiply in each product concentration raised to its stoichiometry + for (std::size_t i = 0; i < product_stoich.size(); ++i) + { + const double stoich = product_stoich[i]; + const std::size_t species_idx = product_state_idx[i]; + + state.ForEachRow([stoich](const double& conc, double& product) + { + product *= std::pow(std::max(0.0, conc), stoich); + }, state.GetConstColumnView(species_idx), product_product); + } + + // Write G = K_eq * [reactants]^stoich - [products]^stoich to the constraint row + state.ForEachRow([](const double& rp, const double& pp, double& forcing_term) + { + forcing_term = rp - pp; + }, reactant_product, product_product, force.GetColumnView(row_idx)); + }, + temp_state_variables, + temp_state_variables + ); + } + + /// @brief Compute Jacobian entries dG/d[species] + /// For reactant R with stoichiometry n: + /// dG/d[R] = K_eq * n * [R]^(n-1) * prod([other_reactants]^stoich) + /// For product P with stoichiometry m: + /// dG/d[P] = -m * [P]^(m-1) * prod([other_products]^stoich) + /// @param info Constraint information including species indices + /// @param state_variable_indices Mapping of state variable names to indices + /// @param jacobian_flat_ids Iterator to this constraint's flat Jacobian indices in shared storage + /// @param jacobian Sparse matrix to store Jacobian values + /// @return Function object that takes (state_variables, jacobian) and computes partials + template + std::function + JacobianFunction( + const ConstraintInfo& info, + const auto& state_variable_indices, + auto jacobian_flat_ids, + SparseMatrixPolicy& jacobian) const + { + DenseMatrixPolicy temp_state_variables{1, state_variable_indices.size(), 0.0}; + + // Pre-compute flat IDs and store them in a vector + // This avoids iterator issues when the lambda is called multiple times (once per block) + std::vector flat_ids; + flat_ids.reserve(reactants_.size() + products_.size()); + for (std::size_t i = 0; i < reactants_.size(); ++i) + { + flat_ids.push_back(*jacobian_flat_ids++); + } + for (std::size_t i = 0; i < products_.size(); ++i) + { + flat_ids.push_back(*jacobian_flat_ids++); + } + + // Copy data to avoid issues when ConstraintSet is moved + double K_eq = this->equilibrium_constant_; + std::vector reactant_stoich; + std::vector reactant_state_idx; + for (std::size_t i = 0; i < this->reactants_.size(); ++i) + { + reactant_stoich.push_back(this->reactants_[i].coefficient_); + reactant_state_idx.push_back(info.state_indices_[reactant_dependency_indices_[i]]); + } + std::vector product_stoich; + std::vector product_state_idx; + for (std::size_t i = 0; i < this->products_.size(); ++i) + { + product_stoich.push_back(this->products_[i].coefficient_); + product_state_idx.push_back(info.state_indices_[product_dependency_indices_[i]]); + } + + return SparseMatrixPolicy::Function( + [K_eq, reactant_stoich, reactant_state_idx, product_stoich, product_state_idx, flat_ids](auto&& state, auto&& jacobian_values) + { + // Create temporary variables for computing partials + auto reactant_product = jacobian_values.GetBlockVariable(); + auto product_product = jacobian_values.GetBlockVariable(); + auto partial_derivative = jacobian_values.GetBlockVariable(); + + jacobian_values.ForEachBlock([K_eq](double& rp, double& pp) + { + rp = K_eq; + pp = 1.0; + }, reactant_product, product_product); + + for (std::size_t i = 0; i < reactant_stoich.size(); ++i) + { + const double stoich = reactant_stoich[i]; + const std::size_t species_idx = reactant_state_idx[i]; + + jacobian_values.ForEachBlock([stoich](const double& conc, double& product) + { + product *= std::pow(std::max(0.0, conc), stoich); + }, state.GetConstColumnView(species_idx), reactant_product); + } + + for (std::size_t i = 0; i < product_stoich.size(); ++i) + { + const double stoich = product_stoich[i]; + const std::size_t species_idx = product_state_idx[i]; + + jacobian_values.ForEachBlock([stoich](const double& conc, double& product) + { + product *= std::pow(std::max(0.0, conc), stoich); + }, state.GetConstColumnView(species_idx), product_product); + } + + // Compute Jacobian entries for each reactant: dG/d[R_i] = K_eq * stoich_i * prod([other_reactants]^stoich) * [R_i]^(stoich_i-1) + for (std::size_t i = 0; i < reactant_stoich.size(); ++i) + { + const double stoich_i = reactant_stoich[i]; + const std::size_t species_idx_i = reactant_state_idx[i]; + + // Compute product of K_eq * all reactants except R_i + auto partial_product = jacobian_values.GetBlockVariable(); + jacobian_values.ForEachBlock([K_eq](double& prod) + { + prod = K_eq; + }, partial_product); + + for (std::size_t j = 0; j < reactant_stoich.size(); ++j) + { + if (j != i) // Skip current species + { + const double stoich_j = reactant_stoich[j]; + const std::size_t species_idx_j = reactant_state_idx[j]; + + jacobian_values.ForEachBlock([stoich_j](const double& conc, double& prod) + { + prod *= std::pow(std::max(0.0, conc), stoich_j); + }, state.GetConstColumnView(species_idx_j), partial_product); + } + } + + // Multiply by stoich_i * [R_i]^(stoich_i-1) + jacobian_values.ForEachBlock( + [stoich_i](const double& conc, const double& prod, double& partial) + { + if (stoich_i == 1.0) + { + partial = prod; + } + else if (conc > 0.0) + { + partial = stoich_i * prod * std::pow(conc, stoich_i - 1.0); + } + else + { + partial = 0.0; + } + }, + state.GetConstColumnView(species_idx_i), + partial_product, + partial_derivative + ); + + // Subtract partial from Jacobian (matching the SubtractJacobianTerms convention) + // Use pre-computed flat ID for this reactant + jacobian_values.ForEachBlock( + [](const double& partial, double& jac) + { + jac -= partial; + }, + partial_derivative, + jacobian_values.GetBlockView(flat_ids[i]) + ); + } + + // Compute Jacobian entries for each product: dG/d[P_i] = -stoich_i * prod([other_products]^stoich) * [P_i]^(stoich_i-1) + for (std::size_t i = 0; i < product_stoich.size(); ++i) + { + const double stoich_i = product_stoich[i]; + const std::size_t species_idx_i = product_state_idx[i]; + + // Compute product of all products except P_i + auto partial_product = jacobian_values.GetBlockVariable(); + jacobian_values.ForEachBlock([](double& prod) + { + prod = 1.0; + }, partial_product); + + for (std::size_t j = 0; j < product_stoich.size(); ++j) + { + if (j != i) // Skip current species + { + const double stoich_j = product_stoich[j]; + const std::size_t species_idx_j = product_state_idx[j]; + + jacobian_values.ForEachBlock([stoich_j](const double& conc, double& prod) + { + prod *= std::pow(std::max(0.0, conc), stoich_j); + }, state.GetConstColumnView(species_idx_j), partial_product); + } + } + + // Multiply by stoich_i * [P_i]^(stoich_i-1) + jacobian_values.ForEachBlock( + [stoich_i](const double& conc, const double& prod, double& partial) + { + if (stoich_i == 1.0) + { + partial = prod; + } + else if (conc > 0.0) + { + partial = stoich_i * prod * std::pow(conc, stoich_i - 1.0); + } + else + { + partial = 0.0; + } + }, + state.GetConstColumnView(species_idx_i), + partial_product, + partial_derivative + ); + + // Add partial to Jacobian (note: G = ... - [products], so derivative gets positive sign after subtraction) + // Use pre-computed flat ID for this product (products come after reactants in flat_ids) + jacobian_values.ForEachBlock( + [](const double& partial, double& jac) + { + jac += partial; + }, + partial_derivative, + jacobian_values.GetBlockView(flat_ids[reactant_stoich.size() + i]) + ); + } + }, + temp_state_variables, + jacobian + ); + } + }; + +} // namespace micm diff --git a/include/micm/constraint/types/linear_constraint.hpp b/include/micm/constraint/types/linear_constraint.hpp new file mode 100644 index 000000000..1a6c1d9d5 --- /dev/null +++ b/include/micm/constraint/types/linear_constraint.hpp @@ -0,0 +1,189 @@ +// Copyright (C) 2023-2026 University Corporation for Atmospheric Research +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace micm +{ + + /// @brief Constraint for linear relationships: sum(coeff[i] * [species[i]]) = constant + /// For example: A + B + C = 1.0 represents a conservation law + /// The linear constraint is: G = c1*[A] + c2*[B] + c3*[C] - constant = 0 + class LinearConstraint + { + public: + /// @brief Name of the constraint + std::string name_; + + /// @brief Names of species this constraint depends on + std::vector species_dependencies_; + + /// @brief Species and their coefficients in the linear sum + std::vector terms_; + + /// @brief The constant value the linear sum should equal + double constant_; + + public: + /// @brief Default constructor + LinearConstraint() = default; + + /// @brief Construct a linear constraint + /// Validates that terms are non-empty + /// Builds species_dependencies_ from terms + /// @param name Constraint identifier + /// @param terms Vector of StoichSpecies (species, coefficient) in the linear sum + /// @param constant The value that sum(coeff[i] * [species[i]]) should equal + LinearConstraint( + const std::string& name, + const std::vector& terms, + double constant) + : name_(name), + terms_(terms), + constant_(constant) + { + if (terms_.empty()) + { + throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_CONSTRAINT, MICM_CONSTRAINT_ERROR_CODE_EMPTY_REACTANTS, ""); + } + for (const auto& term : terms_) + { + species_dependencies_.push_back(term.species_.name_); + } + } + + /// @brief Returns the species whose row should be replaced by this algebraic constraint + /// For linear constraints, the last species is used in the terms list as the algebraic row target. + /// @return Species name of the primary algebraic variable + const std::string& AlgebraicSpecies() const + { + return terms_.back().species_.name_; + } + + /// @brief Create a function to compute the linear constraint residual + /// Returns a reusable function object that evaluates: + /// G = sum(coeff[i] * [species[i]]) - constant + /// @param info Constraint information including species indices and row index + /// @param state_variable_indices Mapping of state variable names to indices + /// @return Function object that takes (state_variables, forcing) and computes residual + template + std::function + ResidualFunction( + const ConstraintInfo& info, + const auto& state_variable_indices) const + { + DenseMatrixPolicy temp_state_variables{1, state_variable_indices.size(), 0.0}; + + // Copy data to avoid issues when ConstraintSet is moved + std::vector coeffs; + std::vector species_indices; + for (std::size_t i = 0; i < this->terms_.size(); ++i) + { + coeffs.push_back(this->terms_[i].coefficient_); + species_indices.push_back(info.state_indices_[i]); + } + double constant = this->constant_; + std::size_t row_idx = info.row_index_; + + return DenseMatrixPolicy::Function( + [coeffs, species_indices, constant, row_idx](auto&& state, auto&& force) + { + // Create a variable for accumulating the linear sum + auto linear_sum = force.GetRowVariable(); + + // Initialize linear_sum to 0.0 before accumulation + state.ForEachRow([](double& sum) + { + sum = 0.0; + }, linear_sum); + + for (std::size_t i = 0; i < coeffs.size(); ++i) + { + const double coeff = coeffs[i]; + const std::size_t species_idx = species_indices[i]; + + state.ForEachRow([coeff](const double& conc, double& sum) + { + sum += coeff * conc; + }, state.GetConstColumnView(species_idx), linear_sum); + } + + // Forcing term = sum(coeff[i] * [species[i]]) - constant + state.ForEachRow([constant](const double& sum_val, double& forcing_term) + { + forcing_term = sum_val - constant; + }, linear_sum, force.GetColumnView(row_idx)); + }, + temp_state_variables, + temp_state_variables + ); + } + + /// @brief Create a function to compute Jacobian entries dG/d[species] + /// For a linear constraint, the Jacobian is simply the coefficients: + /// dG/d[species[i]] = coeff[i] + /// @param info Constraint information including species indices + /// @param state_variable_indices Mapping of state variable names to indices + /// @param jacobian_flat_ids Iterator to this constraint's flat Jacobian indices in shared storage + /// @param jacobian Sparse matrix to store Jacobian values + /// @return Function object that takes (state_variables, jacobian) and computes partials + template + std::function + JacobianFunction( + const ConstraintInfo& info, + const auto& state_variable_indices, + auto jacobian_flat_ids, + SparseMatrixPolicy& jacobian) const + { + DenseMatrixPolicy temp_state_variables{1, state_variable_indices.size(), 0.0}; + + // Pre-compute flat IDs and store them in a vector + // This avoids iterator issues when the lambda is called multiple times + std::vector flat_ids; + flat_ids.reserve(this->terms_.size()); + for (std::size_t i = 0; i < this->terms_.size(); ++i) + { + flat_ids.push_back(*jacobian_flat_ids++); + } + + // Copy data to avoid issues when ConstraintSet is moved + std::vector coeffs; + for (std::size_t i = 0; i < this->terms_.size(); ++i) + { + coeffs.push_back(this->terms_[i].coefficient_); + } + + return SparseMatrixPolicy::Function( + [coeffs, flat_ids](auto&& state, auto&& jacobian_values) + { + // For linear constraints, dG/d[species[i]] = coeff[i] + // We subtract the coefficient from the Jacobian (matching the SubtractJacobianTerms convention) + for (std::size_t i = 0; i < coeffs.size(); ++i) + { + const double coeff = coeffs[i]; + + jacobian_values.ForEachBlock( + [coeff](double& jac) + { + jac -= coeff; + }, + jacobian_values.GetBlockView(flat_ids[i]) + ); + } + }, + temp_state_variables, + jacobian + ); + } + }; + +} // namespace micm diff --git a/include/micm/cuda/solver/cuda_rosenbrock.hpp b/include/micm/cuda/solver/cuda_rosenbrock.hpp index cd820d830..4261b741f 100644 --- a/include/micm/cuda/solver/cuda_rosenbrock.hpp +++ b/include/micm/cuda/solver/cuda_rosenbrock.hpp @@ -14,11 +14,12 @@ namespace micm { struct CudaRosenbrockSolverParameters; - template + template class CudaRosenbrockSolver : public AbstractRosenbrockSolver< RatesPolicy, LinearSolverPolicy, - CudaRosenbrockSolver> + ConstraintSetPolicy, + CudaRosenbrockSolver> { ///@brief Default constructor public: @@ -28,12 +29,12 @@ namespace micm CudaRosenbrockSolver(const CudaRosenbrockSolver&) = delete; CudaRosenbrockSolver& operator=(const CudaRosenbrockSolver&) = delete; CudaRosenbrockSolver(CudaRosenbrockSolver&& other) - : AbstractRosenbrockSolver>( + : AbstractRosenbrockSolver>( std::move(other)){}; CudaRosenbrockSolver& operator=(CudaRosenbrockSolver&& other) { - RosenbrockSolver::operator=(std::move(other)); + RosenbrockSolver::operator=(std::move(other)); return *this; }; @@ -47,8 +48,8 @@ namespace micm CudaRosenbrockSolver( LinearSolverPolicy&& linear_solver, RatesPolicy&& rates, - ConstraintSet&& constraints) - : AbstractRosenbrockSolver>( + ConstraintSetPolicy&& constraints) + : AbstractRosenbrockSolver>( std::move(linear_solver), std::move(rates), std::move(constraints)) {}; diff --git a/include/micm/cuda/solver/cuda_solver_parameters.hpp b/include/micm/cuda/solver/cuda_solver_parameters.hpp index 229ac8608..83dfc1ce4 100644 --- a/include/micm/cuda/solver/cuda_solver_parameters.hpp +++ b/include/micm/cuda/solver/cuda_solver_parameters.hpp @@ -12,8 +12,8 @@ namespace micm /// @brief Parameters for the CUDA Rosenbrock solver struct CudaRosenbrockSolverParameters : public RosenbrockSolverParameters { - template - using SolverType = CudaRosenbrockSolver; + template + using SolverType = CudaRosenbrockSolver; /// @brief Constructor from base class /// @param base diff --git a/include/micm/solver/backward_euler.hpp b/include/micm/solver/backward_euler.hpp index e0cba2b95..5fc27f5d5 100644 --- a/include/micm/solver/backward_euler.hpp +++ b/include/micm/solver/backward_euler.hpp @@ -27,13 +27,13 @@ namespace micm { /// @brief An implementation of the fully implicit backward euler method - template + template class AbstractBackwardEuler { public: LinearSolverPolicy linear_solver_; RatesPolicy rates_; - ConstraintSet constraints_; + ConstraintSetPolicy constraints_; /// @brief Solver parameters typename using ParametersType = BackwardEulerSolverParameters; @@ -45,7 +45,7 @@ namespace micm AbstractBackwardEuler( LinearSolverPolicy&& linear_solver, RatesPolicy&& rates, - ConstraintSet&& constraints) + ConstraintSetPolicy&& constraints) : linear_solver_(std::move(linear_solver)), rates_(std::move(rates)), constraints_(std::move(constraints)) diff --git a/include/micm/solver/backward_euler.inl b/include/micm/solver/backward_euler.inl index 8e5ab6f14..4edc70ccf 100644 --- a/include/micm/solver/backward_euler.inl +++ b/include/micm/solver/backward_euler.inl @@ -3,8 +3,8 @@ namespace micm { - template - inline SolverResult AbstractBackwardEuler::Solve( + template + inline SolverResult AbstractBackwardEuler::Solve( double time_step, auto& state, const BackwardEulerSolverParameters& parameters) const @@ -150,9 +150,9 @@ namespace micm return result; } - template + template template - inline bool AbstractBackwardEuler::IsConverged( + inline bool AbstractBackwardEuler::IsConverged( const BackwardEulerSolverParameters& parameters, const DenseMatrixPolicy& residual, const DenseMatrixPolicy& Yn1, @@ -179,9 +179,9 @@ namespace micm return true; } - template + template template - inline bool AbstractBackwardEuler::IsConverged( + inline bool AbstractBackwardEuler::IsConverged( const BackwardEulerSolverParameters& parameters, const DenseMatrixPolicy& residual, const DenseMatrixPolicy& Yn1, diff --git a/include/micm/solver/backward_euler_solver_parameters.hpp b/include/micm/solver/backward_euler_solver_parameters.hpp index 1a00157fc..641f97d87 100644 --- a/include/micm/solver/backward_euler_solver_parameters.hpp +++ b/include/micm/solver/backward_euler_solver_parameters.hpp @@ -10,14 +10,14 @@ namespace micm { - template + template class AbstractBackwardEuler; - /// @brief BackwardEuler solver parameters + /// @brief Backward Euler solver parameters struct BackwardEulerSolverParameters { - template - using SolverType = AbstractBackwardEuler; + template + using SolverType = AbstractBackwardEuler; double small_{ 1.0e-40 }; double h_start_{ 0.0 }; diff --git a/include/micm/solver/rosenbrock.hpp b/include/micm/solver/rosenbrock.hpp index a7e058d91..193201d95 100644 --- a/include/micm/solver/rosenbrock.hpp +++ b/include/micm/solver/rosenbrock.hpp @@ -37,6 +37,7 @@ namespace micm /// @brief An implementation of the Rosenbrock ODE solver /// @tparam RatesPolicy Calculator of forcing and Jacobian terms /// @tparam LinearSolverPolicy Linear solver + /// @tparam ConstraintSetPolicy Constraint set for algebraic constraints /// @tparam Derived Implementation of the Rosenbock solver /// /// This implements the Curiously Recurring Template Pattern to allow @@ -44,13 +45,13 @@ namespace micm /// in extending classes and called from the base class Solve() function. /// https://en.cppreference.com/w/cpp/language/crtp /// - template + template class AbstractRosenbrockSolver { public: LinearSolverPolicy linear_solver_; RatesPolicy rates_; - ConstraintSet constraints_; + ConstraintSetPolicy constraints_; static constexpr double DEFAULT_H_MIN = 1.0e-15; // Minimum internal time step relative to overall time step static constexpr double DEFAULT_H_START = 1.0e-6; // Default initial time step relative to overall time step @@ -66,7 +67,7 @@ namespace micm AbstractRosenbrockSolver( LinearSolverPolicy&& linear_solver, RatesPolicy&& rates, - ConstraintSet&& constraints) + ConstraintSetPolicy&& constraints) : linear_solver_(std::move(linear_solver)), rates_(std::move(rates)), constraints_(std::move(constraints)) @@ -123,9 +124,9 @@ namespace micm requires(VectorizableDense); }; // end of Abstract Rosenbrock Solver - template + template class RosenbrockSolver - : public AbstractRosenbrockSolver> + : public AbstractRosenbrockSolver> { public: /// @brief Default constructor @@ -137,8 +138,8 @@ namespace micm RosenbrockSolver( LinearSolverPolicy&& linear_solver, RatesPolicy&& rates, - ConstraintSet&& constraints) - : AbstractRosenbrockSolver>( + ConstraintSetPolicy&& constraints) + : AbstractRosenbrockSolver>( std::move(linear_solver), std::move(rates), std::move(constraints)) diff --git a/include/micm/solver/rosenbrock.inl b/include/micm/solver/rosenbrock.inl index 54a93d216..be5643fc3 100644 --- a/include/micm/solver/rosenbrock.inl +++ b/include/micm/solver/rosenbrock.inl @@ -3,8 +3,8 @@ namespace micm { - template - inline SolverResult AbstractRosenbrockSolver::Solve( + template + inline SolverResult AbstractRosenbrockSolver::Solve( double time_step, auto& state, const RosenbrockSolverParameters& parameters) const noexcept @@ -54,11 +54,8 @@ namespace micm initial_forcing.Fill(0); rates_.AddForcingTerms(state, Y, initial_forcing); - // Add constraint residuals to forcing (for DAE systems) - if (has_constraints) - { + if (has_constraints) constraints_.AddForcingTerms(Y, initial_forcing); - } result.stats_.function_calls_ += 1; @@ -66,11 +63,8 @@ namespace micm state.jacobian_.Fill(0); rates_.SubtractJacobianTerms(state, Y, state.jacobian_); - // Add constraint Jacobian terms (for DAE systems) if (has_constraints) - { constraints_.SubtractJacobianTerms(Y, state.jacobian_); - } result.stats_.jacobian_updates_ += 1; @@ -111,7 +105,6 @@ namespace micm } K[stage].Fill(0); rates_.AddForcingTerms(state, Ynew, K[stage]); - // Add constraint residuals for DAE systems if (has_constraints) { constraints_.AddForcingTerms(Ynew, K[stage]); @@ -223,9 +216,7 @@ namespace micm rates_.SubtractJacobianTerms(state, Y, state.jacobian_); // Subtract constraint Jacobian terms (for DAE systems) if (has_constraints) - { constraints_.SubtractJacobianTerms(Y, state.jacobian_); - } result.stats_.jacobian_updates_ += 1; } } @@ -242,9 +233,9 @@ namespace micm return result; } - template + template template - inline void AbstractRosenbrockSolver::AlphaMinusJacobian( + inline void AbstractRosenbrockSolver::AlphaMinusJacobian( auto& state, const double& alpha) const requires(!VectorizableSparse) @@ -262,9 +253,9 @@ namespace micm } } - template + template template - inline void AbstractRosenbrockSolver::AlphaMinusJacobian( + inline void AbstractRosenbrockSolver::AlphaMinusJacobian( auto& state, const double& alpha) const requires(VectorizableSparse) @@ -284,8 +275,8 @@ namespace micm } } - template - inline void AbstractRosenbrockSolver::LinearFactor( + template + inline void AbstractRosenbrockSolver::LinearFactor( const double alpha, SolverStats& stats, auto& state) const @@ -306,9 +297,9 @@ namespace micm stats.decompositions_ += 1; } - template + template template - inline double AbstractRosenbrockSolver::NormalizedError( + inline double AbstractRosenbrockSolver::NormalizedError( const DenseMatrixPolicy& Y, const DenseMatrixPolicy& Ynew, const DenseMatrixPolicy& errors, @@ -349,9 +340,9 @@ namespace micm return std::max(std::sqrt(error / N), error_min); } - template + template template - inline double AbstractRosenbrockSolver::NormalizedError( + inline double AbstractRosenbrockSolver::NormalizedError( const DenseMatrixPolicy& Y, const DenseMatrixPolicy& Ynew, const DenseMatrixPolicy& errors, diff --git a/include/micm/solver/rosenbrock_solver_parameters.hpp b/include/micm/solver/rosenbrock_solver_parameters.hpp index 6c7a952e1..1ed807b91 100644 --- a/include/micm/solver/rosenbrock_solver_parameters.hpp +++ b/include/micm/solver/rosenbrock_solver_parameters.hpp @@ -12,14 +12,14 @@ namespace micm { - template + template class RosenbrockSolver; /// @brief Rosenbrock solver parameters struct RosenbrockSolverParameters { - template - using SolverType = RosenbrockSolver; + template + using SolverType = RosenbrockSolver; std::size_t stages_{}; std::size_t upper_limit_tolerance_{}; diff --git a/include/micm/solver/solver_builder.hpp b/include/micm/solver/solver_builder.hpp index ed0f71b41..6625806be 100644 --- a/include/micm/solver/solver_builder.hpp +++ b/include/micm/solver/solver_builder.hpp @@ -4,7 +4,7 @@ #include #include -#include +#include #include #include #include diff --git a/include/micm/solver/solver_builder.inl b/include/micm/solver/solver_builder.inl index 1345fe4d3..a0982a44a 100644 --- a/include/micm/solver/solver_builder.inl +++ b/include/micm/solver/solver_builder.inl @@ -205,9 +205,9 @@ namespace micm // Include species referenced by constraints (dependencies and algebraic targets) for (const auto& constraint : constraints_) { - for (const auto& dep : constraint.GetSpeciesDependencies()) + for (const auto& dep : constraint.SpeciesDependencies()) used_species.insert(dep); - used_species.insert(constraint.GetAlgebraicSpecies()); + used_species.insert(constraint.AlgebraicSpecies()); } auto available_species = system_.UniqueNames(); @@ -380,7 +380,8 @@ namespace micm throw MicmException(MicmSeverity::Error, MICM_ERROR_CATEGORY_SOLVER, MICM_SOLVER_ERROR_CODE_MISSING_CHEMICAL_SPECIES, "Provided chemical system contains no species."); } - using SolverPolicy = typename SolverParametersPolicy::template SolverType; + using ConstraintSetPolicy = ConstraintSet; + using SolverPolicy = typename SolverParametersPolicy::template SolverType; auto species_map = this->GetSpeciesMap(); auto params_map = this->GetCustomParameterMap(); @@ -390,7 +391,7 @@ namespace micm auto nonzero_elements = rates.NonZeroJacobianElements(); // Create ConstraintSet from stored constraints (if any) - ConstraintSet constraint_set; + ConstraintSetPolicy constraint_set; std::set algebraic_variable_ids; std::size_t number_of_constraints = constraints_.size(); @@ -400,7 +401,7 @@ namespace micm auto constraints_copy = constraints_; // Constraints replace selected species rows in the mass-matrix DAE formulation. // Pass species_map so constraints can resolve dependencies and row targets to species indices. - constraint_set = ConstraintSet(std::move(constraints_copy), species_map); + constraint_set = ConstraintSetPolicy(std::move(constraints_copy), species_map); algebraic_variable_ids = constraint_set.AlgebraicVariableIds(); rates.SetAlgebraicVariableIds(algebraic_variable_ids); @@ -430,10 +431,10 @@ namespace micm rates.SetJacobianFlatIds(jacobian); rates.SetExternalModelFunctions(params_map, species_map, jacobian); - // Set Jacobian flat IDs for constraints if (constraint_set.Size() > 0) { constraint_set.SetJacobianFlatIds(jacobian); + constraint_set.SetConstraintFunctions(species_map, jacobian); } std::vector variable_names{ number_of_species }; diff --git a/include/micm/util/error.hpp b/include/micm/util/error.hpp index 0ed0f8612..056da37aa 100644 --- a/include/micm/util/error.hpp +++ b/include/micm/util/error.hpp @@ -71,3 +71,4 @@ #define MICM_CONSTRAINT_ERROR_CODE_EMPTY_REACTANTS 3 #define MICM_CONSTRAINT_ERROR_CODE_EMPTY_PRODUCTS 4 #define MICM_CONSTRAINT_ERROR_CODE_INVALID_STOICHIOMETRY 5 +#define MICM_CONSTRAINT_ERROR_CODE_DUPLICATE_ALGEBRAIC_SPECIES 6 diff --git a/test/integration/CMakeLists.txt b/test/integration/CMakeLists.txt index 7fc237736..2bab2a7c4 100644 --- a/test/integration/CMakeLists.txt +++ b/test/integration/CMakeLists.txt @@ -8,7 +8,7 @@ create_standard_test(NAME terminator SOURCES test_terminator.cpp) create_standard_test(NAME integrated_reaction_rates SOURCES test_integrated_reaction_rates.cpp) create_standard_test(NAME aerosol_model_backward_euler SOURCES test_aerosol_model_backward_euler.cpp) create_standard_test(NAME aerosol_model_rosenbrock SOURCES test_aerosol_model_rosenbrock.cpp) -create_standard_test(NAME equilibrium SOURCES test_equilibrium.cpp) +create_standard_test(NAME equilibrium_constraint SOURCES test_equilibrium_constraint.cpp) create_standard_test(NAME linear_constraint SOURCES test_linear_constraint.cpp) if(NOT ${MICM_GPU_TYPE} STREQUAL "None") diff --git a/test/integration/test_equilibrium.cpp b/test/integration/test_equilibrium_constraint.cpp similarity index 96% rename from test/integration/test_equilibrium.cpp rename to test/integration/test_equilibrium_constraint.cpp index 298b973c0..dff7f8f1d 100644 --- a/test/integration/test_equilibrium.cpp +++ b/test/integration/test_equilibrium_constraint.cpp @@ -4,7 +4,7 @@ #include #include #include -#include +#include #include @@ -13,45 +13,7 @@ #include #include -/// @brief Test ConstraintSet API directly (unit-level test for DAE infrastructure) -/// Uses replace-state-rows mode: AB's row (index 2) is replaced by the constraint. -TEST(EquilibriumIntegration, ConstraintSetAPITest) -{ - auto A = micm::Species("A"); - auto B = micm::Species("B"); - auto AB = micm::Species("AB"); - - // Create constraint: A + B <-> AB with K_eq = 1000 - std::vector constraints; - constraints.push_back(micm::EquilibriumConstraint( - "A_B_eq", - std::vector{ { A, 1.0 }, { B, 1.0 } }, - std::vector{ { AB, 1.0 } }, - 1000.0)); - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 } - }; - - micm::ConstraintSet set(std::move(constraints), variable_map); - - // Test at equilibrium: [A] = 0.0312, [B] = 0.0312, [AB] = 0.9737 - // K_eq * [A] * [B] = 1000 * 0.0312^2 ≈ 0.9734 - // Should approximately equal [AB] = 0.9737 - micm::Matrix state(1, 3); - state[0][0] = 0.0312; // A - state[0][1] = 0.0312; // B - state[0][2] = 0.9737; // AB - calculated to be at equilibrium - - micm::Matrix forcing(1, 3, 0.0); - set.AddForcingTerms(state, forcing); - - // At equilibrium, residual replaces AB's row (index 2) - // G = K_eq * [A] * [B] - [AB] = 1000 * 0.0312 * 0.0312 - 0.9737 ≈ 0 - EXPECT_NEAR(forcing[0][2], 0.0, 0.01); -} /// @brief Test SetConstraints API integration - verifies solver builds and runs with constraints /// @@ -1073,7 +1035,7 @@ TEST(EquilibriumIntegration, DAEOverlappingSpeciesJacobian) }; // Algebraic species is A (first product), so constraint replaces row 0 - micm::ConstraintSet set(std::move(constraints), variable_map); + micm::ConstraintSet, micm::SparseMatrix> set(std::move(constraints), variable_map); // Build a sparse Jacobian and set flat IDs auto nonzero = set.NonZeroJacobianElements(); diff --git a/test/integration/test_linear_constraint.cpp b/test/integration/test_linear_constraint.cpp index 561090c34..784a58f01 100644 --- a/test/integration/test_linear_constraint.cpp +++ b/test/integration/test_linear_constraint.cpp @@ -4,8 +4,8 @@ #include #include #include -#include -#include +#include +#include #include diff --git a/test/unit/constraint/CMakeLists.txt b/test/unit/constraint/CMakeLists.txt index 51478d578..94277134b 100644 --- a/test/unit/constraint/CMakeLists.txt +++ b/test/unit/constraint/CMakeLists.txt @@ -1,5 +1,5 @@ ################################################################################ # Tests -create_standard_test(NAME equilibrium_constraint SOURCES test_equilibrium_constraint.cpp) -create_standard_test(NAME linear SOURCES test_linear.cpp) create_standard_test(NAME constraint_set SOURCES test_constraint_set.cpp) + +add_subdirectory(type) \ No newline at end of file diff --git a/test/unit/constraint/test_constraint_set.cpp b/test/unit/constraint/test_constraint_set.cpp index 860e1d89e..a75a9e871 100644 --- a/test/unit/constraint/test_constraint_set.cpp +++ b/test/unit/constraint/test_constraint_set.cpp @@ -1,11 +1,9 @@ // Copyright (C) 2023-2026 University Corporation for Atmospheric Research // SPDX-License-Identifier: Apache-2.0 -#include +#include "test_constraint_set_policy.hpp" + #include -#include -#include -#include #include #include #include @@ -14,646 +12,99 @@ #include -#include -#include -#include -#include -#include - using namespace micm; +using StandardSparseMatrix = SparseMatrix; -TEST(ConstraintSet, Construction) -{ - // Create a simple constraint set with one equilibrium constraint - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - 1000.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 } - }; +using Group1VectorMatrix = VectorMatrix; +using Group2VectorMatrix = VectorMatrix; +using Group3VectorMatrix = VectorMatrix; +using Group4VectorMatrix = VectorMatrix; - ConstraintSet set(std::move(constraints), variable_map); +using Group1SparseVectorMatrix = SparseMatrix>; +using Group2SparseVectorMatrix = SparseMatrix>; +using Group3SparseVectorMatrix = SparseMatrix>; +using Group4SparseVectorMatrix = SparseMatrix>; - EXPECT_EQ(set.Size(), 1); +TEST(ConstraintSet, Construction) +{ + testConstruction, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, ReplaceStateRowsMapsToAlgebraicSpecies) { - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "B_C_eq", - std::vector{ StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("C"), 1.0) }, - 10.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "C", 2 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - - EXPECT_EQ(set.Size(), 1); - EXPECT_EQ(set.AlgebraicVariableIds().size(), 1); - EXPECT_TRUE(set.AlgebraicVariableIds().count(2)); // C row is algebraic - - auto non_zero_elements = set.NonZeroJacobianElements(); - EXPECT_EQ(non_zero_elements.size(), 2); - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 1))); // row C, col B - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // row C, col C + testReplaceStateRowsMapsToAlgebraicSpecies, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, NonZeroJacobianElements) { - // Create constraint set - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - 1000.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - - auto non_zero_elements = set.NonZeroJacobianElements(); - - // Algebraic species = AB (index 2), constraint replaces row 2 - // Dependencies: A (col 0), B (col 1), AB (col 2) + diagonal (2,2) - EXPECT_EQ(non_zero_elements.size(), 3); - - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 0))); // dG/dA at row 2, col 0 - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 1))); // dG/dB at row 2, col 1 - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // dG/dAB at row 2, col 2 (+ diagonal) + testNonZeroJacobianElements, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, MultipleConstraints) { - // Create constraint set with two equilibrium constraints - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - 1000.0)); - constraints.push_back(EquilibriumConstraint( - "C_D_eq", - std::vector{ StoichSpecies(Species("C"), 1.0) }, - std::vector{ StoichSpecies(Species("D"), 1.0) }, - 500.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 }, - { "C", 3 }, - { "D", 4 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - - EXPECT_EQ(set.Size(), 2); - - auto non_zero_elements = set.NonZeroJacobianElements(); - - // First constraint: algebraic = AB (index 2), replaces row 2 - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 0))); // dG1/dA - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 1))); // dG1/dB - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // dG1/dAB + diagonal - - // Second constraint: algebraic = D (index 4), replaces row 4 - EXPECT_TRUE(non_zero_elements.count(std::make_pair(4, 3))); // dG2/dC - EXPECT_TRUE(non_zero_elements.count(std::make_pair(4, 4))); // dG2/dD + diagonal - - EXPECT_EQ(non_zero_elements.size(), 5); + testMultipleConstraints, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, AddForcingTerms) { - // Create constraint set - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - 1000.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 } - }; - - std::size_t num_species = 3; - - ConstraintSet set(std::move(constraints), variable_map); - - // State with 2 grid cells - Matrix state(2, num_species); - state[0] = { 0.01, 0.01, 0.001 }; // Away from equilibrium - state[1] = { 0.001, 0.001, 0.001 }; // At equilibrium - - // Forcing vector (same size as state; constraint replaces AB row at index 2) - Matrix forcing(2, num_species, 0.0); - - set.AddForcingTerms(state, forcing); - - // For grid cell 0: G = 1000 * 0.01 * 0.01 - 0.001 = 0.1 - 0.001 = 0.099 - EXPECT_NEAR(forcing[0][2], 0.099, 1e-10); - - // For grid cell 1: G = 1000 * 0.001 * 0.001 - 0.001 = 0.001 - 0.001 = 0.0 - EXPECT_NEAR(forcing[1][2], 0.0, 1e-10); + testAddForcingTerms, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, SubtractJacobianTerms) { - // Create constraint set - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - 1000.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 } - }; - - std::size_t num_species = 3; - - ConstraintSet set(std::move(constraints), variable_map); - - // Get non-zero elements and build sparse Jacobian - auto non_zero_elements = set.NonZeroJacobianElements(); - - // Build a 3x3 sparse Jacobian (constraint replaces AB's row) - auto builder = SparseMatrix::Create(num_species) - .SetNumberOfBlocks(1) - .InitialValue(0.0); - - for (std::size_t i = 0; i < num_species; ++i) - builder = builder.WithElement(i, i); // Diagonals - for (auto& elem : non_zero_elements) - builder = builder.WithElement(elem.first, elem.second); - - SparseMatrix jacobian{ builder }; - - set.SetJacobianFlatIds(jacobian); - - // State with 1 grid cell - Matrix state(1, num_species); - state[0] = { 0.01, 0.02, 0.05 }; - - set.SubtractJacobianTerms(state, jacobian); - - // G = K_eq * [A] * [B] - [AB] - // dG/d[A] = K_eq * [B] = 1000 * 0.02 = 20 - // dG/d[B] = K_eq * [A] = 1000 * 0.01 = 10 - // dG/d[AB] = -1 - - // Jacobian subtracts these values (matching ProcessSet convention) - // Constraint replaces row 2 (AB's row) - EXPECT_NEAR(jacobian[0][2][0], -20.0, 1e-10); // J[2, A] -= dG/dA - EXPECT_NEAR(jacobian[0][2][1], -10.0, 1e-10); // J[2, B] -= dG/dB - EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-10); // J[2, AB] -= dG/dAB = -(-1) = 1 + testSubtractJacobianTerms, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, EmptyConstraintSet) { - // Empty constraint set should be valid and do nothing - ConstraintSet set; - - EXPECT_EQ(set.Size(), 0); - - auto non_zero_elements = set.NonZeroJacobianElements(); - EXPECT_TRUE(non_zero_elements.empty()); - - // AddForcingTerms and SubtractJacobianTerms should do nothing for empty set - Matrix state(1, 3); - state[0] = { 0.1, 0.2, 0.3 }; - - Matrix forcing(1, 4, 1.0); - - set.AddForcingTerms(state, forcing); - - // Forcing should be unchanged - EXPECT_DOUBLE_EQ(forcing[0][0], 1.0); - EXPECT_DOUBLE_EQ(forcing[0][1], 1.0); - EXPECT_DOUBLE_EQ(forcing[0][2], 1.0); - EXPECT_DOUBLE_EQ(forcing[0][3], 1.0); + testEmptyConstraintSet, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, UnknownSpeciesThrows) { - // Creating a constraint with unknown species should throw - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("X"), 1.0), StoichSpecies(Species("Y"), 1.0) }, - std::vector{ StoichSpecies(Species("XY"), 1.0) }, - 1000.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 } - }; - - EXPECT_THROW( - ConstraintSet(std::move(constraints), variable_map), - micm::MicmException); + testUnknownSpeciesThrows, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } -/// @brief Test 3D state (3 species) with 1 constraint -/// -/// System: Species X, Y, Z with constraint X <-> Y (K_eq = 50) -/// Algebraic species = Y (first product), replaces row 1 -/// Jacobian: 3x3 matrix -/// -/// Constraint: G = K_eq * [X] - [Y] = 0 -/// At equilibrium: [Y]/[X] = K_eq = 50 TEST(ConstraintSet, ThreeDStateOneConstraint) { - const double K_eq = 50.0; - const std::size_t num_species = 3; - - // Create constraint: X <-> Y with K_eq = 50 - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "X_Y_eq", - std::vector{ StoichSpecies(Species("X"), 1.0) }, - std::vector{ StoichSpecies(Species("Y"), 1.0) }, - K_eq)); - - std::unordered_map variable_map = { - { "X", 0 }, - { "Y", 1 }, - { "Z", 2 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - - EXPECT_EQ(set.Size(), 1); - - // Check non-zero Jacobian elements - // Algebraic species = Y (index 1), constraint replaces row 1 - auto non_zero_elements = set.NonZeroJacobianElements(); - EXPECT_EQ(non_zero_elements.size(), 2); // (1,0) dG/dX, (1,1) dG/dY + diagonal - EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 0))); // dG/dX at row 1, col 0 - EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 1))); // dG/dY at row 1, col 1 + diagonal - EXPECT_FALSE(non_zero_elements.count(std::make_pair(1, 2))); // Z not involved - - // Build sparse Jacobian (3x3) - auto builder = SparseMatrix::Create(num_species) - .SetNumberOfBlocks(2) // Test with 2 grid cells - .InitialValue(0.0); - - for (std::size_t i = 0; i < num_species; ++i) - builder = builder.WithElement(i, i); - for (auto& elem : non_zero_elements) - builder = builder.WithElement(elem.first, elem.second); - - SparseMatrix jacobian{ builder }; - set.SetJacobianFlatIds(jacobian); - - // State with 2 grid cells - Matrix state(2, num_species); - // Grid cell 0: Away from equilibrium - state[0][0] = 0.1; // X - state[0][1] = 2.0; // Y - state[0][2] = 0.5; // Z (uninvolved) - // Grid cell 1: At equilibrium (Y/X = 50) - state[1][0] = 0.02; // X - state[1][1] = 1.0; // Y = 50 * 0.02 = 1.0 - state[1][2] = 0.3; // Z - - // Test forcing terms - Matrix forcing(2, num_species, 0.0); - set.AddForcingTerms(state, forcing); - - // Constraint replaces row 1 (Y's row) - // Grid cell 0: G = K_eq * [X] - [Y] = 50 * 0.1 - 2.0 = 3.0 - EXPECT_NEAR(forcing[0][1], 3.0, 1e-10); - // Grid cell 1: G = 50 * 0.02 - 1.0 = 0.0 (at equilibrium) - EXPECT_NEAR(forcing[1][1], 0.0, 1e-10); - - // Test Jacobian terms - set.SubtractJacobianTerms(state, jacobian); - - // For constraint G = K_eq * [X] - [Y]: - // dG/dX = K_eq = 50 - // dG/dY = -1 - // Jacobian subtracts at row 1: - - // Grid cell 0 - EXPECT_NEAR(jacobian[0][1][0], -K_eq, 1e-10); // dG/dX - EXPECT_NEAR(jacobian[0][1][1], 1.0, 1e-10); // dG/dY (subtracted -1) - // Grid cell 1 - EXPECT_NEAR(jacobian[1][1][0], -K_eq, 1e-10); - EXPECT_NEAR(jacobian[1][1][1], 1.0, 1e-10); - - // Z row should be unaffected - EXPECT_NEAR(jacobian[0][2][2], 0.0, 1e-10); - EXPECT_NEAR(jacobian[1][2][2], 0.0, 1e-10); + testThreeDStateOneConstraint, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } -/// @brief Test 4D state (4 species) with 2 constraints -/// -/// System: Species A, B, C, D with two constraints: -/// 1. A <-> B with K_eq1 = 10, algebraic species = B (row 1) -/// 2. C + D <-> A with K_eq2 = 100, algebraic species = A (row 0) -/// Jacobian: 4x4 matrix -/// -/// Constraint 1: G1 = K_eq1 * [A] - [B] = 0 -/// Constraint 2: G2 = K_eq2 * [C] * [D] - [A] = 0 TEST(ConstraintSet, FourDStateTwoConstraints) { - const double K_eq1 = 10.0; - const double K_eq2 = 100.0; - const std::size_t num_species = 4; - - // Create two constraints - std::vector constraints; - - // Constraint 1: A <-> B with K_eq1 = 10, algebraic species = B (row 1) - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - K_eq1)); - - // Constraint 2: C + D <-> A with K_eq2 = 100, algebraic species = A (row 0) - constraints.push_back(EquilibriumConstraint( - "CD_A_eq", - std::vector{ StoichSpecies(Species("C"), 1.0), StoichSpecies(Species("D"), 1.0) }, - std::vector{ StoichSpecies(Species("A"), 1.0) }, - K_eq2)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "C", 2 }, - { "D", 3 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - - EXPECT_EQ(set.Size(), 2); - - // Check non-zero Jacobian elements - auto non_zero_elements = set.NonZeroJacobianElements(); - - // Constraint 1 replaces row 1 (B): depends on A (col 0), B (col 1) - EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 0))); // dG1/dA - EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 1))); // dG1/dB + diagonal - - // Constraint 2 replaces row 0 (A): depends on C (col 2), D (col 3), A (col 0) - EXPECT_TRUE(non_zero_elements.count(std::make_pair(0, 2))); // dG2/dC - EXPECT_TRUE(non_zero_elements.count(std::make_pair(0, 3))); // dG2/dD - EXPECT_TRUE(non_zero_elements.count(std::make_pair(0, 0))); // dG2/dA + diagonal - - // Total: 2 (from constraint 1) + 3 (from constraint 2) = 5 - EXPECT_EQ(non_zero_elements.size(), 5); - - // Build sparse Jacobian (4x4) - auto builder = SparseMatrix::Create(num_species) - .SetNumberOfBlocks(3) // Test with 3 grid cells - .InitialValue(0.0); - - for (std::size_t i = 0; i < num_species; ++i) - builder = builder.WithElement(i, i); - for (auto& elem : non_zero_elements) - builder = builder.WithElement(elem.first, elem.second); - - SparseMatrix jacobian{ builder }; - set.SetJacobianFlatIds(jacobian); - - // State with 3 grid cells - Matrix state(3, num_species); - - // Grid cell 0: Both constraints satisfied - state[0][0] = 0.1; // A - state[0][1] = 1.0; // B = 10 * 0.1 - state[0][2] = 0.1; // C - state[0][3] = 0.01; // D, so C*D = 0.001, K_eq2*C*D = 0.1 = A - - // Grid cell 1: First constraint satisfied, second not - state[1][0] = 0.2; // A - state[1][1] = 2.0; // B = 10 * 0.2 (constraint 1 satisfied) - state[1][2] = 0.1; // C - state[1][3] = 0.1; // D, C*D = 0.01, K_eq2*C*D = 1.0 != 0.2 - - // Grid cell 2: Neither constraint satisfied - state[2][0] = 0.5; // A - state[2][1] = 3.0; // B != 10 * 0.5 = 5.0 - state[2][2] = 0.2; // C - state[2][3] = 0.3; // D, C*D = 0.06, K_eq2*C*D = 6.0 != 0.5 - - // Test forcing terms - Matrix forcing(3, num_species, 0.0); - set.AddForcingTerms(state, forcing); - - // Constraint 1 replaces row 1, Constraint 2 replaces row 0 - // Grid cell 0: Both at equilibrium - EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); // G1 = 10 * 0.1 - 1.0 = 0 - EXPECT_NEAR(forcing[0][0], 0.0, 1e-10); // G2 = 100 * 0.1 * 0.01 - 0.1 = 0 - - // Grid cell 1: First satisfied, second not - EXPECT_NEAR(forcing[1][1], 0.0, 1e-10); // G1 = 10 * 0.2 - 2.0 = 0 - EXPECT_NEAR(forcing[1][0], 0.8, 1e-10); // G2 = 100 * 0.1 * 0.1 - 0.2 = 0.8 - - // Grid cell 2: Neither satisfied - EXPECT_NEAR(forcing[2][1], 2.0, 1e-10); // G1 = 10 * 0.5 - 3.0 = 2.0 - EXPECT_NEAR(forcing[2][0], 5.5, 1e-10); // G2 = 100 * 0.2 * 0.3 - 0.5 = 5.5 - - // Test Jacobian terms - set.SubtractJacobianTerms(state, jacobian); - - // Constraint 1 at row 1: dG1/dA = K_eq1, dG1/dB = -1 - // Constraint 2 at row 0: dG2/dC = K_eq2*[D], dG2/dD = K_eq2*[C], dG2/dA = -1 - - // Grid cell 0: - EXPECT_NEAR(jacobian[0][1][0], -K_eq1, 1e-10); - EXPECT_NEAR(jacobian[0][1][1], 1.0, 1e-10); - EXPECT_NEAR(jacobian[0][0][2], -K_eq2 * state[0][3], 1e-10); - EXPECT_NEAR(jacobian[0][0][3], -K_eq2 * state[0][2], 1e-10); - EXPECT_NEAR(jacobian[0][0][0], 1.0, 1e-10); - - // Grid cell 1: - EXPECT_NEAR(jacobian[1][1][0], -K_eq1, 1e-10); - EXPECT_NEAR(jacobian[1][1][1], 1.0, 1e-10); - EXPECT_NEAR(jacobian[1][0][2], -K_eq2 * state[1][3], 1e-10); - EXPECT_NEAR(jacobian[1][0][3], -K_eq2 * state[1][2], 1e-10); - EXPECT_NEAR(jacobian[1][0][0], 1.0, 1e-10); - - // Grid cell 2: - EXPECT_NEAR(jacobian[2][1][0], -K_eq1, 1e-10); - EXPECT_NEAR(jacobian[2][1][1], 1.0, 1e-10); - EXPECT_NEAR(jacobian[2][0][2], -K_eq2 * state[2][3], 1e-10); - EXPECT_NEAR(jacobian[2][0][3], -K_eq2 * state[2][2], 1e-10); - EXPECT_NEAR(jacobian[2][0][0], 1.0, 1e-10); + testFourDStateTwoConstraints, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } -/// @brief Test coupled constraints where constraints share species -/// -/// System: A, B, C with constraints that both involve A: -/// 1. A <-> B (K_eq1 = 5), algebraic = B (row 1) -/// 2. A <-> C (K_eq2 = 20), algebraic = C (row 2) -/// -/// This tests that the Jacobian correctly handles overlapping dependencies TEST(ConstraintSet, CoupledConstraintsSharedSpecies) { - const double K_eq1 = 5.0; - const double K_eq2 = 20.0; - const std::size_t num_species = 3; - - std::vector constraints; - - // Both constraints depend on species A - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - K_eq1)); - - constraints.push_back(EquilibriumConstraint( - "A_C_eq", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("C"), 1.0) }, - K_eq2)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "C", 2 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - - EXPECT_EQ(set.Size(), 2); - - // Check Jacobian structure - both constraints depend on A - auto non_zero_elements = set.NonZeroJacobianElements(); - - // Constraint 1 replaces row 1 (B): dG1/dA, dG1/dB + diagonal - EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 0))); // dG1/dA - EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 1))); // dG1/dB + diagonal - - // Constraint 2 replaces row 2 (C): dG2/dA, dG2/dC + diagonal - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 0))); // dG2/dA - EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // dG2/dC + diagonal - - EXPECT_EQ(non_zero_elements.size(), 4); - - // Build Jacobian (3x3) - auto builder = SparseMatrix::Create(num_species) - .SetNumberOfBlocks(1) - .InitialValue(0.0); - - for (std::size_t i = 0; i < num_species; ++i) - builder = builder.WithElement(i, i); - for (auto& elem : non_zero_elements) - builder = builder.WithElement(elem.first, elem.second); - - SparseMatrix jacobian{ builder }; - set.SetJacobianFlatIds(jacobian); - - // State at dual equilibrium: [B]/[A] = 5, [C]/[A] = 20 - Matrix state(1, num_species); - state[0][0] = 0.1; // A - state[0][1] = 0.5; // B = 5 * 0.1 - state[0][2] = 2.0; // C = 20 * 0.1 - - // Test forcing terms - Matrix forcing(1, num_species, 0.0); - set.AddForcingTerms(state, forcing); - - // Both constraints should be satisfied - EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); // G1 at row 1 - EXPECT_NEAR(forcing[0][2], 0.0, 1e-10); // G2 at row 2 - - // Test Jacobian terms - set.SubtractJacobianTerms(state, jacobian); - - // Constraint 1 at row 1: dG1/dA = 5, dG1/dB = -1 - EXPECT_NEAR(jacobian[0][1][0], -K_eq1, 1e-10); - EXPECT_NEAR(jacobian[0][1][1], 1.0, 1e-10); - - // Constraint 2 at row 2: dG2/dA = 20, dG2/dC = -1 - EXPECT_NEAR(jacobian[0][2][0], -K_eq2, 1e-10); - EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-10); + testCoupledConstraintsSharedSpecies, StandardSparseMatrix, ConstraintSet, StandardSparseMatrix>>(); } TEST(ConstraintSet, VectorizedMatricesRespectGridCellIndexing) { - const std::size_t num_species = 3; - - std::vector constraints; - constraints.push_back(EquilibriumConstraint( - "A_B_eq", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - 1000.0)); - - std::unordered_map variable_map = { - { "A", 0 }, - { "B", 1 }, - { "AB", 2 } - }; - - ConstraintSet set(std::move(constraints), variable_map); - auto non_zero_elements = set.NonZeroJacobianElements(); - - // Constraint replaces AB's row (index 2), Jacobian is 3x3 - auto builder = SparseMatrix>::Create(num_species) - .SetNumberOfBlocks(3) - .InitialValue(0.0); - for (std::size_t i = 0; i < num_species; ++i) - builder = builder.WithElement(i, i); - for (const auto& elem : non_zero_elements) - builder = builder.WithElement(elem.first, elem.second); - - SparseMatrix> jacobian{ builder }; - set.SetJacobianFlatIds(jacobian); - - VectorMatrix state(3, num_species, 0.0); - state[0] = { 0.01, 0.02, 0.05 }; - state[1] = { 0.03, 0.01, 0.2 }; - state[2] = { 0.001, 0.002, 0.004 }; - - VectorMatrix forcing(3, num_species, 0.0); - set.AddForcingTerms(state, forcing); - - // Constraint residual replaces row 2 (AB) - EXPECT_NEAR(forcing[0][2], 1000.0 * 0.01 * 0.02 - 0.05, 1e-12); - EXPECT_NEAR(forcing[1][2], 1000.0 * 0.03 * 0.01 - 0.2, 1e-12); - EXPECT_NEAR(forcing[2][2], 1000.0 * 0.001 * 0.002 - 0.004, 1e-12); - - set.SubtractJacobianTerms(state, jacobian); - - // Jacobian entries at row 2 (AB's row, replaced by constraint) - EXPECT_NEAR(jacobian[0][2][0], -(1000.0 * 0.02), 1e-12); - EXPECT_NEAR(jacobian[0][2][1], -(1000.0 * 0.01), 1e-12); - EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-12); + testVectorizedMatricesRespectGridCellIndexing>(); +} - EXPECT_NEAR(jacobian[1][2][0], -(1000.0 * 0.01), 1e-12); - EXPECT_NEAR(jacobian[1][2][1], -(1000.0 * 0.03), 1e-12); - EXPECT_NEAR(jacobian[1][2][2], 1.0, 1e-12); +TEST(ConstraintSet, VectorMatrix1) +{ + testConstruction>(); + testNonZeroJacobianElements>(); + testAddForcingTerms>(); + testSubtractJacobianTerms>(); +} - EXPECT_NEAR(jacobian[2][2][0], -(1000.0 * 0.002), 1e-12); - EXPECT_NEAR(jacobian[2][2][1], -(1000.0 * 0.001), 1e-12); - EXPECT_NEAR(jacobian[2][2][2], 1.0, 1e-12); +TEST(ConstraintSet, VectorMatrix2) +{ + testConstruction>(); + testNonZeroJacobianElements>(); + testAddForcingTerms>(); + testSubtractJacobianTerms>(); } + +TEST(ConstraintSet, VectorMatrix3) +{ + testConstruction>(); + testNonZeroJacobianElements>(); + testAddForcingTerms>(); + testSubtractJacobianTerms>(); +} \ No newline at end of file diff --git a/test/unit/constraint/test_constraint_set_policy.hpp b/test/unit/constraint/test_constraint_set_policy.hpp new file mode 100644 index 000000000..12a221236 --- /dev/null +++ b/test/unit/constraint/test_constraint_set_policy.hpp @@ -0,0 +1,658 @@ +// Copyright (C) 2023-2026 University Corporation for Atmospheric Research +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +using namespace micm; + +template +void testConstruction() +{ + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + EXPECT_EQ(set.Size(), 1); +} + +template +void testReplaceStateRowsMapsToAlgebraicSpecies() +{ + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "B_C_eq", + std::vector{ StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("C"), 1.0) }, + 10.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "C", 2 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + EXPECT_EQ(set.Size(), 1); + EXPECT_EQ(set.AlgebraicVariableIds().size(), 1); + EXPECT_TRUE(set.AlgebraicVariableIds().count(2)); // C row is algebraic + + auto non_zero_elements = set.NonZeroJacobianElements(); + EXPECT_EQ(non_zero_elements.size(), 2); + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 1))); // row C, col B + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // row C, col C +} + +template +void testNonZeroJacobianElements() +{ + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + auto non_zero_elements = set.NonZeroJacobianElements(); + + // Algebraic species = AB (index 2), constraint replaces row 2 + // Dependencies: A (col 0), B (col 1), AB (col 2) + diagonal (2,2) + EXPECT_EQ(non_zero_elements.size(), 3); + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 0))); // dG/dA at row 2, col 0 + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 1))); // dG/dB at row 2, col 1 + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // dG/dAB at row 2, col 2 (+ diagonal) +} + +template +void testMultipleConstraints() +{ + // Create constraint set with two equilibrium constraints + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + constraints.push_back(EquilibriumConstraint( + "C_D_eq", + std::vector{ StoichSpecies(Species("C"), 1.0) }, + std::vector{ StoichSpecies(Species("D"), 1.0) }, + 500.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 }, + { "C", 3 }, + { "D", 4 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + EXPECT_EQ(set.Size(), 2); + + auto non_zero_elements = set.NonZeroJacobianElements(); + + // First constraint: algebraic = AB (index 2), replaces row 2 + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 0))); // dG1/dA + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 1))); // dG1/dB + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // dG1/dAB + diagonal + + // Second constraint: algebraic = D (index 4), replaces row 4 + EXPECT_TRUE(non_zero_elements.count(std::make_pair(4, 3))); // dG2/dC + EXPECT_TRUE(non_zero_elements.count(std::make_pair(4, 4))); // dG2/dD + diagonal + + EXPECT_EQ(non_zero_elements.size(), 5); +} + +template +void testAddForcingTerms() +{ + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + // Build sparse Jacobian for SetConstraintFunctions + auto non_zero_elements = set.NonZeroJacobianElements(); + auto builder = SparseMatrixPolicy::Create(num_species) + .SetNumberOfBlocks(2) + .InitialValue(0.0); + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + SparseMatrixPolicy jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // State with 2 grid cells + DenseMatrixPolicy state(2, num_species); + state[0] = { 0.2, 0.4, 0.6 }; // Away from equilibrium + state[1] = { 0.1, 0.3, 0.7 }; // Away from equilibrium + + // Forcing vector (same size as state; constraint replaces AB row at index 2) + DenseMatrixPolicy forcing(2, num_species, 0.0); + + set.AddForcingTerms(state, forcing); + + // For grid cell 0: G = 1000 * 0.2 * 0.4 - 0.6 = 80.0 - 0.6 = 79.4 + EXPECT_NEAR(forcing[0][2], 79.4, 1e-10); + + // For grid cell 1: G = 1000 * 0.1 * 0.3 - 0.7 = 30.0 - 0.7 = 29.3 + EXPECT_NEAR(forcing[1][2], 29.3, 1e-10); +} + +template +void testSubtractJacobianTerms() +{ + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + // Get non-zero elements and build sparse Jacobian + auto non_zero_elements = set.NonZeroJacobianElements(); + + // Build a 3x3 sparse Jacobian (constraint replaces AB's row) + auto builder = SparseMatrixPolicy::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); // Diagonals + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + SparseMatrixPolicy jacobian{ builder }; + + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // State with 1 grid cell + DenseMatrixPolicy state(1, num_species); + state[0] = { 0.01, 0.02, 0.05 }; + + set.SubtractJacobianTerms(state, jacobian); + + // G = K_eq * [A] * [B] - [AB] + // dG/d[A] = K_eq * [B] = 1000 * 0.02 = 20 + // dG/d[B] = K_eq * [A] = 1000 * 0.01 = 10 + // dG/d[AB] = -1 + + // Jacobian subtracts these values (matching ProcessSet convention) + // Constraint replaces row 2 (AB's row) + EXPECT_NEAR(jacobian[0][2][0], -20.0, 1e-10); // J[2, A] -= dG/dA + EXPECT_NEAR(jacobian[0][2][1], -10.0, 1e-10); // J[2, B] -= dG/dB + EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-10); // J[2, AB] -= dG/dAB = -(-1) = 1 +} + +template +void testEmptyConstraintSet() +{ + // Empty constraint set should be valid and do nothing + ConstraintSetPolicy set; + + EXPECT_EQ(set.Size(), 0); + + auto non_zero_elements = set.NonZeroJacobianElements(); + EXPECT_TRUE(non_zero_elements.empty()); + + // AddForcingTerms and SubtractJacobianTerms should do nothing for empty set + DenseMatrixPolicy state(1, 3); + state[0] = { 0.1, 0.2, 0.3 }; + + DenseMatrixPolicy forcing(1, 4, 1.0); + + set.AddForcingTerms(state, forcing); + + // Forcing should be unchanged + EXPECT_DOUBLE_EQ(forcing[0][0], 1.0); + EXPECT_DOUBLE_EQ(forcing[0][1], 1.0); + EXPECT_DOUBLE_EQ(forcing[0][2], 1.0); + EXPECT_DOUBLE_EQ(forcing[0][3], 1.0); +} + +template +void testUnknownSpeciesThrows() +{ + // Creating a constraint with unknown species should throw + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("X"), 1.0), StoichSpecies(Species("Y"), 1.0) }, + std::vector{ StoichSpecies(Species("XY"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 } + }; + + EXPECT_THROW((ConstraintSetPolicy(std::move(constraints), variable_map)), micm::MicmException); +} + +/// @brief Test 3D state (3 species) with 1 constraint +template +void testThreeDStateOneConstraint() +{ + const double K_eq = 50.0; + const std::size_t num_species = 3; + + // Create constraint: X <-> Y with K_eq = 50 + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "X_Y_eq", + std::vector{ StoichSpecies(Species("X"), 1.0) }, + std::vector{ StoichSpecies(Species("Y"), 1.0) }, + K_eq)); + + std::unordered_map variable_map = { + { "X", 0 }, + { "Y", 1 }, + { "Z", 2 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + EXPECT_EQ(set.Size(), 1); + + // Check non-zero Jacobian elements + // Algebraic species = Y (index 1), constraint replaces row 1 + auto non_zero_elements = set.NonZeroJacobianElements(); + EXPECT_EQ(non_zero_elements.size(), 2); // (1,0) dG/dX, (1,1) dG/dY + diagonal + EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 0))); // dG/dX at row 1, col 0 + EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 1))); // dG/dY at row 1, col 1 + diagonal + EXPECT_FALSE(non_zero_elements.count(std::make_pair(1, 2))); // Z not involved + + // Build sparse Jacobian (3x3) + auto builder = SparseMatrixPolicy::Create(num_species) + .SetNumberOfBlocks(2) // Test with 2 grid cells + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + SparseMatrixPolicy jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // State with 2 grid cells + DenseMatrixPolicy state(2, num_species); + // Grid cell 0: Away from equilibrium + state[0][0] = 0.1; // X + state[0][1] = 0.3; // Y + state[0][2] = 0.5; // Z (uninvolved) + // Grid cell 1: At equilibrium (Y/X = 50) + state[1][0] = 0.02; // X + state[1][1] = 1.0; // Y = 50 * 0.02 = 1.0 + state[1][2] = 0.3; // Z + + // Test forcing terms + DenseMatrixPolicy forcing(2, num_species, 0.0); + set.AddForcingTerms(state, forcing); + + // Constraint replaces row 1 (Y's row) + // Grid cell 0: G = K_eq * [X] - [Y] = 50 * 0.1 - 0.3 = 5.0 - 0.3 = 4.7 + EXPECT_NEAR(forcing[0][1], 4.7, 1e-10); + // Grid cell 1: G = 50 * 0.02 - 1.0 = 0.0 (at equilibrium) + EXPECT_NEAR(forcing[1][1], 0.0, 1e-10); + + // Test Jacobian terms + set.SubtractJacobianTerms(state, jacobian); + + // For constraint G = K_eq * [X] - [Y]: + // dG/dX = K_eq = 50 + // dG/dY = -1 + // Jacobian subtracts at row 1: + + // Grid cell 0 + EXPECT_NEAR(jacobian[0][1][0], -K_eq, 1e-10); // dG/dX + EXPECT_NEAR(jacobian[0][1][1], 1.0, 1e-10); // dG/dY (subtracted -1) + // Grid cell 1 + EXPECT_NEAR(jacobian[1][1][0], -K_eq, 1e-10); + EXPECT_NEAR(jacobian[1][1][1], 1.0, 1e-10); + + // Z row should be unaffected + EXPECT_NEAR(jacobian[0][2][2], 0.0, 1e-10); + EXPECT_NEAR(jacobian[1][2][2], 0.0, 1e-10); +} + +/// @brief Test 4D state (4 species) with 2 constraints +template +void testFourDStateTwoConstraints() +{ + const double K_eq1 = 10.0; + const double K_eq2 = 100.0; + const std::size_t num_species = 4; + + // Create two constraints + std::vector constraints; + + // Constraint 1: A <-> B with K_eq1 = 10, algebraic species = B (row 1) + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + K_eq1)); + + // Constraint 2: C + D <-> A with K_eq2 = 100, algebraic species = A (row 0) + constraints.push_back(EquilibriumConstraint( + "CD_A_eq", + std::vector{ StoichSpecies(Species("C"), 1.0), StoichSpecies(Species("D"), 1.0) }, + std::vector{ StoichSpecies(Species("A"), 1.0) }, + K_eq2)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "C", 2 }, + { "D", 3 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + EXPECT_EQ(set.Size(), 2); + + // Check non-zero Jacobian elements + auto non_zero_elements = set.NonZeroJacobianElements(); + + // Constraint 1 replaces row 1 (B): depends on A (col 0), B (col 1) + EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 0))); // dG1/dA + EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 1))); // dG1/dB + diagonal + + // Constraint 2 replaces row 0 (A): depends on C (col 2), D (col 3), A (col 0) + EXPECT_TRUE(non_zero_elements.count(std::make_pair(0, 2))); // dG2/dC + EXPECT_TRUE(non_zero_elements.count(std::make_pair(0, 3))); // dG2/dD + EXPECT_TRUE(non_zero_elements.count(std::make_pair(0, 0))); // dG2/dA + diagonal + + // Total: 2 (from constraint 1) + 3 (from constraint 2) = 5 + EXPECT_EQ(non_zero_elements.size(), 5); + + // Build sparse Jacobian (4x4) + auto builder = SparseMatrixPolicy::Create(num_species) + .SetNumberOfBlocks(3) // Test with 3 grid cells + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + SparseMatrixPolicy jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // State with 3 grid cells + DenseMatrixPolicy state(3, num_species); + + // Grid cell 0: Both constraints satisfied + state[0][0] = 0.1; // A + state[0][1] = 1.0; // B = 10 * 0.1 + state[0][2] = 0.1; // C + state[0][3] = 0.01; // D, so C*D = 0.001, K_eq2*C*D = 0.1 = A + + // Grid cell 1: First constraint satisfied, second not + state[1][0] = 0.2; // A + state[1][1] = 2.0; // B = 10 * 0.2 (constraint 1 satisfied) + state[1][2] = 0.1; // C + state[1][3] = 0.1; // D, C*D = 0.01, K_eq2*C*D = 1.0 != 0.2 + + // Grid cell 2: Neither constraint satisfied + state[2][0] = 0.5; // A + state[2][1] = 3.0; // B != 10 * 0.5 = 5.0 + state[2][2] = 0.2; // C + state[2][3] = 0.3; // D, C*D = 0.06, K_eq2*C*D = 6.0 != 0.5 + + // Test forcing terms + DenseMatrixPolicy forcing(3, num_species, 0.0); + set.AddForcingTerms(state, forcing); + + // Constraint 1 replaces row 1, Constraint 2 replaces row 0 + // Grid cell 0: Both at equilibrium + EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); // G1 = 10 * 0.1 - 1.0 = 0 + EXPECT_NEAR(forcing[0][0], 0.0, 1e-10); // G2 = 100 * 0.1 * 0.01 - 0.1 = 0 + + // Grid cell 1: First satisfied, second not + EXPECT_NEAR(forcing[1][1], 0.0, 1e-10); // G1 = 10 * 0.2 - 2.0 = 0 + EXPECT_NEAR(forcing[1][0], 0.8, 1e-10); // G2 = 100 * 0.1 * 0.1 - 0.2 = 0.8 + + // Grid cell 2: Neither satisfied + EXPECT_NEAR(forcing[2][1], 2.0, 1e-10); // G1 = 10 * 0.5 - 3.0 = 2.0 + EXPECT_NEAR(forcing[2][0], 5.5, 1e-10); // G2 = 100 * 0.2 * 0.3 - 0.5 = 5.5 + + // Test Jacobian terms + set.SubtractJacobianTerms(state, jacobian); + + // Constraint 1 at row 1: dG1/dA = K_eq1, dG1/dB = -1 + // Constraint 2 at row 0: dG2/dC = K_eq2*[D], dG2/dD = K_eq2*[C], dG2/dA = -1 + + // Grid cell 0: + EXPECT_NEAR(jacobian[0][1][0], -K_eq1, 1e-10); + EXPECT_NEAR(jacobian[0][1][1], 1.0, 1e-10); + EXPECT_NEAR(jacobian[0][0][2], -K_eq2 * state[0][3], 1e-10); + EXPECT_NEAR(jacobian[0][0][3], -K_eq2 * state[0][2], 1e-10); + EXPECT_NEAR(jacobian[0][0][0], 1.0, 1e-10); + + // Grid cell 1: + EXPECT_NEAR(jacobian[1][1][0], -K_eq1, 1e-10); + EXPECT_NEAR(jacobian[1][1][1], 1.0, 1e-10); + EXPECT_NEAR(jacobian[1][0][2], -K_eq2 * state[1][3], 1e-10); + EXPECT_NEAR(jacobian[1][0][3], -K_eq2 * state[1][2], 1e-10); + EXPECT_NEAR(jacobian[1][0][0], 1.0, 1e-10); + + // Grid cell 2: + EXPECT_NEAR(jacobian[2][1][0], -K_eq1, 1e-10); + EXPECT_NEAR(jacobian[2][1][1], 1.0, 1e-10); + EXPECT_NEAR(jacobian[2][0][2], -K_eq2 * state[2][3], 1e-10); + EXPECT_NEAR(jacobian[2][0][3], -K_eq2 * state[2][2], 1e-10); + EXPECT_NEAR(jacobian[2][0][0], 1.0, 1e-10); +} + +/// @brief Test coupled constraints where constraints share species +template +void testCoupledConstraintsSharedSpecies() +{ + const double K_eq1 = 5.0; + const double K_eq2 = 20.0; + const std::size_t num_species = 3; + + std::vector constraints; + + // Both constraints depend on species A + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + K_eq1)); + + constraints.push_back(EquilibriumConstraint( + "A_C_eq", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("C"), 1.0) }, + K_eq2)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "C", 2 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + + EXPECT_EQ(set.Size(), 2); + + // Check Jacobian structure - both constraints depend on A + auto non_zero_elements = set.NonZeroJacobianElements(); + + // Constraint 1 replaces row 1 (B): dG1/dA, dG1/dB + diagonal + EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 0))); // dG1/dA + EXPECT_TRUE(non_zero_elements.count(std::make_pair(1, 1))); // dG1/dB + diagonal + + // Constraint 2 replaces row 2 (C): dG2/dA, dG2/dC + diagonal + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 0))); // dG2/dA + EXPECT_TRUE(non_zero_elements.count(std::make_pair(2, 2))); // dG2/dC + diagonal + + EXPECT_EQ(non_zero_elements.size(), 4); + + // Build Jacobian (3x3) + auto builder = SparseMatrixPolicy::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + SparseMatrixPolicy jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // State at dual equilibrium: [B]/[A] = 5, [C]/[A] = 20 + DenseMatrixPolicy state(1, num_species); + state[0][0] = 0.1; // A + state[0][1] = 0.5; // B = 5 * 0.1 + state[0][2] = 2.0; // C = 20 * 0.1 + + // Test forcing terms + DenseMatrixPolicy forcing(1, num_species, 0.0); + set.AddForcingTerms(state, forcing); + + // Both constraints should be satisfied + EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); // G1 at row 1 + EXPECT_NEAR(forcing[0][2], 0.0, 1e-10); // G2 at row 2 + + // Test Jacobian terms + set.SubtractJacobianTerms(state, jacobian); + + // Constraint 1 at row 1: dG1/dA = 5, dG1/dB = -1 + EXPECT_NEAR(jacobian[0][1][0], -K_eq1, 1e-10); + EXPECT_NEAR(jacobian[0][1][1], 1.0, 1e-10); + + // Constraint 2 at row 2: dG2/dA = 20, dG2/dC = -1 + EXPECT_NEAR(jacobian[0][2][0], -K_eq2, 1e-10); + EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-10); +} + +template +void testVectorizedMatricesRespectGridCellIndexing() +{ + const std::size_t num_species = 3; + + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_eq", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + ConstraintSetPolicy set(std::move(constraints), variable_map); + auto non_zero_elements = set.NonZeroJacobianElements(); + + // Constraint replaces AB's row (index 2), Jacobian is 3x3 + auto builder = SparseMatrixPolicy::Create(num_species) + .SetNumberOfBlocks(3) + .InitialValue(0.0); + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (const auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + SparseMatrixPolicy jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + DenseMatrixPolicy state(3, num_species, 0.0); + state[0] = { 0.01, 0.02, 0.05 }; + state[1] = { 0.03, 0.01, 0.2 }; + state[2] = { 0.001, 0.002, 0.004 }; + + DenseMatrixPolicy forcing(3, num_species, 0.0); + set.AddForcingTerms(state, forcing); + + // Constraint residual replaces row 2 (AB) + EXPECT_NEAR(forcing[0][2], 1000.0 * 0.01 * 0.02 - 0.05, 1e-12); + EXPECT_NEAR(forcing[1][2], 1000.0 * 0.03 * 0.01 - 0.2, 1e-12); + EXPECT_NEAR(forcing[2][2], 1000.0 * 0.001 * 0.002 - 0.004, 1e-12); + + set.SubtractJacobianTerms(state, jacobian); + + // Jacobian entries at row 2 (AB's row, replaced by constraint) + EXPECT_NEAR(jacobian[0][2][0], -(1000.0 * 0.02), 1e-12); + EXPECT_NEAR(jacobian[0][2][1], -(1000.0 * 0.01), 1e-12); + EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-12); + + EXPECT_NEAR(jacobian[1][2][0], -(1000.0 * 0.01), 1e-12); + EXPECT_NEAR(jacobian[1][2][1], -(1000.0 * 0.03), 1e-12); + EXPECT_NEAR(jacobian[1][2][2], 1.0, 1e-12); + + EXPECT_NEAR(jacobian[2][2][0], -(1000.0 * 0.002), 1e-12); + EXPECT_NEAR(jacobian[2][2][1], -(1000.0 * 0.001), 1e-12); + EXPECT_NEAR(jacobian[2][2][2], 1.0, 1e-12); +} diff --git a/test/unit/constraint/test_equilibrium_constraint.cpp b/test/unit/constraint/test_equilibrium_constraint.cpp deleted file mode 100644 index 4c86b8a72..000000000 --- a/test/unit/constraint/test_equilibrium_constraint.cpp +++ /dev/null @@ -1,212 +0,0 @@ -// Copyright (C) 2023-2026 University Corporation for Atmospheric Research -// SPDX-License-Identifier: Apache-2.0 - -#include -#include -#include -#include - -#include - -#include -#include -#include - -using namespace micm; - -TEST(EquilibriumConstraint, SimpleABEquilibrium) -{ - // Test: A + B <-> AB with K_eq = 1000 - // At equilibrium: [AB] / ([A][B]) = K_eq - // Constraint: G = K_eq * [A] * [B] - [AB] = 0 - - double K_eq = 1000.0; - EquilibriumConstraint constraint( - "A_B_equilibrium", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, // reactants with stoich - std::vector{ StoichSpecies(Species("AB"), 1.0) }, // products with stoich - K_eq); - - EXPECT_EQ(constraint.name_, "A_B_equilibrium"); - EXPECT_EQ(constraint.species_dependencies_.size(), 3); - EXPECT_EQ(constraint.species_dependencies_[0], "A"); - EXPECT_EQ(constraint.species_dependencies_[1], "B"); - EXPECT_EQ(constraint.species_dependencies_[2], "AB"); - - // Test at equilibrium: [A] = 0.001, [B] = 0.001, [AB] = 0.001 - // K_eq * [A] * [B] = 1000 * 0.001 * 0.001 = 0.001 = [AB] - // Residual should be 0 - std::vector concentrations = { 0.001, 0.001, 0.001 }; - std::vector indices = { 0, 1, 2 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - // Test away from equilibrium: [A] = 0.01, [B] = 0.01, [AB] = 0.001 - // K_eq * [A] * [B] = 1000 * 0.01 * 0.01 = 0.1 - // Residual = 0.1 - 0.001 = 0.099 - concentrations = { 0.01, 0.01, 0.001 }; - residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.099, 1e-10); -} - -TEST(EquilibriumConstraint, Jacobian) -{ - // Test Jacobian for A + B <-> AB with K_eq = 1000 - // G = K_eq * [A] * [B] - [AB] - // dG/d[A] = K_eq * [B] - // dG/d[B] = K_eq * [A] - // dG/d[AB] = -1 - - double K_eq = 1000.0; - EquilibriumConstraint constraint( - "A_B_equilibrium", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - std::vector{ StoichSpecies(Species("AB"), 1.0) }, - K_eq); - - std::vector concentrations = { 0.01, 0.02, 0.05 }; - std::vector indices = { 0, 1, 2 }; - std::vector jacobian(3); - - constraint.Jacobian(concentrations.data(), indices.data(), jacobian.data()); - - EXPECT_NEAR(jacobian[0], K_eq * concentrations[1], 1e-10); // dG/d[A] = K_eq * [B] - EXPECT_NEAR(jacobian[1], K_eq * concentrations[0], 1e-10); // dG/d[B] = K_eq * [A] - EXPECT_NEAR(jacobian[2], -1.0, 1e-10); // dG/d[AB] = -1 -} - -TEST(EquilibriumConstraint, SingleReactantSingleProduct) -{ - // Test: A <-> B with K_eq = 10 - // At equilibrium: [B] / [A] = K_eq - // Constraint: G = K_eq * [A] - [B] = 0 - - double K_eq = 10.0; - EquilibriumConstraint constraint( - "A_B_simple", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - K_eq); - - EXPECT_EQ(constraint.species_dependencies_.size(), 2); - - // At equilibrium: [A] = 0.1, [B] = 1.0 - std::vector concentrations = { 0.1, 1.0 }; - std::vector indices = { 0, 1 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - std::vector jacobian(2); - constraint.Jacobian(concentrations.data(), indices.data(), jacobian.data()); - EXPECT_NEAR(jacobian[0], K_eq, 1e-10); // dG/d[A] = K_eq - EXPECT_NEAR(jacobian[1], -1.0, 1e-10); // dG/d[B] = -1 -} - -TEST(EquilibriumConstraint, TwoProductsOneReactant) -{ - // Test: 2A <-> B + C with K_eq = 100 - // At equilibrium: [B][C] / [A]^2 = K_eq - // Constraint: G = K_eq * [A]^2 - [B] * [C] = 0 - - double K_eq = 100.0; - EquilibriumConstraint constraint( - "dissociation", - std::vector{ StoichSpecies(Species("A"), 2.0) }, // A with stoich 2 - std::vector{ StoichSpecies(Species("B"), 1.0), StoichSpecies(Species("C"), 1.0) }, - K_eq); - - // Dependencies should be A, B, C - EXPECT_EQ(constraint.species_dependencies_.size(), 3); - - // At equilibrium: [A] = 0.1, [B] = 0.5, [C] = 2.0 - // K_eq * [A]^2 = 100 * 0.01 = 1.0 - // [B] * [C] = 0.5 * 2.0 = 1.0 - std::vector concentrations = { 0.1, 0.5, 2.0 }; - std::vector indices = { 0, 1, 2 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); -} - -TEST(EquilibriumConstraint, InvalidEquilibriumConstant) -{ - // Test that negative or zero K_eq throws - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - -1.0), - micm::MicmException); - - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - 0.0), - micm::MicmException); -} - -TEST(EquilibriumConstraint, EmptyReactantsThrows) -{ - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{}, // empty reactants - std::vector{ StoichSpecies(Species("B"), 1.0) }, - 1.0), - micm::MicmException); -} - -TEST(EquilibriumConstraint, EmptyProductsThrows) -{ - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{}, // empty products - 1.0), - micm::MicmException); -} - -TEST(EquilibriumConstraint, InvalidStoichiometryThrows) -{ - // Zero stoichiometry for reactant - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), 0.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - 1.0), - micm::MicmException); - - // Negative stoichiometry for reactant - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), -1.0) }, - std::vector{ StoichSpecies(Species("B"), 1.0) }, - 1.0), - micm::MicmException); - - // Zero stoichiometry for product - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), 0.0) }, - 1.0), - micm::MicmException); - - // Negative stoichiometry for product - EXPECT_THROW( - EquilibriumConstraint( - "invalid", - std::vector{ StoichSpecies(Species("A"), 1.0) }, - std::vector{ StoichSpecies(Species("B"), -2.0) }, - 1.0), - micm::MicmException); -} diff --git a/test/unit/constraint/test_linear.cpp b/test/unit/constraint/test_linear.cpp deleted file mode 100644 index 7fed964c6..000000000 --- a/test/unit/constraint/test_linear.cpp +++ /dev/null @@ -1,227 +0,0 @@ -// Copyright (C) 2023-2026 University Corporation for Atmospheric Research -// SPDX-License-Identifier: Apache-2.0 - -#include -#include -#include -#include - -#include - -#include -#include -#include -#include - -using namespace micm; - -TEST(LinearConstraint, SimpleConservation) -{ - // Test: A + B = 1.0 (total concentration conservation) - // Constraint: G = [A] + [B] - 1.0 = 0 - - LinearConstraint constraint( - "A_B_conservation", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - 1.0); - - EXPECT_EQ(constraint.name_, "A_B_conservation"); - EXPECT_EQ(constraint.constant_, 1.0); - EXPECT_EQ(constraint.species_dependencies_.size(), 2); - EXPECT_EQ(constraint.species_dependencies_[0], "A"); - EXPECT_EQ(constraint.species_dependencies_[1], "B"); - EXPECT_EQ(constraint.terms_.size(), 2); - - // Test when constraint is satisfied: [A] = 0.3, [B] = 0.7 - // Residual = 0.3 + 0.7 - 1.0 = 0 - std::vector concentrations = { 0.3, 0.7 }; - std::vector indices = { 0, 1 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - // Test when constraint is not satisfied: [A] = 0.5, [B] = 0.6 - // Residual = 0.5 + 0.6 - 1.0 = 0.1 - concentrations = { 0.5, 0.6 }; - residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.1, 1e-10); -} - -TEST(LinearConstraint, Jacobian) -{ - // Test Jacobian for A + B = 1.0 - // G = [A] + [B] - 1.0 - // dG/d[A] = 1.0 - // dG/d[B] = 1.0 - - LinearConstraint constraint( - "A_B_conservation", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, - 1.0); - - std::vector concentrations = { 0.3, 0.7 }; - std::vector indices = { 0, 1 }; - std::vector jacobian(2); - - constraint.Jacobian(concentrations.data(), indices.data(), jacobian.data()); - - EXPECT_NEAR(jacobian[0], 1.0, 1e-10); // dG/d[A] = 1.0 - EXPECT_NEAR(jacobian[1], 1.0, 1e-10); // dG/d[B] = 1.0 -} - -TEST(LinearConstraint, WeightedSum) -{ - // Test: 2*A + 3*B - C = 5.0 - // Constraint: G = 2*[A] + 3*[B] - [C] - 5.0 = 0 - - LinearConstraint constraint( - "weighted_sum", - std::vector{ - StoichSpecies(Species("A"), 2.0), - StoichSpecies(Species("B"), 3.0), - StoichSpecies(Species("C"), -1.0) }, - 5.0); - - EXPECT_EQ(constraint.species_dependencies_.size(), 3); - EXPECT_EQ(constraint.constant_, 5.0); - - // Test when constraint is satisfied: [A] = 1.0, [B] = 2.0, [C] = 3.0 - // Residual = 2*1.0 + 3*2.0 - 3.0 - 5.0 = 2 + 6 - 3 - 5 = 0 - std::vector concentrations = { 1.0, 2.0, 3.0 }; - std::vector indices = { 0, 1, 2 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - // Test Jacobian - std::vector jacobian(3); - constraint.Jacobian(concentrations.data(), indices.data(), jacobian.data()); - - EXPECT_NEAR(jacobian[0], 2.0, 1e-10); // dG/d[A] = 2.0 - EXPECT_NEAR(jacobian[1], 3.0, 1e-10); // dG/d[B] = 3.0 - EXPECT_NEAR(jacobian[2], -1.0, 1e-10); // dG/d[C] = -1.0 -} - -TEST(LinearConstraint, ThreeSpeciesConservation) -{ - // Test: A + B + C = 10.0 - // Constraint: G = [A] + [B] + [C] - 10.0 = 0 - - LinearConstraint constraint( - "ABC_total", - std::vector{ - StoichSpecies(Species("A"), 1.0), - StoichSpecies(Species("B"), 1.0), - StoichSpecies(Species("C"), 1.0) }, - 10.0); - - EXPECT_EQ(constraint.species_dependencies_.size(), 3); - EXPECT_EQ(constraint.constant_, 10.0); - - // Test when constraint is satisfied: [A] = 2.0, [B] = 3.0, [C] = 5.0 - std::vector concentrations = { 2.0, 3.0, 5.0 }; - std::vector indices = { 0, 1, 2 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - // Test when constraint is violated: [A] = 2.0, [B] = 3.0, [C] = 6.0 - // Residual = 2.0 + 3.0 + 6.0 - 10.0 = 1.0 - concentrations = { 2.0, 3.0, 6.0 }; - residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 1.0, 1e-10); -} - -TEST(LinearConstraint, AlgebraicSpecies) -{ - // Test that AlgebraicSpecies returns the last species in terms list - LinearConstraint constraint( - "A_B_C_conservation", - std::vector{ - StoichSpecies(Species("A"), 1.0), - StoichSpecies(Species("B"), 1.0), - StoichSpecies(Species("C"), 1.0) }, - 10.0); - - EXPECT_EQ(constraint.AlgebraicSpecies(), "C"); -} - -TEST(LinearConstraint, ZeroConstant) -{ - // Test: A - B = 0 (species balance) - // Constraint: G = [A] - [B] = 0 - - LinearConstraint constraint( - "A_equals_B", - std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), -1.0) }, - 0.0); - - EXPECT_EQ(constraint.constant_, 0.0); - - // Test when constraint is satisfied: [A] = 0.5, [B] = 0.5 - std::vector concentrations = { 0.5, 0.5 }; - std::vector indices = { 0, 1 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - // Test when constraint is violated: [A] = 0.6, [B] = 0.5 - // Residual = 0.6 - 0.5 = 0.1 - concentrations = { 0.6, 0.5 }; - residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.1, 1e-10); -} - -TEST(LinearConstraint, FractionalCoefficients) -{ - // Test: 0.5*A + 1.5*B = 2.0 - // Constraint: G = 0.5*[A] + 1.5*[B] - 2.0 = 0 - - LinearConstraint constraint( - "fractional_conservation", - std::vector{ StoichSpecies(Species("A"), 0.5), StoichSpecies(Species("B"), 1.5) }, - 2.0); - - // Test when constraint is satisfied: [A] = 2.0, [B] = 1.0 - // 0.5*2.0 + 1.5*1.0 = 1.0 + 1.5 = 2.5... wait that's 2.5 - // Let's use [A] = 1.0, [B] = 1.0 - // 0.5*1.0 + 1.5*1.0 = 0.5 + 1.5 = 2.0 ✓ - std::vector concentrations = { 1.0, 1.0 }; - std::vector indices = { 0, 1 }; - - double residual = constraint.Residual(concentrations.data(), indices.data()); - EXPECT_NEAR(residual, 0.0, 1e-10); - - // Test Jacobian - std::vector jacobian(2); - constraint.Jacobian(concentrations.data(), indices.data(), jacobian.data()); - - EXPECT_NEAR(jacobian[0], 0.5, 1e-10); // dG/d[A] = 0.5 - EXPECT_NEAR(jacobian[1], 1.5, 1e-10); // dG/d[B] = 1.5 -} - -TEST(LinearConstraint, JacobianIndependentOfConcentrations) -{ - // Test that Jacobian is constant (independent of concentrations) for linear constraints - LinearConstraint constraint( - "A_B_sum", - std::vector{ StoichSpecies(Species("A"), 2.0), StoichSpecies(Species("B"), 3.0) }, - 1.0); - - // Test at two different concentration points - std::vector conc1 = { 0.1, 0.2 }; - std::vector conc2 = { 10.0, 20.0 }; - std::vector indices = { 0, 1 }; - - std::vector jacobian1(2); - std::vector jacobian2(2); - - constraint.Jacobian(conc1.data(), indices.data(), jacobian1.data()); - constraint.Jacobian(conc2.data(), indices.data(), jacobian2.data()); - - // Jacobian should be the same regardless of concentrations - EXPECT_NEAR(jacobian1[0], jacobian2[0], 1e-10); - EXPECT_NEAR(jacobian1[1], jacobian2[1], 1e-10); - EXPECT_NEAR(jacobian1[0], 2.0, 1e-10); - EXPECT_NEAR(jacobian1[1], 3.0, 1e-10); -} diff --git a/test/unit/constraint/type/CMakeLists.txt b/test/unit/constraint/type/CMakeLists.txt new file mode 100644 index 000000000..2a97d4e2e --- /dev/null +++ b/test/unit/constraint/type/CMakeLists.txt @@ -0,0 +1,4 @@ +################################################################################ +# Tests +create_standard_test(NAME equilibrium SOURCES test_equilibrium.cpp) +create_standard_test(NAME linear SOURCES test_linear.cpp) diff --git a/test/unit/constraint/type/test_equilibrium.cpp b/test/unit/constraint/type/test_equilibrium.cpp new file mode 100644 index 000000000..2b3cbd409 --- /dev/null +++ b/test/unit/constraint/type/test_equilibrium.cpp @@ -0,0 +1,378 @@ +// Copyright (C) 2023-2026 University Corporation for Atmospheric Research +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +using namespace micm; +using StandardSparseMatrix = SparseMatrix; + +TEST(EquilibriumConstraint, Construction) +{ + // Test: A + B <-> AB with K_eq = 1000 + // At equilibrium: [AB] / ([A][B]) = K_eq + // Constraint: G = K_eq * [A] * [B] - [AB] = 0 + + double K_eq = 1000.0; + EquilibriumConstraint constraint( + "A_B_equilibrium", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + K_eq); + + EXPECT_EQ(constraint.name_, "A_B_equilibrium"); + EXPECT_EQ(constraint.equilibrium_constant_, K_eq); + EXPECT_EQ(constraint.species_dependencies_.size(), 3); + EXPECT_EQ(constraint.species_dependencies_[0], "A"); + EXPECT_EQ(constraint.species_dependencies_[1], "B"); + EXPECT_EQ(constraint.species_dependencies_[2], "AB"); + EXPECT_EQ(constraint.reactants_.size(), 2); + EXPECT_EQ(constraint.products_.size(), 1); +} + +TEST(EquilibriumConstraint, AlgebraicSpecies) +{ + // Test that AlgebraicSpecies returns the first product species + + double K_eq = 1000.0; + EquilibriumConstraint constraint( + "A_B_equilibrium", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + K_eq); + + EXPECT_EQ(constraint.AlgebraicSpecies(), "AB"); +} + +TEST(EquilibriumConstraint, SingleReactantSingleProduct) +{ + // Test: A <-> B with K_eq = 10 + // At equilibrium: [B] / [A] = K_eq + // Constraint: G = K_eq * [A] - [B] = 0 + + double K_eq = 10.0; + EquilibriumConstraint constraint( + "A_B_simple", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + K_eq); + + EXPECT_EQ(constraint.name_, "A_B_simple"); + EXPECT_EQ(constraint.species_dependencies_.size(), 2); + EXPECT_EQ(constraint.species_dependencies_[0], "A"); + EXPECT_EQ(constraint.species_dependencies_[1], "B"); + EXPECT_EQ(constraint.AlgebraicSpecies(), "B"); +} + +TEST(EquilibriumConstraint, MultipleReactantsAndProducts) +{ + // Test: 2A <-> B + C with K_eq = 100 + // At equilibrium: [B][C] / [A]^2 = K_eq + // Constraint: G = K_eq * [A]^2 - [B] * [C] = 0 + + double K_eq = 100.0; + EquilibriumConstraint constraint( + "dissociation", + std::vector{ StoichSpecies(Species("A"), 2.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0), StoichSpecies(Species("C"), 1.0) }, + K_eq); + + EXPECT_EQ(constraint.name_, "dissociation"); + EXPECT_EQ(constraint.species_dependencies_.size(), 3); + EXPECT_EQ(constraint.species_dependencies_[0], "A"); + EXPECT_EQ(constraint.species_dependencies_[1], "B"); + EXPECT_EQ(constraint.species_dependencies_[2], "C"); + EXPECT_EQ(constraint.reactants_.size(), 1); + EXPECT_EQ(constraint.reactants_[0].coefficient_, 2.0); + EXPECT_EQ(constraint.products_.size(), 2); + EXPECT_EQ(constraint.AlgebraicSpecies(), "B"); +} + +TEST(EquilibriumConstraint, InvalidEquilibriumConstant) +{ + // Test that negative or zero K_eq throws + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + -1.0), + micm::MicmException); + + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + 0.0), + micm::MicmException); +} + +TEST(EquilibriumConstraint, EmptyReactantsThrows) +{ + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{}, // empty reactants + std::vector{ StoichSpecies(Species("B"), 1.0) }, + 1.0), + micm::MicmException); +} + +TEST(EquilibriumConstraint, EmptyProductsThrows) +{ + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{}, // empty products + 1.0), + micm::MicmException); +} + +TEST(EquilibriumConstraint, InvalidStoichiometryThrows) +{ + // Zero stoichiometry for reactant + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), 0.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + 1.0), + micm::MicmException); + + // Negative stoichiometry for reactant + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), -1.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0) }, + 1.0), + micm::MicmException); + + // Zero stoichiometry for product + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), 0.0) }, + 1.0), + micm::MicmException); + + // Negative stoichiometry for product + EXPECT_THROW( + EquilibriumConstraint( + "invalid", + std::vector{ StoichSpecies(Species("A"), 1.0) }, + std::vector{ StoichSpecies(Species("B"), -2.0) }, + 1.0), + micm::MicmException); +} + +// Integration tests using ConstraintSet to test residual and Jacobian computation + +TEST(EquilibriumConstraint, ResidualComputationThroughConstraintSet) +{ + // Test: A + B <-> AB with K_eq = 1000 + // Constraint: G = K_eq * [A] * [B] - [AB] = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_equilibrium", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // Create state matrix with 1 grid cell and 3 species + DenseMatrix state(1, 3); + DenseMatrix forcing(1, 3); + + // Test at equilibrium: [A] = 0.001, [B] = 0.001, [AB] = 0.001 + // K_eq * [A] * [B] = 1000 * 0.001 * 0.001 = 0.001 = [AB] + // Residual should be 0 + state[0][0] = 0.001; // A + state[0][1] = 0.001; // B + state[0][2] = 0.001; // AB + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + // The forcing term for AB (row 2) should be the constraint residual + EXPECT_NEAR(forcing[0][2], 0.0, 1e-10); + + // Test away from equilibrium: [A] = 0.02, [B] = 0.03, [AB] = 0.05 + // K_eq * [A] * [B] = 1000 * 0.02 * 0.03 = 0.6 + // Residual = 0.6 - 0.05 = 0.55 + state[0][0] = 0.02; // A + state[0][1] = 0.03; // B + state[0][2] = 0.05; // AB + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][2], 0.55, 1e-10); +} + +TEST(EquilibriumConstraint, JacobianComputationThroughConstraintSet) +{ + // Test Jacobian for A + B <-> AB with K_eq = 1000 + // G = K_eq * [A] * [B] - [AB] + // dG/d[A] = K_eq * [B] + // dG/d[B] = K_eq * [A] + // dG/d[AB] = -1 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "A_B_equilibrium", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + std::vector{ StoichSpecies(Species("AB"), 1.0) }, + 1000.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "AB", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for Jacobian using builder + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); // Diagonals + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // Create state matrix + DenseMatrix state(1, 3); + state[0][0] = 0.01; // A + state[0][1] = 0.02; // B + state[0][2] = 0.05; // AB + + // Compute Jacobian + set.SubtractJacobianTerms(state, jacobian); + + // The Jacobian computation uses subtraction convention + // Row 2 (AB, the algebraic species): contains dG/d[A], dG/d[B], dG/d[AB] + double dG_dA = 1000.0 * 0.02; // K_eq * [B] = 20.0 + double dG_dB = 1000.0 * 0.01; // K_eq * [A] = 10.0 + double dG_dAB = -1.0; + + // Due to SubtractJacobianTerms convention, the values are negated + EXPECT_NEAR(jacobian[0][2][0], -dG_dA, 1e-8); // -dG/d[A] = -20.0 + EXPECT_NEAR(jacobian[0][2][1], -dG_dB, 1e-8); // -dG/d[B] = -10.0 + EXPECT_NEAR(jacobian[0][2][2], -dG_dAB, 1e-8); // -dG/d[AB] = 1.0 +} + +TEST(EquilibriumConstraint, ComplexStoichiometryResidual) +{ + // Test: 2A <-> B + C with K_eq = 100 + // Constraint: G = K_eq * [A]^2 - [B] * [C] = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(EquilibriumConstraint( + "dissociation", + std::vector{ StoichSpecies(Species("A"), 2.0) }, + std::vector{ StoichSpecies(Species("B"), 1.0), StoichSpecies(Species("C"), 1.0) }, + 100.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "C", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + DenseMatrix state(1, 3); + DenseMatrix forcing(1, 3); + + // At equilibrium: [A] = 0.1, [B] = 0.5, [C] = 2.0 + // K_eq * [A]^2 = 100 * 0.01 = 1.0 + // [B] * [C] = 0.5 * 2.0 = 1.0 + state[0][0] = 0.1; // A + state[0][1] = 0.5; // B + state[0][2] = 2.0; // C + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + // The forcing term for B (row 1, first product) should be the constraint residual + EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); +} + diff --git a/test/unit/constraint/type/test_linear.cpp b/test/unit/constraint/type/test_linear.cpp new file mode 100644 index 000000000..cb35545da --- /dev/null +++ b/test/unit/constraint/type/test_linear.cpp @@ -0,0 +1,556 @@ +// Copyright (C) 2023-2026 University Corporation for Atmospheric Research +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +using namespace micm; +using StandardSparseMatrix = SparseMatrix; + +TEST(LinearConstraint, Construction) +{ + // Test: A + B = 1.0 (total concentration conservation) + // Constraint: G = [A] + [B] - 1.0 = 0 + + LinearConstraint constraint( + "A_B_conservation", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + 1.0); + + EXPECT_EQ(constraint.name_, "A_B_conservation"); + EXPECT_EQ(constraint.constant_, 1.0); + EXPECT_EQ(constraint.species_dependencies_.size(), 2); + EXPECT_EQ(constraint.species_dependencies_[0], "A"); + EXPECT_EQ(constraint.species_dependencies_[1], "B"); + EXPECT_EQ(constraint.terms_.size(), 2); +} + +TEST(LinearConstraint, AlgebraicSpecies) +{ + // Test that AlgebraicSpecies returns the last species in terms list + LinearConstraint constraint( + "A_B_C_conservation", + std::vector{ + StoichSpecies(Species("A"), 1.0), + StoichSpecies(Species("B"), 1.0), + StoichSpecies(Species("C"), 1.0) }, + 10.0); + + EXPECT_EQ(constraint.AlgebraicSpecies(), "C"); +} + +TEST(LinearConstraint, WeightedTerms) +{ + // Test: 2*A + 3*B - C = 5.0 + // Constraint: G = 2*[A] + 3*[B] - [C] - 5.0 = 0 + + LinearConstraint constraint( + "weighted_sum", + std::vector{ + StoichSpecies(Species("A"), 2.0), + StoichSpecies(Species("B"), 3.0), + StoichSpecies(Species("C"), -1.0) }, + 5.0); + + EXPECT_EQ(constraint.name_, "weighted_sum"); + EXPECT_EQ(constraint.species_dependencies_.size(), 3); + EXPECT_EQ(constraint.constant_, 5.0); + EXPECT_EQ(constraint.terms_.size(), 3); + EXPECT_EQ(constraint.terms_[0].coefficient_, 2.0); + EXPECT_EQ(constraint.terms_[1].coefficient_, 3.0); + EXPECT_EQ(constraint.terms_[2].coefficient_, -1.0); +} + +TEST(LinearConstraint, ZeroConstant) +{ + // Test: A - B = 0 (species balance) + // Constraint: G = [A] - [B] = 0 + + LinearConstraint constraint( + "A_equals_B", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), -1.0) }, + 0.0); + + EXPECT_EQ(constraint.name_, "A_equals_B"); + EXPECT_EQ(constraint.constant_, 0.0); + EXPECT_EQ(constraint.species_dependencies_.size(), 2); +} + +TEST(LinearConstraint, FractionalCoefficients) +{ + // Test: 0.5*A + 1.5*B = 2.0 + // Constraint: G = 0.5*[A] + 1.5*[B] - 2.0 = 0 + + LinearConstraint constraint( + "fractional_conservation", + std::vector{ StoichSpecies(Species("A"), 0.5), StoichSpecies(Species("B"), 1.5) }, + 2.0); + + EXPECT_EQ(constraint.name_, "fractional_conservation"); + EXPECT_EQ(constraint.constant_, 2.0); + EXPECT_EQ(constraint.terms_[0].coefficient_, 0.5); + EXPECT_EQ(constraint.terms_[1].coefficient_, 1.5); +} + +// Integration tests using ConstraintSet to test residual and Jacobian computation + +TEST(LinearConstraint, ResidualComputationThroughConstraintSet) +{ + // Test: A + B = 1.0 (total concentration conservation) + // Constraint: G = [A] + [B] - 1.0 = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "A_B_conservation", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + 1.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 } + }; + + std::size_t num_species = 2; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // Create state matrix with 1 grid cell and 2 species + DenseMatrix state(1, 2); + DenseMatrix forcing(1, 2); + + // Test when constraint is satisfied: [A] = 0.3, [B] = 0.7 + // Residual = 0.3 + 0.7 - 1.0 = 0 + state[0][0] = 0.3; // A + state[0][1] = 0.7; // B + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + // The forcing term for B (row 1, last term) should be the constraint residual + EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); + + // Test when constraint is not satisfied: [A] = 0.5, [B] = 0.6 + // Residual = 0.5 + 0.6 - 1.0 = 0.1 + state[0][0] = 0.5; // A + state[0][1] = 0.6; // B + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][1], 0.1, 1e-10); +} + +TEST(LinearConstraint, JacobianComputationThroughConstraintSet) +{ + // Test Jacobian for A + B = 1.0 + // G = [A] + [B] - 1.0 + // dG/d[A] = 1.0 + // dG/d[B] = 1.0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "A_B_conservation", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), 1.0) }, + 1.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 } + }; + + std::size_t num_species = 2; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for Jacobian using builder + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); // Diagonals + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // Create state matrix + DenseMatrix state(1, 2); + state[0][0] = 0.3; // A + state[0][1] = 0.7; // B + + // Compute Jacobian + set.SubtractJacobianTerms(state, jacobian); + + // Due to SubtractJacobianTerms convention, the values are negated + // Row 1 (B, the algebraic species): contains dG/d[A], dG/d[B] + EXPECT_NEAR(jacobian[0][1][0], -1.0, 1e-10); // -dG/d[A] = -1.0 + EXPECT_NEAR(jacobian[0][1][1], -1.0, 1e-10); // -dG/d[B] = -1.0 +} + +TEST(LinearConstraint, WeightedSumResidualAndJacobian) +{ + // Test: 2*A + 3*B - C = 5.0 + // Constraint: G = 2*[A] + 3*[B] - [C] - 5.0 = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "weighted_sum", + std::vector{ + StoichSpecies(Species("A"), 2.0), + StoichSpecies(Species("B"), 3.0), + StoichSpecies(Species("C"), -1.0) }, + 5.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "C", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + DenseMatrix state(1, 3); + DenseMatrix forcing(1, 3); + + // Test when constraint is satisfied: [A] = 1.0, [B] = 2.0, [C] = 3.0 + // Residual = 2*1.0 + 3*2.0 - 3.0 - 5.0 = 2 + 6 - 3 - 5 = 0 + state[0][0] = 1.0; // A + state[0][1] = 2.0; // B + state[0][2] = 3.0; // C + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + // The forcing term for C (row 2, last term) should be the constraint residual + EXPECT_NEAR(forcing[0][2], 0.0, 1e-10); + + // Test Jacobian - reset jacobian values to zero first + for (auto& elem : jacobian.AsVector()) + { + elem = 0.0; + } + + set.SubtractJacobianTerms(state, jacobian); + + // Row 2 (C, the algebraic species): contains dG/d[A], dG/d[B], dG/d[C] + // Due to SubtractJacobianTerms convention, the values are negated + EXPECT_NEAR(jacobian[0][2][0], -2.0, 1e-10); // -dG/d[A] = -2.0 + EXPECT_NEAR(jacobian[0][2][1], -3.0, 1e-10); // -dG/d[B] = -3.0 + EXPECT_NEAR(jacobian[0][2][2], 1.0, 1e-10); // -dG/d[C] = -(-1.0) = 1.0 +} + +TEST(LinearConstraint, ThreeSpeciesConservationResidual) +{ + // Test: A + B + C = 10.0 + // Constraint: G = [A] + [B] + [C] - 10.0 = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "ABC_total", + std::vector{ + StoichSpecies(Species("A"), 1.0), + StoichSpecies(Species("B"), 1.0), + StoichSpecies(Species("C"), 1.0) }, + 10.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 }, + { "C", 2 } + }; + + std::size_t num_species = 3; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + DenseMatrix state(1, 3); + DenseMatrix forcing(1, 3); + + // Test when constraint is satisfied: [A] = 2.0, [B] = 3.0, [C] = 5.0 + state[0][0] = 2.0; // A + state[0][1] = 3.0; // B + state[0][2] = 5.0; // C + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][2], 0.0, 1e-10); + + // Test when constraint is violated: [A] = 2.0, [B] = 3.0, [C] = 6.0 + // Residual = 2.0 + 3.0 + 6.0 - 10.0 = 1.0 + state[0][0] = 2.0; // A + state[0][1] = 3.0; // B + state[0][2] = 6.0; // C + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][2], 1.0, 1e-10); +} + +TEST(LinearConstraint, ZeroConstantResidual) +{ + // Test: A - B = 0 (species balance) + // Constraint: G = [A] - [B] = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "A_equals_B", + std::vector{ StoichSpecies(Species("A"), 1.0), StoichSpecies(Species("B"), -1.0) }, + 0.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 } + }; + + std::size_t num_species = 2; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + DenseMatrix state(1, 2); + DenseMatrix forcing(1, 2); + + // Test when constraint is satisfied: [A] = 0.5, [B] = 0.5 + state[0][0] = 0.5; // A + state[0][1] = 0.5; // B + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); + + // Test when constraint is violated: [A] = 0.6, [B] = 0.5 + // Residual = 0.6 - 0.5 = 0.1 + state[0][0] = 0.6; // A + state[0][1] = 0.5; // B + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][1], 0.1, 1e-10); +} + +TEST(LinearConstraint, FractionalCoefficientsResidualAndJacobian) +{ + // Test: 0.5*A + 1.5*B = 2.0 + // Constraint: G = 0.5*[A] + 1.5*[B] - 2.0 = 0 + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "fractional_conservation", + std::vector{ StoichSpecies(Species("A"), 0.5), StoichSpecies(Species("B"), 1.5) }, + 2.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 } + }; + + std::size_t num_species = 2; + + ConstraintSet set(std::move(constraints), variable_map); + + // Create sparse matrix for constraint setup + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + DenseMatrix state(1, 2); + DenseMatrix forcing(1, 2); + + // Test when constraint is satisfied: [A] = 1.0, [B] = 1.0 + // 0.5*1.0 + 1.5*1.0 = 0.5 + 1.5 = 2.0 ✓ + state[0][0] = 1.0; // A + state[0][1] = 1.0; // B + + forcing.Fill(0.0); + set.AddForcingTerms(state, forcing); + + EXPECT_NEAR(forcing[0][1], 0.0, 1e-10); + + // Test Jacobian computation + set.SubtractJacobianTerms(state, jacobian); + + // Due to SubtractJacobianTerms convention, the values are negated + EXPECT_NEAR(jacobian[0][1][0], -0.5, 1e-10); // -dG/d[A] = -0.5 + EXPECT_NEAR(jacobian[0][1][1], -1.5, 1e-10); // -dG/d[B] = -1.5 +} + +TEST(LinearConstraint, JacobianIndependentOfConcentrations) +{ + // Test that Jacobian is constant (independent of concentrations) for linear constraints + + using DenseMatrix = Matrix; + + std::vector constraints; + constraints.push_back(LinearConstraint( + "A_B_sum", + std::vector{ StoichSpecies(Species("A"), 2.0), StoichSpecies(Species("B"), 3.0) }, + 1.0)); + + std::unordered_map variable_map = { + { "A", 0 }, + { "B", 1 } + }; + + std::size_t num_species = 2; + + ConstraintSet set(std::move(constraints), variable_map); + + auto non_zero_elements = set.NonZeroJacobianElements(); + + auto builder = StandardSparseMatrix::Create(num_species) + .SetNumberOfBlocks(1) + .InitialValue(0.0); + + for (std::size_t i = 0; i < num_species; ++i) + builder = builder.WithElement(i, i); + for (auto& elem : non_zero_elements) + builder = builder.WithElement(elem.first, elem.second); + + StandardSparseMatrix jacobian{ builder }; + + set.SetJacobianFlatIds(jacobian); + set.SetConstraintFunctions(variable_map, jacobian); + + // Test at two different concentration points + DenseMatrix state1(1, 2); + state1[0][0] = 0.1; // A + state1[0][1] = 0.2; // B + + DenseMatrix state2(1, 2); + state2[0][0] = 10.0; // A + state2[0][1] = 20.0; // B + + // Compute Jacobian at first state + set.SubtractJacobianTerms(state1, jacobian); + + double jac1_dA = jacobian[0][1][0]; + double jac1_dB = jacobian[0][1][1]; + + // Reset jacobian and compute at second state + for (auto& elem : jacobian.AsVector()) + { + elem = 0.0; + } + + set.SubtractJacobianTerms(state2, jacobian); + + double jac2_dA = jacobian[0][1][0]; + double jac2_dB = jacobian[0][1][1]; + + // Jacobian should be the same regardless of concentrations + EXPECT_NEAR(jac1_dA, jac2_dA, 1e-10); + EXPECT_NEAR(jac1_dB, jac2_dB, 1e-10); + EXPECT_NEAR(jac1_dA, -2.0, 1e-10); // -dG/d[A] = -2.0 + EXPECT_NEAR(jac1_dB, -3.0, 1e-10); // -dG/d[B] = -3.0 +} \ No newline at end of file diff --git a/test/unit/solver/test_backward_euler.cpp b/test/unit/solver/test_backward_euler.cpp index 388380ce0..76432e8a3 100644 --- a/test/unit/solver/test_backward_euler.cpp +++ b/test/unit/solver/test_backward_euler.cpp @@ -67,7 +67,8 @@ void CheckIsConverged() { using LinearSolverPolicy = micm::LinearSolver; using RatesPolicy = micm::ProcessSet; - using BackwardEuler = micm::AbstractBackwardEuler; + using ConstraintSetPolicy = micm::ConstraintSet; + using BackwardEuler = micm::AbstractBackwardEuler; micm::BackwardEulerSolverParameters parameters; DenseMatrixPolicy residual{ 4, 3, 0.0 }; diff --git a/test/unit/solver/test_rosenbrock.cpp b/test/unit/solver/test_rosenbrock.cpp index 34ab540fc..1a3a3b515 100644 --- a/test/unit/solver/test_rosenbrock.cpp +++ b/test/unit/solver/test_rosenbrock.cpp @@ -1,7 +1,7 @@ #include "test_rosenbrock_solver_policy.hpp" #include -#include +#include #include #include #include