From 828bc4554ab98db95e732e54f9e29eb0b0d13255 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 15:16:20 -0400 Subject: [PATCH 1/4] Seed flaky hetero rectangular gradient test test_rectangular_model_forward_and_gradients asserts a nonzero relation gradient, which about 7.5% of random initializations (15 of 200 seeds) fail through dead ReLU paths on the tiny graph. It inherited RNG state from whatever ran earlier in its xdist worker, so adding unrelated tests could flip it. Pin torch.manual_seed(0). Co-Authored-By: Claude Opus 5.5 --- tests/operators/test_hetero_rectangular.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/operators/test_hetero_rectangular.py b/tests/operators/test_hetero_rectangular.py index 917091b..60d85dd 100644 --- a/tests/operators/test_hetero_rectangular.py +++ b/tests/operators/test_hetero_rectangular.py @@ -477,6 +477,10 @@ def test_relgraph_encode_decode_shapes_match_per_type_widths() -> None: def test_rectangular_model_forward_and_gradients() -> None: """Factory model encode→advance→decode runs with gradients on unequal d_τ.""" + # About 7% of random inits give an all-zero relation gradient (dead ReLU + # paths on this tiny graph), so pin the init instead of inheriting + # whatever RNG state earlier tests in the xdist worker left behind. + torch.manual_seed(0) encoder, decoder = _relgraph_peers() model = GraphKoopmanModel( encoder=encoder, From befd033d2e43ecd2f82f753d384368008d88b6d1 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 15:16:20 -0400 Subject: [PATCH 2/4] Add reaction graph encodings and integrators Refs #65. Fills in koopman_graph.datasets.reaction.graphs and .integrate: - graphs: species_influence_graph (Jacobian-sparsity species graph with linking-reaction counts), reaction_hyperedge_index (library hyperedge layout), reaction_directed_hyperedges (reactant tails / product heads), species_reaction_bipartite (HeteroData with stoichiometric weights). - integrate: integrate_expm (exact augmented matrix exponential for linear systems, including zeroth-order feeds), linear_rate_equations, integrate_rk4 (batched, differentiable), integrate_stiff (SciPy BDF/Radau/LSODA with an exact Jacobian, behind require_scipy and the [reaction] extra), integrate_reaction_system (method="auto"), and mass_action_jacobian. SciPy is imported only inside the stiff path, so import hygiene holds. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 13 + .../datasets/reaction/__init__.py | 36 ++ src/koopman_graph/datasets/reaction/graphs.py | 200 ++++++++ .../datasets/reaction/integrate.py | 468 ++++++++++++++++++ tests/datasets/test_reaction_graphs.py | 129 +++++ tests/datasets/test_reaction_integrate.py | 269 ++++++++++ 6 files changed, 1115 insertions(+) create mode 100644 tests/datasets/test_reaction_graphs.py create mode 100644 tests/datasets/test_reaction_integrate.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 285fcb7..c9cfb3b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (floored `log c`, its inverse, a below-floor mask, and a serializable record of the floor) and `element_balance_violation` for element-total drift from a reference state. +- `koopman_graph.datasets.reaction.graphs`: `species_influence_graph` + (directed species graph matching the mass-action Jacobian's sparsity, + weighted by linking-reaction count), `reaction_hyperedge_index` (one + hyperedge per reaction), `reaction_directed_hyperedges` (reactant tail / + product head incidences), and `species_reaction_bipartite` (PyG + `HeteroData` with stoichiometric edge weights). +- `koopman_graph.datasets.reaction.integrate`: exact augmented + matrix-exponential propagation for linear systems (`integrate_expm`, + `linear_rate_equations`), batched and differentiable torch RK4 + (`integrate_rk4`), SciPy BDF / Radau / LSODA with an exact mass-action + Jacobian behind a lazy `require_scipy()` guard (`integrate_stiff`; + `[reaction]` extra), the `integrate_reaction_system` dispatcher + (`method="auto"`), and `mass_action_jacobian`. - `[reaction]` extra (`scipy>=1.11.0`) for stiff reaction-network solvers and `[cantera]` extra (`cantera>=3.2,<3.3`) for Cantera mechanisms. - `cantera` pytest marker; marked tests are skipped automatically when the diff --git a/src/koopman_graph/datasets/reaction/__init__.py b/src/koopman_graph/datasets/reaction/__init__.py index c09e16d..8aa8629 100644 --- a/src/koopman_graph/datasets/reaction/__init__.py +++ b/src/koopman_graph/datasets/reaction/__init__.py @@ -38,6 +38,26 @@ replacement for a production kinetics solver. """ +from koopman_graph.datasets.reaction.graphs import ( + PRODUCT_RELATION, + REACTANT_RELATION, + REACTION_NODE, + SPECIES_NODE, + reaction_directed_hyperedges, + reaction_hyperedge_index, + species_influence_graph, + species_reaction_bipartite, +) +from koopman_graph.datasets.reaction.integrate import ( + IntegrationMethod, + integrate_expm, + integrate_reaction_system, + integrate_rk4, + integrate_stiff, + linear_rate_equations, + mass_action_jacobian, + require_scipy, +) from koopman_graph.datasets.reaction.system import ( GAS_CONSTANT_J_PER_MOL_K, ArrheniusParams, @@ -54,10 +74,26 @@ __all__ = [ "DEFAULT_LOG_FLOOR", "GAS_CONSTANT_J_PER_MOL_K", + "PRODUCT_RELATION", + "REACTANT_RELATION", + "REACTION_NODE", + "SPECIES_NODE", "ArrheniusParams", "ElementBalanceReport", + "IntegrationMethod", "LogConcentrationTransform", "ReactionSystem", "check_element_balance", "element_balance_violation", + "integrate_expm", + "integrate_reaction_system", + "integrate_rk4", + "integrate_stiff", + "linear_rate_equations", + "mass_action_jacobian", + "reaction_directed_hyperedges", + "reaction_hyperedge_index", + "require_scipy", + "species_influence_graph", + "species_reaction_bipartite", ] diff --git a/src/koopman_graph/datasets/reaction/graphs.py b/src/koopman_graph/datasets/reaction/graphs.py index e69de29..ed277c7 100644 --- a/src/koopman_graph/datasets/reaction/graphs.py +++ b/src/koopman_graph/datasets/reaction/graphs.py @@ -0,0 +1,200 @@ +"""Graph encodings of a :class:`~koopman_graph.datasets.reaction.ReactionSystem`. + +Everything here is derived from stoichiometry alone: + +* :func:`species_influence_graph` builds a directed species graph with an + edge ``s → t`` whenever species ``s`` is a reactant of a reaction that + changes ``t``. This is the sparsity pattern of the mass-action Jacobian + ``∂(dc_t/dt)/∂c_s``. +* :func:`reaction_hyperedge_index` builds an undirected hypergraph with one + hyperedge per reaction, using the library's ``hyperedge_index`` layout + ``(2, nnz)`` (row 0 = species, row 1 = reaction). +* :func:`reaction_directed_hyperedges` returns directed-hypergraph + ``tail_index`` (reactants) and ``head_index`` (products) in the same + layout. +* :func:`species_reaction_bipartite` builds a PyG + :class:`~torch_geometric.data.HeteroData` graph with ``species`` and + ``reaction`` nodes and stoichiometric coefficients as edge weights. + +Pure PyTorch / PyG; importing this module does not import scipy or cantera. +""" + +from __future__ import annotations + +import torch +from torch import Tensor +from torch_geometric.data import HeteroData + +from koopman_graph.datasets.reaction.system import ReactionSystem + +SPECIES_NODE = "species" +"""Node type name for species in :func:`species_reaction_bipartite`.""" + +REACTION_NODE = "reaction" +"""Node type name for reactions in :func:`species_reaction_bipartite`.""" + +REACTANT_RELATION = (SPECIES_NODE, "reactant_of", REACTION_NODE) +"""Edge type from a reactant species to its reaction (weight = ``α``).""" + +PRODUCT_RELATION = (REACTION_NODE, "produces", SPECIES_NODE) +"""Edge type from a reaction to a product species (weight = ``β``).""" + + +def species_influence_graph( + system: ReactionSystem, + *, + self_loops: bool = False, +) -> tuple[Tensor, Tensor]: + """Return the directed species influence graph and its edge weights. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + self_loops : bool, optional + When ``True``, keep ``s → s`` edges (a species that is consumed or + produced by a reaction it drives). Default is ``False``. + + Returns + ------- + edge_index : Tensor + ``long`` tensor with shape ``(2, num_edges)``; row 0 = source species + ``s``, row 1 = target species ``t``. Edges are unique and sorted by + ``(source, target)``. + edge_weight : Tensor + ``float32`` tensor with shape ``(num_edges,)``: the number of + reactions that link ``s`` to ``t``. + + Notes + ----- + An edge exists when some reaction ``j`` has ``α_js > 0`` and + ``N_tj ≠ 0``. A zeroth-order source (``∅ → A``) has no reactants and so + adds no edges. + """ + _require_system(system) + drives = (system.reactant_stoichiometry > 0).to(torch.float64) # (R, S) + changes = (system.stoichiometric_matrix != 0).to(torch.float64) # (S, R) + counts = drives.T @ changes.T # (S_source, S_target) + if not self_loops: + counts.fill_diagonal_(0.0) + source, target = counts.nonzero(as_tuple=True) + edge_index = torch.stack([source, target]).to(dtype=torch.long) + edge_weight = counts[source, target].to(dtype=torch.float32) + return edge_index, edge_weight + + +def reaction_hyperedge_index(system: ReactionSystem) -> Tensor: + """Return an undirected hyperedge per reaction over its participating species. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + + Returns + ------- + Tensor + ``long`` incidence with shape ``(2, nnz)``: row 0 = species index, + row 1 = reaction (hyperedge) index, sorted by ``(reaction, + species)``. A species is incident to reaction ``j`` when it is a + reactant or a product of ``j`` (also when it is a catalyst that + appears on both sides). + """ + _require_system(system) + participates = (system.reactant_stoichiometry > 0) | ( + system.product_stoichiometry > 0 + ) + return _incidence(participates) + + +def reaction_directed_hyperedges(system: ReactionSystem) -> tuple[Tensor, Tensor]: + """Return directed-hypergraph tail (reactant) and head (product) incidences. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + + Returns + ------- + tail_index : Tensor + ``long`` incidence ``(2, nnz_tail)`` from reactant species (row 0) + to reaction (row 1). + head_index : Tensor + ``long`` incidence ``(2, nnz_head)`` from product species (row 0) + to reaction (row 1). + + Notes + ----- + Source reactions have no tail entries and sink reactions no head + entries, so the hyperedge count is ``system.num_reactions``, not + ``max(index[1]) + 1``. Pass it explicitly to consumers that infer it. + """ + _require_system(system) + tail_index = _incidence(system.reactant_stoichiometry > 0) + head_index = _incidence(system.product_stoichiometry > 0) + return tail_index, head_index + + +def species_reaction_bipartite(system: ReactionSystem) -> HeteroData: + """Return the bipartite species-reaction graph as PyG ``HeteroData``. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + + Returns + ------- + HeteroData + Node types :data:`SPECIES_NODE` (``num_species`` nodes) and + :data:`REACTION_NODE` (``num_reactions`` nodes). Edge type + :data:`REACTANT_RELATION` carries ``edge_index`` (species → reaction) + and ``edge_weight`` = reactant coefficient ``α``; + :data:`PRODUCT_RELATION` carries reaction → species edges with + ``edge_weight`` = product coefficient ``β``. The ``reaction`` nodes + also store ``rate_constant`` with shape ``(num_reactions,)``. + """ + _require_system(system) + data = HeteroData() + data[SPECIES_NODE].num_nodes = system.num_species + data[REACTION_NODE].num_nodes = system.num_reactions + data[REACTION_NODE].rate_constant = system.rate_constants.to(torch.float32) + + alpha = system.reactant_stoichiometry + reaction, species = (alpha > 0).nonzero(as_tuple=True) + data[REACTANT_RELATION].edge_index = torch.stack([species, reaction]) + data[REACTANT_RELATION].edge_weight = alpha[reaction, species].to(torch.float32) + + beta = system.product_stoichiometry + reaction, species = (beta > 0).nonzero(as_tuple=True) + data[PRODUCT_RELATION].edge_index = torch.stack([reaction, species]) + data[PRODUCT_RELATION].edge_weight = beta[reaction, species].to(torch.float32) + return data + + +def _incidence(mask: Tensor) -> Tensor: + """Turn a ``(num_reactions, num_species)`` mask into ``(2, nnz)`` incidence. + + Parameters + ---------- + mask + See signature. + + Returns + ------- + See signature.""" + reaction, species = mask.nonzero(as_tuple=True) + return torch.stack([species, reaction]).to(dtype=torch.long) + + +def _require_system(system: object) -> None: + """Raise ``TypeError`` unless ``system`` is a :class:`ReactionSystem`. + + Parameters + ---------- + system + See signature.""" + if not isinstance(system, ReactionSystem): + msg = f"system must be a ReactionSystem, got {type(system).__name__}" + raise TypeError(msg) diff --git a/src/koopman_graph/datasets/reaction/integrate.py b/src/koopman_graph/datasets/reaction/integrate.py index e69de29..a7c04ff 100644 --- a/src/koopman_graph/datasets/reaction/integrate.py +++ b/src/koopman_graph/datasets/reaction/integrate.py @@ -0,0 +1,468 @@ +"""Time integration of mass-action reaction systems. + +Three paths, picked by :func:`integrate_reaction_system` (``method="auto"`` +uses the exact path for linear systems and RK4 otherwise): + +``"expm"`` + Exact propagation for linear systems (every reaction of order at most + one). The affine field ``dc/dt = A c + b`` is folded into one augmented + matrix exponential, so zeroth-order sources are exact too. +``"rk4"`` + Classical fixed-step fourth-order Runge-Kutta in PyTorch for non-stiff + systems. It supports batched initial conditions and autograd. +``"bdf"`` / ``"radau"`` / ``"lsoda"`` + SciPy :func:`scipy.integrate.solve_ivp` stiff solvers with an exact + mass-action Jacobian. They need the ``[reaction]`` extra; SciPy is + imported only when one of these methods runs. + +All paths take output ``times`` (strictly increasing, with ``times[0]`` the +time of ``initial``) and return concentrations of shape +``(len(times), ..., num_species)`` whose first entry is ``initial``. +Importing this module does not import scipy or cantera. +""" + +from __future__ import annotations + +from typing import Any, Literal + +import torch +from torch import Tensor + +from koopman_graph.datasets.reaction.system import ReactionSystem + +IntegrationMethod = Literal["auto", "expm", "rk4", "bdf", "radau", "lsoda"] +"""Accepted ``method`` names for :func:`integrate_reaction_system`.""" + +_STIFF_METHODS = {"bdf": "BDF", "radau": "Radau", "lsoda": "LSODA"} + +_SCIPY_IMPORT_ERROR = ( + "scipy is required for stiff reaction integration (method='bdf', " + "'radau', or 'lsoda'). Install with: pip install koopman-graph[reaction]" +) + + +def require_scipy() -> Any: + """Import and return :mod:`scipy.integrate`, or explain how to install it. + + Returns + ------- + module + The :mod:`scipy.integrate` module. + + Raises + ------ + ImportError + If SciPy is not installed; the message names the ``[reaction]`` + extra. + """ + try: + import scipy.integrate as scipy_integrate + except ImportError as exc: + raise ImportError(_SCIPY_IMPORT_ERROR) from exc + return scipy_integrate + + +def linear_rate_equations(system: ReactionSystem) -> tuple[Tensor, Tensor]: + """Return ``(A, b)`` with ``dc/dt = A c + b`` for a linear system. + + Parameters + ---------- + system : ReactionSystem + System with :attr:`~ReactionSystem.is_linear` true. + + Returns + ------- + A : Tensor + ``float64`` rate matrix with shape ``(num_species, num_species)``. + b : Tensor + ``float64`` constant source term with shape ``(num_species,)``. + + Raises + ------ + ValueError + If any reaction has total order greater than one. + """ + _require_linear(system) + n = system.stoichiometric_matrix # (S, R) + # For order-0 / order-1 reactions r_j = k_j α_j · c + k_j [order_j == 0]. + weighted_alpha = system.rate_constants.unsqueeze(-1) * ( + system.reactant_stoichiometry + ) # (R, S) + a = n @ weighted_alpha + zeroth = (system.reaction_orders == 0).to(torch.float64) + b = n @ (system.rate_constants * zeroth) + return a, b + + +def integrate_expm( + system: ReactionSystem, + initial: Tensor, + times: Tensor, +) -> Tensor: + """Propagate a linear system exactly with an augmented matrix exponential. + + Parameters + ---------- + system : ReactionSystem + System with :attr:`~ReactionSystem.is_linear` true. + initial : Tensor + Floating initial concentrations with shape ``(..., num_species)``. + times : Tensor + Strictly increasing 1-D output times; ``times[0]`` is the time of + ``initial``. + + Returns + ------- + Tensor + Concentrations ``(len(times), ..., num_species)`` in ``float64``. + + Raises + ------ + ValueError + If the system is not linear or inputs are malformed. + """ + a, b = linear_rate_equations(system) + c0 = _validate_initial(system, initial).to(torch.float64) + t = _validate_times(times) + num_species = system.num_species + generator = torch.zeros(num_species + 1, num_species + 1, dtype=torch.float64) + generator[:num_species, :num_species] = a + generator[:num_species, num_species] = b + elapsed = (t - t[0]).to(dtype=torch.float64, device="cpu") + propagators = torch.linalg.matrix_exp(elapsed[:, None, None] * generator) + augmented = torch.cat([c0, torch.ones_like(c0[..., :1])], dim=-1) + # result[t, ..., i] = Σ_j P[t, i, j] y[..., j] + states = torch.einsum("tij,...j->t...i", propagators, augmented) + return states[..., :num_species] + + +def integrate_rk4( + system: ReactionSystem, + initial: Tensor, + times: Tensor, + *, + substeps: int = 10, +) -> Tensor: + """Integrate with classical fixed-step RK4 (non-stiff systems). + + Parameters + ---------- + system : ReactionSystem + Reaction network. + initial : Tensor + Floating initial concentrations with shape ``(..., num_species)``. + The output keeps its dtype and device; gradients flow through. + times : Tensor + Strictly increasing 1-D output times; ``times[0]`` is the time of + ``initial``. + substeps : int, optional + Equal RK4 steps per output interval. Default is ``10``. + + Returns + ------- + Tensor + Concentrations ``(len(times), ..., num_species)``. + + Raises + ------ + ValueError + If ``substeps < 1`` or inputs are malformed. + """ + if isinstance(substeps, bool) or not isinstance(substeps, int) or substeps < 1: + msg = f"substeps must be a positive int, got {substeps!r}" + raise ValueError(msg) + c = _validate_initial(system, initial) + t = _validate_times(times).to(dtype=c.dtype, device=c.device) + states = [c] + for index in range(1, int(t.shape[0])): + h = (t[index] - t[index - 1]) / substeps + for _ in range(substeps): + k1 = system.time_derivative(c) + k2 = system.time_derivative(c + 0.5 * h * k1) + k3 = system.time_derivative(c + 0.5 * h * k2) + k4 = system.time_derivative(c + h * k3) + c = c + (h / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + states.append(c) + return torch.stack(states) + + +def integrate_stiff( + system: ReactionSystem, + initial: Tensor, + times: Tensor, + *, + method: Literal["bdf", "radau", "lsoda"] = "bdf", + rtol: float = 1e-8, + atol: float = 1e-12, +) -> Tensor: + """Integrate a stiff system with SciPy ``solve_ivp`` (``[reaction]`` extra). + + Parameters + ---------- + system : ReactionSystem + Reaction network. + initial : Tensor + Floating initial concentrations with shape ``(num_species,)``. + times : Tensor + Strictly increasing 1-D output times; ``times[0]`` is the time of + ``initial``. + method : {"bdf", "radau", "lsoda"}, optional + SciPy implicit method. Default is ``"bdf"``. + rtol : float, optional + Relative tolerance. Default is ``1e-8``. + atol : float, optional + Absolute tolerance. Default is ``1e-12``. + + Returns + ------- + Tensor + ``float64`` concentrations with shape ``(len(times), num_species)``. + + Raises + ------ + ImportError + If SciPy is not installed. + ValueError + If ``method`` is unknown, ``initial`` is batched, or inputs are + malformed. + RuntimeError + If the solver reports failure. + """ + if method not in _STIFF_METHODS: + msg = f"method must be one of {sorted(_STIFF_METHODS)}, got {method!r}" + raise ValueError(msg) + c0 = _validate_initial(system, initial) + if c0.ndim != 1: + msg = ( + "stiff integration takes a single initial state of shape " + f"({system.num_species},), got {tuple(c0.shape)}" + ) + raise ValueError(msg) + t = _validate_times(times).to(dtype=torch.float64, device="cpu") + scipy_integrate = require_scipy() + + def rhs(_time: float, y: Any) -> Any: + """Evaluate ``dc/dt`` on a NumPy state for ``solve_ivp``. + + Parameters + ---------- + _time + See signature. + y + See signature. + + Returns + ------- + See signature.""" + state = torch.from_numpy(y) + return system.time_derivative(state).numpy() + + def jac(_time: float, y: Any) -> Any: + """Evaluate the mass-action Jacobian on a NumPy state for ``solve_ivp``. + + Parameters + ---------- + _time + See signature. + y + See signature. + + Returns + ------- + See signature.""" + return mass_action_jacobian(system, torch.from_numpy(y)).numpy() + + t_values = t.tolist() + solution = scipy_integrate.solve_ivp( + rhs, + (t_values[0], t_values[-1]), + c0.detach().to(dtype=torch.float64, device="cpu").numpy(), + method=_STIFF_METHODS[method], + t_eval=t_values, + jac=jac, + rtol=float(rtol), + atol=float(atol), + ) + if not solution.success: + msg = f"solve_ivp ({_STIFF_METHODS[method]}) failed: {solution.message}" + raise RuntimeError(msg) + return torch.from_numpy(solution.y.T.copy()) + + +def integrate_reaction_system( + system: ReactionSystem, + initial: Tensor, + times: Tensor, + *, + method: IntegrationMethod = "auto", + **kwargs: Any, +) -> Tensor: + """Integrate a reaction system with the requested (or automatic) method. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + initial : Tensor + Floating initial concentrations with shape ``(..., num_species)`` + (unbatched for stiff methods). + times : Tensor + Strictly increasing 1-D output times; ``times[0]`` is the time of + ``initial``. + method : {"auto", "expm", "rk4", "bdf", "radau", "lsoda"}, optional + ``"auto"`` (default) uses ``"expm"`` when the system is linear and + ``"rk4"`` otherwise. Stiff systems should ask for a SciPy method + explicitly. + **kwargs + Forwarded to :func:`integrate_rk4` (``substeps``) or + :func:`integrate_stiff` (``rtol``, ``atol``). + + Returns + ------- + Tensor + Concentrations ``(len(times), ..., num_species)``. + + Raises + ------ + ValueError + If ``method`` is unknown or ``kwargs`` do not apply to it. + """ + if method == "auto": + method = "expm" if system.is_linear else "rk4" + if method == "expm": + if kwargs: + msg = f"method='expm' takes no options, got {sorted(kwargs)}" + raise ValueError(msg) + return integrate_expm(system, initial, times) + if method == "rk4": + return integrate_rk4(system, initial, times, **kwargs) + if method in _STIFF_METHODS: + return integrate_stiff(system, initial, times, method=method, **kwargs) + msg = ( + "method must be one of ['auto', 'expm', 'rk4', 'bdf', 'radau', " + f"'lsoda'], got {method!r}" + ) + raise ValueError(msg) + + +def mass_action_jacobian(system: ReactionSystem, concentrations: Tensor) -> Tensor: + """Return the exact Jacobian ``∂(dc/dt)/∂c`` of the mass-action field. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + concentrations : Tensor + Floating state with shape ``(num_species,)``. + + Returns + ------- + Tensor + Jacobian ``J[t, s] = Σ_j N_tj ∂r_j/∂c_s`` with shape + ``(num_species, num_species)`` in the dtype of ``concentrations``. + Finite at ``c = 0``. + + Raises + ------ + ValueError + If ``concentrations`` is not a single state. + """ + c = _validate_initial(system, concentrations) + if c.ndim != 1: + msg = ( + f"concentrations must have shape ({system.num_species},), got " + f"{tuple(c.shape)}" + ) + raise ValueError(msg) + alpha = system.reactant_stoichiometry.to(dtype=c.dtype, device=c.device) + k = system.rate_constants.to(dtype=c.dtype, device=c.device) + num_species = system.num_species + # ∂r_j/∂c_s = k_j α_js c_s^(α_js - 1) ∏_{s'≠s} c_s'^α_js'; building the + # product with exponent α - e_s keeps c = 0 finite (no division by c). + eye = torch.eye(num_species, dtype=c.dtype, device=c.device) + shifted = alpha.unsqueeze(1) - eye.unsqueeze(0) # (R, S_s, S_s') + powered = torch.where( + shifted < 0, + torch.ones((), dtype=c.dtype, device=c.device), + torch.pow(c, shifted.clamp(min=0)), + ) + d_rates = k[:, None] * alpha * powered.prod(dim=-1) # (R, S) + n = system.stoichiometric_matrix.to(dtype=c.dtype, device=c.device) + return n @ d_rates + + +def _require_linear(system: ReactionSystem) -> None: + """Raise ``ValueError`` unless every reaction has order at most one. + + Parameters + ---------- + system + See signature.""" + if not isinstance(system, ReactionSystem): + msg = f"system must be a ReactionSystem, got {type(system).__name__}" + raise TypeError(msg) + if not system.is_linear: + orders = system.reaction_orders.tolist() + msg = ( + f"{system.name!r} is not linear (reaction orders {orders}); " + "use method='rk4' or a stiff SciPy method" + ) + raise ValueError(msg) + + +def _validate_initial(system: ReactionSystem, initial: Tensor) -> Tensor: + """Validate an initial state ``(..., num_species)`` of finite floats. + + Parameters + ---------- + system + See signature. + initial + See signature. + + Returns + ------- + See signature.""" + if not isinstance(system, ReactionSystem): + msg = f"system must be a ReactionSystem, got {type(system).__name__}" + raise TypeError(msg) + if not isinstance(initial, Tensor) or not initial.is_floating_point(): + msg = "initial concentrations must be a floating-point Tensor" + raise TypeError(msg) + if initial.ndim == 0 or int(initial.shape[-1]) != system.num_species: + msg = ( + f"initial concentrations must have shape (..., {system.num_species}), " + f"got {tuple(initial.shape)}" + ) + raise ValueError(msg) + if not bool(torch.isfinite(initial).all()): + msg = "initial concentrations must be finite" + raise ValueError(msg) + return initial + + +def _validate_times(times: Tensor) -> Tensor: + """Validate strictly increasing, finite 1-D output times. + + Parameters + ---------- + times + See signature. + + Returns + ------- + See signature.""" + if not isinstance(times, Tensor): + msg = f"times must be a Tensor, got {type(times).__name__}" + raise TypeError(msg) + t = times if times.is_floating_point() else times.to(torch.float64) + if t.ndim != 1 or int(t.shape[0]) < 1: + msg = f"times must be a non-empty 1-D Tensor, got {tuple(times.shape)}" + raise ValueError(msg) + if not bool(torch.isfinite(t).all()): + msg = "times must be finite" + raise ValueError(msg) + if int(t.shape[0]) > 1 and not bool((t[1:] > t[:-1]).all()): + msg = "times must be strictly increasing" + raise ValueError(msg) + return t diff --git a/tests/datasets/test_reaction_graphs.py b/tests/datasets/test_reaction_graphs.py new file mode 100644 index 0000000..f1a20ac --- /dev/null +++ b/tests/datasets/test_reaction_graphs.py @@ -0,0 +1,129 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.graphs`.""" + +from __future__ import annotations + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + PRODUCT_RELATION, + REACTANT_RELATION, + REACTION_NODE, + SPECIES_NODE, + ReactionSystem, + mass_action_jacobian, + reaction_directed_hyperedges, + reaction_hyperedge_index, + species_influence_graph, + species_reaction_bipartite, +) + + +def _michaelis_menten() -> ReactionSystem: + """``E + S <-> ES -> E + P`` plus a feed ``∅ -> S`` (species E, S, ES, P).""" + return ReactionSystem( + species=("E", "S", "ES", "P"), + reactant_stoichiometry=torch.tensor( + [ + [1.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ] + ), + product_stoichiometry=torch.tensor( + [ + [0.0, 0.0, 1.0, 0.0], + [1.0, 1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 0.0], + ] + ), + rate_constants=torch.tensor([2.0, 1.0, 0.5, 0.1]), + ) + + +def test_influence_graph_matches_jacobian_sparsity() -> None: + system = _michaelis_menten() + edge_index, edge_weight = species_influence_graph(system, self_loops=True) + c = torch.tensor([0.3, 0.7, 0.2, 0.1], dtype=torch.float64) + jac = mass_action_jacobian(system, c) + pattern = torch.zeros(4, 4, dtype=torch.bool) + pattern[edge_index[1], edge_index[0]] = True # J[target, source] + assert torch.equal(pattern, jac != 0) + assert edge_index.dtype == torch.long + assert edge_weight.dtype == torch.float32 + assert edge_index.shape[1] == edge_weight.shape[0] + # Edges are unique and sorted by (source, target). + keys = (edge_index[0] * 4 + edge_index[1]).tolist() + assert keys == sorted(set(keys)) + + +def test_influence_graph_weights_and_self_loops() -> None: + system = _michaelis_menten() + edge_index, edge_weight = species_influence_graph(system) + assert not bool((edge_index[0] == edge_index[1]).any()) + pairs = { + (int(s), int(t)): float(w) + for s, t, w in zip(edge_index[0], edge_index[1], edge_weight, strict=True) + } + # ES drives unbinding (-> E, S) and catalysis (-> E, P): E gets two. + es, e, s, p = 2, 0, 1, 3 + assert pairs[(es, e)] == 2.0 + assert pairs[(es, s)] == 1.0 + assert pairs[(es, p)] == 1.0 + # P never drives anything; the zeroth-order feed adds no edges. + assert p not in edge_index[0].tolist() + + +def test_reaction_hyperedge_index_layout() -> None: + system = _michaelis_menten() + hyperedge_index = reaction_hyperedge_index(system) + assert hyperedge_index.dtype == torch.long + assert hyperedge_index.shape[0] == 2 + members = { + j: sorted(hyperedge_index[0, hyperedge_index[1] == j].tolist()) + for j in range(system.num_reactions) + } + assert members == {0: [0, 1, 2], 1: [0, 1, 2], 2: [0, 2, 3], 3: [1]} + + +def test_directed_hyperedges_split_reactants_and_products() -> None: + system = _michaelis_menten() + tail_index, head_index = reaction_directed_hyperedges(system) + tails = {j: tail_index[0, tail_index[1] == j].tolist() for j in range(4)} + heads = {j: head_index[0, head_index[1] == j].tolist() for j in range(4)} + assert tails == {0: [0, 1], 1: [2], 2: [2], 3: []} + assert heads == {0: [2], 1: [0, 1], 2: [0, 3], 3: [1]} + + +def test_species_reaction_bipartite_heterodata() -> None: + system = ReactionSystem( + species=("A", "A2"), + reactant_stoichiometry=torch.tensor([[2.0, 0.0]]), + product_stoichiometry=torch.tensor([[0.0, 1.0]]), + rate_constants=torch.tensor([4.0]), + ) + data = species_reaction_bipartite(system) + assert data[SPECIES_NODE].num_nodes == 2 + assert data[REACTION_NODE].num_nodes == 1 + assert data[REACTION_NODE].rate_constant.tolist() == [4.0] + assert data[REACTANT_RELATION].edge_index.tolist() == [[0], [0]] + assert data[REACTANT_RELATION].edge_weight.tolist() == [2.0] + assert data[PRODUCT_RELATION].edge_index.tolist() == [[0], [1]] + assert data[PRODUCT_RELATION].edge_weight.tolist() == [1.0] + data.validate() + + +@pytest.mark.parametrize( + "fn", + [ + species_influence_graph, + reaction_hyperedge_index, + reaction_directed_hyperedges, + species_reaction_bipartite, + ], +) +def test_graph_builders_require_reaction_system(fn: object) -> None: + with pytest.raises(TypeError, match="must be a ReactionSystem"): + fn("not a system") # type: ignore[operator] diff --git a/tests/datasets/test_reaction_integrate.py b/tests/datasets/test_reaction_integrate.py new file mode 100644 index 0000000..46f567f --- /dev/null +++ b/tests/datasets/test_reaction_integrate.py @@ -0,0 +1,269 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.integrate`.""" + +from __future__ import annotations + +import math +import sys + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + ReactionSystem, + element_balance_violation, + integrate_expm, + integrate_reaction_system, + integrate_rk4, + integrate_stiff, + linear_rate_equations, + mass_action_jacobian, + require_scipy, +) + + +def _chain(*, feed: float = 0.0) -> ReactionSystem: + """``A -> B -> C`` (rates 1.0, 0.3) with optional zeroth-order ``∅ -> A``.""" + return ReactionSystem( + species=("A", "B", "C"), + reactant_stoichiometry=torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]] + ), + product_stoichiometry=torch.tensor( + [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]] + ), + rate_constants=torch.tensor([1.0, 0.3, feed], dtype=torch.float64), + element_matrix=torch.tensor([[1.0, 1.0, 1.0]]), + elements=("M",), + ) + + +def _dimerization() -> ReactionSystem: + """Closed ``2 A <-> A2``; element ``X`` with ``A: 1, A2: 2``.""" + return ReactionSystem( + species=("A", "A2"), + reactant_stoichiometry=torch.tensor([[2.0, 0.0], [0.0, 1.0]]), + product_stoichiometry=torch.tensor([[0.0, 1.0], [2.0, 0.0]]), + rate_constants=torch.tensor([3.0, 0.5]), + element_matrix=torch.tensor([[1.0, 2.0]]), + elements=("X",), + ) + + +def _robertson() -> ReactionSystem: + """Robertson's stiff problem (rates 0.04, 3e7, 1e4).""" + return ReactionSystem( + species=("A", "B", "C"), + reactant_stoichiometry=torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 1.0, 1.0]] + ), + product_stoichiometry=torch.tensor( + [[0.0, 1.0, 0.0], [0.0, 1.0, 1.0], [1.0, 0.0, 1.0]] + ), + rate_constants=torch.tensor([0.04, 3.0e7, 1.0e4]), + element_matrix=torch.tensor([[1.0, 1.0, 1.0]]), + elements=("M",), + ) + + +def test_linear_rate_equations_for_chain_with_feed() -> None: + a, b = linear_rate_equations(_chain(feed=0.2)) + expected_a = torch.tensor( + [[-1.0, 0.0, 0.0], [1.0, -0.3, 0.0], [0.0, 0.3, 0.0]], dtype=torch.float64 + ) + assert torch.allclose(a, expected_a) + assert torch.allclose(b, torch.tensor([0.2, 0.0, 0.0], dtype=torch.float64)) + with pytest.raises(ValueError, match="is not linear"): + linear_rate_equations(_dimerization()) + with pytest.raises(TypeError, match="must be a ReactionSystem"): + linear_rate_equations("chain") # type: ignore[arg-type] + + +def test_expm_matches_closed_form_chain() -> None: + system = _chain() + times = torch.tensor([0.0, 0.5, 2.0, 7.0], dtype=torch.float64) + c0 = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float64) + states = integrate_expm(system, c0, times) + assert states.shape == (4, 3) + assert torch.equal(states[0], c0) + k1, k2 = 1.0, 0.3 + for t, state in zip(times.tolist(), states, strict=True): + a = math.exp(-k1 * t) + b = k1 / (k2 - k1) * (math.exp(-k1 * t) - math.exp(-k2 * t)) + assert math.isclose(float(state[0]), a, rel_tol=1e-10, abs_tol=1e-12) + assert math.isclose(float(state[1]), b, rel_tol=1e-10, abs_tol=1e-12) + assert math.isclose(float(state.sum()), 1.0, rel_tol=1e-12) + + +def test_expm_feed_is_exact_and_rk4_agrees() -> None: + system = _chain(feed=0.2) + times = torch.linspace(1.0, 6.0, 6, dtype=torch.float64) + c0 = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.5, 0.5]], dtype=torch.float64) + exact = integrate_expm(system, c0, times) + assert exact.shape == (6, 2, 3) + # With feed f, total mass grows linearly: M(t) = M(t0) + f (t - t0). + totals = exact.sum(dim=-1) + expected = c0.sum(dim=-1) + 0.2 * (times - times[0]).unsqueeze(-1) + assert torch.allclose(totals, expected, atol=1e-12) + approx = integrate_rk4(system, c0, times, substeps=20) + assert torch.allclose(approx, exact, atol=1e-7) + + +def test_auto_picks_expm_for_linear_and_rk4_otherwise() -> None: + times = torch.tensor([0.0, 1.0], dtype=torch.float64) + linear = integrate_reaction_system( + _chain(), torch.tensor([1.0, 0.0, 0.0], dtype=torch.float64), times + ) + assert torch.allclose( + linear, integrate_expm(_chain(), torch.tensor([1.0, 0.0, 0.0]), times) + ) + nonlinear = integrate_reaction_system( + _dimerization(), torch.tensor([1.0, 0.0], dtype=torch.float64), times + ) + assert torch.allclose( + nonlinear, + integrate_rk4( + _dimerization(), torch.tensor([1.0, 0.0], dtype=torch.float64), times + ), + ) + + +def test_rk4_conserves_elements_and_keeps_dtype() -> None: + system = _dimerization() + c0 = torch.tensor([1.0, 0.2], dtype=torch.float32) + times = torch.linspace(0.0, 3.0, 13) + states = integrate_rk4(system, c0, times, substeps=8) + assert states.dtype == torch.float32 + violation = element_balance_violation(states, c0, system) + assert float(violation.abs().max()) < 1e-5 + # Reaches the dimerization equilibrium k_f A^2 = k_r A2. + a, a2 = states[-1].tolist() + assert math.isclose(3.0 * a * a, 0.5 * a2, rel_tol=1e-3) + + +def test_rk4_is_differentiable_in_initial_state() -> None: + c0 = torch.tensor([1.0, 0.0], dtype=torch.float64, requires_grad=True) + states = integrate_rk4(_dimerization(), c0, torch.tensor([0.0, 0.5])) + states[-1, 1].backward() + assert c0.grad is not None + assert bool(torch.isfinite(c0.grad).all()) + + +def test_mass_action_jacobian_matches_autograd() -> None: + system = _robertson() + for c in ( + torch.tensor([0.7, 2e-5, 0.3], dtype=torch.float64), + torch.zeros(3, dtype=torch.float64), + ): + expected = torch.autograd.functional.jacobian(system.time_derivative, c) + jac = mass_action_jacobian(system, c) + assert torch.allclose(jac, expected) + assert bool(torch.isfinite(jac).all()) + with pytest.raises(ValueError, match=r"shape \(3,\)"): + mass_action_jacobian(system, torch.zeros(2, 3, dtype=torch.float64)) + + +def test_stiff_robertson_conserves_mass() -> None: + pytest.importorskip("scipy") + system = _robertson() + times = torch.tensor([0.0, 1.0, 10.0, 100.0], dtype=torch.float64) + c0 = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float64) + for method in ("bdf", "radau", "lsoda"): + states = integrate_reaction_system(system, c0, times, method=method) + assert states.shape == (4, 3) + assert states.dtype == torch.float64 + assert torch.allclose(states[0], c0) + violation = element_balance_violation(states, c0, system) + assert float(violation.abs().max()) < 1e-6 + # Reference value at t = 1 (Hairer & Wanner): A ≈ 0.9664597. + bdf = integrate_stiff(system, c0, times) + assert math.isclose(float(bdf[1, 0]), 0.9664597, rel_tol=1e-5) + + +def test_stiff_solver_failure_raises(monkeypatch: pytest.MonkeyPatch) -> None: + scipy_integrate = pytest.importorskip("scipy.integrate") + + class _Failed: + success = False + message = "step size too small" + + monkeypatch.setattr(scipy_integrate, "solve_ivp", lambda *a, **k: _Failed()) + with pytest.raises(RuntimeError, match="step size too small"): + integrate_stiff( + _robertson(), torch.tensor([1.0, 0.0, 0.0]), torch.tensor([0.0, 1.0]) + ) + + +def test_require_scipy_names_reaction_extra(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem(sys.modules, "scipy", None) + monkeypatch.setitem(sys.modules, "scipy.integrate", None) + with pytest.raises(ImportError, match=r"koopman-graph\[reaction\]"): + require_scipy() + with pytest.raises(ImportError, match=r"koopman-graph\[reaction\]"): + integrate_reaction_system( + _robertson(), + torch.tensor([1.0, 0.0, 0.0]), + torch.tensor([0.0, 1.0]), + method="bdf", + ) + + +def test_single_time_returns_initial_state() -> None: + c0 = torch.tensor([0.4, 0.6, 0.0], dtype=torch.float64) + times = torch.tensor([2.0], dtype=torch.float64) + assert torch.equal(integrate_expm(_chain(), c0, times), c0.unsqueeze(0)) + assert torch.equal(integrate_rk4(_chain(), c0, times), c0.unsqueeze(0)) + + +@pytest.mark.parametrize( + ("times", "error", "match"), + [ + ([0.0, 1.0], TypeError, "times must be a Tensor"), + (torch.zeros(0), ValueError, "non-empty 1-D"), + (torch.zeros(2, 2), ValueError, "non-empty 1-D"), + (torch.tensor([0.0, math.inf]), ValueError, "must be finite"), + (torch.tensor([0.0, 1.0, 1.0]), ValueError, "strictly increasing"), + ], +) +def test_time_validation(times: object, error: type[Exception], match: str) -> None: + with pytest.raises(error, match=match): + integrate_rk4(_chain(), torch.ones(3), times) # type: ignore[arg-type] + + +def test_integer_times_are_accepted() -> None: + states = integrate_rk4(_chain(), torch.ones(3), torch.tensor([0, 1, 2])) + assert states.shape == (3, 3) + + +@pytest.mark.parametrize( + ("initial", "error", "match"), + [ + ([1.0, 0.0, 0.0], TypeError, "floating-point Tensor"), + (torch.ones(3, dtype=torch.long), TypeError, "floating-point Tensor"), + (torch.ones(2), ValueError, r"shape \(\.\.\., 3\)"), + (torch.tensor(1.0), ValueError, r"shape \(\.\.\., 3\)"), + (torch.tensor([1.0, math.nan, 0.0]), ValueError, "must be finite"), + ], +) +def test_initial_validation( + initial: object, error: type[Exception], match: str +) -> None: + with pytest.raises(error, match=match): + integrate_expm(_chain(), initial, torch.tensor([0.0, 1.0])) # type: ignore[arg-type] + + +def test_option_validation() -> None: + c0 = torch.ones(3) + times = torch.tensor([0.0, 1.0]) + for substeps in (0, 1.5, True): + with pytest.raises(ValueError, match="substeps must be a positive int"): + integrate_rk4(_chain(), c0, times, substeps=substeps) # type: ignore[arg-type] + with pytest.raises(ValueError, match="takes no options"): + integrate_reaction_system(_chain(), c0, times, method="expm", substeps=4) + with pytest.raises(ValueError, match="method must be one of"): + integrate_reaction_system(_chain(), c0, times, method="euler") # type: ignore[arg-type] + with pytest.raises(ValueError, match="method must be one of"): + integrate_stiff(_chain(), c0, times, method="rk45") # type: ignore[arg-type] + with pytest.raises(ValueError, match="single initial state"): + integrate_stiff(_chain(), torch.ones(2, 3), times) + with pytest.raises(TypeError, match="must be a ReactionSystem"): + integrate_rk4("chain", c0, times) # type: ignore[arg-type] From 63db4e212b172234c4e69238b5c6893a9bd9dcb0 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 15:34:57 -0400 Subject: [PATCH 3/4] Address review findings on reaction graphs and integrators - integrate_reaction_system: method="auto" accepts substeps for linear and nonlinear systems alike (ignored on the exact path), rejects other options up front, and type-checks the system before dispatch. - integrate_expm: build the generator and propagators on the input's device instead of always on CPU. - species_reaction_bipartite: keep rate_constant in float64 so tiny or huge rate constants do not under/overflow. - mass_action_jacobian: evaluate over padded reactant slots, so memory is O(R * width^2) instead of a dense (R, S, S) tensor. - integrate_stiff: precompute stoichiometry and Jacobian slots once rather than revalidating and recasting on every solver callback. - Share one _require_system helper and reuse ReactionSystem._validate_concentrations; finiteness errors name the argument. Co-Authored-By: Claude Opus 5.5 --- src/koopman_graph/datasets/reaction/graphs.py | 20 +-- .../datasets/reaction/integrate.py | 169 +++++++++++++----- src/koopman_graph/datasets/reaction/system.py | 12 ++ tests/datasets/test_reaction_graphs.py | 12 ++ tests/datasets/test_reaction_integrate.py | 61 ++++++- 5 files changed, 210 insertions(+), 64 deletions(-) diff --git a/src/koopman_graph/datasets/reaction/graphs.py b/src/koopman_graph/datasets/reaction/graphs.py index ed277c7..07e8e34 100644 --- a/src/koopman_graph/datasets/reaction/graphs.py +++ b/src/koopman_graph/datasets/reaction/graphs.py @@ -25,7 +25,7 @@ from torch import Tensor from torch_geometric.data import HeteroData -from koopman_graph.datasets.reaction.system import ReactionSystem +from koopman_graph.datasets.reaction.system import ReactionSystem, _require_system SPECIES_NODE = "species" """Node type name for species in :func:`species_reaction_bipartite`.""" @@ -153,13 +153,15 @@ def species_reaction_bipartite(system: ReactionSystem) -> HeteroData: and ``edge_weight`` = reactant coefficient ``α``; :data:`PRODUCT_RELATION` carries reaction → species edges with ``edge_weight`` = product coefficient ``β``. The ``reaction`` nodes - also store ``rate_constant`` with shape ``(num_reactions,)``. + also store ``float64`` ``rate_constant`` with shape + ``(num_reactions,)``; it stays ``float64`` because rate constants + span many decades and would under- or overflow in ``float32``. """ _require_system(system) data = HeteroData() data[SPECIES_NODE].num_nodes = system.num_species data[REACTION_NODE].num_nodes = system.num_reactions - data[REACTION_NODE].rate_constant = system.rate_constants.to(torch.float32) + data[REACTION_NODE].rate_constant = system.rate_constants.clone() alpha = system.reactant_stoichiometry reaction, species = (alpha > 0).nonzero(as_tuple=True) @@ -186,15 +188,3 @@ def _incidence(mask: Tensor) -> Tensor: See signature.""" reaction, species = mask.nonzero(as_tuple=True) return torch.stack([species, reaction]).to(dtype=torch.long) - - -def _require_system(system: object) -> None: - """Raise ``TypeError`` unless ``system`` is a :class:`ReactionSystem`. - - Parameters - ---------- - system - See signature.""" - if not isinstance(system, ReactionSystem): - msg = f"system must be a ReactionSystem, got {type(system).__name__}" - raise TypeError(msg) diff --git a/src/koopman_graph/datasets/reaction/integrate.py b/src/koopman_graph/datasets/reaction/integrate.py index a7c04ff..cc87f77 100644 --- a/src/koopman_graph/datasets/reaction/integrate.py +++ b/src/koopman_graph/datasets/reaction/integrate.py @@ -28,7 +28,7 @@ import torch from torch import Tensor -from koopman_graph.datasets.reaction.system import ReactionSystem +from koopman_graph.datasets.reaction.system import ReactionSystem, _require_system IntegrationMethod = Literal["auto", "expm", "rk4", "bdf", "radau", "lsoda"] """Accepted ``method`` names for :func:`integrate_reaction_system`.""" @@ -114,7 +114,8 @@ def integrate_expm( Returns ------- Tensor - Concentrations ``(len(times), ..., num_species)`` in ``float64``. + Concentrations ``(len(times), ..., num_species)`` in ``float64`` on + the device of ``initial``. Raises ------ @@ -122,13 +123,17 @@ def integrate_expm( If the system is not linear or inputs are malformed. """ a, b = linear_rate_equations(system) - c0 = _validate_initial(system, initial).to(torch.float64) + c0 = _validate_state(system, initial, "initial").to(torch.float64) t = _validate_times(times) num_species = system.num_species - generator = torch.zeros(num_species + 1, num_species + 1, dtype=torch.float64) - generator[:num_species, :num_species] = a - generator[:num_species, num_species] = b - elapsed = (t - t[0]).to(dtype=torch.float64, device="cpu") + device = c0.device + generator = torch.zeros( + num_species + 1, num_species + 1, dtype=torch.float64, device=device + ) + generator[:num_species, :num_species] = a.to(device) + generator[:num_species, num_species] = b.to(device) + t = t.to(dtype=torch.float64, device=device) + elapsed = t - t[0] propagators = torch.linalg.matrix_exp(elapsed[:, None, None] * generator) augmented = torch.cat([c0, torch.ones_like(c0[..., :1])], dim=-1) # result[t, ..., i] = Σ_j P[t, i, j] y[..., j] @@ -171,7 +176,7 @@ def integrate_rk4( if isinstance(substeps, bool) or not isinstance(substeps, int) or substeps < 1: msg = f"substeps must be a positive int, got {substeps!r}" raise ValueError(msg) - c = _validate_initial(system, initial) + c = _validate_state(system, initial, "initial") t = _validate_times(times).to(dtype=c.dtype, device=c.device) states = [c] for index in range(1, int(t.shape[0])): @@ -231,7 +236,7 @@ def integrate_stiff( if method not in _STIFF_METHODS: msg = f"method must be one of {sorted(_STIFF_METHODS)}, got {method!r}" raise ValueError(msg) - c0 = _validate_initial(system, initial) + c0 = _validate_state(system, initial, "initial") if c0.ndim != 1: msg = ( "stiff integration takes a single initial state of shape " @@ -240,6 +245,11 @@ def integrate_stiff( raise ValueError(msg) t = _validate_times(times).to(dtype=torch.float64, device="cpu") scipy_integrate = require_scipy() + # Built once; solve_ivp calls rhs / jac many times per solve. + alpha = system.reactant_stoichiometry + k = system.rate_constants + n = system.stoichiometric_matrix + terms = _jacobian_terms(system, dtype=torch.float64, device=torch.device("cpu")) def rhs(_time: float, y: Any) -> Any: """Evaluate ``dc/dt`` on a NumPy state for ``solve_ivp``. @@ -254,8 +264,8 @@ def rhs(_time: float, y: Any) -> Any: Returns ------- See signature.""" - state = torch.from_numpy(y) - return system.time_derivative(state).numpy() + rates = k * torch.pow(torch.from_numpy(y), alpha).prod(dim=-1) + return (n @ rates).numpy() def jac(_time: float, y: Any) -> Any: """Evaluate the mass-action Jacobian on a NumPy state for ``solve_ivp``. @@ -270,7 +280,7 @@ def jac(_time: float, y: Any) -> Any: Returns ------- See signature.""" - return mass_action_jacobian(system, torch.from_numpy(y)).numpy() + return _jacobian_from_terms(torch.from_numpy(y), *terms).numpy() t_values = t.tolist() solution = scipy_integrate.solve_ivp( @@ -315,7 +325,8 @@ def integrate_reaction_system( explicitly. **kwargs Forwarded to :func:`integrate_rk4` (``substeps``) or - :func:`integrate_stiff` (``rtol``, ``atol``). + :func:`integrate_stiff` (``rtol``, ``atol``). With ``"auto"`` only + ``substeps`` is accepted; it is ignored when the exact path is used. Returns ------- @@ -324,11 +335,20 @@ def integrate_reaction_system( Raises ------ + TypeError + If ``system`` is not a :class:`ReactionSystem`. ValueError If ``method`` is unknown or ``kwargs`` do not apply to it. """ + _require_system(system) if method == "auto": - method = "expm" if system.is_linear else "rk4" + unknown = sorted(set(kwargs) - {"substeps"}) + if unknown: + msg = f"method='auto' only accepts 'substeps', got {unknown}" + raise ValueError(msg) + if system.is_linear: + return integrate_expm(system, initial, times) + method = "rk4" if method == "expm": if kwargs: msg = f"method='expm' takes no options, got {sorted(kwargs)}" @@ -367,27 +387,94 @@ def mass_action_jacobian(system: ReactionSystem, concentrations: Tensor) -> Tens ValueError If ``concentrations`` is not a single state. """ - c = _validate_initial(system, concentrations) + c = _validate_state(system, concentrations, "concentrations") if c.ndim != 1: msg = ( f"concentrations must have shape ({system.num_species},), got " f"{tuple(c.shape)}" ) raise ValueError(msg) - alpha = system.reactant_stoichiometry.to(dtype=c.dtype, device=c.device) - k = system.rate_constants.to(dtype=c.dtype, device=c.device) - num_species = system.num_species - # ∂r_j/∂c_s = k_j α_js c_s^(α_js - 1) ∏_{s'≠s} c_s'^α_js'; building the - # product with exponent α - e_s keeps c = 0 finite (no division by c). - eye = torch.eye(num_species, dtype=c.dtype, device=c.device) - shifted = alpha.unsqueeze(1) - eye.unsqueeze(0) # (R, S_s, S_s') + terms = _jacobian_terms(system, dtype=c.dtype, device=c.device) + return _jacobian_from_terms(c, *terms) + + +def _jacobian_terms( + system: ReactionSystem, + *, + dtype: torch.dtype, + device: torch.device, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Pack each reaction's reactants into padded slots for the Jacobian. + + Parameters + ---------- + system + See signature. + dtype + See signature. + device + See signature. + + Returns + ------- + ``(slot_species, slot_alpha, k, n)``: reactant species indices and + coefficients with shape ``(num_reactions, width)`` (padding has + ``α = 0``), rate constants, and the stoichiometric matrix.""" + alpha = system.reactant_stoichiometry + is_reactant = alpha > 0 + width = max(int(is_reactant.sum(dim=1).max()), 1) + # A stable sort moves each reaction's reactant columns to the front. + order = torch.sort((~is_reactant).to(torch.uint8), dim=1, stable=True).indices + slot_species = order[:, :width] + slot_alpha = alpha.gather(1, slot_species) + return ( + slot_species.to(device=device), + slot_alpha.to(dtype=dtype, device=device), + system.rate_constants.to(dtype=dtype, device=device), + system.stoichiometric_matrix.to(dtype=dtype, device=device), + ) + + +def _jacobian_from_terms( + c: Tensor, + slot_species: Tensor, + slot_alpha: Tensor, + k: Tensor, + n: Tensor, +) -> Tensor: + """Evaluate the mass-action Jacobian from :func:`_jacobian_terms` slots. + + Parameters + ---------- + c + See signature. + slot_species + See signature. + slot_alpha + See signature. + k + See signature. + n + See signature. + + Returns + ------- + See signature.""" + # ∂r_j/∂c_s = k_j α_js c_s^(α_js - 1) ∏_{s'≠s} c_s'^α_js', taken over the + # reactant slots only (memory O(R·width²), not O(R·S²)). Lowering the + # exponent by one instead of dividing by c keeps c = 0 finite. + values = c[slot_species].unsqueeze(1) # (R, 1, width) + eye = torch.eye(slot_alpha.shape[1], dtype=c.dtype, device=c.device) + shifted = slot_alpha.unsqueeze(1) - eye # (R, width_s, width_s') powered = torch.where( shifted < 0, torch.ones((), dtype=c.dtype, device=c.device), - torch.pow(c, shifted.clamp(min=0)), + torch.pow(values, shifted.clamp(min=0)), ) - d_rates = k[:, None] * alpha * powered.prod(dim=-1) # (R, S) - n = system.stoichiometric_matrix.to(dtype=c.dtype, device=c.device) + slot_grad = k[:, None] * slot_alpha * powered.prod(dim=-1) # (R, width) + d_rates = torch.zeros( + k.shape[0], n.shape[0], dtype=c.dtype, device=c.device + ).scatter_add(1, slot_species, slot_grad) # (R, S) return n @ d_rates @@ -398,9 +485,7 @@ def _require_linear(system: ReactionSystem) -> None: ---------- system See signature.""" - if not isinstance(system, ReactionSystem): - msg = f"system must be a ReactionSystem, got {type(system).__name__}" - raise TypeError(msg) + _require_system(system) if not system.is_linear: orders = system.reaction_orders.tolist() msg = ( @@ -410,35 +495,27 @@ def _require_linear(system: ReactionSystem) -> None: raise ValueError(msg) -def _validate_initial(system: ReactionSystem, initial: Tensor) -> Tensor: - """Validate an initial state ``(..., num_species)`` of finite floats. +def _validate_state(system: ReactionSystem, state: Tensor, name: str) -> Tensor: + """Validate a finite floating state ``(..., num_species)``. Parameters ---------- system See signature. - initial + state See signature. + name + Argument name used in the finiteness error. Returns ------- See signature.""" - if not isinstance(system, ReactionSystem): - msg = f"system must be a ReactionSystem, got {type(system).__name__}" - raise TypeError(msg) - if not isinstance(initial, Tensor) or not initial.is_floating_point(): - msg = "initial concentrations must be a floating-point Tensor" - raise TypeError(msg) - if initial.ndim == 0 or int(initial.shape[-1]) != system.num_species: - msg = ( - f"initial concentrations must have shape (..., {system.num_species}), " - f"got {tuple(initial.shape)}" - ) - raise ValueError(msg) - if not bool(torch.isfinite(initial).all()): - msg = "initial concentrations must be finite" + _require_system(system) + state = system._validate_concentrations(state) + if not bool(torch.isfinite(state).all()): + msg = f"{name} must be finite" raise ValueError(msg) - return initial + return state def _validate_times(times: Tensor) -> Tensor: diff --git a/src/koopman_graph/datasets/reaction/system.py b/src/koopman_graph/datasets/reaction/system.py index d96fe3b..36715a7 100644 --- a/src/koopman_graph/datasets/reaction/system.py +++ b/src/koopman_graph/datasets/reaction/system.py @@ -626,6 +626,18 @@ def check_element_balance( ) +def _require_system(system: object) -> None: + """Raise ``TypeError`` unless ``system`` is a :class:`ReactionSystem`. + + Parameters + ---------- + system + See signature.""" + if not isinstance(system, ReactionSystem): + msg = f"system must be a ReactionSystem, got {type(system).__name__}" + raise TypeError(msg) + + def _finite_float(value: Any, name: str) -> float: """Coerce ``value`` to a finite float or raise ``ValueError``. diff --git a/tests/datasets/test_reaction_graphs.py b/tests/datasets/test_reaction_graphs.py index f1a20ac..11a9b9f 100644 --- a/tests/datasets/test_reaction_graphs.py +++ b/tests/datasets/test_reaction_graphs.py @@ -108,6 +108,7 @@ def test_species_reaction_bipartite_heterodata() -> None: assert data[SPECIES_NODE].num_nodes == 2 assert data[REACTION_NODE].num_nodes == 1 assert data[REACTION_NODE].rate_constant.tolist() == [4.0] + assert data[REACTION_NODE].rate_constant.dtype == torch.float64 assert data[REACTANT_RELATION].edge_index.tolist() == [[0], [0]] assert data[REACTANT_RELATION].edge_weight.tolist() == [2.0] assert data[PRODUCT_RELATION].edge_index.tolist() == [[0], [1]] @@ -115,6 +116,17 @@ def test_species_reaction_bipartite_heterodata() -> None: data.validate() +def test_bipartite_keeps_tiny_rate_constants() -> None: + system = ReactionSystem( + species=("A", "B"), + reactant_stoichiometry=torch.tensor([[1.0, 0.0]]), + product_stoichiometry=torch.tensor([[0.0, 1.0]]), + rate_constants=torch.tensor([1e-50], dtype=torch.float64), + ) + rate = species_reaction_bipartite(system)[REACTION_NODE].rate_constant + assert rate.tolist() == [1e-50] + + @pytest.mark.parametrize( "fn", [ diff --git a/tests/datasets/test_reaction_integrate.py b/tests/datasets/test_reaction_integrate.py index 46f567f..761e5b8 100644 --- a/tests/datasets/test_reaction_integrate.py +++ b/tests/datasets/test_reaction_integrate.py @@ -127,6 +127,29 @@ def test_auto_picks_expm_for_linear_and_rk4_otherwise() -> None: ) +def test_auto_accepts_substeps_for_linear_and_nonlinear() -> None: + times = torch.tensor([0.0, 1.0], dtype=torch.float64) + c0 = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float64) + linear = integrate_reaction_system(_chain(), c0, times, substeps=20) + assert torch.allclose(linear, integrate_expm(_chain(), c0, times)) + explicit = integrate_reaction_system(_chain(), c0, times, method="expm") + assert torch.allclose(explicit, linear) + dimer0 = torch.tensor([1.0, 0.0], dtype=torch.float64) + nonlinear = integrate_reaction_system(_dimerization(), dimer0, times, substeps=20) + assert torch.allclose( + nonlinear, integrate_rk4(_dimerization(), dimer0, times, substeps=20) + ) + with pytest.raises(ValueError, match="only accepts 'substeps'"): + integrate_reaction_system(_chain(), c0, times, rtol=1e-6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_expm_keeps_input_device() -> None: + c0 = torch.tensor([1.0, 0.0, 0.0], device="cuda") + states = integrate_expm(_chain(), c0, torch.tensor([0.0, 1.0])) + assert states.device == c0.device + + def test_rk4_conserves_elements_and_keeps_dtype() -> None: system = _dimerization() c0 = torch.tensor([1.0, 0.2], dtype=torch.float32) @@ -160,6 +183,36 @@ def test_mass_action_jacobian_matches_autograd() -> None: assert bool(torch.isfinite(jac).all()) with pytest.raises(ValueError, match=r"shape \(3,\)"): mass_action_jacobian(system, torch.zeros(2, 3, dtype=torch.float64)) + with pytest.raises(ValueError, match="concentrations must be finite"): + mass_action_jacobian(system, torch.tensor([math.nan, 0.0, 0.0])) + + +def test_mass_action_jacobian_with_wide_reactant_sets() -> None: + """Three-reactant and zeroth-order reactions over unused species.""" + system = ReactionSystem( + species=("A", "B", "C", "D", "E"), + reactant_stoichiometry=torch.tensor( + [ + [2.0, 1.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 3.0, 0.0, 0.0], + ] + ), + product_stoichiometry=torch.tensor( + [ + [0.0, 0.0, 1.0, 1.0, 0.0], + [1.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0, 1.0], + ] + ), + rate_constants=torch.tensor([1.5, 0.2, 0.7]), + ) + for c in ( + torch.tensor([0.4, 1.3, 0.8, 0.6, 2.0], dtype=torch.float64), + torch.tensor([0.0, 1.3, 0.0, 0.6, 0.0], dtype=torch.float64), + ): + expected = torch.autograd.functional.jacobian(system.time_derivative, c) + assert torch.allclose(mass_action_jacobian(system, c), expected) def test_stiff_robertson_conserves_mass() -> None: @@ -237,11 +290,11 @@ def test_integer_times_are_accepted() -> None: @pytest.mark.parametrize( ("initial", "error", "match"), [ - ([1.0, 0.0, 0.0], TypeError, "floating-point Tensor"), - (torch.ones(3, dtype=torch.long), TypeError, "floating-point Tensor"), + ([1.0, 0.0, 0.0], TypeError, "must be a Tensor"), + (torch.ones(3, dtype=torch.long), TypeError, "must be floating point"), (torch.ones(2), ValueError, r"shape \(\.\.\., 3\)"), (torch.tensor(1.0), ValueError, r"shape \(\.\.\., 3\)"), - (torch.tensor([1.0, math.nan, 0.0]), ValueError, "must be finite"), + (torch.tensor([1.0, math.nan, 0.0]), ValueError, "initial must be finite"), ], ) def test_initial_validation( @@ -267,3 +320,5 @@ def test_option_validation() -> None: integrate_stiff(_chain(), torch.ones(2, 3), times) with pytest.raises(TypeError, match="must be a ReactionSystem"): integrate_rk4("chain", c0, times) # type: ignore[arg-type] + with pytest.raises(TypeError, match="must be a ReactionSystem"): + integrate_reaction_system(None, c0, times) # type: ignore[arg-type] From 32704d8306cc226f41937177988b6f179eddb8b4 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 16:50:48 -0400 Subject: [PATCH 4/4] Seed flaky Sinkhorn doubly-stochastic test test_doubly_stochastic_row_and_column_sums drew unseeded torch.randn logits and asserted row/column sums within 1e-4. The default 8 Sinkhorn sweeps are documented as approximate: 63 of 500 random seeds leave a row sum off by more than 1e-4 (max ~2.6e-3), so the test failed about 13% of the time (seen while propagating the reaction stack). Pin torch.manual_seed(0) (error ~5e-7) instead of loosening the tolerance. Co-Authored-By: Claude Opus 5.5 --- tests/operators/test_parameterizations.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/operators/test_parameterizations.py b/tests/operators/test_parameterizations.py index 46d9a46..7df9175 100644 --- a/tests/operators/test_parameterizations.py +++ b/tests/operators/test_parameterizations.py @@ -53,6 +53,10 @@ def test_row_stochastic_conserves_latent_mass() -> None: def test_doubly_stochastic_row_and_column_sums() -> None: """Sinkhorn assembly is approximately doubly stochastic.""" + # The default 8 Sinkhorn sweeps are approximate: about 13% of random 5x5 + # logits leave a row sum off by more than 1e-4 (up to ~3e-3). Pin the + # draw instead of inheriting whatever RNG state earlier tests left. + torch.manual_seed(0) raw = torch.randn(5, 5) matrix = assemble_doubly_stochastic_matrix(raw) ones = torch.ones(5)