From 7567f85ccd308d669392e1898a5b8d15fb425ff1 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Thu, 9 Jul 2026 21:17:06 -0400 Subject: [PATCH 1/8] wip --- .gitignore | 3 + README.md | 23 ++ pyproject.toml | 2 + python/__init__.py | 1 + python/csrc/kintera.cpp | 2 + python/csrc/pyequilibrium.cpp | 50 ++++ python/equilibrium/__init__.py | 23 ++ python/equilibrium/schlichting_young_2022.py | 236 ++++++++++++++++++ .../equilibrium/schlichting_young_2022.yaml | 102 ++++++++ python/kintera.pyi | 40 ++- src/CMakeLists.txt | 2 + src/equilibrium/equilibrate.h | 220 ++++++++++++++++ src/equilibrium/equilibrium.cpp | 208 +++++++++++++++ src/equilibrium/equilibrium.hpp | 76 ++++++ src/equilibrium/equilibrium_dispatch.cpp | 52 ++++ src/equilibrium/equilibrium_dispatch.cu | 74 ++++++ src/equilibrium/equilibrium_dispatch.hpp | 17 ++ src/equilibrium/equilibrium_options.cpp | 136 ++++++++++ src/math/leastsq_kkt.h | 2 +- tests/CMakeLists.txt | 1 + tests/test_equilibrium.cpp | 111 ++++++++ tests/test_equilibrium.py | 89 +++++++ 22 files changed, 1468 insertions(+), 2 deletions(-) create mode 100644 python/csrc/pyequilibrium.cpp create mode 100644 python/equilibrium/__init__.py create mode 100644 python/equilibrium/schlichting_young_2022.py create mode 100644 python/equilibrium/schlichting_young_2022.yaml create mode 100644 src/equilibrium/equilibrate.h create mode 100644 src/equilibrium/equilibrium.cpp create mode 100644 src/equilibrium/equilibrium.hpp create mode 100644 src/equilibrium/equilibrium_dispatch.cpp create mode 100644 src/equilibrium/equilibrium_dispatch.cu create mode 100644 src/equilibrium/equilibrium_dispatch.hpp create mode 100644 src/equilibrium/equilibrium_options.cpp create mode 100644 tests/test_equilibrium.cpp create mode 100644 tests/test_equilibrium.py 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..44b1b742 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,29 @@ KINTERA provides efficient implementations of: - Phase equilibrium computations - Atmospheric chemistry models +### Multiphase Equilibrium + +`Equilibrium` is a fixed-temperature, fixed-pressure constrained chemistry +solver. The C++/CUDA core accepts component moles, phase membership, +stoichiometry, and precomputed logarithmic equilibrium constants. Case-specific +physics remains in Python under `kintera.equilibrium`; the +`SchlichtingYoung2022` driver supplies the published core-mantle-atmosphere +topology and iterates its atmosphere mass-pressure relation around the core +solve. Thermodynamic coefficients referenced by that paper but not published +in it are supplied through the driver's required `log_k_model` callback. + +Equilibrium networks use the repository's top-level `phases`, `species`, and +`reactions` YAML layout. Phase species determine component ordering, species +compositions generate the elemental matrix, and reactions with +`type: equilibrium` generate stoichiometry: + +```python +from kintera import Equilibrium, EquilibriumOptions + +options = EquilibriumOptions.from_yaml("equilibrium.yaml") +solver = Equilibrium(options) +``` + 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/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..2a5b02c8 100644 --- a/python/csrc/kintera.cpp +++ b/python/csrc/kintera.cpp @@ -33,6 +33,7 @@ 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"; @@ -219,6 +220,7 @@ PYBIND11_MODULE(kintera, m) { bind_photolysis(m); bind_diffusion(m); bind_sparse_solver(m); + bind_equilibrium(m); m.def("species_names", []() -> const std::vector& { return kintera::species_names; diff --git a/python/csrc/pyequilibrium.cpp b/python/csrc/pyequilibrium.cpp new file mode 100644 index 00000000..84145352 --- /dev/null +++ b/python/csrc/pyequilibrium.cpp @@ -0,0 +1,50 @@ +#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, + elements) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, + phases) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, + reactions) + .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, phase_ids) + .ADD_OPTION(kintera::Matrix, kintera::EquilibriumOptionsImpl, stoich) + .ADD_OPTION(kintera::Matrix, kintera::EquilibriumOptionsImpl, + element_matrix) + .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(Equilibrium, EquilibriumOptions, + &kintera::EquilibriumImpl::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..88ea1555 --- /dev/null +++ b/python/equilibrium/__init__.py @@ -0,0 +1,23 @@ +from .schlichting_young_2022 import ( + COMPONENTS, + ELEMENTS, + PHASES, + REACTIONS, + SchlichtingYoung2022, + SchlichtingYoungResult, + initial_moles, + make_options, + surface_pressure, +) + +__all__ = [ + "COMPONENTS", + "ELEMENTS", + "PHASES", + "REACTIONS", + "SchlichtingYoung2022", + "SchlichtingYoungResult", + "initial_moles", + "make_options", + "surface_pressure", +] diff --git a/python/equilibrium/schlichting_young_2022.py b/python/equilibrium/schlichting_young_2022.py new file mode 100644 index 00000000..2f2615c6 --- /dev/null +++ b/python/equilibrium/schlichting_young_2022.py @@ -0,0 +1,236 @@ +"""Schlichting & Young (2022) core-mantle-atmosphere equilibrium case. + +The paper specifies the reaction topology and pressure relation completely, +but refers several standard-state Gibbs energies to external NIST, MAGMA, and +experimental datasets. Consequently this module requires a ``log_k_model`` +callable rather than silently substituting thermodynamic data. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Literal + +import torch + +from ..kintera import Equilibrium, EquilibriumOptions + +_CONFIG = Path(__file__).with_name("schlichting_young_2022.yaml") +_BASE_OPTIONS = EquilibriumOptions.from_yaml(str(_CONFIG)) +COMPONENTS = tuple(_BASE_OPTIONS.components()) +ELEMENTS = tuple(_BASE_OPTIONS.elements()) +PHASES = tuple(_BASE_OPTIONS.phases()) +REACTIONS = tuple(_BASE_OPTIONS.reactions()) +_PHASE_IDS = tuple(_BASE_OPTIONS.phase_ids()) +_GAS_PHASE = _BASE_OPTIONS.gas_phase() +_COMPONENT_INDEX = {name: i for i, name in enumerate(COMPONENTS)} + +_ATOMIC_MASS = { + "C": 0.012011, + "Fe": 0.055845, + "H": 0.001008, + "Mg": 0.024305, + "Na": 0.022990, + "O": 0.015999, + "Si": 0.028085, +} +_ELEMENT_MATRIX = _BASE_OPTIONS.element_matrix() +_MOLAR_MASS = tuple( + sum( + _ELEMENT_MATRIX[e][i] * _ATOMIC_MASS[element] + for e, element in enumerate(ELEMENTS) + ) + for i in range(len(COMPONENTS)) +) + + +def make_options( + scenario: Literal["reactive", "isolated-core"] = "reactive", + *, + max_iter: int = 100, + ftol: float = 5e-3, +) -> EquilibriumOptions: + """Create the generic core options for the paper's reaction system.""" + options = EquilibriumOptions.from_yaml(str(_CONFIG)) + if scenario == "isolated-core": + keep = [i for i in range(18) if i not in (1, 3, 4, 6)] + options.stoich([[row[j] for j in keep] for row in options.stoich()]) + options.reactions([options.reactions()[j] for j in keep]) + elif scenario != "reactive": + raise ValueError(f"unknown scenario: {scenario}") + options.max_iter(max_iter).ftol(ftol) + options.validate() + return options + + +def initial_moles( + hydrogen_mass_fraction: float = 0.04, + metal_mass_fraction: float = 0.25, + *, + dtype: torch.dtype = torch.float64, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Construct the paper's nominal composition on a one-kilogram basis.""" + if not 0 < hydrogen_mass_fraction < 1: + raise ValueError("hydrogen_mass_fraction must lie between zero and one") + if not 0 <= metal_mass_fraction < 1 - hydrogen_mass_fraction: + raise ValueError("metal_mass_fraction leaves no silicate inventory") + + mass = torch.tensor(_MOLAR_MASS, dtype=dtype, device=device) + result = torch.full((len(COMPONENTS),), 1.0e-20, dtype=dtype, device=device) + fe_metal = _COMPONENT_INDEX["Fe[metal]"] + result[fe_metal] = metal_mass_fraction / mass[fe_metal] + + silicate_mass = 1.0 - metal_mass_fraction - hydrogen_mass_fraction + mgsio3 = _COMPONENT_INDEX["MgSiO3[silicate]"] + mgo = _COMPONENT_INDEX["MgO[silicate]"] + fesio3 = _COMPONENT_INDEX["FeSiO3[silicate]"] + na2o = _COMPONENT_INDEX["Na2O[silicate]"] + silicate_moles = silicate_mass / ( + 0.921 * mass[mgsio3] + + 0.032 * mass[mgo] + + 0.035 * mass[fesio3] + + 0.007 * mass[na2o] + ) + result[mgsio3] = 0.921 * silicate_moles + result[mgo] = 0.032 * silicate_moles + result[fesio3] = 0.035 * silicate_moles + result[na2o] = 0.007 * silicate_moles + + # The primary atmosphere is 99.9 mol% H2 and 0.1 mol% CO. + h2_gas = _COMPONENT_INDEX["H2[gas]"] + co_gas = _COMPONENT_INDEX["CO[gas]"] + gas_moles = hydrogen_mass_fraction / ( + 0.999 * mass[h2_gas] + 0.001 * mass[co_gas] + ) + result[h2_gas] = 0.999 * gas_moles + result[co_gas] = 0.001 * gas_moles + return result + + +def surface_pressure( + moles: torch.Tensor, planet_mass_earth: float | torch.Tensor = 4.0 +) -> torch.Tensor: + """Evaluate Equation 8 using the atmosphere mass implied by ``moles``.""" + mass = torch.as_tensor(_MOLAR_MASS, dtype=moles.dtype, device=moles.device) + component_mass = moles * mass + total_mass = component_mass.sum(dim=-1) + gas = torch.tensor( + [i for i, phase in enumerate(_PHASE_IDS) if phase == _GAS_PHASE], + dtype=torch.long, + device=moles.device, + ) + atmosphere_fraction = ( + component_mass.index_select(-1, gas).sum(dim=-1) / total_mass + ) + planet_mass = torch.as_tensor( + planet_mass_earth, dtype=moles.dtype, device=moles.device + ) + return 1.1e11 * atmosphere_fraction * planet_mass.pow(2.0 / 3.0) + + +@dataclass(frozen=True) +class SchlichtingYoungResult: + moles: torch.Tensor + pressure: torch.Tensor + phase_totals: torch.Tensor + phase_fractions: torch.Tensor + gain: torch.Tensor + diagnostics: torch.Tensor + pressure_error: torch.Tensor + outer_iterations: int + + +class SchlichtingYoung2022: + """Python driver coupling the generic core to the paper's pressure law.""" + + def __init__( + self, + log_k_model: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + scenario: Literal["reactive", "isolated-core"] = "reactive", + *, + max_outer_iter: int = 30, + pressure_tolerance: float = 0.1, + pressure_damping: float = 0.5, + ) -> None: + if log_k_model is None: + raise TypeError( + "log_k_model is required because the paper refers some Gibbs " + "energies to external datasets" + ) + if max_outer_iter <= 0: + raise ValueError("max_outer_iter must be positive") + if pressure_tolerance <= 0: + raise ValueError("pressure_tolerance must be positive") + if not 0 < pressure_damping <= 1: + raise ValueError("pressure_damping must lie in (0, 1]") + self.scenario = scenario + self.log_k_model = log_k_model + self.max_outer_iter = max_outer_iter + self.pressure_tolerance = pressure_tolerance + self.pressure_damping = pressure_damping + self.options = make_options(scenario) + self.solver = Equilibrium(self.options) + + def solve( + self, + temp: torch.Tensor, + moles: torch.Tensor, + pressure: torch.Tensor | None = None, + *, + planet_mass_earth: float = 4.0, + ) -> SchlichtingYoungResult: + self.solver.to(moles.device, moles.dtype) + pressure = ( + surface_pressure(moles, planet_mass_earth) + if pressure is None + else pressure.to(device=moles.device, dtype=moles.dtype) + ) + state = moles + gain = diagnostics = None + pressure_error = torch.full_like(pressure, float("inf")) + + for _iteration in range(1, self.max_outer_iter + 1): + log_k = self.log_k_model(temp, pressure) + if self.scenario == "isolated-core" and log_k.size(-1) == 18: + keep = [i for i in range(18) if i not in (1, 3, 4, 6)] + log_k = log_k[..., keep] + state, gain, diagnostics = self.solver(temp, pressure, state, log_k) + target = surface_pressure(state, planet_mass_earth) + pressure_error = (target - pressure).abs() / target.clamp_min( + torch.finfo(target.dtype).tiny + ) + chemistry_ok = diagnostics[..., 0].eq(0).all() + if chemistry_ok and pressure_error.max() <= self.pressure_tolerance: + pressure = target + break + pressure = torch.exp( + (1.0 - self.pressure_damping) * torch.log(pressure) + + self.pressure_damping * torch.log(target) + ) + + phase_totals = torch.stack( + [ + state[..., [i for i, p in enumerate(_PHASE_IDS) if p == phase]].sum( + -1 + ) + for phase in range(len(PHASES)) + ], + dim=-1, + ) + phase_fractions = torch.empty_like(state) + for phase in range(len(PHASES)): + ids = [i for i, value in enumerate(_PHASE_IDS) if value == phase] + phase_fractions[..., ids] = ( + state[..., ids] / phase_totals[..., phase, None] + ) + + return SchlichtingYoungResult( + state, + pressure, + phase_totals, + phase_fractions, + gain, + diagnostics, + pressure_error, + _iteration, + ) diff --git a/python/equilibrium/schlichting_young_2022.yaml b/python/equilibrium/schlichting_young_2022.yaml new file mode 100644 index 00000000..7ea1cf34 --- /dev/null +++ b/python/equilibrium/schlichting_young_2022.yaml @@ -0,0 +1,102 @@ +equilibrium: + standard-pressure: 1.0e5 + max-iter: 100 + ftol: 5.0e-3 + mole-floor: 1.0e-30 + +phases: + - name: silicate + thermo: ideal-solution + species: + - "MgO[silicate]" + - "SiO2[silicate]" + - "MgSiO3[silicate]" + - "FeO[silicate]" + - "FeSiO3[silicate]" + - "Na2O[silicate]" + - "Na2SiO3[silicate]" + - "H2[silicate]" + - "H2O[silicate]" + - "CO[silicate]" + - "CO2[silicate]" + - name: metal + thermo: ideal-solution + species: ["Fe[metal]", "Si[metal]", "O[metal]", "H[metal]"] + - name: atmosphere + thermo: ideal-gas + species: + - "H2[gas]" + - "CO[gas]" + - "CO2[gas]" + - "CH4[gas]" + - "O2[gas]" + - "H2O[gas]" + - "Fe[gas]" + - "Mg[gas]" + - "SiO[gas]" + - "Na[gas]" + +species: + - {name: "MgO[silicate]", composition: {Mg: 1, O: 1}} + - {name: "SiO2[silicate]", composition: {Si: 1, O: 2}} + - {name: "MgSiO3[silicate]", composition: {Mg: 1, Si: 1, O: 3}} + - {name: "FeO[silicate]", composition: {Fe: 1, O: 1}} + - {name: "FeSiO3[silicate]", composition: {Fe: 1, Si: 1, O: 3}} + - {name: "Na2O[silicate]", composition: {Na: 2, O: 1}} + - {name: "Na2SiO3[silicate]", composition: {Na: 2, Si: 1, O: 3}} + - {name: "H2[silicate]", composition: {H: 2}} + - {name: "H2O[silicate]", composition: {H: 2, O: 1}} + - {name: "CO[silicate]", composition: {C: 1, O: 1}} + - {name: "CO2[silicate]", composition: {C: 1, O: 2}} + - {name: "Fe[metal]", composition: {Fe: 1}} + - {name: "Si[metal]", composition: {Si: 1}} + - {name: "O[metal]", composition: {O: 1}} + - {name: "H[metal]", composition: {H: 1}} + - {name: "H2[gas]", composition: {H: 2}} + - {name: "CO[gas]", composition: {C: 1, O: 1}} + - {name: "CO2[gas]", composition: {C: 1, O: 2}} + - {name: "CH4[gas]", composition: {C: 1, H: 4}} + - {name: "O2[gas]", composition: {O: 2}} + - {name: "H2O[gas]", composition: {H: 2, O: 1}} + - {name: "Fe[gas]", composition: {Fe: 1}} + - {name: "Mg[gas]", composition: {Mg: 1}} + - {name: "SiO[gas]", composition: {Si: 1, O: 1}} + - {name: "Na[gas]", composition: {Na: 1}} + +reactions: + - type: equilibrium + equation: "Na2SiO3[silicate] <=> Na2O[silicate] + SiO2[silicate]" + - type: equilibrium + equation: "0.5 SiO2[silicate] + Fe[metal] <=> FeO[silicate] + 0.5 Si[metal]" + - type: equilibrium + equation: "MgSiO3[silicate] <=> MgO[silicate] + SiO2[silicate]" + - type: equilibrium + equation: "O[metal] + 0.5 Si[metal] <=> 0.5 SiO2[silicate]" + - type: equilibrium + equation: "2 H[metal] <=> H2[silicate]" + - type: equilibrium + equation: "FeSiO3[silicate] <=> FeO[silicate] + SiO2[silicate]" + - type: equilibrium + equation: "2 H2O[silicate] + Si[metal] <=> SiO2[silicate] + 2 H2[silicate]" + - type: equilibrium + equation: "CO[gas] + 0.5 O2[gas] <=> CO2[gas]" + - type: equilibrium + equation: "CH4[gas] + 0.5 O2[gas] <=> 2 H2[gas] + CO[gas]" + - type: equilibrium + equation: "H2[gas] + 0.5 O2[gas] <=> H2O[gas]" + - type: equilibrium + equation: "FeO[silicate] <=> Fe[gas] + 0.5 O2[gas]" + - type: equilibrium + equation: "MgO[silicate] <=> Mg[gas] + 0.5 O2[gas]" + - type: equilibrium + equation: "SiO2[silicate] <=> SiO[gas] + 0.5 O2[gas]" + - type: equilibrium + equation: "Na2O[silicate] <=> 2 Na[gas] + 0.5 O2[gas]" + - type: equilibrium + equation: "H2[gas] <=> H2[silicate]" + - type: equilibrium + equation: "H2O[gas] <=> H2O[silicate]" + - type: equilibrium + equation: "CO[gas] <=> CO[silicate]" + - type: equilibrium + equation: "CO2[gas] <=> CO2[silicate]" diff --git a/python/kintera.pyi b/python/kintera.pyi index 3f732668..b583ab3f 100644 --- a/python/kintera.pyi +++ b/python/kintera.pyi @@ -6,7 +6,7 @@ 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 @@ -456,6 +456,44 @@ 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 elements(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 stoich(self, value: List[List[float]] = ...) -> Any: ... + def element_matrix(self, value: List[List[float]] = ...) -> 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 Equilibrium: + 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/equilibrate.h b/src/equilibrium/equilibrate.h new file mode 100644 index 00000000..a53e2510 --- /dev/null +++ b/src/equilibrium/equilibrate.h @@ -0,0 +1,220 @@ +#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 +equilibrate(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, + T const *element_matrix, int nspecies, int nreaction, int nphase, + int nelement, 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, *initial_elements; + 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)); + initial_elements = (T *)malloc(nelement * 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); + initial_elements = alloc_from(work, nelement); + } + + 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]; + } + } + for (int e = 0; e < nelement; ++e) { + initial_elements[e] = 0.; + for (int i = 0; i < nspecies; ++i) { + initial_elements[e] += element_matrix[e * nspecies + i] * in_moles[i]; + } + } + + 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)); + + T element_error = 0.; + for (int e = 0; e < nelement; ++e) { + T final_value = 0.; + for (int i = 0; i < nspecies; ++i) { + final_value += element_matrix[e * nspecies + i] * out_moles[i]; + } + T denom = fabs(initial_elements[e]); + T error = denom > 0. ? fabs(final_value - initial_elements[e]) / denom + : fabs(final_value - initial_elements[e]); + if (error > element_error) + element_error = error; + } + + diag[0] = static_cast(status); + diag[1] = static_cast(iter); + diag[2] = exp(max_error) - 1.; + diag[3] = element_error; + + if (own_work) { + free(phase_totals); + free(phase_stoich); + free(residual); + free(jac); + free(constraints); + free(bounds); + free(step); + free(trial); + free(initial_elements); + } + return status; +} + +} // namespace kintera diff --git a/src/equilibrium/equilibrium.cpp b/src/equilibrium/equilibrium.cpp new file mode 100644 index 00000000..c9e1e315 --- /dev/null +++ b/src/equilibrium/equilibrium.cpp @@ -0,0 +1,208 @@ +#include +#include +#include + +#include "equilibrium.hpp" +#include "equilibrium_dispatch.hpp" + +namespace kintera { + +void EquilibriumOptionsImpl::report(std::ostream &os) const { + os << "-- equilibrium options --\n" + << "* components = " << components().size() << "\n" + << "* elements = " << elements().size() << "\n" + << "* phases = " << phases().size() << "\n" + << "* reactions = " << (stoich().empty() ? 0 : stoich()[0].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(stoich().size() == components().size(), + "stoich must have one row per component"); + TORCH_CHECK(!stoich()[0].empty(), "Equilibrium requires reactions"); + if (!reactions().empty()) { + TORCH_CHECK(reactions().size() == stoich()[0].size(), + "reactions and stoich column counts must match"); + } + if (!elements().empty()) { + TORCH_CHECK(elements().size() == element_matrix().size(), + "elements and element_matrix row counts must match"); + } + TORCH_CHECK(components().size() <= 64, + "the active-set solver currently supports at most 64 components"); + auto nreaction = stoich()[0].size(); + for (auto const &row : stoich()) { + TORCH_CHECK(row.size() == nreaction, "stoich rows must have equal length"); + } + + // Independent reaction columns are required by the square Newton system. + auto work = stoich(); + size_t rank = 0; + for (size_t column = 0; column < nreaction && rank < work.size(); ++column) { + size_t pivot = rank; + for (size_t row = rank + 1; row < work.size(); ++row) { + if (std::abs(work[row][column]) > std::abs(work[pivot][column])) + pivot = row; + } + if (std::abs(work[pivot][column]) <= 1.e-12) + continue; + std::swap(work[rank], work[pivot]); + for (size_t row = rank + 1; row < work.size(); ++row) { + double factor = work[row][column] / work[rank][column]; + for (size_t j = column; j < nreaction; ++j) { + work[row][j] -= factor * work[rank][j]; + } + } + ++rank; + } + TORCH_CHECK(rank == nreaction, "stoich reaction columns must be independent"); + 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"); + + if (!element_matrix().empty()) { + for (auto const &row : element_matrix()) { + TORCH_CHECK(row.size() == components().size(), + "element_matrix rows must match component count"); + for (size_t j = 0; j < nreaction; ++j) { + double balance = 0.; + for (size_t i = 0; i < components().size(); ++i) { + balance += row[i] * stoich()[i][j]; + } + TORCH_CHECK(std::abs(balance) <= 1.e-10, "reaction ", j, + " does not conserve an element"); + } + } + } +} + +EquilibriumImpl::EquilibriumImpl(EquilibriumOptions const &options_) + : options(options_) { + reset(); +} + +void EquilibriumImpl::reset() { + options->validate(); + int nspecies = static_cast(options->components().size()); + int nreaction = static_cast(options->stoich()[0].size()); + + std::vector stoich_flat; + stoich_flat.reserve(nspecies * nreaction); + for (auto const &row : options->stoich()) { + stoich_flat.insert(stoich_flat.end(), row.begin(), row.end()); + } + stoich = register_buffer( + "stoich", + torch::tensor(stoich_flat, torch::kFloat64).view({nspecies, nreaction})); + phase_ids = register_buffer("phase_ids", + torch::tensor(options->phase_ids(), torch::kInt)); + + int nelement = static_cast(options->element_matrix().size()); + std::vector element_flat; + if (nelement == 0) { + // A zero row keeps the device kernel and diagnostic layout uniform. + nelement = 1; + element_flat.resize(nspecies, 0.); + } else { + element_flat.reserve(nelement * nspecies); + for (auto const &row : options->element_matrix()) { + element_flat.insert(element_flat.end(), row.begin(), row.end()); + } + } + element_matrix = register_buffer( + "element_matrix", + torch::tensor(element_flat, torch::kFloat64).view({nelement, nspecies})); +} + +void EquilibriumImpl::pretty_print(std::ostream &os) const { + os << "Equilibrium(components=" << options->components().size() + << ", phases=" << options->phases().size() << ")"; +} + +std::tuple +EquilibriumImpl::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(4); + 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(); + auto e = element_matrix.to(moles.options()).contiguous(); + at::native::call_equilibrium(moles.device().type(), iter, s, p, e, + 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..bfba951d --- /dev/null +++ b/src/equilibrium/equilibrium.hpp @@ -0,0 +1,76 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include +#include + +#include + +namespace kintera { + +using Matrix = std::vector>; + +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, elements); + ADD_ARG(std::vector, phases); + ADD_ARG(std::vector, reactions); + ADD_ARG(std::vector, phase_ids); + ADD_ARG(Matrix, stoich); + ADD_ARG(Matrix, element_matrix); + 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 EquilibriumImpl : public torch::nn::Cloneable { +public: + EquilibriumOptions options; + + torch::Tensor stoich; + torch::Tensor phase_ids; + torch::Tensor element_matrix; + + EquilibriumImpl() : options(EquilibriumOptionsImpl::create()) {} + explicit EquilibriumImpl(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 (..., 4): status, iterations, equilibrium error, and element + * conservation error. + */ + std::tuple + forward(torch::Tensor temp, torch::Tensor pres, torch::Tensor moles, + torch::Tensor log_k, bool warm_start = false); +}; + +TORCH_MODULE(Equilibrium); + +} // namespace kintera diff --git a/src/equilibrium/equilibrium_dispatch.cpp b/src/equilibrium/equilibrium_dispatch.cpp new file mode 100644 index 00000000..3655abc6 --- /dev/null +++ b/src/equilibrium/equilibrium_dispatch.cpp @@ -0,0 +1,52 @@ +#include +#include +#include + +#include "equilibrate.h" +#include "equilibrium_dispatch.hpp" + +namespace kintera { + +void call_equilibrium_cpu(at::TensorIterator &iter, at::Tensor const &stoich, + at::Tensor const &phase_ids, + at::Tensor const &element_matrix, 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); + int nelement = element_matrix.size(0); + auto stoich_ptr = stoich.data_ptr(); + auto phase_ptr = phase_ids.data_ptr(); + auto element_ptr = element_matrix.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]); + equilibrate(gain, diag, out, *temp, *pres, moles, log_k, stoich_ptr, + phase_ptr, element_ptr, nspecies, nreaction, nphase, + nelement, 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..49296836 --- /dev/null +++ b/src/equilibrium/equilibrium_dispatch.cu @@ -0,0 +1,74 @@ +#include +#include + +#include + +#include "equilibrate.h" +#include "equilibrium_dispatch.hpp" + +namespace kintera { + +template +size_t equilibrium_space(int nspecies, int nreaction, int nphase, + int nelement) { + 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)); + bump(alignof(T), nelement * 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, + at::Tensor const &element_matrix, 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); + int nelement = element_matrix.size(0); + auto stoich_ptr = stoich.data_ptr(); + auto phase_ptr = phase_ids.data_ptr(); + auto element_ptr = element_matrix.data_ptr(); + int work_size = + equilibrium_space(nspecies, nreaction, nphase, nelement); + // The KKT workspace scales quadratically with reaction count. One cell + // per block keeps the 25-component paper case within shared-memory limits. + native::gpu_mem_kernel<1, 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]); + equilibrate(gain, diag, out, *temp, *pres, moles, log_k, stoich_ptr, + phase_ptr, element_ptr, nspecies, nreaction, nphase, + nelement, 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..bdddb293 --- /dev/null +++ b/src/equilibrium/equilibrium_dispatch.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include +#include + +namespace at::native { + +using equilibrium_fn = void (*)(at::TensorIterator &iter, + at::Tensor const &stoich, + at::Tensor const &phase_ids, + at::Tensor const &element_matrix, 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..da13b849 --- /dev/null +++ b/src/equilibrium/equilibrium_options.cpp @@ -0,0 +1,136 @@ +#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); + } + + options->elements().assign(element_set.begin(), element_set.end()); + options->element_matrix().resize(options->elements().size()); + for (size_t e = 0; e < options->elements().size(); ++e) { + auto &row = options->element_matrix()[e]; + row.resize(options->components().size(), 0.); + for (size_t i = 0; i < options->components().size(); ++i) { + auto const &composition = compositions.at(options->components()[i]); + auto found = composition.find(options->elements()[e]); + if (found != composition.end()) + row[i] = found->second; + } + } + + 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"); + + options->stoich().assign(options->components().size(), + std::vector(reactions.size(), 0.)); + 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); + options->stoich()[found - options->components().begin()][j] -= + 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); + options->stoich()[found - options->components().begin()][j] += + coefficient; + } + } + + 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/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/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..122d3d3f --- /dev/null +++ b/tests/test_equilibrium.cpp @@ -0,0 +1,111 @@ +#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}) + .stoich({{-1.}, {1.}}) + .element_matrix({{1., 1.}}) + .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(); + Equilibrium 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_LT(diag.select(-1, 3).max().item(), 1.e-6); + EXPECT_TRUE( + torch::allclose(moles, torch::tensor({{0.8, 0.2}}, tensor_options))); +} + +TEST(EquilibriumOptions, RejectsUnbalancedReaction) { + auto op = make_options(); + op->element_matrix({{1., 2.}}); + EXPECT_THROW(op->validate(), c10::Error); +} + +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: {X: 1}}\n" + << " - {name: B, composition: {X: 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->elements(), std::vector({"X"})); + EXPECT_EQ(op->reactions(), std::vector({"A <=> B"})); + EXPECT_EQ(op->phase_ids(), std::vector({0, 0})); + EXPECT_DOUBLE_EQ(op->stoich()[0][0], -1.); + EXPECT_DOUBLE_EQ(op->stoich()[1][0], 1.); + EXPECT_DOUBLE_EQ(op->standard_pressure(), 2.e5); + EXPECT_EQ(op->max_iter(), 12); + 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..10407571 --- /dev/null +++ b/tests/test_equilibrium.py @@ -0,0 +1,89 @@ +import pytest +import torch +from kintera import Equilibrium, EquilibriumOptions +from kintera.equilibrium import ( + COMPONENTS, + SchlichtingYoung2022, + initial_moles, + make_options, + surface_pressure, +) + + +def test_python_core_binding_is_functional(): + options = ( + EquilibriumOptions() + .components(["A", "B"]) + .phases(["gas"]) + .phase_ids([0, 0]) + .stoich([[-1.0], [1.0]]) + .element_matrix([[1.0, 1.0]]) + .gas_phase(0) + ) + solver = Equilibrium(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_paper_topology_and_pressure_relation(): + options = make_options() + options.validate() + assert len(options.components()) == 25 + assert len(options.elements()) == 7 + assert len(options.reactions()) == 18 + assert len(options.stoich()[0]) == 18 + + moles = initial_moles() + pressure = surface_pressure(moles) + assert moles.shape == (len(COMPONENTS),) + assert pressure > 0 + + +def test_paper_driver_accepts_external_thermodynamics(): + moles = initial_moles() + options = make_options() + stoich = torch.tensor(options.stoich(), dtype=moles.dtype) + phase_ids = options.phase_ids() + + def equilibrium_at_initial_state(temp, pres): + del temp + log_activity = torch.empty_like(moles) + for phase in range(3): + ids = [i for i, value in enumerate(phase_ids) if value == phase] + log_activity[ids] = torch.log(moles[ids] / moles[ids].sum()) + gas = [i for i, value in enumerate(phase_ids) if value == 2] + log_activity[gas] += torch.log(pres / 1.0e5) + return log_activity @ stoich + + model = SchlichtingYoung2022(equilibrium_at_initial_state, max_outer_iter=2) + result = model.solve(torch.tensor(4500.0, dtype=moles.dtype), moles) + assert model.options.gas_phase() == 2 + assert result.diagnostics[0] == 0 + assert result.pressure_error < 0.1 + for phase in range(3): + ids = [i for i, value in enumerate(model.options.phase_ids()) if value == phase] + torch.testing.assert_close( + result.phase_fractions[ids].sum(), torch.tensor(1.0, dtype=moles.dtype) + ) + + +def test_paper_driver_rejects_invalid_outer_solve_options(): + def model(temp, pres): + del pres + return torch.zeros(18, dtype=temp.dtype, device=temp.device) + + with pytest.raises(ValueError, match="max_outer_iter"): + SchlichtingYoung2022(model, max_outer_iter=0) + with pytest.raises(ValueError, match="pressure_damping"): + SchlichtingYoung2022(model, pressure_damping=0.0) From 4ba829c6adf568a26696a831068415614e640330 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Thu, 9 Jul 2026 21:55:02 -0400 Subject: [PATCH 2/8] wip --- README.md | 15 +-- python/csrc/kintera.cpp | 61 +++++---- python/csrc/pyequilibrium.cpp | 9 +- python/equilibrium/__init__.py | 2 - python/equilibrium/schlichting_young_2022.py | 24 +--- python/kintera.pyi | 12 +- src/equilibrium/equilibrate.h | 33 +---- src/equilibrium/equilibrium.cpp | 126 ++++++------------- src/equilibrium/equilibrium.hpp | 17 +-- src/equilibrium/equilibrium_dispatch.cpp | 8 +- src/equilibrium/equilibrium_dispatch.cu | 15 +-- src/equilibrium/equilibrium_dispatch.hpp | 3 +- src/equilibrium/equilibrium_options.cpp | 37 +++--- src/species.cpp | 124 +++++++++--------- src/utils/molar_mass.cpp | 85 +++++++++++++ src/utils/molar_mass.hpp | 24 ++++ tests/test_equilibrium.cpp | 45 +++++-- tests/test_equilibrium.py | 34 +++-- 18 files changed, 364 insertions(+), 310 deletions(-) create mode 100644 src/utils/molar_mass.cpp create mode 100644 src/utils/molar_mass.hpp diff --git a/README.md b/README.md index 44b1b742..134e7449 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,10 @@ KINTERA provides efficient implementations of: ### Multiphase Equilibrium -`Equilibrium` is a fixed-temperature, fixed-pressure constrained chemistry -solver. The C++/CUDA core accepts component moles, phase membership, -stoichiometry, and precomputed logarithmic equilibrium constants. Case-specific +`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 physics remains in Python under `kintera.equilibrium`; the `SchlichtingYoung2022` driver supplies the published core-mantle-atmosphere topology and iterates its atmosphere mass-pressure relation around the core @@ -36,14 +37,14 @@ in it are supplied through the driver's required `log_k_model` callback. Equilibrium networks use the repository's top-level `phases`, `species`, and `reactions` YAML layout. Phase species determine component ordering, species -compositions generate the elemental matrix, and reactions with -`type: equilibrium` generate stoichiometry: +compositions validate elemental balance, and reactions with `type: equilibrium` +generate the module's stoichiometric buffer: ```python -from kintera import Equilibrium, EquilibriumOptions +from kintera import EquilibriumOptions, EquilibriumTP options = EquilibriumOptions.from_yaml("equilibrium.yaml") -solver = Equilibrium(options) +solver = EquilibriumTP(options) ``` The library is written in C++17 with Python bindings, leveraging PyTorch for tensor operations and providing GPU acceleration support via CUDA. diff --git a/python/csrc/kintera.cpp b/python/csrc/kintera.cpp index 2a5b02c8..91d49bc9 100644 --- a/python/csrc/kintera.cpp +++ b/python/csrc/kintera.cpp @@ -11,6 +11,7 @@ #include #include #include +#include // python #include "pyoptions.hpp" @@ -25,27 +26,34 @@ extern std::vector species_cref_R; extern std::vector species_uref_R; extern std::vector species_sref_R; -} // namespace kintera +} // 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_equilibrium(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")); + 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) @@ -60,9 +68,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) @@ -210,7 +218,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")); @@ -222,44 +230,47 @@ PYBIND11_MODULE(kintera, 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 index 84145352..d53aa295 100644 --- a/python/csrc/pyequilibrium.cpp +++ b/python/csrc/pyequilibrium.cpp @@ -26,16 +26,11 @@ void bind_equilibrium(py::module &m) { py::arg("filename"), py::arg("verbose") = false) .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, components) - .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, - elements) .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, phases) .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, reactions) .ADD_OPTION(std::vector, kintera::EquilibriumOptionsImpl, phase_ids) - .ADD_OPTION(kintera::Matrix, kintera::EquilibriumOptionsImpl, stoich) - .ADD_OPTION(kintera::Matrix, kintera::EquilibriumOptionsImpl, - element_matrix) .ADD_OPTION(int, kintera::EquilibriumOptionsImpl, gas_phase) .ADD_OPTION(double, kintera::EquilibriumOptionsImpl, standard_pressure) .ADD_OPTION(int, kintera::EquilibriumOptionsImpl, max_iter) @@ -43,8 +38,8 @@ void bind_equilibrium(py::module &m) { .ADD_OPTION(double, kintera::EquilibriumOptionsImpl, mole_floor) .def("validate", &kintera::EquilibriumOptionsImpl::validate); - ADD_KINTERA_MODULE(Equilibrium, EquilibriumOptions, - &kintera::EquilibriumImpl::forward, py::arg("temp"), + 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 index 88ea1555..dbc1fb7e 100644 --- a/python/equilibrium/__init__.py +++ b/python/equilibrium/__init__.py @@ -1,6 +1,5 @@ from .schlichting_young_2022 import ( COMPONENTS, - ELEMENTS, PHASES, REACTIONS, SchlichtingYoung2022, @@ -12,7 +11,6 @@ __all__ = [ "COMPONENTS", - "ELEMENTS", "PHASES", "REACTIONS", "SchlichtingYoung2022", diff --git a/python/equilibrium/schlichting_young_2022.py b/python/equilibrium/schlichting_young_2022.py index 2f2615c6..2fc54424 100644 --- a/python/equilibrium/schlichting_young_2022.py +++ b/python/equilibrium/schlichting_young_2022.py @@ -12,35 +12,18 @@ import torch -from ..kintera import Equilibrium, EquilibriumOptions +from ..kintera import EquilibriumOptions, EquilibriumTP, molar_masses_from_yaml _CONFIG = Path(__file__).with_name("schlichting_young_2022.yaml") _BASE_OPTIONS = EquilibriumOptions.from_yaml(str(_CONFIG)) COMPONENTS = tuple(_BASE_OPTIONS.components()) -ELEMENTS = tuple(_BASE_OPTIONS.elements()) PHASES = tuple(_BASE_OPTIONS.phases()) REACTIONS = tuple(_BASE_OPTIONS.reactions()) _PHASE_IDS = tuple(_BASE_OPTIONS.phase_ids()) _GAS_PHASE = _BASE_OPTIONS.gas_phase() _COMPONENT_INDEX = {name: i for i, name in enumerate(COMPONENTS)} -_ATOMIC_MASS = { - "C": 0.012011, - "Fe": 0.055845, - "H": 0.001008, - "Mg": 0.024305, - "Na": 0.022990, - "O": 0.015999, - "Si": 0.028085, -} -_ELEMENT_MATRIX = _BASE_OPTIONS.element_matrix() -_MOLAR_MASS = tuple( - sum( - _ELEMENT_MATRIX[e][i] * _ATOMIC_MASS[element] - for e, element in enumerate(ELEMENTS) - ) - for i in range(len(COMPONENTS)) -) +_MOLAR_MASS = tuple(molar_masses_from_yaml(str(_CONFIG))) def make_options( @@ -53,7 +36,6 @@ def make_options( options = EquilibriumOptions.from_yaml(str(_CONFIG)) if scenario == "isolated-core": keep = [i for i in range(18) if i not in (1, 3, 4, 6)] - options.stoich([[row[j] for j in keep] for row in options.stoich()]) options.reactions([options.reactions()[j] for j in keep]) elif scenario != "reactive": raise ValueError(f"unknown scenario: {scenario}") @@ -169,7 +151,7 @@ def __init__( self.pressure_tolerance = pressure_tolerance self.pressure_damping = pressure_damping self.options = make_options(scenario) - self.solver = Equilibrium(self.options) + self.solver = EquilibriumTP(self.options) def solve( self, diff --git a/python/kintera.pyi b/python/kintera.pyi index b583ab3f..9a3b9b47 100644 --- a/python/kintera.pyi +++ b/python/kintera.pyi @@ -12,6 +12,13 @@ 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]: ... + class SpeciesThermo: """ Species thermodynamics configuration. @@ -461,12 +468,9 @@ class EquilibriumOptions: @staticmethod def from_yaml(filename: str, verbose: bool = False) -> EquilibriumOptions: ... def components(self, value: List[str] = ...) -> Any: ... - def elements(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 stoich(self, value: List[List[float]] = ...) -> Any: ... - def element_matrix(self, value: List[List[float]] = ...) -> Any: ... def gas_phase(self, value: int = ...) -> Any: ... def standard_pressure(self, value: float = ...) -> Any: ... def max_iter(self, value: int = ...) -> Any: ... @@ -474,7 +478,7 @@ class EquilibriumOptions: def mole_floor(self, value: float = ...) -> Any: ... def validate(self) -> None: ... -class Equilibrium: +class EquilibriumTP: options: EquilibriumOptions def __init__(self, options: EquilibriumOptions) -> None: ... def forward( diff --git a/src/equilibrium/equilibrate.h b/src/equilibrium/equilibrate.h index a53e2510..49a0def5 100644 --- a/src/equilibrium/equilibrate.h +++ b/src/equilibrium/equilibrate.h @@ -48,10 +48,9 @@ DISPATCH_MACRO T equilibrium_max_error(T const *moles, T pres, template DISPATCH_MACRO int equilibrate(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, - T const *element_matrix, int nspecies, int nreaction, int nphase, - int nelement, int gas_phase, T standard_pressure, T ftol, - T mole_floor, int max_iter, char *work = nullptr) { + 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.; @@ -65,7 +64,7 @@ equilibrate(T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, } T *phase_totals, *phase_stoich, *residual, *jac, *constraints, *bounds, *step; - T *trial, *initial_elements; + T *trial; bool own_work = work == nullptr; if (own_work) { phase_totals = (T *)malloc(nphase * sizeof(T)); @@ -76,7 +75,6 @@ equilibrate(T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, bounds = (T *)malloc(nspecies * sizeof(T)); step = (T *)malloc(nreaction * sizeof(T)); trial = (T *)malloc(nspecies * sizeof(T)); - initial_elements = (T *)malloc(nelement * sizeof(T)); } else { phase_totals = alloc_from(work, nphase); phase_stoich = alloc_from(work, nphase * nreaction); @@ -86,7 +84,6 @@ equilibrate(T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, bounds = alloc_from(work, nspecies); step = alloc_from(work, nreaction); trial = alloc_from(work, nspecies); - initial_elements = alloc_from(work, nelement); } memcpy(out_moles, in_moles, nspecies * sizeof(T)); @@ -97,13 +94,6 @@ equilibrate(T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, phase_stoich[phase_ids[i] * nreaction + j] += stoich[i * nreaction + j]; } } - for (int e = 0; e < nelement; ++e) { - initial_elements[e] = 0.; - for (int i = 0; i < nspecies; ++i) { - initial_elements[e] += element_matrix[e * nspecies + i] * in_moles[i]; - } - } - T target = log(1. + ftol); T max_error = 0.; int status = 2; @@ -185,23 +175,9 @@ equilibrate(T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, nphase, gas_phase, phase_totals, residual); memcpy(gain, jac, nreaction * nreaction * sizeof(T)); - T element_error = 0.; - for (int e = 0; e < nelement; ++e) { - T final_value = 0.; - for (int i = 0; i < nspecies; ++i) { - final_value += element_matrix[e * nspecies + i] * out_moles[i]; - } - T denom = fabs(initial_elements[e]); - T error = denom > 0. ? fabs(final_value - initial_elements[e]) / denom - : fabs(final_value - initial_elements[e]); - if (error > element_error) - element_error = error; - } - diag[0] = static_cast(status); diag[1] = static_cast(iter); diag[2] = exp(max_error) - 1.; - diag[3] = element_error; if (own_work) { free(phase_totals); @@ -212,7 +188,6 @@ equilibrate(T *gain, T *diag, T *out_moles, T temp, T pres, T const *in_moles, free(bounds); free(step); free(trial); - free(initial_elements); } return status; } diff --git a/src/equilibrium/equilibrium.cpp b/src/equilibrium/equilibrium.cpp index c9e1e315..45359e69 100644 --- a/src/equilibrium/equilibrium.cpp +++ b/src/equilibrium/equilibrium.cpp @@ -1,18 +1,21 @@ #include + #include #include #include "equilibrium.hpp" #include "equilibrium_dispatch.hpp" +#include +#include + namespace kintera { void EquilibriumOptionsImpl::report(std::ostream &os) const { os << "-- equilibrium options --\n" << "* components = " << components().size() << "\n" - << "* elements = " << elements().size() << "\n" << "* phases = " << phases().size() << "\n" - << "* reactions = " << (stoich().empty() ? 0 : stoich()[0].size()) << "\n" + << "* reactions = " << reactions().size() << "\n" << "* gas_phase = " << gas_phase() << "\n" << "* standard_pressure = " << standard_pressure() << " Pa\n" << "* max_iter = " << max_iter() << "\n" @@ -24,45 +27,9 @@ void EquilibriumOptionsImpl::validate() const { TORCH_CHECK(!phases().empty(), "Equilibrium requires phases"); TORCH_CHECK(phase_ids().size() == components().size(), "phase_ids must contain one entry per component"); - TORCH_CHECK(stoich().size() == components().size(), - "stoich must have one row per component"); - TORCH_CHECK(!stoich()[0].empty(), "Equilibrium requires reactions"); - if (!reactions().empty()) { - TORCH_CHECK(reactions().size() == stoich()[0].size(), - "reactions and stoich column counts must match"); - } - if (!elements().empty()) { - TORCH_CHECK(elements().size() == element_matrix().size(), - "elements and element_matrix row counts must match"); - } + TORCH_CHECK(!reactions().empty(), "Equilibrium requires reactions"); TORCH_CHECK(components().size() <= 64, "the active-set solver currently supports at most 64 components"); - auto nreaction = stoich()[0].size(); - for (auto const &row : stoich()) { - TORCH_CHECK(row.size() == nreaction, "stoich rows must have equal length"); - } - - // Independent reaction columns are required by the square Newton system. - auto work = stoich(); - size_t rank = 0; - for (size_t column = 0; column < nreaction && rank < work.size(); ++column) { - size_t pivot = rank; - for (size_t row = rank + 1; row < work.size(); ++row) { - if (std::abs(work[row][column]) > std::abs(work[pivot][column])) - pivot = row; - } - if (std::abs(work[pivot][column]) <= 1.e-12) - continue; - std::swap(work[rank], work[pivot]); - for (size_t row = rank + 1; row < work.size(); ++row) { - double factor = work[row][column] / work[rank][column]; - for (size_t j = column; j < nreaction; ++j) { - work[row][j] -= factor * work[rank][j]; - } - } - ++rank; - } - TORCH_CHECK(rank == nreaction, "stoich reaction columns must be independent"); for (int phase : phase_ids()) { TORCH_CHECK(phase >= 0 && phase < static_cast(phases().size()), "phase id out of range: ", phase); @@ -75,69 +42,55 @@ void EquilibriumOptionsImpl::validate() const { TORCH_CHECK(ftol() > 0., "ftol must be positive"); TORCH_CHECK(mole_floor() >= 0., "mole_floor must be nonnegative"); - if (!element_matrix().empty()) { - for (auto const &row : element_matrix()) { - TORCH_CHECK(row.size() == components().size(), - "element_matrix rows must match component count"); - for (size_t j = 0; j < nreaction; ++j) { - double balance = 0.; - for (size_t i = 0; i < components().size(); ++i) { - balance += row[i] * stoich()[i][j]; - } - TORCH_CHECK(std::abs(balance) <= 1.e-10, "reaction ", j, - " does not conserve an element"); - } + 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); } } } -EquilibriumImpl::EquilibriumImpl(EquilibriumOptions const &options_) +EquilibriumTPImpl::EquilibriumTPImpl(EquilibriumOptions const &options_) : options(options_) { reset(); } -void EquilibriumImpl::reset() { +void EquilibriumTPImpl::reset() { options->validate(); - int nspecies = static_cast(options->components().size()); - int nreaction = static_cast(options->stoich()[0].size()); - - std::vector stoich_flat; - stoich_flat.reserve(nspecies * nreaction); - for (auto const &row : options->stoich()) { - stoich_flat.insert(stoich_flat.end(), row.begin(), row.end()); - } + std::vector reactions; + reactions.reserve(options->reactions().size()); + for (auto const &equation : options->reactions()) + reactions.emplace_back(equation); stoich = register_buffer( - "stoich", - torch::tensor(stoich_flat, torch::kFloat64).view({nspecies, nreaction})); + "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)); - - int nelement = static_cast(options->element_matrix().size()); - std::vector element_flat; - if (nelement == 0) { - // A zero row keeps the device kernel and diagnostic layout uniform. - nelement = 1; - element_flat.resize(nspecies, 0.); - } else { - element_flat.reserve(nelement * nspecies); - for (auto const &row : options->element_matrix()) { - element_flat.insert(element_flat.end(), row.begin(), row.end()); - } - } - element_matrix = register_buffer( - "element_matrix", - torch::tensor(element_flat, torch::kFloat64).view({nelement, nspecies})); } -void EquilibriumImpl::pretty_print(std::ostream &os) const { - os << "Equilibrium(components=" << options->components().size() +void EquilibriumTPImpl::pretty_print(std::ostream &os) const { + os << "EquilibriumTP(components=" << options->components().size() << ", phases=" << options->phases().size() << ")"; } std::tuple -EquilibriumImpl::forward(torch::Tensor temp, torch::Tensor pres, - torch::Tensor moles, torch::Tensor log_k, - bool warm_start) { +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"); @@ -174,7 +127,7 @@ EquilibriumImpl::forward(torch::Tensor temp, torch::Tensor pres, gain_shape.push_back(nreaction * nreaction); auto gain = torch::zeros(gain_shape, moles.options()); auto diag_shape = batch; - diag_shape.push_back(4); + diag_shape.push_back(3); auto diag = torch::zeros(diag_shape, moles.options()); auto iter = at::TensorIteratorConfig() @@ -193,8 +146,7 @@ EquilibriumImpl::forward(torch::Tensor temp, torch::Tensor pres, auto s = stoich.to(moles.options()).contiguous(); auto p = phase_ids.to(moles.device(), torch::kInt).contiguous(); - auto e = element_matrix.to(moles.options()).contiguous(); - at::native::call_equilibrium(moles.device().type(), iter, s, p, e, + 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()); diff --git a/src/equilibrium/equilibrium.hpp b/src/equilibrium/equilibrium.hpp index bfba951d..81fbc09c 100644 --- a/src/equilibrium/equilibrium.hpp +++ b/src/equilibrium/equilibrium.hpp @@ -13,8 +13,6 @@ namespace kintera { -using Matrix = std::vector>; - struct EquilibriumOptionsImpl final { static std::shared_ptr create() { return std::make_shared(); @@ -27,12 +25,9 @@ struct EquilibriumOptionsImpl final { void validate() const; ADD_ARG(std::vector, components); - ADD_ARG(std::vector, elements); ADD_ARG(std::vector, phases); ADD_ARG(std::vector, reactions); ADD_ARG(std::vector, phase_ids); - ADD_ARG(Matrix, stoich); - ADD_ARG(Matrix, element_matrix); ADD_ARG(int, gas_phase) = 0; ADD_ARG(double, standard_pressure) = 1.e5; ADD_ARG(int, max_iter) = 50; @@ -41,16 +36,15 @@ struct EquilibriumOptionsImpl final { }; using EquilibriumOptions = std::shared_ptr; -class EquilibriumImpl : public torch::nn::Cloneable { +class EquilibriumTPImpl : public torch::nn::Cloneable { public: EquilibriumOptions options; torch::Tensor stoich; torch::Tensor phase_ids; - torch::Tensor element_matrix; - EquilibriumImpl() : options(EquilibriumOptionsImpl::create()) {} - explicit EquilibriumImpl(EquilibriumOptions const &options_); + EquilibriumTPImpl() : options(EquilibriumOptionsImpl::create()) {} + explicit EquilibriumTPImpl(EquilibriumOptions const &options_); void reset() override; void pretty_print(std::ostream &os) const override; @@ -63,14 +57,13 @@ class EquilibriumImpl : public torch::nn::Cloneable { * log_k: (..., nreaction) * * Returns equilibrium moles, the final reaction Jacobian/gain matrix, and - * diagnostics (..., 4): status, iterations, equilibrium error, and element - * conservation error. + * 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(Equilibrium); +TORCH_MODULE(EquilibriumTP); } // namespace kintera diff --git a/src/equilibrium/equilibrium_dispatch.cpp b/src/equilibrium/equilibrium_dispatch.cpp index 3655abc6..b82cc4a4 100644 --- a/src/equilibrium/equilibrium_dispatch.cpp +++ b/src/equilibrium/equilibrium_dispatch.cpp @@ -8,18 +8,15 @@ namespace kintera { void call_equilibrium_cpu(at::TensorIterator &iter, at::Tensor const &stoich, - at::Tensor const &phase_ids, - at::Tensor const &element_matrix, int nphase, + 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); - int nelement = element_matrix.size(0); auto stoich_ptr = stoich.data_ptr(); auto phase_ptr = phase_ids.data_ptr(); - auto element_ptr = element_matrix.data_ptr(); iter.for_each( [&](char **data, const int64_t *strides, int64_t n) { for (int64_t i = 0; i < n; ++i) { @@ -31,8 +28,7 @@ void call_equilibrium_cpu(at::TensorIterator &iter, at::Tensor const &stoich, auto moles = reinterpret_cast(data[5] + i * strides[5]); auto log_k = reinterpret_cast(data[6] + i * strides[6]); equilibrate(gain, diag, out, *temp, *pres, moles, log_k, stoich_ptr, - phase_ptr, element_ptr, nspecies, nreaction, nphase, - nelement, gas_phase, + phase_ptr, nspecies, nreaction, nphase, gas_phase, static_cast(standard_pressure), static_cast(ftol), static_cast(mole_floor), max_iter); diff --git a/src/equilibrium/equilibrium_dispatch.cu b/src/equilibrium/equilibrium_dispatch.cu index 49296836..b9af0f55 100644 --- a/src/equilibrium/equilibrium_dispatch.cu +++ b/src/equilibrium/equilibrium_dispatch.cu @@ -9,8 +9,7 @@ namespace kintera { template -size_t equilibrium_space(int nspecies, int nreaction, int nphase, - int nelement) { +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; @@ -23,25 +22,20 @@ size_t equilibrium_space(int nspecies, int nreaction, int nphase, bump(alignof(T), nspecies * sizeof(T)); bump(alignof(T), nreaction * sizeof(T)); bump(alignof(T), nspecies * sizeof(T)); - bump(alignof(T), nelement * 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, - at::Tensor const &element_matrix, int nphase, + 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); - int nelement = element_matrix.size(0); auto stoich_ptr = stoich.data_ptr(); auto phase_ptr = phase_ids.data_ptr(); - auto element_ptr = element_matrix.data_ptr(); - int work_size = - equilibrium_space(nspecies, nreaction, nphase, nelement); + int work_size = equilibrium_space(nspecies, nreaction, nphase); // The KKT workspace scales quadratically with reaction count. One cell // per block keeps the 25-component paper case within shared-memory limits. native::gpu_mem_kernel<1, 7>( @@ -56,8 +50,7 @@ void call_equilibrium_cuda(at::TensorIterator &iter, at::Tensor const &stoich, auto moles = reinterpret_cast(data[5] + strides[5]); auto log_k = reinterpret_cast(data[6] + strides[6]); equilibrate(gain, diag, out, *temp, *pres, moles, log_k, stoich_ptr, - phase_ptr, element_ptr, nspecies, nreaction, nphase, - nelement, gas_phase, + phase_ptr, nspecies, nreaction, nphase, gas_phase, static_cast(standard_pressure), static_cast(ftol), static_cast(mole_floor), max_iter, work); diff --git a/src/equilibrium/equilibrium_dispatch.hpp b/src/equilibrium/equilibrium_dispatch.hpp index bdddb293..f04f1473 100644 --- a/src/equilibrium/equilibrium_dispatch.hpp +++ b/src/equilibrium/equilibrium_dispatch.hpp @@ -7,8 +7,7 @@ namespace at::native { using equilibrium_fn = void (*)(at::TensorIterator &iter, at::Tensor const &stoich, - at::Tensor const &phase_ids, - at::Tensor const &element_matrix, int nphase, + at::Tensor const &phase_ids, int nphase, int gas_phase, double standard_pressure, double ftol, double mole_floor, int max_iter); diff --git a/src/equilibrium/equilibrium_options.cpp b/src/equilibrium/equilibrium_options.cpp index da13b849..8f614cb6 100644 --- a/src/equilibrium/equilibrium_options.cpp +++ b/src/equilibrium/equilibrium_options.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -71,19 +72,6 @@ EquilibriumOptionsImpl::from_yaml(std::string const &filename, bool verbose) { "phase component has no species definition: ", component); } - options->elements().assign(element_set.begin(), element_set.end()); - options->element_matrix().resize(options->elements().size()); - for (size_t e = 0; e < options->elements().size(); ++e) { - auto &row = options->element_matrix()[e]; - row.resize(options->components().size(), 0.); - for (size_t i = 0; i < options->components().size(); ++i) { - auto const &composition = compositions.at(options->components()[i]); - auto found = composition.find(options->elements()[e]); - if (found != composition.end()) - row[i] = found->second; - } - } - std::vector reactions; for (auto const &reaction_node : config["reactions"]) { if (reaction_node["type"] && @@ -97,24 +85,35 @@ EquilibriumOptionsImpl::from_yaml(std::string const &filename, bool verbose) { } TORCH_CHECK(!reactions.empty(), "equilibrium YAML contains no reactions"); - options->stoich().assign(options->components().size(), - std::vector(reactions.size(), 0.)); 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); - options->stoich()[found - options->components().begin()][j] -= - coefficient; + (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); - options->stoich()[found - options->components().begin()][j] += - coefficient; + (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); } } diff --git a/src/species.cpp b/src/species.cpp index a19af559..31946640 100644 --- a/src/species.cpp +++ b/src/species.cpp @@ -14,13 +14,11 @@ // torch #include -// harp -#include - // kintera #include #include +#include #include #include "species.hpp" @@ -53,16 +51,17 @@ void clear_species_registry() { species_nasa9_Tmid.clear(); } -} // namespace +} // 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; + if (!db.empty()) + return db; std::string path; try { path = find_resource("nasa9.dat"); - } catch (std::exception const& e) { + } catch (std::exception const &e) { TORCH_CHECK(false, e.what()); } std::ifstream ifs(path); @@ -70,12 +69,14 @@ static std::unordered_map& get_nasa9_db() { std::string line; while (std::getline(ifs, line)) { - if (line.empty() || line[0] == '#') continue; + 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(); + while (!name.empty() && std::isspace(name.back())) + name.pop_back(); // read 4 lines of 5 values = 20 coefficients double vals[20]; @@ -83,21 +84,25 @@ static std::unordered_map& get_nasa9_db() { 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; + while (iss >> v && idx < 20) + vals[idx++] = v; } - if (idx < 20) continue; + 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 + 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 + 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; } @@ -110,23 +115,23 @@ void init_species_from_yaml(std::string 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 +156,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 +174,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 +203,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 +222,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,22 +243,24 @@ 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; + 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; + if (v != 0.0) + return true; } } return false; } -static at::Tensor nasa9_coeffs_to_tensor( - std::vector> const& coeffs, - c10::TensorOptions const& options) { +static at::Tensor +nasa9_coeffs_to_tensor(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 +275,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 { +at::Tensor +SpeciesThermoImpl::nasa9_Tmid_tensor(c10::TensorOptions const &options) const { auto tensor = torch::empty({static_cast(nasa9_Tmid().size())}, torch::dtype(torch::kFloat64)); if (!nasa9_Tmid().empty()) { @@ -329,7 +336,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 +387,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 +396,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 +517,8 @@ 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); @@ -531,4 +539,4 @@ SpeciesThermo merge_thermo(SpeciesThermo const& thermo1, return merged; } -} // namespace kintera +} // namespace kintera diff --git a/src/utils/molar_mass.cpp b/src/utils/molar_mass.cpp new file mode 100644 index 00000000..c88b679b --- /dev/null +++ b/src/utils/molar_mass.cpp @@ -0,0 +1,85 @@ +#include +#include + +#include + +#include + +#include + +#include "molar_mass.hpp" + +#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..58b77671 --- /dev/null +++ b/src/utils/molar_mass.hpp @@ -0,0 +1,24 @@ +#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/test_equilibrium.cpp b/tests/test_equilibrium.cpp index 122d3d3f..20e9e795 100644 --- a/tests/test_equilibrium.cpp +++ b/tests/test_equilibrium.cpp @@ -4,6 +4,7 @@ #include #include +#include #define DEVICE_TESTING_SKIP_DEFAULT_INSTANTIATION #include "device_testing.hpp" @@ -17,8 +18,7 @@ EquilibriumOptions make_options() { op->components({"A", "B"}) .phases({"gas"}) .phase_ids({0, 0}) - .stoich({{-1.}, {1.}}) - .element_matrix({{1., 1.}}) + .reactions({"A <=> B"}) .gas_phase(0) .standard_pressure(1.e5) .max_iter(50) @@ -34,7 +34,7 @@ TEST_P(EquilibriumDeviceTest, IdealGasReaction) { if (device.type() == torch::kMPS) GTEST_SKIP(); auto op = make_options(); - Equilibrium eq(op); + EquilibriumTP eq(op); eq->to(device, dtype); auto tensor_options = torch::device(device).dtype(dtype); @@ -52,15 +52,13 @@ TEST_P(EquilibriumDeviceTest, IdealGasReaction) { 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_LT(diag.select(-1, 3).max().item(), 1.e-6); EXPECT_TRUE( torch::allclose(moles, torch::tensor({{0.8, 0.2}}, tensor_options))); } -TEST(EquilibriumOptions, RejectsUnbalancedReaction) { - auto op = make_options(); - op->element_matrix({{1., 2.}}); - EXPECT_THROW(op->validate(), c10::Error); +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(EquilibriumOptions, ReadsSpeciesAndReactionsFromYaml) { @@ -73,8 +71,8 @@ TEST(EquilibriumOptions, ReadsSpeciesAndReactionsFromYaml) { << " thermo: ideal-gas\n" << " species: [A, B]\n" << "species:\n" - << " - {name: A, composition: {X: 1}}\n" - << " - {name: B, composition: {X: 1}}\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" @@ -84,13 +82,34 @@ TEST(EquilibriumOptions, ReadsSpeciesAndReactionsFromYaml) { auto op = EquilibriumOptionsImpl::from_yaml(path.string()); EXPECT_EQ(op->components(), std::vector({"A", "B"})); - EXPECT_EQ(op->elements(), std::vector({"X"})); EXPECT_EQ(op->reactions(), std::vector({"A <=> B"})); EXPECT_EQ(op->phase_ids(), std::vector({0, 0})); - EXPECT_DOUBLE_EQ(op->stoich()[0][0], -1.); - EXPECT_DOUBLE_EQ(op->stoich()[1][0], 1.); 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); } diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index 10407571..416290a1 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -1,6 +1,14 @@ +from pathlib import Path + import pytest import torch -from kintera import Equilibrium, EquilibriumOptions +from kintera import ( + EquilibriumOptions, + EquilibriumTP, + atomic_mass, + molar_mass, + molar_masses_from_yaml, +) from kintera.equilibrium import ( COMPONENTS, SchlichtingYoung2022, @@ -16,11 +24,10 @@ def test_python_core_binding_is_functional(): .components(["A", "B"]) .phases(["gas"]) .phase_ids([0, 0]) - .stoich([[-1.0], [1.0]]) - .element_matrix([[1.0, 1.0]]) + .reactions(["A <=> B"]) .gas_phase(0) ) - solver = Equilibrium(options) + solver = EquilibriumTP(options) moles = torch.tensor([0.8, 0.2], dtype=torch.float64) result, gain, diagnostics = solver( torch.tensor(1000.0, dtype=torch.float64), @@ -36,13 +43,26 @@ def test_python_core_binding_is_functional(): 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_paper_topology_and_pressure_relation(): options = make_options() options.validate() assert len(options.components()) == 25 - assert len(options.elements()) == 7 assert len(options.reactions()) == 18 - assert len(options.stoich()[0]) == 18 + config = ( + Path(__file__).parents[1] + / "python" + / "equilibrium" + / "schlichting_young_2022.yaml" + ) + assert len(molar_masses_from_yaml(str(config))) == 25 + assert EquilibriumTP(options).buffer("stoich").shape == (25, 18) + isolated = make_options("isolated-core") + assert EquilibriumTP(isolated).buffer("stoich").shape == (25, 14) moles = initial_moles() pressure = surface_pressure(moles) @@ -53,7 +73,7 @@ def test_paper_topology_and_pressure_relation(): def test_paper_driver_accepts_external_thermodynamics(): moles = initial_moles() options = make_options() - stoich = torch.tensor(options.stoich(), dtype=moles.dtype) + stoich = EquilibriumTP(options).buffer("stoich").to(dtype=moles.dtype) phase_ids = options.phase_ids() def equilibrium_at_initial_state(temp, pres): From 7d662313b74e4f3869a5938cd6afdd5dbdf799b0 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Thu, 9 Jul 2026 22:35:52 -0400 Subject: [PATCH 3/8] wip --- README.md | 16 +- examples/equilibrium_nasa9.py | 31 +++ examples/equilibrium_nasa9.yaml | 24 ++ python/csrc/kintera.cpp | 2 + python/equilibrium/__init__.py | 22 +- python/equilibrium/nasa9.py | 28 +++ python/equilibrium/schlichting_young_2022.py | 218 ------------------ .../equilibrium/schlichting_young_2022.yaml | 102 -------- python/kintera.pyi | 1 + src/equilibrium/equilibrium_dispatch.cu | 7 +- src/loops.cuh | 107 ++++++--- src/species.cpp | 32 +++ src/species.hpp | 36 +-- tests/test_equilibrium.py | 88 ++----- 14 files changed, 253 insertions(+), 461 deletions(-) create mode 100644 examples/equilibrium_nasa9.py create mode 100644 examples/equilibrium_nasa9.yaml create mode 100644 python/equilibrium/nasa9.py delete mode 100644 python/equilibrium/schlichting_young_2022.py delete mode 100644 python/equilibrium/schlichting_young_2022.yaml diff --git a/README.md b/README.md index 134e7449..eefc9789 100644 --- a/README.md +++ b/README.md @@ -28,12 +28,8 @@ KINTERA provides efficient implementations of: `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 -physics remains in Python under `kintera.equilibrium`; the -`SchlichtingYoung2022` driver supplies the published core-mantle-atmosphere -topology and iterates its atmosphere mass-pressure relation around the core -solve. Thermodynamic coefficients referenced by that paper but not published -in it are supplied through the driver's required `log_k_model` callback. +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 @@ -47,6 +43,14 @@ 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/python/csrc/kintera.cpp b/python/csrc/kintera.cpp index 91d49bc9..4518f8de 100644 --- a/python/csrc/kintera.cpp +++ b/python/csrc/kintera.cpp @@ -46,6 +46,8 @@ PYBIND11_MODULE(kintera, m) { 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_( diff --git a/python/equilibrium/__init__.py b/python/equilibrium/__init__.py index dbc1fb7e..244c2cc5 100644 --- a/python/equilibrium/__init__.py +++ b/python/equilibrium/__init__.py @@ -1,21 +1,3 @@ -from .schlichting_young_2022 import ( - COMPONENTS, - PHASES, - REACTIONS, - SchlichtingYoung2022, - SchlichtingYoungResult, - initial_moles, - make_options, - surface_pressure, -) +from .nasa9 import Nasa9LogK -__all__ = [ - "COMPONENTS", - "PHASES", - "REACTIONS", - "SchlichtingYoung2022", - "SchlichtingYoungResult", - "initial_moles", - "make_options", - "surface_pressure", -] +__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/equilibrium/schlichting_young_2022.py b/python/equilibrium/schlichting_young_2022.py deleted file mode 100644 index 2fc54424..00000000 --- a/python/equilibrium/schlichting_young_2022.py +++ /dev/null @@ -1,218 +0,0 @@ -"""Schlichting & Young (2022) core-mantle-atmosphere equilibrium case. - -The paper specifies the reaction topology and pressure relation completely, -but refers several standard-state Gibbs energies to external NIST, MAGMA, and -experimental datasets. Consequently this module requires a ``log_k_model`` -callable rather than silently substituting thermodynamic data. -""" - -from dataclasses import dataclass -from pathlib import Path -from typing import Callable, Literal - -import torch - -from ..kintera import EquilibriumOptions, EquilibriumTP, molar_masses_from_yaml - -_CONFIG = Path(__file__).with_name("schlichting_young_2022.yaml") -_BASE_OPTIONS = EquilibriumOptions.from_yaml(str(_CONFIG)) -COMPONENTS = tuple(_BASE_OPTIONS.components()) -PHASES = tuple(_BASE_OPTIONS.phases()) -REACTIONS = tuple(_BASE_OPTIONS.reactions()) -_PHASE_IDS = tuple(_BASE_OPTIONS.phase_ids()) -_GAS_PHASE = _BASE_OPTIONS.gas_phase() -_COMPONENT_INDEX = {name: i for i, name in enumerate(COMPONENTS)} - -_MOLAR_MASS = tuple(molar_masses_from_yaml(str(_CONFIG))) - - -def make_options( - scenario: Literal["reactive", "isolated-core"] = "reactive", - *, - max_iter: int = 100, - ftol: float = 5e-3, -) -> EquilibriumOptions: - """Create the generic core options for the paper's reaction system.""" - options = EquilibriumOptions.from_yaml(str(_CONFIG)) - if scenario == "isolated-core": - keep = [i for i in range(18) if i not in (1, 3, 4, 6)] - options.reactions([options.reactions()[j] for j in keep]) - elif scenario != "reactive": - raise ValueError(f"unknown scenario: {scenario}") - options.max_iter(max_iter).ftol(ftol) - options.validate() - return options - - -def initial_moles( - hydrogen_mass_fraction: float = 0.04, - metal_mass_fraction: float = 0.25, - *, - dtype: torch.dtype = torch.float64, - device: torch.device | str | None = None, -) -> torch.Tensor: - """Construct the paper's nominal composition on a one-kilogram basis.""" - if not 0 < hydrogen_mass_fraction < 1: - raise ValueError("hydrogen_mass_fraction must lie between zero and one") - if not 0 <= metal_mass_fraction < 1 - hydrogen_mass_fraction: - raise ValueError("metal_mass_fraction leaves no silicate inventory") - - mass = torch.tensor(_MOLAR_MASS, dtype=dtype, device=device) - result = torch.full((len(COMPONENTS),), 1.0e-20, dtype=dtype, device=device) - fe_metal = _COMPONENT_INDEX["Fe[metal]"] - result[fe_metal] = metal_mass_fraction / mass[fe_metal] - - silicate_mass = 1.0 - metal_mass_fraction - hydrogen_mass_fraction - mgsio3 = _COMPONENT_INDEX["MgSiO3[silicate]"] - mgo = _COMPONENT_INDEX["MgO[silicate]"] - fesio3 = _COMPONENT_INDEX["FeSiO3[silicate]"] - na2o = _COMPONENT_INDEX["Na2O[silicate]"] - silicate_moles = silicate_mass / ( - 0.921 * mass[mgsio3] - + 0.032 * mass[mgo] - + 0.035 * mass[fesio3] - + 0.007 * mass[na2o] - ) - result[mgsio3] = 0.921 * silicate_moles - result[mgo] = 0.032 * silicate_moles - result[fesio3] = 0.035 * silicate_moles - result[na2o] = 0.007 * silicate_moles - - # The primary atmosphere is 99.9 mol% H2 and 0.1 mol% CO. - h2_gas = _COMPONENT_INDEX["H2[gas]"] - co_gas = _COMPONENT_INDEX["CO[gas]"] - gas_moles = hydrogen_mass_fraction / ( - 0.999 * mass[h2_gas] + 0.001 * mass[co_gas] - ) - result[h2_gas] = 0.999 * gas_moles - result[co_gas] = 0.001 * gas_moles - return result - - -def surface_pressure( - moles: torch.Tensor, planet_mass_earth: float | torch.Tensor = 4.0 -) -> torch.Tensor: - """Evaluate Equation 8 using the atmosphere mass implied by ``moles``.""" - mass = torch.as_tensor(_MOLAR_MASS, dtype=moles.dtype, device=moles.device) - component_mass = moles * mass - total_mass = component_mass.sum(dim=-1) - gas = torch.tensor( - [i for i, phase in enumerate(_PHASE_IDS) if phase == _GAS_PHASE], - dtype=torch.long, - device=moles.device, - ) - atmosphere_fraction = ( - component_mass.index_select(-1, gas).sum(dim=-1) / total_mass - ) - planet_mass = torch.as_tensor( - planet_mass_earth, dtype=moles.dtype, device=moles.device - ) - return 1.1e11 * atmosphere_fraction * planet_mass.pow(2.0 / 3.0) - - -@dataclass(frozen=True) -class SchlichtingYoungResult: - moles: torch.Tensor - pressure: torch.Tensor - phase_totals: torch.Tensor - phase_fractions: torch.Tensor - gain: torch.Tensor - diagnostics: torch.Tensor - pressure_error: torch.Tensor - outer_iterations: int - - -class SchlichtingYoung2022: - """Python driver coupling the generic core to the paper's pressure law.""" - - def __init__( - self, - log_k_model: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], - scenario: Literal["reactive", "isolated-core"] = "reactive", - *, - max_outer_iter: int = 30, - pressure_tolerance: float = 0.1, - pressure_damping: float = 0.5, - ) -> None: - if log_k_model is None: - raise TypeError( - "log_k_model is required because the paper refers some Gibbs " - "energies to external datasets" - ) - if max_outer_iter <= 0: - raise ValueError("max_outer_iter must be positive") - if pressure_tolerance <= 0: - raise ValueError("pressure_tolerance must be positive") - if not 0 < pressure_damping <= 1: - raise ValueError("pressure_damping must lie in (0, 1]") - self.scenario = scenario - self.log_k_model = log_k_model - self.max_outer_iter = max_outer_iter - self.pressure_tolerance = pressure_tolerance - self.pressure_damping = pressure_damping - self.options = make_options(scenario) - self.solver = EquilibriumTP(self.options) - - def solve( - self, - temp: torch.Tensor, - moles: torch.Tensor, - pressure: torch.Tensor | None = None, - *, - planet_mass_earth: float = 4.0, - ) -> SchlichtingYoungResult: - self.solver.to(moles.device, moles.dtype) - pressure = ( - surface_pressure(moles, planet_mass_earth) - if pressure is None - else pressure.to(device=moles.device, dtype=moles.dtype) - ) - state = moles - gain = diagnostics = None - pressure_error = torch.full_like(pressure, float("inf")) - - for _iteration in range(1, self.max_outer_iter + 1): - log_k = self.log_k_model(temp, pressure) - if self.scenario == "isolated-core" and log_k.size(-1) == 18: - keep = [i for i in range(18) if i not in (1, 3, 4, 6)] - log_k = log_k[..., keep] - state, gain, diagnostics = self.solver(temp, pressure, state, log_k) - target = surface_pressure(state, planet_mass_earth) - pressure_error = (target - pressure).abs() / target.clamp_min( - torch.finfo(target.dtype).tiny - ) - chemistry_ok = diagnostics[..., 0].eq(0).all() - if chemistry_ok and pressure_error.max() <= self.pressure_tolerance: - pressure = target - break - pressure = torch.exp( - (1.0 - self.pressure_damping) * torch.log(pressure) - + self.pressure_damping * torch.log(target) - ) - - phase_totals = torch.stack( - [ - state[..., [i for i, p in enumerate(_PHASE_IDS) if p == phase]].sum( - -1 - ) - for phase in range(len(PHASES)) - ], - dim=-1, - ) - phase_fractions = torch.empty_like(state) - for phase in range(len(PHASES)): - ids = [i for i, value in enumerate(_PHASE_IDS) if value == phase] - phase_fractions[..., ids] = ( - state[..., ids] / phase_totals[..., phase, None] - ) - - return SchlichtingYoungResult( - state, - pressure, - phase_totals, - phase_fractions, - gain, - diagnostics, - pressure_error, - _iteration, - ) diff --git a/python/equilibrium/schlichting_young_2022.yaml b/python/equilibrium/schlichting_young_2022.yaml deleted file mode 100644 index 7ea1cf34..00000000 --- a/python/equilibrium/schlichting_young_2022.yaml +++ /dev/null @@ -1,102 +0,0 @@ -equilibrium: - standard-pressure: 1.0e5 - max-iter: 100 - ftol: 5.0e-3 - mole-floor: 1.0e-30 - -phases: - - name: silicate - thermo: ideal-solution - species: - - "MgO[silicate]" - - "SiO2[silicate]" - - "MgSiO3[silicate]" - - "FeO[silicate]" - - "FeSiO3[silicate]" - - "Na2O[silicate]" - - "Na2SiO3[silicate]" - - "H2[silicate]" - - "H2O[silicate]" - - "CO[silicate]" - - "CO2[silicate]" - - name: metal - thermo: ideal-solution - species: ["Fe[metal]", "Si[metal]", "O[metal]", "H[metal]"] - - name: atmosphere - thermo: ideal-gas - species: - - "H2[gas]" - - "CO[gas]" - - "CO2[gas]" - - "CH4[gas]" - - "O2[gas]" - - "H2O[gas]" - - "Fe[gas]" - - "Mg[gas]" - - "SiO[gas]" - - "Na[gas]" - -species: - - {name: "MgO[silicate]", composition: {Mg: 1, O: 1}} - - {name: "SiO2[silicate]", composition: {Si: 1, O: 2}} - - {name: "MgSiO3[silicate]", composition: {Mg: 1, Si: 1, O: 3}} - - {name: "FeO[silicate]", composition: {Fe: 1, O: 1}} - - {name: "FeSiO3[silicate]", composition: {Fe: 1, Si: 1, O: 3}} - - {name: "Na2O[silicate]", composition: {Na: 2, O: 1}} - - {name: "Na2SiO3[silicate]", composition: {Na: 2, Si: 1, O: 3}} - - {name: "H2[silicate]", composition: {H: 2}} - - {name: "H2O[silicate]", composition: {H: 2, O: 1}} - - {name: "CO[silicate]", composition: {C: 1, O: 1}} - - {name: "CO2[silicate]", composition: {C: 1, O: 2}} - - {name: "Fe[metal]", composition: {Fe: 1}} - - {name: "Si[metal]", composition: {Si: 1}} - - {name: "O[metal]", composition: {O: 1}} - - {name: "H[metal]", composition: {H: 1}} - - {name: "H2[gas]", composition: {H: 2}} - - {name: "CO[gas]", composition: {C: 1, O: 1}} - - {name: "CO2[gas]", composition: {C: 1, O: 2}} - - {name: "CH4[gas]", composition: {C: 1, H: 4}} - - {name: "O2[gas]", composition: {O: 2}} - - {name: "H2O[gas]", composition: {H: 2, O: 1}} - - {name: "Fe[gas]", composition: {Fe: 1}} - - {name: "Mg[gas]", composition: {Mg: 1}} - - {name: "SiO[gas]", composition: {Si: 1, O: 1}} - - {name: "Na[gas]", composition: {Na: 1}} - -reactions: - - type: equilibrium - equation: "Na2SiO3[silicate] <=> Na2O[silicate] + SiO2[silicate]" - - type: equilibrium - equation: "0.5 SiO2[silicate] + Fe[metal] <=> FeO[silicate] + 0.5 Si[metal]" - - type: equilibrium - equation: "MgSiO3[silicate] <=> MgO[silicate] + SiO2[silicate]" - - type: equilibrium - equation: "O[metal] + 0.5 Si[metal] <=> 0.5 SiO2[silicate]" - - type: equilibrium - equation: "2 H[metal] <=> H2[silicate]" - - type: equilibrium - equation: "FeSiO3[silicate] <=> FeO[silicate] + SiO2[silicate]" - - type: equilibrium - equation: "2 H2O[silicate] + Si[metal] <=> SiO2[silicate] + 2 H2[silicate]" - - type: equilibrium - equation: "CO[gas] + 0.5 O2[gas] <=> CO2[gas]" - - type: equilibrium - equation: "CH4[gas] + 0.5 O2[gas] <=> 2 H2[gas] + CO[gas]" - - type: equilibrium - equation: "H2[gas] + 0.5 O2[gas] <=> H2O[gas]" - - type: equilibrium - equation: "FeO[silicate] <=> Fe[gas] + 0.5 O2[gas]" - - type: equilibrium - equation: "MgO[silicate] <=> Mg[gas] + 0.5 O2[gas]" - - type: equilibrium - equation: "SiO2[silicate] <=> SiO[gas] + 0.5 O2[gas]" - - type: equilibrium - equation: "Na2O[silicate] <=> 2 Na[gas] + 0.5 O2[gas]" - - type: equilibrium - equation: "H2[gas] <=> H2[silicate]" - - type: equilibrium - equation: "H2O[gas] <=> H2O[silicate]" - - type: equilibrium - equation: "CO[gas] <=> CO[silicate]" - - type: equilibrium - equation: "CO2[gas] <=> CO2[silicate]" diff --git a/python/kintera.pyi b/python/kintera.pyi index 9a3b9b47..e2675e5e 100644 --- a/python/kintera.pyi +++ b/python/kintera.pyi @@ -18,6 +18,7 @@ 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: """ diff --git a/src/equilibrium/equilibrium_dispatch.cu b/src/equilibrium/equilibrium_dispatch.cu index b9af0f55..c6231884 100644 --- a/src/equilibrium/equilibrium_dispatch.cu +++ b/src/equilibrium/equilibrium_dispatch.cu @@ -35,10 +35,9 @@ void call_equilibrium_cuda(at::TensorIterator &iter, at::Tensor const &stoich, int nreaction = stoich.size(1); auto stoich_ptr = stoich.data_ptr(); auto phase_ptr = phase_ids.data_ptr(); - int work_size = equilibrium_space(nspecies, nreaction, nphase); - // The KKT workspace scales quadratically with reaction count. One cell - // per block keeps the 25-component paper case within shared-memory limits. - native::gpu_mem_kernel<1, 7>( + 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) { 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/species.cpp b/src/species.cpp index 31946640..d2c26615 100644 --- a/src/species.cpp +++ b/src/species.cpp @@ -21,6 +21,8 @@ #include #include +#include + #include "species.hpp" namespace kintera { @@ -110,6 +112,36 @@ static std::unordered_map &get_nasa9_db() { 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"); + 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); + } + + auto tensor_options = temp.options(); + auto low_tensor = + torch::empty({static_cast(species.size()), 9}, tensor_options); + auto high_tensor = torch::empty_like(low_tensor); + for (size_t i = 0; i < species.size(); ++i) { + for (int j = 0; j < 9; ++j) { + low_tensor.index_put_({static_cast(i), j}, low[i][j]); + high_tensor.index_put_({static_cast(i), j}, high[i][j]); + } + } + auto midpoint = torch::full({static_cast(species.size())}, 1000., + tensor_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); diff --git a/src/species.hpp b/src/species.hpp index aefda416..2aa84c75 100644 --- a/src/species.hpp +++ b/src/species.hpp @@ -16,11 +16,11 @@ namespace at { class Tensor; -} // namespace at +} // namespace at namespace YAML { class Node; -} // namespace YAML +} // namespace YAML namespace kintera { @@ -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; @@ -104,6 +108,6 @@ extern std::vector species_uref_R; extern std::vector species_sref_R; extern bool species_initialized; -} // namespace kintera +} // namespace kintera #undef ADD_ARG diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index 416290a1..1ed6314e 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -7,15 +7,8 @@ EquilibriumTP, atomic_mass, molar_mass, - molar_masses_from_yaml, -) -from kintera.equilibrium import ( - COMPONENTS, - SchlichtingYoung2022, - initial_moles, - make_options, - surface_pressure, ) +from kintera.equilibrium import Nasa9LogK def test_python_core_binding_is_functional(): @@ -48,62 +41,27 @@ def test_standalone_molar_mass_utilities(): assert molar_mass({"H": 2.0, "O": 1.0}) == pytest.approx(18.015e-3) -def test_paper_topology_and_pressure_relation(): - options = make_options() - options.validate() - assert len(options.components()) == 25 - assert len(options.reactions()) == 18 - config = ( - Path(__file__).parents[1] - / "python" - / "equilibrium" - / "schlichting_young_2022.yaml" - ) - assert len(molar_masses_from_yaml(str(config))) == 25 - assert EquilibriumTP(options).buffer("stoich").shape == (25, 18) - isolated = make_options("isolated-core") - assert EquilibriumTP(isolated).buffer("stoich").shape == (25, 14) - - moles = initial_moles() - pressure = surface_pressure(moles) - assert moles.shape == (len(COMPONENTS),) - assert pressure > 0 - - -def test_paper_driver_accepts_external_thermodynamics(): - moles = initial_moles() - options = make_options() - stoich = EquilibriumTP(options).buffer("stoich").to(dtype=moles.dtype) - phase_ids = options.phase_ids() - - def equilibrium_at_initial_state(temp, pres): - del temp - log_activity = torch.empty_like(moles) - for phase in range(3): - ids = [i for i, value in enumerate(phase_ids) if value == phase] - log_activity[ids] = torch.log(moles[ids] / moles[ids].sum()) - gas = [i for i, value in enumerate(phase_ids) if value == 2] - log_activity[gas] += torch.log(pres / 1.0e5) - return log_activity @ stoich - - model = SchlichtingYoung2022(equilibrium_at_initial_state, max_outer_iter=2) - result = model.solve(torch.tensor(4500.0, dtype=moles.dtype), moles) - assert model.options.gas_phase() == 2 - assert result.diagnostics[0] == 0 - assert result.pressure_error < 0.1 - for phase in range(3): - ids = [i for i, value in enumerate(model.options.phase_ids()) if value == phase] - torch.testing.assert_close( - result.phase_fractions[ids].sum(), torch.tensor(1.0, dtype=moles.dtype) - ) - +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) -def test_paper_driver_rejects_invalid_outer_solve_options(): - def model(temp, pres): - del pres - return torch.zeros(18, dtype=temp.dtype, device=temp.device) + log_k = model(temp, pressure) + result, _, diagnostics = solver(temp, pressure, initial, log_k) - with pytest.raises(ValueError, match="max_outer_iter"): - SchlichtingYoung2022(model, max_outer_iter=0) - with pytest.raises(ValueError, match="pressure_damping"): - SchlichtingYoung2022(model, pressure_damping=0.0) + 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), + ) From af73f607de7bbaf18883935cf56469c5fa7dc93c Mon Sep 17 00:00:00 2001 From: mac/cli Date: Thu, 9 Jul 2026 22:42:38 -0400 Subject: [PATCH 4/8] wip --- src/equilibrium/equilibrium_dispatch.cpp | 13 +++++++------ src/equilibrium/equilibrium_dispatch.cu | 13 +++++++------ .../{equilibrate.h => phase_equilibrate_tp.h} | 9 +++++---- 3 files changed, 19 insertions(+), 16 deletions(-) rename src/equilibrium/{equilibrate.h => phase_equilibrate_tp.h} (94%) diff --git a/src/equilibrium/equilibrium_dispatch.cpp b/src/equilibrium/equilibrium_dispatch.cpp index b82cc4a4..a22227f0 100644 --- a/src/equilibrium/equilibrium_dispatch.cpp +++ b/src/equilibrium/equilibrium_dispatch.cpp @@ -2,7 +2,7 @@ #include #include -#include "equilibrate.h" +#include "phase_equilibrate_tp.h" #include "equilibrium_dispatch.hpp" namespace kintera { @@ -27,11 +27,12 @@ void call_equilibrium_cpu(at::TensorIterator &iter, at::Tensor const &stoich, 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]); - equilibrate(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); + 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); diff --git a/src/equilibrium/equilibrium_dispatch.cu b/src/equilibrium/equilibrium_dispatch.cu index c6231884..c14465f1 100644 --- a/src/equilibrium/equilibrium_dispatch.cu +++ b/src/equilibrium/equilibrium_dispatch.cu @@ -3,7 +3,7 @@ #include -#include "equilibrate.h" +#include "phase_equilibrate_tp.h" #include "equilibrium_dispatch.hpp" namespace kintera { @@ -48,11 +48,12 @@ void call_equilibrium_cuda(at::TensorIterator &iter, at::Tensor const &stoich, 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]); - equilibrate(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); + 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); }); }); } diff --git a/src/equilibrium/equilibrate.h b/src/equilibrium/phase_equilibrate_tp.h similarity index 94% rename from src/equilibrium/equilibrate.h rename to src/equilibrium/phase_equilibrate_tp.h index 49a0def5..544e1e4c 100644 --- a/src/equilibrium/equilibrate.h +++ b/src/equilibrium/phase_equilibrate_tp.h @@ -47,10 +47,11 @@ DISPATCH_MACRO T equilibrium_max_error(T const *moles, T pres, template DISPATCH_MACRO int -equilibrate(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) { +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.; From 6c089b72750ea3bad69fce8350eae073b9ca947c Mon Sep 17 00:00:00 2001 From: mac/cli Date: Thu, 9 Jul 2026 22:48:52 -0400 Subject: [PATCH 5/8] wip --- .github/workflows/release.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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: From 1a426c650f78e8daec4badb84f15bdd0832d4e1e Mon Sep 17 00:00:00 2001 From: Cheng Li Date: Thu, 9 Jul 2026 23:00:57 -0400 Subject: [PATCH 6/8] wip --- python/csrc/kintera.cpp | 2 +- python/csrc/pyequilibrium.cpp | 3 +- src/equilibrium/equilibrium.cpp | 16 +++---- src/equilibrium/equilibrium.hpp | 25 ++++++----- src/equilibrium/equilibrium_dispatch.cpp | 7 +-- src/equilibrium/equilibrium_dispatch.hpp | 2 +- src/equilibrium/equilibrium_options.cpp | 18 ++++---- src/equilibrium/phase_equilibrate_tp.h | 34 ++++++--------- src/species.cpp | 55 ++++++++++-------------- src/species.hpp | 6 +-- src/utils/molar_mass.cpp | 24 +++++------ src/utils/molar_mass.hpp | 11 +++-- tests/test_equilibrium.cpp | 6 +-- 13 files changed, 91 insertions(+), 118 deletions(-) diff --git a/python/csrc/kintera.cpp b/python/csrc/kintera.cpp index 4518f8de..0f85581f 100644 --- a/python/csrc/kintera.cpp +++ b/python/csrc/kintera.cpp @@ -26,7 +26,7 @@ extern std::vector species_cref_R; extern std::vector species_uref_R; extern std::vector species_sref_R; -} // namespace kintera +} // namespace kintera void bind_thermo(py::module &m); void bind_constants(py::module &m); diff --git a/python/csrc/pyequilibrium.cpp b/python/csrc/pyequilibrium.cpp index d53aa295..18e53f57 100644 --- a/python/csrc/pyequilibrium.cpp +++ b/python/csrc/pyequilibrium.cpp @@ -1,10 +1,9 @@ -#include - #include #include #include #include +#include #include "pyoptions.hpp" diff --git a/src/equilibrium/equilibrium.cpp b/src/equilibrium/equilibrium.cpp index 45359e69..68143636 100644 --- a/src/equilibrium/equilibrium.cpp +++ b/src/equilibrium/equilibrium.cpp @@ -1,14 +1,14 @@ +#include "equilibrium.hpp" + #include #include #include - -#include "equilibrium.hpp" -#include "equilibrium_dispatch.hpp" - #include #include +#include "equilibrium_dispatch.hpp" + namespace kintera { void EquilibriumOptionsImpl::report(std::ostream &os) const { @@ -34,9 +34,9 @@ void EquilibriumOptionsImpl::validate() const { 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( + 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"); @@ -157,4 +157,4 @@ EquilibriumTPImpl::forward(torch::Tensor temp, torch::Tensor pres, return {out_moles, gain.view(gain_shape), diag}; } -} // namespace kintera +} // namespace kintera diff --git a/src/equilibrium/equilibrium.hpp b/src/equilibrium/equilibrium.hpp index 81fbc09c..57200d82 100644 --- a/src/equilibrium/equilibrium.hpp +++ b/src/equilibrium/equilibrium.hpp @@ -1,15 +1,14 @@ #pragma once -#include -#include -#include -#include - +#include #include #include #include -#include +#include +#include +#include +#include namespace kintera { @@ -18,8 +17,8 @@ struct EquilibriumOptionsImpl final { return std::make_shared(); } - static std::shared_ptr - from_yaml(std::string const &filename, bool verbose = false); + static std::shared_ptr from_yaml( + std::string const &filename, bool verbose = false); void report(std::ostream &os) const; void validate() const; @@ -37,7 +36,7 @@ struct EquilibriumOptionsImpl final { using EquilibriumOptions = std::shared_ptr; class EquilibriumTPImpl : public torch::nn::Cloneable { -public: + public: EquilibriumOptions options; torch::Tensor stoich; @@ -59,11 +58,11 @@ class EquilibriumTPImpl : public torch::nn::Cloneable { * 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); + 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 +} // namespace kintera diff --git a/src/equilibrium/equilibrium_dispatch.cpp b/src/equilibrium/equilibrium_dispatch.cpp index a22227f0..167a33e1 100644 --- a/src/equilibrium/equilibrium_dispatch.cpp +++ b/src/equilibrium/equilibrium_dispatch.cpp @@ -1,9 +1,10 @@ +#include "equilibrium_dispatch.hpp" + #include #include #include #include "phase_equilibrate_tp.h" -#include "equilibrium_dispatch.hpp" namespace kintera { @@ -39,11 +40,11 @@ void call_equilibrium_cpu(at::TensorIterator &iter, at::Tensor const &stoich, }); } -} // namespace kintera +} // namespace kintera namespace at::native { DEFINE_DISPATCH(call_equilibrium); REGISTER_ALL_CPU_DISPATCH(call_equilibrium, &kintera::call_equilibrium_cpu); -} // namespace at::native +} // namespace at::native diff --git a/src/equilibrium/equilibrium_dispatch.hpp b/src/equilibrium/equilibrium_dispatch.hpp index f04f1473..30423ede 100644 --- a/src/equilibrium/equilibrium_dispatch.hpp +++ b/src/equilibrium/equilibrium_dispatch.hpp @@ -13,4 +13,4 @@ using equilibrium_fn = void (*)(at::TensorIterator &iter, DECLARE_DISPATCH(equilibrium_fn, call_equilibrium); -} // namespace at::native +} // namespace at::native diff --git a/src/equilibrium/equilibrium_options.cpp b/src/equilibrium/equilibrium_options.cpp index 8f614cb6..a2a7bd33 100644 --- a/src/equilibrium/equilibrium_options.cpp +++ b/src/equilibrium/equilibrium_options.cpp @@ -1,20 +1,19 @@ +#include + #include #include #include -#include -#include - -#include - #include #include +#include +#include #include "equilibrium.hpp" namespace kintera { -EquilibriumOptions -EquilibriumOptionsImpl::from_yaml(std::string const &filename, bool verbose) { +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'"); @@ -127,9 +126,8 @@ EquilibriumOptionsImpl::from_yaml(std::string const &filename, bool verbose) { } options->validate(); - if (verbose) - options->report(std::cout); + if (verbose) options->report(std::cout); return options; } -} // namespace kintera +} // namespace kintera diff --git a/src/equilibrium/phase_equilibrate_tp.h b/src/equilibrium/phase_equilibrate_tp.h index 544e1e4c..f2ba672c 100644 --- a/src/equilibrium/phase_equilibrate_tp.h +++ b/src/equilibrium/phase_equilibrate_tp.h @@ -1,14 +1,13 @@ #pragma once -#include -#include -#include - #include - #include #include +#include +#include +#include + namespace kintera { template @@ -18,14 +17,12 @@ DISPATCH_MACRO T equilibrium_max_error(T const *moles, T pres, 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 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 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) { @@ -39,19 +36,17 @@ DISPATCH_MACRO T equilibrium_max_error(T const *moles, T pres, T max_error = 0.; for (int j = 0; j < nreaction; ++j) { T value = fabs(residual[j]); - if (value > max_error) - max_error = value; + 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) { +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.; @@ -150,8 +145,7 @@ phase_equilibrate_tp(T *gain, T *diag, T *out_moles, T temp, T pres, delta += stoich[i * nreaction + j] * step[j]; } trial[i] = out_moles[i] + scale * delta; - if (!(trial[i] > mole_floor)) - positive = false; + if (!(trial[i] > mole_floor)) positive = false; } if (positive) { T trial_error = equilibrium_max_error( @@ -193,4 +187,4 @@ phase_equilibrate_tp(T *gain, T *diag, T *out_moles, T temp, T pres, return status; } -} // namespace kintera +} // namespace kintera diff --git a/src/species.cpp b/src/species.cpp index d2c26615..444a1136 100644 --- a/src/species.cpp +++ b/src/species.cpp @@ -17,12 +17,11 @@ // kintera #include +#include #include #include #include -#include - #include "species.hpp" namespace kintera { @@ -53,12 +52,11 @@ void clear_species_registry() { species_nasa9_Tmid.clear(); } -} // namespace +} // namespace static std::unordered_map &get_nasa9_db() { static std::unordered_map db; - if (!db.empty()) - return db; + if (!db.empty()) return db; std::string path; try { @@ -71,14 +69,12 @@ static std::unordered_map &get_nasa9_db() { std::string line; while (std::getline(ifs, line)) { - if (line.empty() || line[0] == '#') - continue; + 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(); + while (!name.empty() && std::isspace(name.back())) name.pop_back(); // read 4 lines of 5 values = 20 coefficients double vals[20]; @@ -86,25 +82,21 @@ static std::unordered_map &get_nasa9_db() { 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; + while (iss >> v && idx < 20) vals[idx++] = v; } - if (idx < 20) - continue; + 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 + 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 + 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; } @@ -277,22 +269,20 @@ void SpeciesThermoImpl::accumulate(at::Tensor &data, bool SpeciesThermoImpl::has_nasa9() const { for (auto const &coeffs : nasa9_low()) { for (double v : coeffs) { - if (v != 0.0) - return true; + if (v != 0.0) return true; } } for (auto const &coeffs : nasa9_high()) { for (double v : coeffs) { - if (v != 0.0) - return true; + if (v != 0.0) return true; } } return false; } -static at::Tensor -nasa9_coeffs_to_tensor(std::vector> const &coeffs, - c10::TensorOptions const &options) { +static at::Tensor nasa9_coeffs_to_tensor( + std::vector> const &coeffs, + c10::TensorOptions const &options) { auto tensor = torch::empty({static_cast(coeffs.size()), 9}, torch::dtype(torch::kFloat64)); if (!coeffs.empty()) { @@ -316,8 +306,8 @@ at::Tensor SpeciesThermoImpl::nasa9_coeffs_high_tensor( return nasa9_coeffs_to_tensor(nasa9_high(), options); } -at::Tensor -SpeciesThermoImpl::nasa9_Tmid_tensor(c10::TensorOptions const &options) const { +at::Tensor SpeciesThermoImpl::nasa9_Tmid_tensor( + c10::TensorOptions const &options) const { auto tensor = torch::empty({static_cast(nasa9_Tmid().size())}, torch::dtype(torch::kFloat64)); if (!nasa9_Tmid().empty()) { @@ -549,8 +539,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); @@ -571,4 +560,4 @@ SpeciesThermo merge_thermo(SpeciesThermo const &thermo1, return merged; } -} // namespace kintera +} // namespace kintera diff --git a/src/species.hpp b/src/species.hpp index 2aa84c75..00bf8b64 100644 --- a/src/species.hpp +++ b/src/species.hpp @@ -16,11 +16,11 @@ namespace at { class Tensor; -} // namespace at +} // namespace at namespace YAML { class Node; -} // namespace YAML +} // namespace YAML namespace kintera { @@ -108,6 +108,6 @@ extern std::vector species_uref_R; extern std::vector species_sref_R; extern bool species_initialized; -} // namespace kintera +} // namespace kintera #undef ADD_ARG diff --git a/src/utils/molar_mass.cpp b/src/utils/molar_mass.cpp index c88b679b..d096d3b4 100644 --- a/src/utils/molar_mass.cpp +++ b/src/utils/molar_mass.cpp @@ -1,15 +1,12 @@ -#include -#include - -#include +#include "molar_mass.hpp" #include +#include -#include - -#include "molar_mass.hpp" - +#include +#include #include +#include namespace kintera { @@ -21,13 +18,12 @@ 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) { +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 {}; + if (elements.empty()) return {}; auto ncomponent = element_matrix[0].size(); for (auto const &row : element_matrix) { @@ -82,4 +78,4 @@ std::vector molar_masses_from_yaml(std::string const &filename) { return result; } -} // namespace kintera +} // namespace kintera diff --git a/src/utils/molar_mass.hpp b/src/utils/molar_mass.hpp index 58b77671..6531496c 100644 --- a/src/utils/molar_mass.hpp +++ b/src/utils/molar_mass.hpp @@ -1,10 +1,9 @@ #pragma once +#include #include #include -#include - namespace kintera { //! Atomic molar mass [kg/mol]. @@ -14,11 +13,11 @@ double atomic_mass(std::string const &element); 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); +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 +} // namespace kintera diff --git a/tests/test_equilibrium.cpp b/tests/test_equilibrium.cpp index 20e9e795..14937a6f 100644 --- a/tests/test_equilibrium.cpp +++ b/tests/test_equilibrium.cpp @@ -2,7 +2,6 @@ #include #include - #include #include @@ -31,8 +30,7 @@ class EquilibriumDeviceTest : public DeviceTest {}; GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(DeviceTest); TEST_P(EquilibriumDeviceTest, IdealGasReaction) { - if (device.type() == torch::kMPS) - GTEST_SKIP(); + if (device.type() == torch::kMPS) GTEST_SKIP(); auto op = make_options(); EquilibriumTP eq(op); eq->to(device, dtype); @@ -127,4 +125,4 @@ INSTANTIATE_TEST_SUITE_P( return name; }); -} // namespace +} // namespace From 48c6a8d3c08c92b4cbe769ac819b2bd8962f5e94 Mon Sep 17 00:00:00 2001 From: mac/cli Date: Thu, 9 Jul 2026 23:41:40 -0400 Subject: [PATCH 7/8] wip --- src/species.cpp | 132 +++++++++++++++++++------------------ tests/test_equilibrium.cpp | 23 ++++++- 2 files changed, 89 insertions(+), 66 deletions(-) diff --git a/src/species.cpp b/src/species.cpp index 444a1136..d29cd543 100644 --- a/src/species.cpp +++ b/src/species.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -52,61 +53,62 @@ void clear_species_registry() { species_nasa9_Tmid.clear(); } -} // namespace +} // namespace 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; @@ -119,18 +121,17 @@ at::Tensor nasa9_gibbs_rt(at::Tensor temp, high.push_back(found->second.high); } - auto tensor_options = temp.options(); - auto low_tensor = - torch::empty({static_cast(species.size()), 9}, tensor_options); - auto high_tensor = torch::empty_like(low_tensor); - for (size_t i = 0; i < species.size(); ++i) { - for (int j = 0; j < 9; ++j) { - low_tensor.index_put_({static_cast(i), j}, low[i][j]); - high_tensor.index_put_({static_cast(i), j}, high[i][j]); - } - } + 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., - tensor_options); + temp.options()); return nasa9_gibbs_RT(temp, low_tensor, high_tensor, midpoint); } @@ -269,20 +270,22 @@ void SpeciesThermoImpl::accumulate(at::Tensor &data, bool SpeciesThermoImpl::has_nasa9() const { for (auto const &coeffs : nasa9_low()) { for (double v : coeffs) { - if (v != 0.0) return true; + if (v != 0.0) + return true; } } for (auto const &coeffs : nasa9_high()) { for (double v : coeffs) { - if (v != 0.0) return true; + if (v != 0.0) + return true; } } return false; } -static at::Tensor nasa9_coeffs_to_tensor( - std::vector> const &coeffs, - c10::TensorOptions const &options) { +static at::Tensor +nasa9_coeffs_to_tensor(std::vector> const &coeffs, + c10::TensorOptions const &options) { auto tensor = torch::empty({static_cast(coeffs.size()), 9}, torch::dtype(torch::kFloat64)); if (!coeffs.empty()) { @@ -306,8 +309,8 @@ at::Tensor SpeciesThermoImpl::nasa9_coeffs_high_tensor( return nasa9_coeffs_to_tensor(nasa9_high(), options); } -at::Tensor SpeciesThermoImpl::nasa9_Tmid_tensor( - c10::TensorOptions const &options) const { +at::Tensor +SpeciesThermoImpl::nasa9_Tmid_tensor(c10::TensorOptions const &options) const { auto tensor = torch::empty({static_cast(nasa9_Tmid().size())}, torch::dtype(torch::kFloat64)); if (!nasa9_Tmid().empty()) { @@ -539,7 +542,8 @@ 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); @@ -560,4 +564,4 @@ SpeciesThermo merge_thermo(SpeciesThermo const &thermo1, return merged; } -} // namespace kintera +} // namespace kintera diff --git a/tests/test_equilibrium.cpp b/tests/test_equilibrium.cpp index 14937a6f..47eb97b6 100644 --- a/tests/test_equilibrium.cpp +++ b/tests/test_equilibrium.cpp @@ -2,7 +2,9 @@ #include #include +#include #include +#include #include #define DEVICE_TESTING_SKIP_DEFAULT_INSTANTIATION @@ -30,7 +32,8 @@ class EquilibriumDeviceTest : public DeviceTest {}; GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(DeviceTest); TEST_P(EquilibriumDeviceTest, IdealGasReaction) { - if (device.type() == torch::kMPS) GTEST_SKIP(); + if (device.type() == torch::kMPS) + GTEST_SKIP(); auto op = make_options(); EquilibriumTP eq(op); eq->to(device, dtype); @@ -59,6 +62,22 @@ TEST(MolarMass, StandaloneUtilities) { 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"; @@ -125,4 +144,4 @@ INSTANTIATE_TEST_SUITE_P( return name; }); -} // namespace +} // namespace From bde0addf5b84da50a4702b80d3d2264e51234a94 Mon Sep 17 00:00:00 2001 From: Cheng Li Date: Thu, 9 Jul 2026 23:43:41 -0400 Subject: [PATCH 8/8] wip --- src/species.cpp | 41 +++++++++++++++----------------------- tests/test_equilibrium.cpp | 5 ++--- 2 files changed, 18 insertions(+), 28 deletions(-) diff --git a/src/species.cpp b/src/species.cpp index d29cd543..62415768 100644 --- a/src/species.cpp +++ b/src/species.cpp @@ -53,7 +53,7 @@ void clear_species_registry() { species_nasa9_Tmid.clear(); } -} // namespace +} // namespace static std::unordered_map &get_nasa9_db() { static std::unordered_map db; @@ -70,32 +70,26 @@ static std::unordered_map &get_nasa9_db() { std::string line; while (std::getline(ifs, line)) { - if (line.empty() || line[0] == '#') - continue; + 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(); + 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; + while (iss >> value && idx < 20) vals[idx++] = value; } - if (idx < 20) - continue; + if (idx < 20) continue; Nasa9Entry entry; - for (int k = 0; k < 7; ++k) - entry.low[k] = vals[k]; + 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]; + 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; @@ -270,22 +264,20 @@ void SpeciesThermoImpl::accumulate(at::Tensor &data, bool SpeciesThermoImpl::has_nasa9() const { for (auto const &coeffs : nasa9_low()) { for (double v : coeffs) { - if (v != 0.0) - return true; + if (v != 0.0) return true; } } for (auto const &coeffs : nasa9_high()) { for (double v : coeffs) { - if (v != 0.0) - return true; + if (v != 0.0) return true; } } return false; } -static at::Tensor -nasa9_coeffs_to_tensor(std::vector> const &coeffs, - c10::TensorOptions const &options) { +static at::Tensor nasa9_coeffs_to_tensor( + std::vector> const &coeffs, + c10::TensorOptions const &options) { auto tensor = torch::empty({static_cast(coeffs.size()), 9}, torch::dtype(torch::kFloat64)); if (!coeffs.empty()) { @@ -309,8 +301,8 @@ at::Tensor SpeciesThermoImpl::nasa9_coeffs_high_tensor( return nasa9_coeffs_to_tensor(nasa9_high(), options); } -at::Tensor -SpeciesThermoImpl::nasa9_Tmid_tensor(c10::TensorOptions const &options) const { +at::Tensor SpeciesThermoImpl::nasa9_Tmid_tensor( + c10::TensorOptions const &options) const { auto tensor = torch::empty({static_cast(nasa9_Tmid().size())}, torch::dtype(torch::kFloat64)); if (!nasa9_Tmid().empty()) { @@ -542,8 +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); @@ -564,4 +555,4 @@ SpeciesThermo merge_thermo(SpeciesThermo const &thermo1, return merged; } -} // namespace kintera +} // namespace kintera diff --git a/tests/test_equilibrium.cpp b/tests/test_equilibrium.cpp index 47eb97b6..d32e6618 100644 --- a/tests/test_equilibrium.cpp +++ b/tests/test_equilibrium.cpp @@ -32,8 +32,7 @@ class EquilibriumDeviceTest : public DeviceTest {}; GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(DeviceTest); TEST_P(EquilibriumDeviceTest, IdealGasReaction) { - if (device.type() == torch::kMPS) - GTEST_SKIP(); + if (device.type() == torch::kMPS) GTEST_SKIP(); auto op = make_options(); EquilibriumTP eq(op); eq->to(device, dtype); @@ -144,4 +143,4 @@ INSTANTIATE_TEST_SUITE_P( return name; }); -} // namespace +} // namespace