diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 407c8654..ecc60156 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -34,7 +34,7 @@ jobs: build-macos: name: Build macOS Wheels if: ${{ github.event.inputs.build_os == 'Both' || github.event.inputs.build_os == 'MacOS' }} - runs-on: macos-latest + runs-on: macos-15 strategy: fail-fast: true matrix: diff --git a/.gitignore b/.gitignore index 0c0ee65d..1820c25a 100644 --- a/.gitignore +++ b/.gitignore @@ -81,3 +81,6 @@ IMPLICIT_INTEGRATION_GAP_ANALYSIS.md ION_CHEMISTRY_SUPPORT_PLAN.md KINETICS_BASE_TITAN_GAP_STATUS.md python/kinetics_base/titan/ARCHITECTURE.md + +# pdfs +*.pdf diff --git a/README.md b/README.md index b1e9a4ff..eefc9789 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,34 @@ KINTERA provides efficient implementations of: - Phase equilibrium computations - Atmospheric chemistry models +### Multiphase Equilibrium + +`EquilibriumTP` is a fixed-temperature, fixed-pressure constrained chemistry +solver. The C++/CUDA core accepts component moles and precomputed logarithmic +equilibrium constants; the module derives phase membership and stoichiometry +from its options. Case-specific thermodynamics remains in Python under +`kintera.equilibrium`. + +Equilibrium networks use the repository's top-level `phases`, `species`, and +`reactions` YAML layout. Phase species determine component ordering, species +compositions validate elemental balance, and reactions with `type: equilibrium` +generate the module's stoichiometric buffer: + +```python +from kintera import EquilibriumOptions, EquilibriumTP + +options = EquilibriumOptions.from_yaml("equilibrium.yaml") +solver = EquilibriumTP(options) +``` + +`Nasa9LogK` evaluates ideal-gas equilibrium constants from the bundled NASA-9 +database. See `examples/equilibrium_nasa9.yaml` and +`examples/equilibrium_nasa9.py` for a complete YAML-defined sample: + +```bash +python examples/equilibrium_nasa9.py +``` + The library is written in C++17 with Python bindings, leveraging PyTorch for tensor operations and providing GPU acceleration support via CUDA. ## Features diff --git a/examples/equilibrium_nasa9.py b/examples/equilibrium_nasa9.py new file mode 100644 index 00000000..f45159b9 --- /dev/null +++ b/examples/equilibrium_nasa9.py @@ -0,0 +1,31 @@ +"""Gas equilibrium sample using the standard NASA-9 database.""" + +from pathlib import Path + +import torch +from kintera import EquilibriumOptions, EquilibriumTP +from kintera.equilibrium import Nasa9LogK + + +def main() -> None: + config = Path(__file__).with_suffix(".yaml") + options = EquilibriumOptions.from_yaml(str(config)) + species = options.components() + solver = EquilibriumTP(options) + log_k_model = Nasa9LogK(options) + + temp = torch.tensor(4000.0, dtype=torch.float64) + pressure = torch.tensor(1.0e5, dtype=torch.float64) + initial_moles = torch.tensor([0.60, 0.20, 0.20], dtype=torch.float64) + log_k = log_k_model(temp, pressure) + moles, _, diagnostics = solver(temp, pressure, initial_moles, log_k) + + print(f"T = {temp.item():.0f} K, P = {pressure.item():.0f} Pa") + print("log(K):", dict(zip(options.reactions(), log_k.tolist()))) + print("initial moles:", dict(zip(species, initial_moles.tolist()))) + print("equilibrium moles:", dict(zip(species, moles.tolist()))) + print("diagnostics [status, iterations, error]:", diagnostics.tolist()) + + +if __name__ == "__main__": + main() diff --git a/examples/equilibrium_nasa9.yaml b/examples/equilibrium_nasa9.yaml new file mode 100644 index 00000000..98723a41 --- /dev/null +++ b/examples/equilibrium_nasa9.yaml @@ -0,0 +1,24 @@ +phases: + - name: gas + thermo: ideal-gas + species: [H2, O2, H2O] + +species: + - name: H2 + composition: {H: 2} + + - name: O2 + composition: {O: 2} + + - name: H2O + composition: {H: 2, O: 1} + +reactions: + - equation: "H2 + 0.5 O2 <=> H2O" + type: equilibrium + +equilibrium: + standard-pressure: 1.0e5 + max-iter: 100 + ftol: 1.0e-8 + mole-floor: 1.0e-30 diff --git a/pyproject.toml b/pyproject.toml index 9e8b47d1..ee4a75ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ packages = [ "kintera.atm2d.conservation", "kintera.atm2d.newton", "kintera.atm2d.sources", + "kintera.equilibrium", "kintera.kinetics_base", "kintera.kinetics_base.titan", "kintera.kinetics_base_titan", @@ -64,6 +65,7 @@ include-package-data = false "kintera" = [ "kintera/**/*", "data/*.dat", + "equilibrium/*.yaml", "lib/*.so", "lib/*.dylib", "*.pyi", diff --git a/python/__init__.py b/python/__init__.py index 23ef16a7..98647cf6 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -18,6 +18,7 @@ def _add_packaged_resource_directory() -> None: _add_packaged_resource_directory() from .atm2d import * +from .equilibrium import * from .kinetics_base_titan import * try: diff --git a/python/csrc/kintera.cpp b/python/csrc/kintera.cpp index 7c04c246..0f85581f 100644 --- a/python/csrc/kintera.cpp +++ b/python/csrc/kintera.cpp @@ -11,6 +11,7 @@ #include #include #include +#include // python #include "pyoptions.hpp" @@ -27,24 +28,34 @@ extern std::vector species_sref_R; } // namespace kintera -void bind_thermo(py::module& m); -void bind_constants(py::module& m); -void bind_kinetics(py::module& m); -void bind_photolysis(py::module& m); -void bind_diffusion(py::module& m); -void bind_sparse_solver(py::module& m); +void bind_thermo(py::module &m); +void bind_constants(py::module &m); +void bind_kinetics(py::module &m); +void bind_photolysis(py::module &m); +void bind_diffusion(py::module &m); +void bind_sparse_solver(py::module &m); +void bind_equilibrium(py::module &m); PYBIND11_MODULE(kintera, m) { m.attr("__name__") = "kintera"; m.doc() = R"(Atmospheric Thermodynamics and Chemistry Library)"; + m.def("atomic_mass", &kintera::atomic_mass, py::arg("element")); + m.def("molar_mass", &kintera::molar_mass, py::arg("composition")); + m.def("molar_masses", &kintera::molar_masses, py::arg("elements"), + py::arg("element_matrix")); + m.def("molar_masses_from_yaml", &kintera::molar_masses_from_yaml, + py::arg("filename")); + m.def("nasa9_gibbs_rt", &kintera::nasa9_gibbs_rt, py::arg("temp"), + py::arg("species")); + auto pySpeciesThermo = py::class_( m, "SpeciesThermo"); pySpeciesThermo.def(py::init<>()) .def("__repr__", - [](const kintera::SpeciesThermo& self) { + [](const kintera::SpeciesThermo &self) { return fmt::format("SpeciesThermo({})", self); }) .def("species", &kintera::SpeciesThermoImpl::species) @@ -59,9 +70,9 @@ PYBIND11_MODULE(kintera, m) { auto pyReaction = py::class_(m, "Reaction"); pyReaction.def(py::init<>()) - .def(py::init()) + .def(py::init()) .def("__repr__", - [](const kintera::Reaction& self) { + [](const kintera::Reaction &self) { return fmt::format("Reaction({})", self); }) .def("equation", &kintera::Reaction::equation) @@ -209,7 +220,7 @@ PYBIND11_MODULE(kintera, m) { &kintera::KBTitanReactionReport::unsupported_reaction_ids); m.def("classify_kinetics_base_titan_reactions", - py::overload_cast( + py::overload_cast( &kintera::classify_kinetics_base_titan_reactions), py::arg("pun_path"), py::arg("run_input_path")); @@ -219,45 +230,49 @@ PYBIND11_MODULE(kintera, m) { bind_photolysis(m); bind_diffusion(m); bind_sparse_solver(m); + bind_equilibrium(m); - m.def("species_names", []() -> const std::vector& { + m.def("species_names", []() -> const std::vector & { return kintera::species_names; }); - m.def("set_species_names", [](const std::vector& names) { + m.def("set_species_names", [](const std::vector &names) { kintera::species_names = names; return kintera::species_names; }); - m.def("species_weights", []() -> const std::vector& { + m.def("species_weights", []() -> const std::vector & { return kintera::species_weights; }); - m.def("set_species_weights", [](const std::vector& weights) { + m.def("set_species_weights", [](const std::vector &weights) { kintera::species_weights = weights; return kintera::species_weights; }); - m.def("species_cref_R", - []() -> const std::vector& { return kintera::species_cref_R; }); + m.def("species_cref_R", []() -> const std::vector & { + return kintera::species_cref_R; + }); - m.def("set_species_cref_R", [](const std::vector& cref_R) { + m.def("set_species_cref_R", [](const std::vector &cref_R) { kintera::species_cref_R = cref_R; return kintera::species_cref_R; }); - m.def("species_uref_R", - []() -> const std::vector& { return kintera::species_uref_R; }); + m.def("species_uref_R", []() -> const std::vector & { + return kintera::species_uref_R; + }); - m.def("set_species_uref_R", [](const std::vector& uref_R) { + m.def("set_species_uref_R", [](const std::vector &uref_R) { kintera::species_uref_R = uref_R; return kintera::species_uref_R; }); - m.def("species_sref_R", - []() -> const std::vector& { return kintera::species_sref_R; }); + m.def("species_sref_R", []() -> const std::vector & { + return kintera::species_sref_R; + }); - m.def("set_species_sref_R", [](const std::vector& sref_R) { + m.def("set_species_sref_R", [](const std::vector &sref_R) { kintera::species_sref_R = sref_R; return kintera::species_sref_R; }); diff --git a/python/csrc/pyequilibrium.cpp b/python/csrc/pyequilibrium.cpp new file mode 100644 index 00000000..18e53f57 --- /dev/null +++ b/python/csrc/pyequilibrium.cpp @@ -0,0 +1,44 @@ +#include +#include +#include + +#include +#include + +#include "pyoptions.hpp" + +namespace py = pybind11; + +void bind_equilibrium(py::module &m) { + auto pyOptions = + py::class_( + m, "EquilibriumOptions"); + + pyOptions.def(py::init<>()) + .def("__repr__", + [](kintera::EquilibriumOptions const &self) { + std::stringstream ss; + self->report(ss); + return ss.str(); + }) + .def_static("from_yaml", &kintera::EquilibriumOptionsImpl::from_yaml, + py::arg("filename"), py::arg("verbose") = false) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, + components) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, + phases) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, + reactions) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, phase_ids) + .ADD_OPTION(int, kintera::EquilibriumOptionsImpl, gas_phase) + .ADD_OPTION(double, kintera::EquilibriumOptionsImpl, standard_pressure) + .ADD_OPTION(int, kintera::EquilibriumOptionsImpl, max_iter) + .ADD_OPTION(double, kintera::EquilibriumOptionsImpl, ftol) + .ADD_OPTION(double, kintera::EquilibriumOptionsImpl, mole_floor) + .def("validate", &kintera::EquilibriumOptionsImpl::validate); + + ADD_KINTERA_MODULE(EquilibriumTP, EquilibriumOptions, + &kintera::EquilibriumTPImpl::forward, py::arg("temp"), + py::arg("pres"), py::arg("moles"), py::arg("log_k"), + py::arg("warm_start") = false); +} diff --git a/python/equilibrium/__init__.py b/python/equilibrium/__init__.py new file mode 100644 index 00000000..244c2cc5 --- /dev/null +++ b/python/equilibrium/__init__.py @@ -0,0 +1,3 @@ +from .nasa9 import Nasa9LogK + +__all__ = ["Nasa9LogK"] diff --git a/python/equilibrium/nasa9.py b/python/equilibrium/nasa9.py new file mode 100644 index 00000000..c7c024d3 --- /dev/null +++ b/python/equilibrium/nasa9.py @@ -0,0 +1,28 @@ +"""NASA-9 equilibrium-constant models.""" + +from collections.abc import Sequence + +import torch + +from ..kintera import EquilibriumOptions, EquilibriumTP, nasa9_gibbs_rt + + +class Nasa9LogK: + """Compute reaction log(K) from bundled ideal-gas NASA-9 data.""" + + def __init__( + self, + options: EquilibriumOptions, + nasa9_species: Sequence[str] | None = None, + ) -> None: + self.species = list(nasa9_species or options.components()) + if len(self.species) != len(options.components()): + raise ValueError("nasa9_species must contain one name per component") + self.stoich = EquilibriumTP(options).buffer("stoich") + + def __call__(self, temp: torch.Tensor, pressure: torch.Tensor) -> torch.Tensor: + """Return natural-log equilibrium constants in option reaction order.""" + del pressure + gibbs_rt = nasa9_gibbs_rt(temp, self.species) + stoich = self.stoich.to(device=temp.device, dtype=temp.dtype) + return -(gibbs_rt @ stoich) diff --git a/python/kintera.pyi b/python/kintera.pyi index 3f732668..e2675e5e 100644 --- a/python/kintera.pyi +++ b/python/kintera.pyi @@ -6,12 +6,20 @@ offering efficient implementations of chemical kinetics calculations, thermodynamic equation of state, and phase equilibrium computations. """ -from typing import Dict, List, Optional, Union, overload +from typing import Any, Dict, List, Optional, Tuple, Union, overload import torch # Type aliases Composition = Dict[str, float] +def atomic_mass(element: str) -> float: ... +def molar_mass(composition: Composition) -> float: ... +def molar_masses( + elements: List[str], element_matrix: List[List[float]] +) -> List[float]: ... +def molar_masses_from_yaml(filename: str) -> List[float]: ... +def nasa9_gibbs_rt(temp: torch.Tensor, species: List[str]) -> torch.Tensor: ... + class SpeciesThermo: """ Species thermodynamics configuration. @@ -456,6 +464,41 @@ class NucleationOptions: """ ... +class EquilibriumOptions: + def __init__(self) -> None: ... + @staticmethod + def from_yaml(filename: str, verbose: bool = False) -> EquilibriumOptions: ... + def components(self, value: List[str] = ...) -> Any: ... + def phases(self, value: List[str] = ...) -> Any: ... + def reactions(self, value: List[str] = ...) -> Any: ... + def phase_ids(self, value: List[int] = ...) -> Any: ... + def gas_phase(self, value: int = ...) -> Any: ... + def standard_pressure(self, value: float = ...) -> Any: ... + def max_iter(self, value: int = ...) -> Any: ... + def ftol(self, value: float = ...) -> Any: ... + def mole_floor(self, value: float = ...) -> Any: ... + def validate(self) -> None: ... + +class EquilibriumTP: + options: EquilibriumOptions + def __init__(self, options: EquilibriumOptions) -> None: ... + def forward( + self, + temp: torch.Tensor, + pres: torch.Tensor, + moles: torch.Tensor, + log_k: torch.Tensor, + warm_start: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... + def __call__( + self, + temp: torch.Tensor, + pres: torch.Tensor, + moles: torch.Tensor, + log_k: torch.Tensor, + warm_start: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... + class ThermoOptions(SpeciesThermo): """ Configuration options for thermodynamic calculations. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c0893286..ef50cbec 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -37,6 +37,7 @@ file(GLOB src_files CONFIGURE_DEPENDS math/*.cpp units/*.cpp diffusion/*.cpp + equilibrium/*.cpp ) add_library(${namel}_${buildl} @@ -70,6 +71,7 @@ if (CUDAToolkit_FOUND) thermo/*.cu math/*.cu kinetics/*.cu + equilibrium/*.cu ) add_library(${namel}_cuda_${buildl} diff --git a/src/equilibrium/equilibrium.cpp b/src/equilibrium/equilibrium.cpp new file mode 100644 index 00000000..68143636 --- /dev/null +++ b/src/equilibrium/equilibrium.cpp @@ -0,0 +1,160 @@ +#include "equilibrium.hpp" + +#include + +#include +#include +#include +#include + +#include "equilibrium_dispatch.hpp" + +namespace kintera { + +void EquilibriumOptionsImpl::report(std::ostream &os) const { + os << "-- equilibrium options --\n" + << "* components = " << components().size() << "\n" + << "* phases = " << phases().size() << "\n" + << "* reactions = " << reactions().size() << "\n" + << "* gas_phase = " << gas_phase() << "\n" + << "* standard_pressure = " << standard_pressure() << " Pa\n" + << "* max_iter = " << max_iter() << "\n" + << "* ftol = " << ftol() << "\n"; +} + +void EquilibriumOptionsImpl::validate() const { + TORCH_CHECK(!components().empty(), "Equilibrium requires components"); + TORCH_CHECK(!phases().empty(), "Equilibrium requires phases"); + TORCH_CHECK(phase_ids().size() == components().size(), + "phase_ids must contain one entry per component"); + TORCH_CHECK(!reactions().empty(), "Equilibrium requires reactions"); + TORCH_CHECK(components().size() <= 64, + "the active-set solver currently supports at most 64 components"); + for (int phase : phase_ids()) { + TORCH_CHECK(phase >= 0 && phase < static_cast(phases().size()), + "phase id out of range: ", phase); + } + TORCH_CHECK( + gas_phase() >= 0 && gas_phase() < static_cast(phases().size()), + "gas_phase is out of range"); + TORCH_CHECK(standard_pressure() > 0., "standard_pressure must be positive"); + TORCH_CHECK(max_iter() > 0, "max_iter must be positive"); + TORCH_CHECK(ftol() > 0., "ftol must be positive"); + TORCH_CHECK(mole_floor() >= 0., "mole_floor must be nonnegative"); + + for (auto const &equation : reactions()) { + Reaction reaction(equation); + for (auto const &[name, coefficient] : reaction.reactants()) { + (void)coefficient; + TORCH_CHECK(std::find(components().begin(), components().end(), name) != + components().end(), + "reaction references unknown component: ", name); + } + for (auto const &[name, coefficient] : reaction.products()) { + (void)coefficient; + TORCH_CHECK(std::find(components().begin(), components().end(), name) != + components().end(), + "reaction references unknown component: ", name); + } + } +} + +EquilibriumTPImpl::EquilibriumTPImpl(EquilibriumOptions const &options_) + : options(options_) { + reset(); +} + +void EquilibriumTPImpl::reset() { + options->validate(); + std::vector reactions; + reactions.reserve(options->reactions().size()); + for (auto const &equation : options->reactions()) + reactions.emplace_back(equation); + stoich = register_buffer( + "stoich", generate_stoichiometry_matrix(reactions, options->components()) + .transpose(0, 1) + .contiguous() + .to(torch::kFloat64)); + auto rank = torch::linalg_matrix_rank(stoich).item(); + TORCH_CHECK(rank == stoich.size(1), + "stoich reaction columns must be independent"); + phase_ids = register_buffer("phase_ids", + torch::tensor(options->phase_ids(), torch::kInt)); +} + +void EquilibriumTPImpl::pretty_print(std::ostream &os) const { + os << "EquilibriumTP(components=" << options->components().size() + << ", phases=" << options->phases().size() << ")"; +} + +std::tuple +EquilibriumTPImpl::forward(torch::Tensor temp, torch::Tensor pres, + torch::Tensor moles, torch::Tensor log_k, + bool warm_start) { + (void)warm_start; + TORCH_CHECK(moles.dim() >= 1, "moles must have a component dimension"); + TORCH_CHECK(log_k.dim() >= 1, "log_k must have a reaction dimension"); + int nspecies = stoich.size(0); + int nreaction = stoich.size(1); + TORCH_CHECK(moles.size(-1) == nspecies, "moles last dimension must be ", + nspecies); + TORCH_CHECK(log_k.size(-1) == nreaction, "log_k last dimension must be ", + nreaction); + TORCH_CHECK(moles.is_floating_point() && log_k.is_floating_point(), + "moles and log_k must be floating-point tensors"); + TORCH_CHECK(moles.device() == log_k.device() && + moles.device() == temp.device() && + moles.device() == pres.device(), + "all equilibrium inputs must be on the same device"); + TORCH_CHECK(moles.scalar_type() == log_k.scalar_type() && + moles.scalar_type() == temp.scalar_type() && + moles.scalar_type() == pres.scalar_type(), + "all equilibrium inputs must have the same dtype"); + + auto moles_scalar = moles.select(-1, 0); + auto logk_scalar = log_k.select(-1, 0); + auto broadcast = + torch::broadcast_tensors({temp, pres, moles_scalar, logk_scalar}); + auto batch = broadcast[0].sizes().vec(); + auto component_shape = batch; + component_shape.push_back(nspecies); + auto reaction_shape = batch; + reaction_shape.push_back(nreaction); + + auto out_moles = moles.expand(component_shape).contiguous().clone(); + auto logk_expanded = log_k.expand(reaction_shape).contiguous(); + auto gain_shape = batch; + gain_shape.push_back(nreaction * nreaction); + auto gain = torch::zeros(gain_shape, moles.options()); + auto diag_shape = batch; + diag_shape.push_back(3); + auto diag = torch::zeros(diag_shape, moles.options()); + + auto iter = at::TensorIteratorConfig() + .resize_outputs(false) + .check_all_same_dtype(false) + .declare_static_shape(component_shape, + {static_cast(batch.size())}) + .add_output(gain) + .add_output(diag) + .add_output(out_moles) + .add_owned_input(broadcast[0].unsqueeze(-1)) + .add_owned_input(broadcast[1].unsqueeze(-1)) + .add_owned_input(moles.expand(component_shape)) + .add_input(logk_expanded) + .build(); + + auto s = stoich.to(moles.options()).contiguous(); + auto p = phase_ids.to(moles.device(), torch::kInt).contiguous(); + at::native::call_equilibrium(moles.device().type(), iter, s, p, + options->phases().size(), options->gas_phase(), + options->standard_pressure(), options->ftol(), + options->mole_floor(), options->max_iter()); + + gain_shape.pop_back(); + gain_shape.push_back(nreaction); + gain_shape.push_back(nreaction); + return {out_moles, gain.view(gain_shape), diag}; +} + +} // namespace kintera diff --git a/src/equilibrium/equilibrium.hpp b/src/equilibrium/equilibrium.hpp new file mode 100644 index 00000000..57200d82 --- /dev/null +++ b/src/equilibrium/equilibrium.hpp @@ -0,0 +1,68 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace kintera { + +struct EquilibriumOptionsImpl final { + static std::shared_ptr create() { + return std::make_shared(); + } + + static std::shared_ptr from_yaml( + std::string const &filename, bool verbose = false); + + void report(std::ostream &os) const; + void validate() const; + + ADD_ARG(std::vector, components); + ADD_ARG(std::vector, phases); + ADD_ARG(std::vector, reactions); + ADD_ARG(std::vector, phase_ids); + ADD_ARG(int, gas_phase) = 0; + ADD_ARG(double, standard_pressure) = 1.e5; + ADD_ARG(int, max_iter) = 50; + ADD_ARG(double, ftol) = 1.e-8; + ADD_ARG(double, mole_floor) = 1.e-30; +}; +using EquilibriumOptions = std::shared_ptr; + +class EquilibriumTPImpl : public torch::nn::Cloneable { + public: + EquilibriumOptions options; + + torch::Tensor stoich; + torch::Tensor phase_ids; + + EquilibriumTPImpl() : options(EquilibriumOptionsImpl::create()) {} + explicit EquilibriumTPImpl(EquilibriumOptions const &options_); + + void reset() override; + void pretty_print(std::ostream &os) const override; + + //! Solve fixed-temperature, fixed-pressure multiphase equilibrium. + /*! + * Shapes use arbitrary broadcast batch dimensions: + * temp, pres: (...) + * moles: (..., ncomponent) + * log_k: (..., nreaction) + * + * Returns equilibrium moles, the final reaction Jacobian/gain matrix, and + * diagnostics (..., 3): status, iterations, and equilibrium error. + */ + std::tuple forward( + torch::Tensor temp, torch::Tensor pres, torch::Tensor moles, + torch::Tensor log_k, bool warm_start = false); +}; + +TORCH_MODULE(EquilibriumTP); + +} // namespace kintera diff --git a/src/equilibrium/equilibrium_dispatch.cpp b/src/equilibrium/equilibrium_dispatch.cpp new file mode 100644 index 00000000..167a33e1 --- /dev/null +++ b/src/equilibrium/equilibrium_dispatch.cpp @@ -0,0 +1,50 @@ +#include "equilibrium_dispatch.hpp" + +#include +#include +#include + +#include "phase_equilibrate_tp.h" + +namespace kintera { + +void call_equilibrium_cpu(at::TensorIterator &iter, at::Tensor const &stoich, + at::Tensor const &phase_ids, int nphase, + int gas_phase, double standard_pressure, double ftol, + double mole_floor, int max_iter) { + int grain_size = std::max(1, iter.numel() / at::get_num_threads()); + AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "call_equilibrium_cpu", [&] { + int nspecies = stoich.size(0); + int nreaction = stoich.size(1); + auto stoich_ptr = stoich.data_ptr(); + auto phase_ptr = phase_ids.data_ptr(); + iter.for_each( + [&](char **data, const int64_t *strides, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + auto gain = reinterpret_cast(data[0] + i * strides[0]); + auto diag = reinterpret_cast(data[1] + i * strides[1]); + auto out = reinterpret_cast(data[2] + i * strides[2]); + auto temp = reinterpret_cast(data[3] + i * strides[3]); + auto pres = reinterpret_cast(data[4] + i * strides[4]); + auto moles = reinterpret_cast(data[5] + i * strides[5]); + auto log_k = reinterpret_cast(data[6] + i * strides[6]); + phase_equilibrate_tp(gain, diag, out, *temp, *pres, moles, log_k, + stoich_ptr, phase_ptr, nspecies, nreaction, + nphase, gas_phase, + static_cast(standard_pressure), + static_cast(ftol), + static_cast(mole_floor), max_iter); + } + }, + grain_size); + }); +} + +} // namespace kintera + +namespace at::native { + +DEFINE_DISPATCH(call_equilibrium); +REGISTER_ALL_CPU_DISPATCH(call_equilibrium, &kintera::call_equilibrium_cpu); + +} // namespace at::native diff --git a/src/equilibrium/equilibrium_dispatch.cu b/src/equilibrium/equilibrium_dispatch.cu new file mode 100644 index 00000000..c14465f1 --- /dev/null +++ b/src/equilibrium/equilibrium_dispatch.cu @@ -0,0 +1,67 @@ +#include +#include + +#include + +#include "phase_equilibrate_tp.h" +#include "equilibrium_dispatch.hpp" + +namespace kintera { + +template +size_t equilibrium_space(int nspecies, int nreaction, int nphase) { + size_t bytes = 0; + auto bump = [&](size_t align, size_t nbytes) { + bytes = static_cast(align_up(bytes, align)) + nbytes; + }; + bump(alignof(T), nphase * sizeof(T)); + bump(alignof(T), nphase * nreaction * sizeof(T)); + bump(alignof(T), nreaction * sizeof(T)); + bump(alignof(T), nreaction * nreaction * sizeof(T)); + bump(alignof(T), nspecies * nreaction * sizeof(T)); + bump(alignof(T), nspecies * sizeof(T)); + bump(alignof(T), nreaction * sizeof(T)); + bump(alignof(T), nspecies * sizeof(T)); + return bytes + leastsq_kkt_space(nreaction, nspecies); +} + +void call_equilibrium_cuda(at::TensorIterator &iter, at::Tensor const &stoich, + at::Tensor const &phase_ids, int nphase, + int gas_phase, double standard_pressure, double ftol, + double mole_floor, int max_iter) { + at::cuda::CUDAGuard device_guard(iter.device()); + AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "call_equilibrium_cuda", [&] { + int nspecies = stoich.size(0); + int nreaction = stoich.size(1); + auto stoich_ptr = stoich.data_ptr(); + auto phase_ptr = phase_ids.data_ptr(); + size_t work_size = equilibrium_space(nspecies, nreaction, nphase); + // Each equilibrium cell receives an independent global-memory workspace. + native::gpu_global_mem_kernel<32, 7>( + iter, work_size, + [=] GPU_LAMBDA(char *const data[7], unsigned int strides[7], + char *work) { + auto gain = reinterpret_cast(data[0] + strides[0]); + auto diag = reinterpret_cast(data[1] + strides[1]); + auto out = reinterpret_cast(data[2] + strides[2]); + auto temp = reinterpret_cast(data[3] + strides[3]); + auto pres = reinterpret_cast(data[4] + strides[4]); + auto moles = reinterpret_cast(data[5] + strides[5]); + auto log_k = reinterpret_cast(data[6] + strides[6]); + phase_equilibrate_tp( + gain, diag, out, *temp, *pres, moles, log_k, stoich_ptr, + phase_ptr, nspecies, nreaction, nphase, gas_phase, + static_cast(standard_pressure), + static_cast(ftol), static_cast(mole_floor), + max_iter, work); + }); + }); +} + +} // namespace kintera + +namespace at::native { + +REGISTER_CUDA_DISPATCH(call_equilibrium, &kintera::call_equilibrium_cuda); + +} // namespace at::native diff --git a/src/equilibrium/equilibrium_dispatch.hpp b/src/equilibrium/equilibrium_dispatch.hpp new file mode 100644 index 00000000..30423ede --- /dev/null +++ b/src/equilibrium/equilibrium_dispatch.hpp @@ -0,0 +1,16 @@ +#pragma once + +#include +#include + +namespace at::native { + +using equilibrium_fn = void (*)(at::TensorIterator &iter, + at::Tensor const &stoich, + at::Tensor const &phase_ids, int nphase, + int gas_phase, double standard_pressure, + double ftol, double mole_floor, int max_iter); + +DECLARE_DISPATCH(equilibrium_fn, call_equilibrium); + +} // namespace at::native diff --git a/src/equilibrium/equilibrium_options.cpp b/src/equilibrium/equilibrium_options.cpp new file mode 100644 index 00000000..a2a7bd33 --- /dev/null +++ b/src/equilibrium/equilibrium_options.cpp @@ -0,0 +1,133 @@ +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "equilibrium.hpp" + +namespace kintera { + +EquilibriumOptions EquilibriumOptionsImpl::from_yaml( + std::string const &filename, bool verbose) { + auto config = YAML::LoadFile(find_resource(filename)); + TORCH_CHECK(config["phases"], "equilibrium YAML requires 'phases'"); + TORCH_CHECK(config["species"], "equilibrium YAML requires 'species'"); + TORCH_CHECK(config["reactions"], "equilibrium YAML requires 'reactions'"); + + auto options = EquilibriumOptionsImpl::create(); + std::set seen_components; + int gas_phase_count = 0; + + for (auto const &phase_node : config["phases"]) { + TORCH_CHECK(phase_node["name"], "phase is missing 'name'"); + TORCH_CHECK(phase_node["species"], "phase is missing 'species'"); + int phase_id = static_cast(options->phases().size()); + options->phases().push_back(phase_node["name"].as()); + + std::string model; + if (phase_node["thermo"]) { + model = phase_node["thermo"].as(); + } else if (phase_node["type"]) { + model = phase_node["type"].as(); + } + if (model == "ideal-gas") { + options->gas_phase(phase_id); + ++gas_phase_count; + } + + for (auto const &component_node : phase_node["species"]) { + auto component = component_node.as(); + TORCH_CHECK(seen_components.insert(component).second, + "component appears in more than one phase: ", component); + options->components().push_back(component); + options->phase_ids().push_back(phase_id); + } + } + TORCH_CHECK(gas_phase_count == 1, + "equilibrium YAML requires exactly one ideal-gas phase"); + + std::map> compositions; + std::set element_set; + for (auto const &species_node : config["species"]) { + TORCH_CHECK(species_node["name"], "species is missing 'name'"); + TORCH_CHECK(species_node["composition"], + "species is missing 'composition'"); + auto name = species_node["name"].as(); + TORCH_CHECK(compositions.find(name) == compositions.end(), + "duplicate species definition: ", name); + for (auto const &entry : species_node["composition"]) { + auto element = entry.first.as(); + compositions[name][element] = entry.second.as(); + element_set.insert(element); + } + } + for (auto const &component : options->components()) { + TORCH_CHECK(compositions.find(component) != compositions.end(), + "phase component has no species definition: ", component); + } + + std::vector reactions; + for (auto const &reaction_node : config["reactions"]) { + if (reaction_node["type"] && + reaction_node["type"].as() != "equilibrium") { + continue; + } + TORCH_CHECK(reaction_node["equation"], "reaction is missing 'equation'"); + auto equation = reaction_node["equation"].as(); + options->reactions().push_back(equation); + reactions.emplace_back(equation); + } + TORCH_CHECK(!reactions.empty(), "equilibrium YAML contains no reactions"); + + for (size_t j = 0; j < reactions.size(); ++j) { + for (auto const &[name, coefficient] : reactions[j].reactants()) { + auto found = std::find(options->components().begin(), + options->components().end(), name); + TORCH_CHECK(found != options->components().end(), + "reaction references unknown component: ", name); + (void)coefficient; + } + for (auto const &[name, coefficient] : reactions[j].products()) { + auto found = std::find(options->components().begin(), + options->components().end(), name); + TORCH_CHECK(found != options->components().end(), + "reaction references unknown component: ", name); + (void)coefficient; + } + for (auto const &element : element_set) { + double balance = 0.; + for (auto const &[name, coefficient] : reactions[j].reactants()) { + auto found = compositions.at(name).find(element); + if (found != compositions.at(name).end()) + balance -= coefficient * found->second; + } + for (auto const &[name, coefficient] : reactions[j].products()) { + auto found = compositions.at(name).find(element); + if (found != compositions.at(name).end()) + balance += coefficient * found->second; + } + TORCH_CHECK(std::abs(balance) <= 1.e-10, "reaction ", j, + " does not conserve element ", element); + } + } + + if (config["equilibrium"]) { + auto equilibrium = config["equilibrium"]; + options->standard_pressure( + equilibrium["standard-pressure"].as(1.e5)); + options->max_iter(equilibrium["max-iter"].as(50)); + options->ftol(equilibrium["ftol"].as(1.e-8)); + options->mole_floor(equilibrium["mole-floor"].as(1.e-30)); + } + + options->validate(); + if (verbose) options->report(std::cout); + return options; +} + +} // namespace kintera diff --git a/src/equilibrium/phase_equilibrate_tp.h b/src/equilibrium/phase_equilibrate_tp.h new file mode 100644 index 00000000..f2ba672c --- /dev/null +++ b/src/equilibrium/phase_equilibrate_tp.h @@ -0,0 +1,190 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace kintera { + +template +DISPATCH_MACRO T equilibrium_max_error(T const *moles, T pres, + T standard_pressure, T const *log_k, + T const *stoich, int const *phase_ids, + int nspecies, int nreaction, int nphase, + int gas_phase, T *phase_totals, + T *residual) { + for (int p = 0; p < nphase; ++p) phase_totals[p] = 0.; + for (int i = 0; i < nspecies; ++i) { + phase_totals[phase_ids[i]] += moles[i]; + } + + for (int j = 0; j < nreaction; ++j) residual[j] = -log_k[j]; + for (int i = 0; i < nspecies; ++i) { + T log_activity = log(moles[i] / phase_totals[phase_ids[i]]); + if (phase_ids[i] == gas_phase) { + log_activity += log(pres / standard_pressure); + } + for (int j = 0; j < nreaction; ++j) { + residual[j] += stoich[i * nreaction + j] * log_activity; + } + } + + T max_error = 0.; + for (int j = 0; j < nreaction; ++j) { + T value = fabs(residual[j]); + if (value > max_error) max_error = value; + } + return max_error; +} + +template +DISPATCH_MACRO int phase_equilibrate_tp( + T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, + T const *log_k, T const *stoich, int const *phase_ids, int nspecies, + int nreaction, int nphase, int gas_phase, T standard_pressure, T ftol, + T mole_floor, int max_iter, char *work = nullptr) { + if (!(temp > 0.) || !(pres > 0.) || nspecies <= 0 || nreaction <= 0 || + nphase <= 0 || gas_phase < 0 || gas_phase >= nphase) { + diag[0] = 1.; + return 1; + } + for (int i = 0; i < nspecies; ++i) { + if (!(in_moles[i] > mole_floor)) { + diag[0] = 1.; + return 1; + } + } + + T *phase_totals, *phase_stoich, *residual, *jac, *constraints, *bounds, *step; + T *trial; + bool own_work = work == nullptr; + if (own_work) { + phase_totals = (T *)malloc(nphase * sizeof(T)); + phase_stoich = (T *)malloc(nphase * nreaction * sizeof(T)); + residual = (T *)malloc(nreaction * sizeof(T)); + jac = (T *)malloc(nreaction * nreaction * sizeof(T)); + constraints = (T *)malloc(nspecies * nreaction * sizeof(T)); + bounds = (T *)malloc(nspecies * sizeof(T)); + step = (T *)malloc(nreaction * sizeof(T)); + trial = (T *)malloc(nspecies * sizeof(T)); + } else { + phase_totals = alloc_from(work, nphase); + phase_stoich = alloc_from(work, nphase * nreaction); + residual = alloc_from(work, nreaction); + jac = alloc_from(work, nreaction * nreaction); + constraints = alloc_from(work, nspecies * nreaction); + bounds = alloc_from(work, nspecies); + step = alloc_from(work, nreaction); + trial = alloc_from(work, nspecies); + } + + memcpy(out_moles, in_moles, nspecies * sizeof(T)); + memset(jac, 0, nreaction * nreaction * sizeof(T)); + memset(phase_stoich, 0, nphase * nreaction * sizeof(T)); + for (int i = 0; i < nspecies; ++i) { + for (int j = 0; j < nreaction; ++j) { + phase_stoich[phase_ids[i] * nreaction + j] += stoich[i * nreaction + j]; + } + } + T target = log(1. + ftol); + T max_error = 0.; + int status = 2; + int iter = 0; + for (; iter < max_iter; ++iter) { + max_error = equilibrium_max_error( + out_moles, pres, standard_pressure, log_k, stoich, phase_ids, nspecies, + nreaction, nphase, gas_phase, phase_totals, residual); + if (max_error <= target) { + status = 0; + break; + } + + // Jacobian d(residual)/d(reaction extent). + for (int j = 0; j < nreaction; ++j) { + for (int k = 0; k < nreaction; ++k) { + T value = 0.; + for (int i = 0; i < nspecies; ++i) { + int phase = phase_ids[i]; + T phase_delta = phase_stoich[phase * nreaction + k]; + value += stoich[i * nreaction + j] * + (stoich[i * nreaction + k] / out_moles[i] - + phase_delta / phase_totals[phase]); + } + jac[j * nreaction + k] = value; + } + step[j] = -residual[j]; + } + + // -S dx <= n - floor is equivalent to n + S dx >= floor. + for (int i = 0; i < nspecies; ++i) { + bounds[i] = out_moles[i] - mole_floor; + for (int j = 0; j < nreaction; ++j) { + constraints[i * nreaction + j] = -stoich[i * nreaction + j]; + } + } + + int kkt_iter = max_iter; + int err = leastsq_kkt(step, jac, constraints, bounds, nreaction, nreaction, + nspecies, 0, &kkt_iter, 1.e-12, work); + if (err != 0) { + status = 3; + break; + } + + T scale = 1.; + bool accepted = false; + while (scale >= 1.e-8) { + bool positive = true; + for (int i = 0; i < nspecies; ++i) { + T delta = 0.; + for (int j = 0; j < nreaction; ++j) { + delta += stoich[i * nreaction + j] * step[j]; + } + trial[i] = out_moles[i] + scale * delta; + if (!(trial[i] > mole_floor)) positive = false; + } + if (positive) { + T trial_error = equilibrium_max_error( + trial, pres, standard_pressure, log_k, stoich, phase_ids, nspecies, + nreaction, nphase, gas_phase, phase_totals, residual); + if (trial_error < max_error) { + memcpy(out_moles, trial, nspecies * sizeof(T)); + accepted = true; + break; + } + } + scale *= .5; + } + if (!accepted) { + status = 4; + break; + } + } + + max_error = equilibrium_max_error(out_moles, pres, standard_pressure, log_k, + stoich, phase_ids, nspecies, nreaction, + nphase, gas_phase, phase_totals, residual); + memcpy(gain, jac, nreaction * nreaction * sizeof(T)); + + diag[0] = static_cast(status); + diag[1] = static_cast(iter); + diag[2] = exp(max_error) - 1.; + + if (own_work) { + free(phase_totals); + free(phase_stoich); + free(residual); + free(jac); + free(constraints); + free(bounds); + free(step); + free(trial); + } + return status; +} + +} // namespace kintera diff --git a/src/loops.cuh b/src/loops.cuh index b10d22fc..6c5c3bd7 100644 --- a/src/loops.cuh +++ b/src/loops.cuh @@ -1,11 +1,13 @@ #pragma once #include +#include #include #include #include // torch +#include #include #include #include @@ -23,7 +25,7 @@ inline size_t get_max_dynamic_shared_memory(int device) { max_dynamic_smem_by_device.resize(device + 1, 0); } - auto& cached = max_dynamic_smem_by_device[device]; + auto &cached = max_dynamic_smem_by_device[device]; if (cached != 0) { return cached; } @@ -32,7 +34,7 @@ inline size_t get_max_dynamic_shared_memory(int device) { // report sharedMemPerBlockOptin as 1 even when opt-in dynamic shared memory // is available, so prefer the CUDA device attribute and never let a bad // opt-in value reduce the ordinary per-block limit. - auto* prop = at::cuda::getDeviceProperties(device); + auto *prop = at::cuda::getDeviceProperties(device); cached = prop->sharedMemPerBlock; int attr_optin_smem = 0; @@ -42,7 +44,8 @@ inline size_t get_max_dynamic_shared_memory(int device) { cached = std::max(cached, static_cast(attr_optin_smem)); } else { C10_CUDA_CLEAR_ERROR(); - cached = std::max(cached, static_cast(prop->sharedMemPerBlockOptin)); + cached = + std::max(cached, static_cast(prop->sharedMemPerBlockOptin)); } return cached; @@ -55,7 +58,7 @@ __global__ void element_kernel(int64_t numel, func_t f) { // Shared memory allocation extern __shared__ unsigned char memory[]; - char* smem = reinterpret_cast(memory); + char *smem = reinterpret_cast(memory); if (idx < numel) { f(idx, smem); @@ -63,31 +66,30 @@ __global__ void element_kernel(int64_t numel, func_t f) { } template -void gpu_kernel(at::TensorIterator& iter, const func_t& f) { +void gpu_kernel(at::TensorIterator &iter, const func_t &f) { TORCH_CHECK(iter.ninputs() + iter.noutputs() == Arity); - std::array data; + std::array data; for (int i = 0; i < Arity; i++) { - data[i] = reinterpret_cast(iter.data_ptr(i)); + data[i] = reinterpret_cast(iter.data_ptr(i)); } auto offset_calc = ::make_offset_calculator(iter); int64_t numel = iter.numel(); - at::native::launch_legacy_kernel<128, 1>(numel, - [=] __device__(int idx) { - auto offsets = offset_calc.get(idx); - f(data.data(), offsets.data()); - }); + at::native::launch_legacy_kernel<128, 1>(numel, [=] __device__(int idx) { + auto offsets = offset_calc.get(idx); + f(data.data(), offsets.data()); + }); } template -void gpu_mem_kernel(at::TensorIterator& iter, int work_size, const func_t& f) { +void gpu_mem_kernel(at::TensorIterator &iter, int work_size, const func_t &f) { TORCH_CHECK(iter.ninputs() + iter.noutputs() == Arity); - std::array data; + std::array data; for (int i = 0; i < Arity; i++) { - data[i] = reinterpret_cast(iter.data_ptr(i)); + data[i] = reinterpret_cast(iter.data_ptr(i)); } auto offset_calc = ::make_offset_calculator(iter); @@ -100,34 +102,34 @@ void gpu_mem_kernel(at::TensorIterator& iter, int work_size, const func_t& f) { int device = -1; C10_CUDA_CHECK(cudaGetDevice(&device)); - auto* prop = at::cuda::getDeviceProperties(device); + auto *prop = at::cuda::getDeviceProperties(device); size_t max_dynamic_smem = get_max_dynamic_shared_memory(device); - //printf("max_dynamic_smem = %zu\n", max_dynamic_smem); + // printf("max_dynamic_smem = %zu\n", max_dynamic_smem); - auto device_lambda = [=] __device__(int idx, char* smem) { - auto offsets = offset_calc.get(idx); - int tid = threadIdx.x; - f(data.data(), offsets.data(), smem + tid * work_size); - }; + auto device_lambda = [=] __device__(int idx, char *smem) { + auto offsets = offset_calc.get(idx); + int tid = threadIdx.x; + f(data.data(), offsets.data(), smem + tid * work_size); + }; // request the full size auto kernelPtr = element_kernel; if (shared > (size_t)max_dynamic_smem) { TORCH_CHECK(false, "Requested shared memory (", shared, - " bytes) exceeds device maximum (", - max_dynamic_smem, " bytes)."); + " bytes) exceeds device maximum (", max_dynamic_smem, + " bytes)."); } if (shared > prop->sharedMemPerBlock) { static std::mutex attr_mutex; static std::vector> attr_once_by_device; - std::once_flag* attr_once = nullptr; + std::once_flag *attr_once = nullptr; { std::lock_guard guard(attr_mutex); if (device >= static_cast(attr_once_by_device.size())) { attr_once_by_device.resize(device + 1); } - auto& flag = attr_once_by_device[device]; + auto &flag = attr_once_by_device[device]; if (!flag) { flag = std::make_unique(); } @@ -136,8 +138,7 @@ void gpu_mem_kernel(at::TensorIterator& iter, int work_size, const func_t& f) { std::call_once(*attr_once, [=] { C10_CUDA_CHECK(cudaFuncSetAttribute( - kernelPtr, - cudaFuncAttributeMaxDynamicSharedMemorySize, + kernelPtr, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(max_dynamic_smem))); }); } @@ -153,5 +154,51 @@ void gpu_mem_kernel(at::TensorIterator& iter, int work_size, const func_t& f) { C10_CUDA_KERNEL_LAUNCH_CHECK(); } -} // namespace native -} // namespace kintera +template +__global__ void global_mem_element_kernel(int64_t numel, char *workspace, + size_t work_size, func_t f) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < numel) { + f(idx, workspace + idx * work_size); + } +} + +template +void gpu_global_mem_kernel(at::TensorIterator &iter, size_t work_size, + const func_t &f) { + TORCH_CHECK(iter.ninputs() + iter.noutputs() == Arity); + TORCH_CHECK(work_size > 0, "CUDA workspace size must be positive"); + + std::array data; + for (int i = 0; i < Arity; ++i) { + data[i] = reinterpret_cast(iter.data_ptr(i)); + } + + int64_t numel = iter.numel(); + if (numel == 0) { + return; + } + TORCH_CHECK(work_size <= + static_cast(std::numeric_limits::max()), + "CUDA workspace per cell is too large"); + auto workspace = + at::empty({numel, static_cast(work_size)}, + at::TensorOptions().device(iter.device()).dtype(at::kByte)); + auto *workspace_ptr = reinterpret_cast(workspace.data_ptr()); + auto offset_calc = ::make_offset_calculator(iter); + + auto device_lambda = [=] __device__(int64_t idx, char *work) { + auto offsets = offset_calc.get(idx); + f(data.data(), offsets.data(), work); + }; + + dim3 block(Threads); + dim3 grid((numel + block.x - 1) / block.x); + auto stream = at::cuda::getCurrentCUDAStream(); + global_mem_element_kernel<<>>( + numel, workspace_ptr, work_size, device_lambda); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace native +} // namespace kintera diff --git a/src/math/leastsq_kkt.h b/src/math/leastsq_kkt.h index da5f8da0..2dbb4ee1 100644 --- a/src/math/leastsq_kkt.h +++ b/src/math/leastsq_kkt.h @@ -20,7 +20,7 @@ namespace kintera { // Compute bitmask hash for a set of integers [0..n-1] -DISPATCH_MACRO uint64_t hash_set(const int* arr, int size, int n) { +DISPATCH_MACRO inline uint64_t hash_set(const int* arr, int size, int n) { uint64_t mask = 0; for (int i = 0; i < size; i++) { int x = arr[i]; diff --git a/src/species.cpp b/src/species.cpp index a19af559..62415768 100644 --- a/src/species.cpp +++ b/src/species.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -14,13 +15,12 @@ // torch #include -// harp -#include - // kintera #include +#include #include +#include #include #include "species.hpp" @@ -55,78 +55,102 @@ void clear_species_registry() { } // namespace -static std::unordered_map& get_nasa9_db() { +static std::unordered_map &get_nasa9_db() { static std::unordered_map db; - if (!db.empty()) return db; - - std::string path; - try { - path = find_resource("nasa9.dat"); - } catch (std::exception const& e) { - TORCH_CHECK(false, e.what()); - } - std::ifstream ifs(path); - TORCH_CHECK(ifs.good(), "Cannot open NASA-9 data file: ", path); - - std::string line; - while (std::getline(ifs, line)) { - if (line.empty() || line[0] == '#') continue; - // species name line (non-numeric first character) - if (!std::isdigit(line[0]) && line[0] != '-' && line[0] != ' ') { - std::string name = line; - // trim whitespace - while (!name.empty() && std::isspace(name.back())) name.pop_back(); - - // read 4 lines of 5 values = 20 coefficients - double vals[20]; - int idx = 0; - for (int row = 0; row < 4 && std::getline(ifs, line); ++row) { - std::istringstream iss(line); - double v; - while (iss >> v && idx < 20) vals[idx++] = v; + static std::once_flag initialized; + std::call_once(initialized, [&] { + std::string path; + try { + path = find_resource("nasa9.dat"); + } catch (std::exception const &e) { + TORCH_CHECK(false, e.what()); + } + std::ifstream ifs(path); + TORCH_CHECK(ifs.good(), "Cannot open NASA-9 data file: ", path); + + std::string line; + while (std::getline(ifs, line)) { + if (line.empty() || line[0] == '#') continue; + // Species name lines begin with a non-numeric, non-whitespace character. + if (!std::isdigit(line[0]) && line[0] != '-' && line[0] != ' ') { + std::string name = line; + while (!name.empty() && std::isspace(name.back())) name.pop_back(); + + double vals[20]; + int idx = 0; + for (int row = 0; row < 4 && std::getline(ifs, line); ++row) { + std::istringstream iss(line); + double value; + while (iss >> value && idx < 20) vals[idx++] = value; + } + if (idx < 20) continue; + + Nasa9Entry entry; + for (int k = 0; k < 7; ++k) entry.low[k] = vals[k]; + entry.low[7] = vals[8]; + entry.low[8] = vals[9]; + for (int k = 0; k < 7; ++k) entry.high[k] = vals[10 + k]; + entry.high[7] = vals[18]; + entry.high[8] = vals[19]; + db[name] = entry; } - if (idx < 20) continue; - - Nasa9Entry e; - // low-T range (vals 0..9): a0-a6, a7(=0), a8, a9 - // store 9 coefficients: a0-a6, a8, a9 (skip a7) - for (int k = 0; k < 7; ++k) e.low[k] = vals[k]; - e.low[7] = vals[8]; // a8 - e.low[8] = vals[9]; // a9 - - // high-T range (vals 10..19) - for (int k = 0; k < 7; ++k) e.high[k] = vals[10 + k]; - e.high[7] = vals[18]; // a8 - e.high[8] = vals[19]; // a9 - - db[name] = e; } - } + }); return db; } +at::Tensor nasa9_gibbs_rt(at::Tensor temp, + std::vector const &species) { + TORCH_CHECK(temp.is_floating_point(), "temp must be a floating-point tensor"); + TORCH_CHECK(!species.empty(), "NASA-9 species list must not be empty"); + auto const &database = get_nasa9_db(); + Nasa9CoeffTable low; + Nasa9CoeffTable high; + low.reserve(species.size()); + high.reserve(species.size()); + for (auto const &name : species) { + auto found = database.find(name); + TORCH_CHECK(found != database.end(), "NASA-9 species not found: ", name); + low.push_back(found->second.low); + high.push_back(found->second.high); + } + + static_assert(sizeof(Nasa9CoeffArray) == 9 * sizeof(double)); + std::array shape = {static_cast(species.size()), 9}; + auto cpu_options = torch::TensorOptions().dtype(torch::kFloat64); + auto low_tensor = torch::from_blob(low.data(), shape, cpu_options) + .clone() + .to(temp.options()); + auto high_tensor = torch::from_blob(high.data(), shape, cpu_options) + .clone() + .to(temp.options()); + auto midpoint = torch::full({static_cast(species.size())}, 1000., + temp.options()); + return nasa9_gibbs_RT(temp, low_tensor, high_tensor, midpoint); +} + void init_species_from_yaml(std::string filename) { auto config = YAML::LoadFile(filename); init_species_from_yaml(config); } -void init_species_from_yaml(YAML::Node const& config) { +void init_species_from_yaml(YAML::Node const &config) { // check if species are defined TORCH_CHECK(config["species"], "'species' is not defined in the kintera configuration file"); clear_species_registry(); - for (const auto& sp : config["species"]) { + for (const auto &sp : config["species"]) { species_names.push_back(sp["name"].as()); std::map comp; - for (const auto& it : sp["composition"]) { + for (const auto &it : sp["composition"]) { std::string key = it.first.as(); double value = it.second.as(); comp[key] = value; } - species_weights.push_back(harp::get_compound_weight(comp)); + species_weights.push_back(molar_mass(comp)); if (sp["cv_R"]) { species_cref_R.push_back(sp["cv_R"].as()); @@ -151,7 +175,7 @@ void init_species_from_yaml(YAML::Node const& config) { std::array high_coeffs = {}; double Tmid = 1000.0; - auto& nasa9_db = get_nasa9_db(); + auto &nasa9_db = get_nasa9_db(); auto name = sp["name"].as(); auto it = nasa9_db.find(name); if (it != nasa9_db.end()) { @@ -169,13 +193,13 @@ void init_species_from_yaml(YAML::Node const& config) { species_initialized = true; } -void ensure_species_initialized(std::string const& filename) { +void ensure_species_initialized(std::string const &filename) { if (!species_initialized) { init_species_from_yaml(filename); } } -void ensure_species_initialized(YAML::Node const& config) { +void ensure_species_initialized(YAML::Node const &config) { if (!species_initialized) { init_species_from_yaml(config); } @@ -198,7 +222,7 @@ std::vector SpeciesThermoImpl::species() const { } at::Tensor SpeciesThermoImpl::narrow_copy(at::Tensor data, - SpeciesThermo const& other) const { + SpeciesThermo const &other) const { auto source_ids = merge_vectors(vapor_ids(), cloud_ids()); auto other_ids = merge_vectors(other->vapor_ids(), other->cloud_ids()); std::vector indices; @@ -217,9 +241,9 @@ at::Tensor SpeciesThermoImpl::narrow_copy(at::Tensor data, return data.index_select(-1, id); } -void SpeciesThermoImpl::accumulate(at::Tensor& data, - at::Tensor const& other_data, - SpeciesThermo const& other) const { +void SpeciesThermoImpl::accumulate(at::Tensor &data, + at::Tensor const &other_data, + SpeciesThermo const &other) const { auto source_ids = merge_vectors(vapor_ids(), cloud_ids()); auto other_ids = merge_vectors(other->vapor_ids(), other->cloud_ids()); std::vector indices; @@ -238,12 +262,12 @@ void SpeciesThermoImpl::accumulate(at::Tensor& data, } bool SpeciesThermoImpl::has_nasa9() const { - for (auto const& coeffs : nasa9_low()) { + for (auto const &coeffs : nasa9_low()) { for (double v : coeffs) { if (v != 0.0) return true; } } - for (auto const& coeffs : nasa9_high()) { + for (auto const &coeffs : nasa9_high()) { for (double v : coeffs) { if (v != 0.0) return true; } @@ -252,8 +276,8 @@ bool SpeciesThermoImpl::has_nasa9() const { } static at::Tensor nasa9_coeffs_to_tensor( - std::vector> const& coeffs, - c10::TensorOptions const& options) { + std::vector> const &coeffs, + c10::TensorOptions const &options) { auto tensor = torch::empty({static_cast(coeffs.size()), 9}, torch::dtype(torch::kFloat64)); if (!coeffs.empty()) { @@ -268,17 +292,17 @@ static at::Tensor nasa9_coeffs_to_tensor( } at::Tensor SpeciesThermoImpl::nasa9_coeffs_low_tensor( - c10::TensorOptions const& options) const { + c10::TensorOptions const &options) const { return nasa9_coeffs_to_tensor(nasa9_low(), options); } at::Tensor SpeciesThermoImpl::nasa9_coeffs_high_tensor( - c10::TensorOptions const& options) const { + c10::TensorOptions const &options) const { return nasa9_coeffs_to_tensor(nasa9_high(), options); } at::Tensor SpeciesThermoImpl::nasa9_Tmid_tensor( - c10::TensorOptions const& options) const { + c10::TensorOptions const &options) const { auto tensor = torch::empty({static_cast(nasa9_Tmid().size())}, torch::dtype(torch::kFloat64)); if (!nasa9_Tmid().empty()) { @@ -329,7 +353,7 @@ void populate_thermo(SpeciesThermo thermo) { } } -void check_dimensions(SpeciesThermo const& thermo) { +void check_dimensions(SpeciesThermo const &thermo) { int nspecies = thermo->vapor_ids().size() + thermo->cloud_ids().size(); TORCH_CHECK(thermo->cref_R().size() == nspecies, @@ -380,8 +404,8 @@ void check_dimensions(SpeciesThermo const& thermo) { ". Expected = ", nspecies); } -SpeciesThermo merge_thermo(SpeciesThermo const& thermo1, - SpeciesThermo const& thermo2) { +SpeciesThermo merge_thermo(SpeciesThermo const &thermo1, + SpeciesThermo const &thermo2) { // check dimensions check_dimensions(thermo1); check_dimensions(thermo2); @@ -389,20 +413,20 @@ SpeciesThermo merge_thermo(SpeciesThermo const& thermo1, // return a new SpeciesThermo object with merged data auto merged = SpeciesThermoImpl::create(); - auto& vapor_ids = merged->vapor_ids(); - auto& cloud_ids = merged->cloud_ids(); - - auto& cref_R = merged->cref_R(); - auto& uref_R = merged->uref_R(); - auto& sref_R = merged->sref_R(); - auto& intEng_R_extra = merged->intEng_R_extra(); - auto& cp_R_extra = merged->cp_R_extra(); - auto& entropy_R_extra = merged->entropy_R_extra(); - auto& czh = merged->czh(); - auto& czh_ddC = merged->czh_ddC(); - auto& nasa9_low = merged->nasa9_low(); - auto& nasa9_high = merged->nasa9_high(); - auto& nasa9_Tmid = merged->nasa9_Tmid(); + auto &vapor_ids = merged->vapor_ids(); + auto &cloud_ids = merged->cloud_ids(); + + auto &cref_R = merged->cref_R(); + auto &uref_R = merged->uref_R(); + auto &sref_R = merged->sref_R(); + auto &intEng_R_extra = merged->intEng_R_extra(); + auto &cp_R_extra = merged->cp_R_extra(); + auto &entropy_R_extra = merged->entropy_R_extra(); + auto &czh = merged->czh(); + auto &czh_ddC = merged->czh_ddC(); + auto &nasa9_low = merged->nasa9_low(); + auto &nasa9_high = merged->nasa9_high(); + auto &nasa9_Tmid = merged->nasa9_Tmid(); // concatenate fields int nvapor1 = thermo1->vapor_ids().size(); @@ -510,7 +534,7 @@ SpeciesThermo merge_thermo(SpeciesThermo const& thermo1, cloud_ids = sort_vectors(cloud_ids, cidx); // add nvapor to cidx - for (auto& idx : cidx) idx += nvapor; + for (auto &idx : cidx) idx += nvapor; auto sorted = merge_vectors(vidx, cidx); diff --git a/src/species.hpp b/src/species.hpp index aefda416..00bf8b64 100644 --- a/src/species.hpp +++ b/src/species.hpp @@ -28,11 +28,15 @@ using Nasa9CoeffArray = std::array; using Nasa9CoeffTable = std::vector; void init_species_from_yaml(std::string filename); -void init_species_from_yaml(YAML::Node const& config); +void init_species_from_yaml(YAML::Node const &config); //! Initialize species and thermo data from a KINETICS-base master input file. -void init_species_from_kinetics_base(std::string const& master_input_path); -void ensure_species_initialized(std::string const& filename); -void ensure_species_initialized(YAML::Node const& config); +void init_species_from_kinetics_base(std::string const &master_input_path); +void ensure_species_initialized(std::string const &filename); +void ensure_species_initialized(YAML::Node const &config); + +//! Evaluate standard-state Gibbs energy g/RT from the bundled NASA-9 database. +at::Tensor nasa9_gibbs_rt(at::Tensor temp, + std::vector const &species); struct SpeciesThermoImpl { static std::shared_ptr create() { @@ -45,16 +49,16 @@ struct SpeciesThermoImpl { std::vector species() const; at::Tensor narrow_copy(at::Tensor data, - std::shared_ptr const& other) const; - void accumulate(at::Tensor& data, at::Tensor const& other_data, - std::shared_ptr const& other) const; + std::shared_ptr const &other) const; + void accumulate(at::Tensor &data, at::Tensor const &other_data, + std::shared_ptr const &other) const; bool has_nasa9() const; at::Tensor nasa9_coeffs_low_tensor( - c10::TensorOptions const& options = c10::TensorOptions()) const; + c10::TensorOptions const &options = c10::TensorOptions()) const; at::Tensor nasa9_coeffs_high_tensor( - c10::TensorOptions const& options = c10::TensorOptions()) const; + c10::TensorOptions const &options = c10::TensorOptions()) const; at::Tensor nasa9_Tmid_tensor( - c10::TensorOptions const& options = c10::TensorOptions()) const; + c10::TensorOptions const &options = c10::TensorOptions()) const; ADD_ARG(std::vector, vapor_ids); ADD_ARG(std::vector, cloud_ids); @@ -92,10 +96,10 @@ using SpeciesThermo = std::shared_ptr; void populate_thermo(SpeciesThermo thermo); -void check_dimensions(SpeciesThermo const& thermo); +void check_dimensions(SpeciesThermo const &thermo); -SpeciesThermo merge_thermo(SpeciesThermo const& thermo1, - SpeciesThermo const& thermo2); +SpeciesThermo merge_thermo(SpeciesThermo const &thermo1, + SpeciesThermo const &thermo2); extern std::vector species_names; extern std::vector species_weights; diff --git a/src/utils/molar_mass.cpp b/src/utils/molar_mass.cpp new file mode 100644 index 00000000..d096d3b4 --- /dev/null +++ b/src/utils/molar_mass.cpp @@ -0,0 +1,81 @@ +#include "molar_mass.hpp" + +#include +#include + +#include +#include +#include +#include + +namespace kintera { + +double atomic_mass(std::string const &element) { + return harp::get_element_weight(element) * 1.e-3; +} + +double molar_mass(Composition const &composition) { + return harp::get_compound_weight(composition); +} + +std::vector molar_masses( + std::vector const &elements, + std::vector> const &element_matrix) { + TORCH_CHECK(elements.size() == element_matrix.size(), + "elements and element_matrix row counts must match"); + if (elements.empty()) return {}; + + auto ncomponent = element_matrix[0].size(); + for (auto const &row : element_matrix) { + TORCH_CHECK(row.size() == ncomponent, + "element_matrix rows must have equal length"); + } + + std::vector result(ncomponent, 0.); + for (size_t i = 0; i < ncomponent; ++i) { + Composition composition; + for (size_t e = 0; e < elements.size(); ++e) { + if (element_matrix[e][i] != 0.) { + composition[elements[e]] = element_matrix[e][i]; + } + } + result[i] = molar_mass(composition); + TORCH_CHECK(result[i] > 0., "component ", i, " has no positive molar mass"); + } + return result; +} + +std::vector molar_masses_from_yaml(std::string const &filename) { + auto config = YAML::LoadFile(find_resource(filename)); + TORCH_CHECK(config["phases"], "chemistry YAML requires 'phases'"); + TORCH_CHECK(config["species"], "chemistry YAML requires 'species'"); + + std::map compositions; + for (auto const &species_node : config["species"]) { + TORCH_CHECK(species_node["name"], "species is missing 'name'"); + TORCH_CHECK(species_node["composition"], + "species is missing 'composition'"); + auto name = species_node["name"].as(); + TORCH_CHECK(compositions.emplace(name, Composition{}).second, + "duplicate species definition: ", name); + for (auto const &entry : species_node["composition"]) { + compositions.at(name)[entry.first.as()] = + entry.second.as(); + } + } + + std::vector result; + for (auto const &phase_node : config["phases"]) { + TORCH_CHECK(phase_node["species"], "phase is missing 'species'"); + for (auto const &component_node : phase_node["species"]) { + auto component = component_node.as(); + auto found = compositions.find(component); + TORCH_CHECK(found != compositions.end(), + "phase component has no species definition: ", component); + result.push_back(molar_mass(found->second)); + } + } + return result; +} + +} // namespace kintera diff --git a/src/utils/molar_mass.hpp b/src/utils/molar_mass.hpp new file mode 100644 index 00000000..6531496c --- /dev/null +++ b/src/utils/molar_mass.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include +#include +#include + +namespace kintera { + +//! Atomic molar mass [kg/mol]. +double atomic_mass(std::string const &element); + +//! Compound molar mass [kg/mol]. +double molar_mass(Composition const &composition); + +//! Component molar masses [kg/mol] from an element-by-component matrix. +std::vector molar_masses( + std::vector const &elements, + std::vector> const &element_matrix); + +//! Component molar masses [kg/mol] in phase order from a chemistry YAML file. +std::vector molar_masses_from_yaml(std::string const &filename); + +} // namespace kintera diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5b0e1a60..2e2db0f1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -19,6 +19,7 @@ setup_test(test_ch4_photolysis) setup_test(test_chapman_cycle) setup_test(test_falloff) setup_test(test_diffusion) +setup_test(test_equilibrium) setup_test(test_kinetics_base) file(GLOB inputs *.inp *.dat *.yaml) diff --git a/tests/test_equilibrium.cpp b/tests/test_equilibrium.cpp new file mode 100644 index 00000000..d32e6618 --- /dev/null +++ b/tests/test_equilibrium.cpp @@ -0,0 +1,146 @@ +#include + +#include +#include +#include +#include +#include +#include + +#define DEVICE_TESTING_SKIP_DEFAULT_INSTANTIATION +#include "device_testing.hpp" + +using namespace kintera; + +namespace { + +EquilibriumOptions make_options() { + auto op = EquilibriumOptionsImpl::create(); + op->components({"A", "B"}) + .phases({"gas"}) + .phase_ids({0, 0}) + .reactions({"A <=> B"}) + .gas_phase(0) + .standard_pressure(1.e5) + .max_iter(50) + .ftol(1.e-7) + .mole_floor(1.e-20); + return op; +} + +class EquilibriumDeviceTest : public DeviceTest {}; +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(DeviceTest); + +TEST_P(EquilibriumDeviceTest, IdealGasReaction) { + if (device.type() == torch::kMPS) GTEST_SKIP(); + auto op = make_options(); + EquilibriumTP eq(op); + eq->to(device, dtype); + + auto tensor_options = torch::device(device).dtype(dtype); + auto temp = torch::tensor({1000., 1500.}, tensor_options); + auto pres = torch::tensor(1.e5, tensor_options); + auto moles = torch::tensor({{0.8, 0.2}}, tensor_options); + auto log_k = torch::zeros({2, 1}, tensor_options); + + auto [result, gain, diag] = eq->forward(temp, pres, moles, log_k); + EXPECT_TRUE(torch::allclose(result.select(-1, 0), + torch::full({2}, .5, tensor_options), 1.e-5, + 1.e-6)); + EXPECT_TRUE(torch::allclose(result.sum(-1), torch::ones({2}, tensor_options), + 1.e-6, 1.e-6)); + EXPECT_EQ(gain.sizes(), torch::IntArrayRef({2, 1, 1})); + EXPECT_TRUE(torch::all(diag.select(-1, 0) == 0).item()); + EXPECT_LT(diag.select(-1, 2).max().item(), 1.e-5); + EXPECT_TRUE( + torch::allclose(moles, torch::tensor({{0.8, 0.2}}, tensor_options))); +} + +TEST(MolarMass, StandaloneUtilities) { + EXPECT_NEAR(atomic_mass("H"), 1.008e-3, 1.e-8); + EXPECT_NEAR(molar_mass({{"H", 2.}, {"O", 1.}}), 18.015e-3, 1.e-8); +} + +TEST(Nasa9, ConcurrentFirstDatabaseAccess) { + std::vector> futures; + for (int i = 0; i < 8; ++i) { + futures.push_back(std::async(std::launch::async, [] { + return nasa9_gibbs_rt(torch::tensor(1500., torch::kFloat64), + {"H2", "O2", "H2O"}); + })); + } + auto expected = futures.front().get(); + EXPECT_EQ(expected.sizes(), torch::IntArrayRef({3})); + EXPECT_TRUE(torch::isfinite(expected).all().item()); + for (size_t i = 1; i < futures.size(); ++i) { + EXPECT_TRUE(torch::allclose(futures[i].get(), expected)); + } +} + +TEST(EquilibriumOptions, ReadsSpeciesAndReactionsFromYaml) { + auto path = std::filesystem::temp_directory_path() / + "kintera_equilibrium_options_test.yaml"; + { + std::ofstream yaml(path); + yaml << "phases:\n" + << " - name: gas\n" + << " thermo: ideal-gas\n" + << " species: [A, B]\n" + << "species:\n" + << " - {name: A, composition: {H: 1}}\n" + << " - {name: B, composition: {H: 1}}\n" + << "reactions:\n" + << " - {type: equilibrium, equation: 'A <=> B'}\n" + << " - {type: arrhenius, equation: 'B => A'}\n" + << "equilibrium: {standard-pressure: 200000, max-iter: 12, " + "ftol: 1.e-7}\n"; + } + + auto op = EquilibriumOptionsImpl::from_yaml(path.string()); + EXPECT_EQ(op->components(), std::vector({"A", "B"})); + EXPECT_EQ(op->reactions(), std::vector({"A <=> B"})); + EXPECT_EQ(op->phase_ids(), std::vector({0, 0})); + EXPECT_DOUBLE_EQ(op->standard_pressure(), 2.e5); + EXPECT_EQ(op->max_iter(), 12); + EquilibriumTP equilibrium(op); + EXPECT_DOUBLE_EQ(equilibrium->stoich[0][0].item(), -1.); + EXPECT_DOUBLE_EQ(equilibrium->stoich[1][0].item(), 1.); + auto masses = molar_masses_from_yaml(path.string()); + ASSERT_EQ(masses.size(), 2u); + EXPECT_NEAR(masses[0], 1.008e-3, 1.e-8); + EXPECT_NEAR(masses[1], 1.008e-3, 1.e-8); + std::filesystem::remove(path); +} + +TEST(EquilibriumOptions, RejectsUnbalancedYamlReaction) { + auto path = std::filesystem::temp_directory_path() / + "kintera_unbalanced_equilibrium_test.yaml"; + { + std::ofstream yaml(path); + yaml << "phases:\n" + << " - {name: gas, thermo: ideal-gas, species: [A, B]}\n" + << "species:\n" + << " - {name: A, composition: {H: 1}}\n" + << " - {name: B, composition: {H: 2}}\n" + << "reactions:\n" + << " - {type: equilibrium, equation: 'A <=> B'}\n"; + } + EXPECT_THROW(EquilibriumOptionsImpl::from_yaml(path.string()), c10::Error); + std::filesystem::remove(path); +} + +INSTANTIATE_TEST_SUITE_P( + DeviceAndDtype, EquilibriumDeviceTest, + testing::Values(Parameters{torch::kCPU, torch::kFloat32}, + Parameters{torch::kCPU, torch::kFloat64}, + Parameters{torch::kCUDA, torch::kFloat32}, + Parameters{torch::kCUDA, torch::kFloat64}), + [](const testing::TestParamInfo &info) { + std::string name = torch::Device(info.param.device_type).str(); + name += "_"; + name += torch::toString(info.param.dtype); + std::replace(name.begin(), name.end(), '.', '_'); + return name; + }); + +} // namespace diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py new file mode 100644 index 00000000..1ed6314e --- /dev/null +++ b/tests/test_equilibrium.py @@ -0,0 +1,67 @@ +from pathlib import Path + +import pytest +import torch +from kintera import ( + EquilibriumOptions, + EquilibriumTP, + atomic_mass, + molar_mass, +) +from kintera.equilibrium import Nasa9LogK + + +def test_python_core_binding_is_functional(): + options = ( + EquilibriumOptions() + .components(["A", "B"]) + .phases(["gas"]) + .phase_ids([0, 0]) + .reactions(["A <=> B"]) + .gas_phase(0) + ) + solver = EquilibriumTP(options) + moles = torch.tensor([0.8, 0.2], dtype=torch.float64) + result, gain, diagnostics = solver( + torch.tensor(1000.0, dtype=torch.float64), + torch.tensor(1.0e5, dtype=torch.float64), + moles, + torch.zeros(1, dtype=torch.float64), + ) + torch.testing.assert_close( + result, torch.tensor([0.5, 0.5], dtype=torch.float64), atol=1e-7, rtol=1e-7 + ) + torch.testing.assert_close(moles, torch.tensor([0.8, 0.2], dtype=torch.float64)) + assert gain.shape == (1, 1) + assert diagnostics[0] == 0 + + +def test_standalone_molar_mass_utilities(): + assert atomic_mass("H") == pytest.approx(1.008e-3) + assert molar_mass({"H": 2.0, "O": 1.0}) == pytest.approx(18.015e-3) + + +def test_nasa9_yaml_equilibrium_case(): + config = Path(__file__).parents[1] / "examples" / "equilibrium_nasa9.yaml" + options = EquilibriumOptions.from_yaml(str(config)) + solver = EquilibriumTP(options) + model = Nasa9LogK(options) + temp = torch.tensor(4000.0, dtype=torch.float64) + pressure = torch.tensor(1.0e5, dtype=torch.float64) + initial = torch.tensor([0.6, 0.2, 0.2], dtype=torch.float64) + + log_k = model(temp, pressure) + result, _, diagnostics = solver(temp, pressure, initial, log_k) + + assert log_k.shape == (1,) + assert torch.isfinite(log_k).all() + assert diagnostics[0] == 0 + assert diagnostics[2] < 1.0e-7 + torch.testing.assert_close( + result @ torch.tensor([2.0, 0.0, 2.0], dtype=result.dtype), + initial @ torch.tensor([2.0, 0.0, 2.0], dtype=initial.dtype), + ) + torch.testing.assert_close( + result @ torch.tensor([0.0, 2.0, 1.0], dtype=result.dtype), + initial @ torch.tensor([0.0, 2.0, 1.0], dtype=initial.dtype), + )