From c6fa500a139f02e6a79fd3c57e48dd30bf13e138 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 15:28:11 -0400 Subject: [PATCH 1/2] Add reaction catalog and oracle equilibria, spectra, timescales MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refs #65. Fills in koopman_graph.datasets.reaction.catalog and .oracles: - catalog: first_order_chain, reversible_dimerization, michaelis_menten (closed, element-balanced), brusselator and lotka_volterra (open), with frozen default parameters exposed as module constants. - oracles: conservation_laws (orthonormal left nullspace of N), stoichiometric_rank, equilibrium (exact for linear systems; damped non-negative Newton within the compatibility class otherwise; reports missing steady states and non-convergence), and oracle_spectrum / oracle_timescales / OracleSpectrum (Jacobian, dynamic eigenvalues with conservation zero modes removed, -1/Re(λ) timescales, exp(λ Δt) flow-map eigenvalues). Pure PyTorch; reuses mass_action_jacobian and linear_rate_equations. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 11 + .../datasets/reaction/__init__.py | 38 ++ .../datasets/reaction/catalog.py | 279 +++++++++++++ .../datasets/reaction/oracles.py | 378 ++++++++++++++++++ tests/datasets/test_reaction_catalog.py | 108 +++++ tests/datasets/test_reaction_oracles.py | 187 +++++++++ 6 files changed, 1001 insertions(+) create mode 100644 tests/datasets/test_reaction_catalog.py create mode 100644 tests/datasets/test_reaction_oracles.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c9cfb3b..101f26f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Jacobian behind a lazy `require_scipy()` guard (`integrate_stiff`; `[reaction]` extra), the `integrate_reaction_system` dispatcher (`method="auto"`), and `mass_action_jacobian`. +- `koopman_graph.datasets.reaction.catalog`: reference systems with frozen + default parameters (`first_order_chain`, `reversible_dimerization`, + `michaelis_menten`, `brusselator`, `lotka_volterra`). Closed systems carry + an element matrix. +- `koopman_graph.datasets.reaction.oracles`: `conservation_laws` (orthonormal + left nullspace of `N`), `stoichiometric_rank`, `equilibrium` (exact for + linear systems, damped non-negative Newton within the compatibility + class otherwise), and `oracle_spectrum` / `oracle_timescales` / + `OracleSpectrum`. These give the Jacobian, dynamic eigenvalues with the + conservation zero modes removed, timescales `-1/Re(λ)`, and flow-map + eigenvalues `exp(λ Δt)` to compare with a learned discrete operator. - `[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 8aa8629..4f58cfb 100644 --- a/src/koopman_graph/datasets/reaction/__init__.py +++ b/src/koopman_graph/datasets/reaction/__init__.py @@ -38,6 +38,19 @@ replacement for a production kinetics solver. """ +from koopman_graph.datasets.reaction.catalog import ( + DEFAULT_BRUSSELATOR_PARAMS, + DEFAULT_CHAIN_RATE, + DEFAULT_CHAIN_RATIO, + DEFAULT_DIMERIZATION_RATES, + DEFAULT_LOTKA_VOLTERRA_RATES, + DEFAULT_MICHAELIS_MENTEN_RATES, + brusselator, + first_order_chain, + lotka_volterra, + michaelis_menten, + reversible_dimerization, +) from koopman_graph.datasets.reaction.graphs import ( PRODUCT_RELATION, REACTANT_RELATION, @@ -58,6 +71,14 @@ mass_action_jacobian, require_scipy, ) +from koopman_graph.datasets.reaction.oracles import ( + OracleSpectrum, + conservation_laws, + equilibrium, + oracle_spectrum, + oracle_timescales, + stoichiometric_rank, +) from koopman_graph.datasets.reaction.system import ( GAS_CONSTANT_J_PER_MOL_K, ArrheniusParams, @@ -72,7 +93,13 @@ ) __all__ = [ + "DEFAULT_BRUSSELATOR_PARAMS", + "DEFAULT_CHAIN_RATE", + "DEFAULT_CHAIN_RATIO", + "DEFAULT_DIMERIZATION_RATES", "DEFAULT_LOG_FLOOR", + "DEFAULT_LOTKA_VOLTERRA_RATES", + "DEFAULT_MICHAELIS_MENTEN_RATES", "GAS_CONSTANT_J_PER_MOL_K", "PRODUCT_RELATION", "REACTANT_RELATION", @@ -82,18 +109,29 @@ "ElementBalanceReport", "IntegrationMethod", "LogConcentrationTransform", + "OracleSpectrum", "ReactionSystem", + "brusselator", "check_element_balance", + "conservation_laws", "element_balance_violation", + "equilibrium", + "first_order_chain", "integrate_expm", "integrate_reaction_system", "integrate_rk4", "integrate_stiff", "linear_rate_equations", + "lotka_volterra", "mass_action_jacobian", + "michaelis_menten", + "oracle_spectrum", + "oracle_timescales", "reaction_directed_hyperedges", "reaction_hyperedge_index", "require_scipy", + "reversible_dimerization", "species_influence_graph", "species_reaction_bipartite", + "stoichiometric_rank", ] diff --git a/src/koopman_graph/datasets/reaction/catalog.py b/src/koopman_graph/datasets/reaction/catalog.py index e69de29..1dfc3e3 100644 --- a/src/koopman_graph/datasets/reaction/catalog.py +++ b/src/koopman_graph/datasets/reaction/catalog.py @@ -0,0 +1,279 @@ +"""Small reference reaction systems with frozen defaults. + +Each constructor returns a validated +:class:`~koopman_graph.datasets.reaction.ReactionSystem`. The default +parameters are part of the dataset contract: changing them changes every +dataset generated from the catalog, so they are module constants. Closed +systems carry an element matrix; open systems (sources, sinks, or a +reservoir held constant) do not, and :func:`check_element_balance` is not +meaningful for them. + +* :func:`first_order_chain`: closed, linear ``S0 → S1 → … → S(n-1)``. +* :func:`reversible_dimerization`: closed, nonlinear ``2 A ⇌ A2``. +* :func:`michaelis_menten`: closed ``E + S ⇌ ES → E + P``. +* :func:`brusselator`: open; limit cycle when ``b > 1 + a²``. +* :func:`lotka_volterra`: open; neutral centre (purely imaginary pair). + +Pure PyTorch; importing this module does not import scipy or cantera. +""" + +from __future__ import annotations + +import math + +import torch + +from koopman_graph.datasets.reaction.system import ReactionSystem + +DEFAULT_CHAIN_RATE = 1.0 +"""Rate of the first step of :func:`first_order_chain`.""" + +DEFAULT_CHAIN_RATIO = 0.5 +"""Geometric ratio between successive :func:`first_order_chain` rates.""" + +DEFAULT_DIMERIZATION_RATES = (3.0, 0.5) +"""Default ``(k_forward, k_reverse)`` for :func:`reversible_dimerization`.""" + +DEFAULT_MICHAELIS_MENTEN_RATES = (10.0, 1.0, 0.5) +"""Default ``(k_on, k_off, k_cat)`` for :func:`michaelis_menten`.""" + +DEFAULT_BRUSSELATOR_PARAMS = (1.0, 3.0) +"""Default ``(a, b)`` for :func:`brusselator` (oscillatory: ``3 > 1 + 1``).""" + +DEFAULT_LOTKA_VOLTERRA_RATES = (1.0, 1.0, 1.0) +"""Default ``(alpha, beta, gamma)`` for :func:`lotka_volterra`.""" + +_MAX_CHAIN_SPECIES = 64 + + +def first_order_chain( + num_species: int = 3, + *, + rate_constants: tuple[float, ...] | None = None, +) -> ReactionSystem: + """Return the linear chain ``S0 → S1 → … → S(n-1)``. + + Parameters + ---------- + num_species : int, optional + Chain length ``n`` in ``[2, 64]``. Default is ``3``. + rate_constants : tuple of float or None, optional + ``n - 1`` step rates. Default is the geometric sequence + ``DEFAULT_CHAIN_RATE * DEFAULT_CHAIN_RATIO**i`` (``1.0, 0.5, …``), + which keeps every Jacobian eigenvalue distinct. + + Returns + ------- + ReactionSystem + Closed, linear system with one element ``M`` (total mass). The + Jacobian eigenvalues are ``-k_i`` and one conservation zero. + + Raises + ------ + ValueError + If ``num_species`` is out of range or the rate count is wrong. + """ + if ( + isinstance(num_species, bool) + or not isinstance(num_species, int) + or not 2 <= num_species <= _MAX_CHAIN_SPECIES + ): + msg = ( + f"num_species must be an int in [2, {_MAX_CHAIN_SPECIES}], got " + f"{num_species!r}" + ) + raise ValueError(msg) + num_reactions = num_species - 1 + if rate_constants is None: + rates = tuple( + DEFAULT_CHAIN_RATE * DEFAULT_CHAIN_RATIO**i for i in range(num_reactions) + ) + else: + rates = tuple(float(k) for k in rate_constants) + if len(rates) != num_reactions: + msg = ( + f"rate_constants must have {num_reactions} entries for " + f"{num_species} species, got {len(rates)}" + ) + raise ValueError(msg) + reactants = torch.zeros(num_reactions, num_species, dtype=torch.float64) + products = torch.zeros(num_reactions, num_species, dtype=torch.float64) + index = torch.arange(num_reactions) + reactants[index, index] = 1.0 + products[index, index + 1] = 1.0 + species = tuple(f"S{i}" for i in range(num_species)) + return ReactionSystem( + species=species, + reactant_stoichiometry=reactants, + product_stoichiometry=products, + rate_constants=torch.tensor(rates, dtype=torch.float64), + element_matrix=torch.ones(1, num_species, dtype=torch.float64), + elements=("M",), + reaction_names=tuple( + f"{species[i]} -> {species[i + 1]}" for i in range(num_reactions) + ), + name=f"first_order_chain_{num_species}", + ) + + +def reversible_dimerization( + k_forward: float = DEFAULT_DIMERIZATION_RATES[0], + k_reverse: float = DEFAULT_DIMERIZATION_RATES[1], +) -> ReactionSystem: + """Return ``2 A ⇌ A2``. + + Parameters + ---------- + k_forward : float, optional + Rate constant of ``2 A → A2``. Default is ``3.0``. + k_reverse : float, optional + Rate constant of ``A2 → 2 A``. Default is ``0.5``. + + Returns + ------- + ReactionSystem + Closed, nonlinear system; element ``X`` counts monomer units + (``A: 1``, ``A2: 2``). Equilibrium satisfies + ``k_forward A² = k_reverse A2``. + """ + return ReactionSystem( + species=("A", "A2"), + reactant_stoichiometry=torch.tensor( + [[2.0, 0.0], [0.0, 1.0]], dtype=torch.float64 + ), + product_stoichiometry=torch.tensor( + [[0.0, 1.0], [2.0, 0.0]], dtype=torch.float64 + ), + rate_constants=torch.tensor([k_forward, k_reverse], dtype=torch.float64), + element_matrix=torch.tensor([[1.0, 2.0]], dtype=torch.float64), + elements=("X",), + reaction_names=("2 A -> A2", "A2 -> 2 A"), + name="reversible_dimerization", + ) + + +def michaelis_menten( + k_on: float = DEFAULT_MICHAELIS_MENTEN_RATES[0], + k_off: float = DEFAULT_MICHAELIS_MENTEN_RATES[1], + k_cat: float = DEFAULT_MICHAELIS_MENTEN_RATES[2], +) -> ReactionSystem: + """Return the mass-action Michaelis-Menten mechanism ``E + S ⇌ ES → E + P``. + + Parameters + ---------- + k_on : float, optional + Binding rate constant. Default is ``10.0``. + k_off : float, optional + Unbinding rate constant. Default is ``1.0``. + k_cat : float, optional + Catalytic rate constant. Default is ``0.5``. + + Returns + ------- + ReactionSystem + Closed system over species ``(E, S, ES, P)`` with elements ``Enz`` + (enzyme units) and ``Sub`` (substrate backbone). All substrate ends + up as ``P``, a boundary equilibrium with ``S = ES = 0``. + """ + 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]], + dtype=torch.float64, + ), + 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]], + dtype=torch.float64, + ), + rate_constants=torch.tensor([k_on, k_off, k_cat], dtype=torch.float64), + element_matrix=torch.tensor( + [[1.0, 0.0, 1.0, 0.0], [0.0, 1.0, 1.0, 1.0]], dtype=torch.float64 + ), + elements=("Enz", "Sub"), + reaction_names=("E + S -> ES", "ES -> E + S", "ES -> E + P"), + name="michaelis_menten", + ) + + +def brusselator( + a: float = DEFAULT_BRUSSELATOR_PARAMS[0], + b: float = DEFAULT_BRUSSELATOR_PARAMS[1], +) -> ReactionSystem: + """Return the Brusselator with reservoir species folded into rates. + + Parameters + ---------- + a : float, optional + Feed rate of ``∅ → X`` (reservoir ``A`` held constant). Default is + ``1.0``. + b : float, optional + Rate of ``X → Y`` (reservoir ``B`` held constant). Default is + ``3.0``. + + Returns + ------- + ReactionSystem + Open system over ``(X, Y)``: ``∅ → X`` (``a``), ``X → Y`` (``b``), + ``2 X + Y → 3 X`` (``1``), ``X → ∅`` (``1``). The unique fixed point + ``(a, b / a)`` is unstable, with a limit cycle around it, when + ``b > 1 + a²``. + + Raises + ------ + ValueError + If ``a`` or ``b`` is not finite and positive. + """ + for name, value in (("a", a), ("b", b)): + if not math.isfinite(float(value)) or float(value) <= 0.0: + msg = f"{name} must be finite and > 0, got {value!r}" + raise ValueError(msg) + return ReactionSystem( + species=("X", "Y"), + reactant_stoichiometry=torch.tensor( + [[0.0, 0.0], [1.0, 0.0], [2.0, 1.0], [1.0, 0.0]], dtype=torch.float64 + ), + product_stoichiometry=torch.tensor( + [[1.0, 0.0], [0.0, 1.0], [3.0, 0.0], [0.0, 0.0]], dtype=torch.float64 + ), + rate_constants=torch.tensor([a, b, 1.0, 1.0], dtype=torch.float64), + reaction_names=("-> X", "X -> Y", "2 X + Y -> 3 X", "X ->"), + name="brusselator", + ) + + +def lotka_volterra( + alpha: float = DEFAULT_LOTKA_VOLTERRA_RATES[0], + beta: float = DEFAULT_LOTKA_VOLTERRA_RATES[1], + gamma: float = DEFAULT_LOTKA_VOLTERRA_RATES[2], +) -> ReactionSystem: + """Return the mass-action Lotka-Volterra predator-prey network. + + Parameters + ---------- + alpha : float, optional + Prey birth ``X → 2 X``. Default is ``1.0``. + beta : float, optional + Predation ``X + Y → 2 Y``. Default is ``1.0``. + gamma : float, optional + Predator death ``Y → ∅``. Default is ``1.0``. + + Returns + ------- + ReactionSystem + Open system over ``(X, Y)``. The interior fixed point + ``(gamma / beta, alpha / beta)`` is a neutral centre with Jacobian + eigenvalues ``± i sqrt(alpha gamma)``, so trajectories are closed + orbits and the oracle timescale is infinite. + """ + return ReactionSystem( + species=("X", "Y"), + reactant_stoichiometry=torch.tensor( + [[1.0, 0.0], [1.0, 1.0], [0.0, 1.0]], dtype=torch.float64 + ), + product_stoichiometry=torch.tensor( + [[2.0, 0.0], [0.0, 2.0], [0.0, 0.0]], dtype=torch.float64 + ), + rate_constants=torch.tensor([alpha, beta, gamma], dtype=torch.float64), + reaction_names=("X -> 2 X", "X + Y -> 2 Y", "Y ->"), + name="lotka_volterra", + ) diff --git a/src/koopman_graph/datasets/reaction/oracles.py b/src/koopman_graph/datasets/reaction/oracles.py index e69de29..58c8044 100644 --- a/src/koopman_graph/datasets/reaction/oracles.py +++ b/src/koopman_graph/datasets/reaction/oracles.py @@ -0,0 +1,378 @@ +"""Exact references for reaction systems: equilibria, spectra, timescales. + +These are the oracle quantities that a learned Koopman operator can be +checked against: + +* :func:`conservation_laws` returns an orthonormal basis of the left + nullspace of the stoichiometric matrix: every ``w`` with ``wᵀ N = 0`` + gives a conserved total ``wᵀ c``. This covers element balance and any + extra stoichiometric conservation. +* :func:`equilibrium` finds a steady state in the stoichiometric + compatibility class of an initial state, ``{c ≥ 0 : W c = W c0}``. It is + exact for linear systems and uses damped Newton otherwise. +* :func:`oracle_spectrum` gives the Jacobian at a state, its eigenvalues, + and relaxation timescales ``-1 / Re(λ)``. It drops the + ``num_species - rank(N)`` zero modes that conservation laws force. + +Pure PyTorch; importing this module does not import scipy or cantera. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from torch import Tensor + +from koopman_graph.datasets.reaction.integrate import ( + linear_rate_equations, + mass_action_jacobian, +) +from koopman_graph.datasets.reaction.system import ReactionSystem + + +@dataclass(frozen=True) +class OracleSpectrum: + """Linearization of a reaction system at one state. + + Attributes + ---------- + state : Tensor + ``float64`` state ``(num_species,)`` the Jacobian is evaluated at. + jacobian : Tensor + ``float64`` Jacobian ``∂(dc/dt)/∂c`` with shape + ``(num_species, num_species)``. + dynamic_eigenvalues : Tensor + ``complex128`` eigenvalues after removing conservation zero modes, + sorted by real part, largest (slowest-decaying) first. + conservation_eigenvalues : Tensor + ``complex128`` eigenvalues removed as conservation zero modes (the + ``num_conservation_laws`` eigenvalues of smallest modulus). + num_conservation_laws : int + ``num_species - rank(N)``. + """ + + state: Tensor + jacobian: Tensor + dynamic_eigenvalues: Tensor + conservation_eigenvalues: Tensor + num_conservation_laws: int + + @property + def timescales(self) -> Tensor: + """Return ``-1 / Re(λ)`` for each dynamic eigenvalue. + + Returns + ------- + Tensor + ``float64`` tensor aligned with :attr:`dynamic_eigenvalues`: + positive for decaying modes, ``inf`` when ``Re(λ) = 0`` (e.g. a + neutral centre), and negative for growing modes. + """ + real = self.dynamic_eigenvalues.real + return torch.where( + real == 0, + torch.full_like(real, math.inf), + -1.0 / real, + ) + + @property + def is_stable(self) -> bool: + """Return whether every dynamic mode decays. + + Returns + ------- + bool + ``True`` when all dynamic eigenvalues have ``Re(λ) < 0``. + """ + return bool((self.dynamic_eigenvalues.real < 0).all()) + + def discrete_eigenvalues(self, time_step: float) -> Tensor: + """Return the eigenvalues of the flow map over ``time_step``. + + Parameters + ---------- + time_step : float + Sampling interval ``Δt`` (finite, > 0). + + Returns + ------- + Tensor + ``complex128`` values ``exp(λ Δt)`` aligned with + :attr:`dynamic_eigenvalues`, the ones to compare with a + discrete-time learned Koopman operator. + + Raises + ------ + ValueError + If ``time_step`` is not finite and positive. + """ + dt = float(time_step) + if not math.isfinite(dt) or dt <= 0.0: + msg = f"time_step must be finite and > 0, got {time_step!r}" + raise ValueError(msg) + return torch.exp(self.dynamic_eigenvalues * dt) + + +def stoichiometric_rank(system: ReactionSystem) -> int: + """Return the rank of the stoichiometric matrix ``N``. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + + Returns + ------- + int + Numerical rank of ``N`` (``num_species × num_reactions``). + """ + _require_system(system) + return int(torch.linalg.matrix_rank(system.stoichiometric_matrix).item()) + + +def conservation_laws(system: ReactionSystem) -> Tensor: + """Return an orthonormal basis of conserved linear combinations. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + + Returns + ------- + Tensor + ``float64`` matrix ``W`` with shape ``(num_species - rank(N), + num_species)`` and orthonormal rows, ``W N = 0``. It has zero rows + when nothing is conserved (e.g. an open system with independent + sources). + """ + rank = stoichiometric_rank(system) + u, _s, _vh = torch.linalg.svd(system.stoichiometric_matrix, full_matrices=True) + return u[:, rank:].T.contiguous() + + +def equilibrium( + system: ReactionSystem, + initial: Tensor, + *, + guess: Tensor | None = None, + tol: float = 1e-12, + max_iter: int = 100, +) -> Tensor: + """Return a steady state in the compatibility class of ``initial``. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + initial : Tensor + Floating state ``(num_species,)`` that fixes the conserved totals + ``W c = W initial``. + guess : Tensor or None, optional + Newton starting point for nonlinear systems. Default is + ``initial``. Systems with several steady states (e.g. Lotka-Volterra, + whose origin is also a fixed point) return the one Newton reaches + from here. + tol : float, optional + Relative step tolerance for Newton convergence. Default is + ``1e-12``. + max_iter : int, optional + Newton iteration cap. Default is ``100``. + + Returns + ------- + Tensor + ``float64`` steady state ``c*`` with ``dc/dt(c*) ≈ 0`` and + ``W c* = W initial``. Nonlinear solves are kept non-negative. + + Raises + ------ + ValueError + If a linear system has no steady state in the class (e.g. a feed + accumulating into a species with no outflow), or arguments are + invalid. + RuntimeError + If Newton does not converge within ``max_iter`` iterations. + """ + c0 = _state(system, initial, "initial") + if not math.isfinite(float(tol)) or float(tol) <= 0.0: + msg = f"tol must be finite and > 0, got {tol!r}" + raise ValueError(msg) + if isinstance(max_iter, bool) or not isinstance(max_iter, int) or max_iter < 1: + msg = f"max_iter must be a positive int, got {max_iter!r}" + raise ValueError(msg) + w = conservation_laws(system) + if system.is_linear: + return _linear_equilibrium(system, c0, w) + + rank = int(system.num_species - w.shape[0]) + u, _s, _vh = torch.linalg.svd(system.stoichiometric_matrix, full_matrices=True) + q = u[:, :rank] # orthonormal basis of range(N); dc/dt always lies in it + targets = w @ c0 + + def residual(c: Tensor) -> Tensor: + """Stack the projected rate equations and conservation constraints. + + Parameters + ---------- + c + See signature. + + Returns + ------- + See signature.""" + return torch.cat([q.T @ system.time_derivative(c), w @ c - targets]) + + start = c0 if guess is None else _state(system, guess, "guess") + c = start.clamp(min=0) + f = residual(c) + for _ in range(max_iter): + jac = torch.cat([q.T @ mass_action_jacobian(system, c), w]) + step = torch.linalg.lstsq(jac, -f.unsqueeze(-1)).solution.squeeze(-1) + alpha = 1.0 + norm = float(f.norm()) + for _ in range(40): + candidate = (c + alpha * step).clamp(min=0) + f_candidate = residual(candidate) + if float(f_candidate.norm()) < norm or norm == 0.0: + break + alpha *= 0.5 + moved = float((candidate - c).abs().max()) + c, f = candidate, f_candidate + if moved <= tol * (1.0 + float(c.abs().max())): + return c + msg = ( + f"Newton did not converge for {system.name!r} in {max_iter} iterations " + f"(residual norm {float(f.norm()):.3e})" + ) + raise RuntimeError(msg) + + +def oracle_spectrum(system: ReactionSystem, state: Tensor) -> OracleSpectrum: + """Linearize the mass-action field at ``state`` and split its spectrum. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + state : Tensor + Floating state ``(num_species,)``, typically an :func:`equilibrium`. + + Returns + ------- + OracleSpectrum + Jacobian, dynamic and conservation eigenvalues, and timescales. + + Notes + ----- + ``J = N ∂r/∂c`` has rank at most ``rank(N)``, so at least + ``num_species - rank(N)`` eigenvalues are exactly zero. These are the + ones removed. Extra zeros (e.g. at a boundary state) stay among the + dynamic eigenvalues and show up as ``inf`` timescales. + """ + c = _state(system, state, "state") + jac = mass_action_jacobian(system, c) + eigenvalues = torch.linalg.eigvals(jac) + num_laws = system.num_species - stoichiometric_rank(system) + by_modulus = torch.argsort(eigenvalues.abs(), stable=True) + conservation = eigenvalues[by_modulus[:num_laws]] + dynamic = eigenvalues[by_modulus[num_laws:]] + dynamic = dynamic[torch.argsort(dynamic.real, descending=True, stable=True)] + return OracleSpectrum( + state=c, + jacobian=jac, + dynamic_eigenvalues=dynamic, + conservation_eigenvalues=conservation, + num_conservation_laws=num_laws, + ) + + +def oracle_timescales(system: ReactionSystem, state: Tensor) -> Tensor: + """Return relaxation timescales at ``state``, slowest first. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + state : Tensor + Floating state ``(num_species,)``. + + Returns + ------- + Tensor + ``OracleSpectrum.timescales`` of :func:`oracle_spectrum`. + """ + return oracle_spectrum(system, state).timescales + + +def _linear_equilibrium(system: ReactionSystem, c0: Tensor, w: Tensor) -> Tensor: + """Solve ``A c + b = 0`` with ``W c = W c0`` for a linear system. + + Parameters + ---------- + system + See signature. + c0 + See signature. + w + See signature. + + Returns + ------- + See signature.""" + a, b = linear_rate_equations(system) + lhs = torch.cat([a, w]) + rhs = torch.cat([-b, w @ c0]) + solution = torch.linalg.lstsq(lhs, rhs.unsqueeze(-1)).solution.squeeze(-1) + scale = 1.0 + float(rhs.abs().max()) + if float((lhs @ solution - rhs).abs().max()) > 1e-9 * scale: + msg = ( + f"{system.name!r} has no steady state for this initial state " + "(a source feeds a species with no outflow)" + ) + raise ValueError(msg) + return solution + + +def _state(system: ReactionSystem, value: Tensor, name: str) -> Tensor: + """Validate one finite state ``(num_species,)`` and return it as float64. + + Parameters + ---------- + system + See signature. + value + See signature. + name + See signature. + + Returns + ------- + See signature.""" + _require_system(system) + if not isinstance(value, Tensor) or not value.is_floating_point(): + msg = f"{name} must be a floating-point Tensor" + raise TypeError(msg) + if value.shape != (system.num_species,): + msg = ( + f"{name} must have shape ({system.num_species},), got {tuple(value.shape)}" + ) + raise ValueError(msg) + if not bool(torch.isfinite(value).all()): + msg = f"{name} must be finite" + raise ValueError(msg) + return value.detach().to(dtype=torch.float64, device="cpu") + + +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/tests/datasets/test_reaction_catalog.py b/tests/datasets/test_reaction_catalog.py new file mode 100644 index 0000000..9bfb435 --- /dev/null +++ b/tests/datasets/test_reaction_catalog.py @@ -0,0 +1,108 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.catalog`.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + DEFAULT_BRUSSELATOR_PARAMS, + DEFAULT_DIMERIZATION_RATES, + DEFAULT_LOTKA_VOLTERRA_RATES, + DEFAULT_MICHAELIS_MENTEN_RATES, + ReactionSystem, + brusselator, + check_element_balance, + first_order_chain, + integrate_rk4, + lotka_volterra, + michaelis_menten, + reversible_dimerization, +) + +_CLOSED = [first_order_chain, reversible_dimerization, michaelis_menten] +_OPEN = [brusselator, lotka_volterra] + + +@pytest.mark.parametrize("factory", _CLOSED + _OPEN) +def test_catalog_returns_named_validated_systems(factory: object) -> None: + system = factory() # type: ignore[operator] + assert isinstance(system, ReactionSystem) + assert system.reaction_names is not None + assert len(system.reaction_names) == system.num_reactions + assert system.rate_constants.dtype == torch.float64 + # Frozen defaults: two calls give the same system. + assert system.to_dict() == factory().to_dict() # type: ignore[operator] + + +@pytest.mark.parametrize("factory", _CLOSED) +def test_closed_systems_are_element_balanced(factory: object) -> None: + system = factory() # type: ignore[operator] + report = check_element_balance(system) + assert report.balanced, report.unbalanced_reactions + + +@pytest.mark.parametrize("factory", _OPEN) +def test_open_systems_have_no_element_matrix(factory: object) -> None: + assert factory().element_matrix is None # type: ignore[operator] + + +def test_first_order_chain_defaults_and_options() -> None: + system = first_order_chain() + assert system.species == ("S0", "S1", "S2") + assert system.rate_constants.tolist() == [1.0, 0.5] + assert system.is_linear + assert system.reaction_names == ("S0 -> S1", "S1 -> S2") + long = first_order_chain(5, rate_constants=(4.0, 3.0, 2.0, 1.0)) + assert long.num_species == 5 + assert long.rate_constants.tolist() == [4.0, 3.0, 2.0, 1.0] + assert long.name == "first_order_chain_5" + for bad in (1, 65, 2.5, True): + with pytest.raises(ValueError, match="num_species must be an int"): + first_order_chain(bad) # type: ignore[arg-type] + with pytest.raises(ValueError, match="must have 2 entries"): + first_order_chain(3, rate_constants=(1.0,)) + + +def test_default_parameters_are_used() -> None: + assert reversible_dimerization().rate_constants.tolist() == list( + DEFAULT_DIMERIZATION_RATES + ) + assert michaelis_menten().rate_constants.tolist() == list( + DEFAULT_MICHAELIS_MENTEN_RATES + ) + a, b = DEFAULT_BRUSSELATOR_PARAMS + assert brusselator().rate_constants.tolist() == [a, b, 1.0, 1.0] + assert b > 1.0 + a * a # default Brusselator is in the oscillatory regime + assert lotka_volterra().rate_constants.tolist() == list( + DEFAULT_LOTKA_VOLTERRA_RATES + ) + assert not michaelis_menten().is_linear + + +@pytest.mark.parametrize(("a", "b"), [(0.0, 1.0), (1.0, -2.0), (math.inf, 1.0)]) +def test_brusselator_rejects_bad_params(a: float, b: float) -> None: + with pytest.raises(ValueError, match="must be finite and > 0"): + brusselator(a, b) + + +def test_lotka_volterra_orbit_is_closed() -> None: + """Conserved quantity ``V = βx - γ ln x + βy - α ln y`` along RK4 orbits.""" + system = lotka_volterra() + c0 = torch.tensor([2.0, 1.0], dtype=torch.float64) + states = integrate_rk4(system, c0, torch.linspace(0.0, 10.0, 101), substeps=20) + x, y = states[:, 0], states[:, 1] + invariant = x - torch.log(x) + y - torch.log(y) + assert float((invariant - invariant[0]).abs().max()) < 1e-6 + + +def test_michaelis_menten_converts_all_substrate() -> None: + system = michaelis_menten() + c0 = torch.tensor([0.1, 1.0, 0.0, 0.0], dtype=torch.float64) + states = integrate_rk4(system, c0, torch.tensor([0.0, 80.0]), substeps=4000) + final = states[-1] + assert torch.allclose( + final, torch.tensor([0.1, 0.0, 0.0, 1.0], dtype=torch.float64), atol=1e-6 + ) diff --git a/tests/datasets/test_reaction_oracles.py b/tests/datasets/test_reaction_oracles.py new file mode 100644 index 0000000..16a89ad --- /dev/null +++ b/tests/datasets/test_reaction_oracles.py @@ -0,0 +1,187 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.oracles`.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + ReactionSystem, + brusselator, + conservation_laws, + equilibrium, + first_order_chain, + integrate_rk4, + lotka_volterra, + michaelis_menten, + oracle_spectrum, + oracle_timescales, + reversible_dimerization, + stoichiometric_rank, +) + +_F64 = torch.float64 + + +def _fed_chain(*, outflow: bool) -> ReactionSystem: + """``∅ → A → B`` (feed 0.5), optionally with ``B → ∅`` (rate 0.25).""" + alpha = [[0.0, 0.0], [1.0, 0.0]] + beta = [[1.0, 0.0], [0.0, 1.0]] + rates = [0.5, 1.0] + if outflow: + alpha.append([0.0, 1.0]) + beta.append([0.0, 0.0]) + rates.append(0.25) + return ReactionSystem( + species=("A", "B"), + reactant_stoichiometry=torch.tensor(alpha, dtype=_F64), + product_stoichiometry=torch.tensor(beta, dtype=_F64), + rate_constants=torch.tensor(rates, dtype=_F64), + ) + + +def test_conservation_laws_span_left_nullspace() -> None: + system = michaelis_menten() + w = conservation_laws(system) + assert stoichiometric_rank(system) == 2 + assert w.shape == (2, 4) + assert torch.allclose( + w @ system.stoichiometric_matrix, torch.zeros(2, 3, dtype=_F64), atol=1e-12 + ) + assert torch.allclose(w @ w.T, torch.eye(2, dtype=_F64), atol=1e-12) + # The element matrix rows lie in the same space. + assert system.element_matrix is not None + projected = system.element_matrix @ w.T @ w + assert torch.allclose(projected, system.element_matrix, atol=1e-12) + assert conservation_laws(brusselator()).shape == (0, 2) + + +def test_chain_spectrum_is_closed_form() -> None: + system = first_order_chain(4, rate_constants=(2.0, 1.0, 0.25)) + c0 = torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=_F64) + c_eq = equilibrium(system, c0) + assert torch.allclose( + c_eq, torch.tensor([0.0, 0.0, 0.0, 1.0], dtype=_F64), atol=1e-12 + ) + spectrum = oracle_spectrum(system, c_eq) + assert spectrum.num_conservation_laws == 1 + assert torch.allclose( + spectrum.dynamic_eigenvalues.real, torch.tensor([-0.25, -1.0, -2.0], dtype=_F64) + ) + assert torch.allclose(spectrum.dynamic_eigenvalues.imag, torch.zeros(3, dtype=_F64)) + assert spectrum.conservation_eigenvalues.abs().max() < 1e-12 + assert torch.allclose( + spectrum.timescales, torch.tensor([4.0, 1.0, 0.5], dtype=_F64) + ) + assert spectrum.is_stable + discrete = spectrum.discrete_eigenvalues(0.1) + assert torch.allclose( + discrete.real, torch.exp(torch.tensor([-0.025, -0.1, -0.2], dtype=_F64)) + ) + assert torch.allclose(oracle_timescales(system, c_eq), spectrum.timescales) + + +def test_linear_open_equilibrium_and_missing_steady_state() -> None: + system = _fed_chain(outflow=True) + c_eq = equilibrium(system, torch.zeros(2, dtype=_F64)) + # A* = feed / k1 = 0.5, B* = k1 A* / k_out = 2.0. + assert torch.allclose(c_eq, torch.tensor([0.5, 2.0], dtype=_F64)) + with pytest.raises(ValueError, match="no steady state"): + equilibrium(_fed_chain(outflow=False), torch.zeros(2, dtype=_F64)) + + +def test_dimerization_equilibrium_matches_closed_form_and_rk4() -> None: + system = reversible_dimerization() + c0 = torch.tensor([1.0, 0.0], dtype=_F64) + c_eq = equilibrium(system, c0) + # A + 2 A2 = 1 and 3 A^2 = 0.5 A2 => A = 1/4, A2 = 3/8. + assert torch.allclose(c_eq, torch.tensor([0.25, 0.375], dtype=_F64), atol=1e-12) + assert float(system.time_derivative(c_eq).abs().max()) < 1e-12 + spectrum = oracle_spectrum(system, c_eq) + # Linearized along the reaction extent: -(4 k_f A* + k_r) = -3.5. + assert torch.allclose( + spectrum.dynamic_eigenvalues, torch.tensor([-3.5 + 0j], dtype=torch.complex128) + ) + late = integrate_rk4(system, c0, torch.tensor([0.0, 20.0]), substeps=400)[-1] + assert torch.allclose(late, c_eq, atol=1e-10) + + +def test_michaelis_menten_boundary_equilibrium_spectrum() -> None: + system = michaelis_menten() + c0 = torch.tensor([0.1, 1.0, 0.0, 0.0], dtype=_F64) + c_eq = equilibrium(system, c0) + assert torch.allclose( + c_eq, torch.tensor([0.1, 0.0, 0.0, 1.0], dtype=_F64), atol=1e-12 + ) + spectrum = oracle_spectrum(system, c_eq) + assert spectrum.num_conservation_laws == 2 + # Free-substrate block [[-k_on E0, k_off], [k_on E0, -(k_off + k_cat)]]. + trace, det = -2.5, 0.5 + root = math.sqrt(trace * trace - 4 * det) + expected = torch.tensor([(trace + root) / 2, (trace - root) / 2], dtype=_F64) + assert torch.allclose(spectrum.dynamic_eigenvalues.real, expected) + late = integrate_rk4(system, c0, torch.tensor([0.0, 80.0]), substeps=4000)[-1] + assert torch.allclose(late, c_eq, atol=1e-6) + + +def test_brusselator_unstable_focus() -> None: + system = brusselator() + c_eq = equilibrium(system, torch.tensor([0.5, 0.5], dtype=_F64)) + assert torch.allclose(c_eq, torch.tensor([1.0, 3.0], dtype=_F64), atol=1e-12) + spectrum = oracle_spectrum(system, c_eq) + assert spectrum.num_conservation_laws == 0 + assert spectrum.conservation_eigenvalues.numel() == 0 + assert torch.allclose( + spectrum.dynamic_eigenvalues.real, torch.full((2,), 0.5, dtype=_F64) + ) + assert torch.allclose( + spectrum.dynamic_eigenvalues.imag.abs(), + torch.full((2,), math.sqrt(3) / 2, dtype=_F64), + ) + assert not spectrum.is_stable + assert torch.allclose(spectrum.timescales, torch.full((2,), -2.0, dtype=_F64)) + + +def test_lotka_volterra_centre_has_infinite_timescale() -> None: + system = lotka_volterra() + c_eq = equilibrium( + system, + torch.tensor([2.0, 2.0], dtype=_F64), + guess=torch.tensor([1.5, 0.8], dtype=_F64), + ) + assert torch.allclose(c_eq, torch.ones(2, dtype=_F64), atol=1e-12) + spectrum = oracle_spectrum(system, c_eq) + assert torch.allclose(spectrum.dynamic_eigenvalues.real, torch.zeros(2, dtype=_F64)) + assert torch.all(torch.isinf(spectrum.timescales)) + assert not spectrum.is_stable + discrete = spectrum.discrete_eigenvalues(0.5) + assert torch.allclose(discrete.abs(), torch.ones(2, dtype=_F64)) + + +def test_equilibrium_validation() -> None: + system = reversible_dimerization() + c0 = torch.tensor([1.0, 0.0], dtype=_F64) + with pytest.raises(TypeError, match="floating-point Tensor"): + equilibrium(system, [1.0, 0.0]) # type: ignore[arg-type] + with pytest.raises(ValueError, match=r"shape \(2,\)"): + equilibrium(system, torch.ones(3, dtype=_F64)) + with pytest.raises(ValueError, match="must be finite"): + equilibrium(system, torch.tensor([math.nan, 0.0])) + with pytest.raises(ValueError, match=r"guess must have shape"): + equilibrium(system, c0, guess=torch.ones(1, dtype=_F64)) + for tol in (0.0, math.nan): + with pytest.raises(ValueError, match="tol must be"): + equilibrium(system, c0, tol=tol) + for max_iter in (0, 1.5, True): + with pytest.raises(ValueError, match="max_iter must be"): + equilibrium(system, c0, max_iter=max_iter) # type: ignore[arg-type] + with pytest.raises(RuntimeError, match="did not converge"): + equilibrium(system, c0, max_iter=1) + with pytest.raises(TypeError, match="must be a ReactionSystem"): + oracle_spectrum("dimer", c0) # type: ignore[arg-type] + spectrum = oracle_spectrum(system, c0) + for dt in (0.0, -1.0, math.inf): + with pytest.raises(ValueError, match="time_step must be"): + spectrum.discrete_eigenvalues(dt) From 0df08cb5dc95da23510f51efa1ba0a3f586050bc Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 15:43:39 -0400 Subject: [PATCH 2/2] Harden reaction oracles from review - Linear equilibrium returns c_p + P0 (c0 - c_p), the t -> inf limit, so non-unique steady states (A -> B, A -> C) split by rate instead of falling back to lstsq's minimum-norm point; a defective zero mode raises. Rounding-level values are snapped to zero instead of going negative. - Newton also requires a small residual to converge and raises when the line search stalls, instead of reporting a tiny clamped step as a root. - OracleSpectrum gains zero_tolerance (1e-10 ||J||_F from oracle_spectrum) so near-zero real parts read as neutral in timescales and is_stable. - One SVD of N gives rank, range basis and conservation laws. - Every catalog constructor rejects non-finite or non-positive rates. Co-Authored-By: Claude Opus 5.5 --- .../datasets/reaction/catalog.py | 46 +++++- .../datasets/reaction/oracles.py | 149 ++++++++++++++---- tests/datasets/test_reaction_catalog.py | 14 ++ tests/datasets/test_reaction_oracles.py | 60 +++++++ 4 files changed, 230 insertions(+), 39 deletions(-) diff --git a/src/koopman_graph/datasets/reaction/catalog.py b/src/koopman_graph/datasets/reaction/catalog.py index 1dfc3e3..b82278b 100644 --- a/src/koopman_graph/datasets/reaction/catalog.py +++ b/src/koopman_graph/datasets/reaction/catalog.py @@ -71,7 +71,8 @@ def first_order_chain( Raises ------ ValueError - If ``num_species`` is out of range or the rate count is wrong. + If ``num_species`` is out of range, the rate count is wrong, or a + rate is not finite and positive. """ if ( isinstance(num_species, bool) @@ -96,6 +97,9 @@ def first_order_chain( f"{num_species} species, got {len(rates)}" ) raise ValueError(msg) + _require_positive( + *((f"rate_constants[{i}]", k) for i, k in enumerate(rate_constants)) + ) reactants = torch.zeros(num_reactions, num_species, dtype=torch.float64) products = torch.zeros(num_reactions, num_species, dtype=torch.float64) index = torch.arange(num_reactions) @@ -135,7 +139,13 @@ def reversible_dimerization( Closed, nonlinear system; element ``X`` counts monomer units (``A: 1``, ``A2: 2``). Equilibrium satisfies ``k_forward A² = k_reverse A2``. + + Raises + ------ + ValueError + If a rate constant is not finite and positive. """ + _require_positive(("k_forward", k_forward), ("k_reverse", k_reverse)) return ReactionSystem( species=("A", "A2"), reactant_stoichiometry=torch.tensor( @@ -174,7 +184,13 @@ def michaelis_menten( Closed system over species ``(E, S, ES, P)`` with elements ``Enz`` (enzyme units) and ``Sub`` (substrate backbone). All substrate ends up as ``P``, a boundary equilibrium with ``S = ES = 0``. + + Raises + ------ + ValueError + If a rate constant is not finite and positive. """ + _require_positive(("k_on", k_on), ("k_off", k_off), ("k_cat", k_cat)) return ReactionSystem( species=("E", "S", "ES", "P"), reactant_stoichiometry=torch.tensor( @@ -223,10 +239,7 @@ def brusselator( ValueError If ``a`` or ``b`` is not finite and positive. """ - for name, value in (("a", a), ("b", b)): - if not math.isfinite(float(value)) or float(value) <= 0.0: - msg = f"{name} must be finite and > 0, got {value!r}" - raise ValueError(msg) + _require_positive(("a", a), ("b", b)) return ReactionSystem( species=("X", "Y"), reactant_stoichiometry=torch.tensor( @@ -264,7 +277,13 @@ def lotka_volterra( ``(gamma / beta, alpha / beta)`` is a neutral centre with Jacobian eigenvalues ``± i sqrt(alpha gamma)``, so trajectories are closed orbits and the oracle timescale is infinite. + + Raises + ------ + ValueError + If a rate constant is not finite and positive. """ + _require_positive(("alpha", alpha), ("beta", beta), ("gamma", gamma)) return ReactionSystem( species=("X", "Y"), reactant_stoichiometry=torch.tensor( @@ -277,3 +296,20 @@ def lotka_volterra( reaction_names=("X -> 2 X", "X + Y -> 2 Y", "Y ->"), name="lotka_volterra", ) + + +def _require_positive(*params: tuple[str, float]) -> None: + """Raise unless every ``(name, value)`` is finite and positive. + + Parameters + ---------- + *params + See signature. + + Raises + ------ + See signature.""" + for name, value in params: + if not math.isfinite(float(value)) or float(value) <= 0.0: + msg = f"{name} must be finite and > 0, got {value!r}" + raise ValueError(msg) diff --git a/src/koopman_graph/datasets/reaction/oracles.py b/src/koopman_graph/datasets/reaction/oracles.py index 58c8044..ac11162 100644 --- a/src/koopman_graph/datasets/reaction/oracles.py +++ b/src/koopman_graph/datasets/reaction/oracles.py @@ -31,6 +31,11 @@ ) from koopman_graph.datasets.reaction.system import ReactionSystem +_ZERO_REAL_RTOL = 1e-10 +"""Relative tolerance, in units of ``‖J‖``, below which ``Re(λ)`` counts as 0.""" + +_MAX_BACKTRACKS = 40 + @dataclass(frozen=True) class OracleSpectrum: @@ -51,6 +56,11 @@ class OracleSpectrum: ``num_conservation_laws`` eigenvalues of smallest modulus). num_conservation_laws : int ``num_species - rank(N)``. + zero_tolerance : float + Absolute tolerance below which ``|Re(λ)|`` is treated as zero, so + rounding in ``eigvals`` cannot turn a neutral mode into a growing or + decaying one. Default is ``0.0`` (exact comparison); + :func:`oracle_spectrum` sets it relative to ``‖J‖``. """ state: Tensor @@ -58,6 +68,7 @@ class OracleSpectrum: dynamic_eigenvalues: Tensor conservation_eigenvalues: Tensor num_conservation_laws: int + zero_tolerance: float = 0.0 @property def timescales(self) -> Tensor: @@ -67,12 +78,13 @@ def timescales(self) -> Tensor: ------- Tensor ``float64`` tensor aligned with :attr:`dynamic_eigenvalues`: - positive for decaying modes, ``inf`` when ``Re(λ) = 0`` (e.g. a - neutral centre), and negative for growing modes. + positive for decaying modes, ``inf`` when ``|Re(λ)|`` is within + :attr:`zero_tolerance` (e.g. a neutral centre), and negative for + growing modes. """ real = self.dynamic_eigenvalues.real return torch.where( - real == 0, + real.abs() <= self.zero_tolerance, torch.full_like(real, math.inf), -1.0 / real, ) @@ -84,9 +96,10 @@ def is_stable(self) -> bool: Returns ------- bool - ``True`` when all dynamic eigenvalues have ``Re(λ) < 0``. + ``True`` when all dynamic eigenvalues have + ``Re(λ) < -zero_tolerance``. """ - return bool((self.dynamic_eigenvalues.real < 0).all()) + return bool((self.dynamic_eigenvalues.real < -self.zero_tolerance).all()) def discrete_eigenvalues(self, time_step: float) -> Tensor: """Return the eigenvalues of the flow map over ``time_step``. @@ -128,8 +141,7 @@ def stoichiometric_rank(system: ReactionSystem) -> int: int Numerical rank of ``N`` (``num_species × num_reactions``). """ - _require_system(system) - return int(torch.linalg.matrix_rank(system.stoichiometric_matrix).item()) + return _stoichiometric_basis(system)[0] def conservation_laws(system: ReactionSystem) -> Tensor: @@ -148,9 +160,7 @@ def conservation_laws(system: ReactionSystem) -> Tensor: when nothing is conserved (e.g. an open system with independent sources). """ - rank = stoichiometric_rank(system) - u, _s, _vh = torch.linalg.svd(system.stoichiometric_matrix, full_matrices=True) - return u[:, rank:].T.contiguous() + return _stoichiometric_basis(system)[2] def equilibrium( @@ -185,16 +195,23 @@ def equilibrium( ------- Tensor ``float64`` steady state ``c*`` with ``dc/dt(c*) ≈ 0`` and - ``W c* = W initial``. Nonlinear solves are kept non-negative. + ``W c* = W initial``. For a linear system this is + ``c_p + P0 (initial - c_p)``, where ``c_p`` is any steady state and + ``P0`` the spectral projector onto the kernel of ``A``. That is the + ``t → ∞`` limit from ``initial`` whenever the nonzero modes decay, + also when the steady state in the class is not unique (e.g. + ``A → B`` and ``A → C`` split ``A`` in the ratio of their rates). + Entries within rounding of zero are set to zero. Raises ------ ValueError - If a linear system has no steady state in the class (e.g. a feed - accumulating into a species with no outflow), or arguments are - invalid. + If a linear system has no steady state (e.g. a feed accumulating + into a species with no outflow) or no limiting one (a defective + zero eigenvalue), or arguments are invalid. RuntimeError - If Newton does not converge within ``max_iter`` iterations. + If Newton does not reach a steady state within ``max_iter`` + iterations, or its line search stalls away from one. """ c0 = _state(system, initial, "initial") if not math.isfinite(float(tol)) or float(tol) <= 0.0: @@ -203,14 +220,15 @@ def equilibrium( if isinstance(max_iter, bool) or not isinstance(max_iter, int) or max_iter < 1: msg = f"max_iter must be a positive int, got {max_iter!r}" raise ValueError(msg) - w = conservation_laws(system) if system.is_linear: - return _linear_equilibrium(system, c0, w) + return _linear_equilibrium(system, c0) - rank = int(system.num_species - w.shape[0]) - u, _s, _vh = torch.linalg.svd(system.stoichiometric_matrix, full_matrices=True) - q = u[:, :rank] # orthonormal basis of range(N); dc/dt always lies in it + # q is an orthonormal basis of range(N); dc/dt always lies in it. + _rank, q, w = _stoichiometric_basis(system) targets = w @ c0 + # Step convergence alone also fires when the line search or the c >= 0 + # clamp stalls, so the residual must be small as well. + residual_tol = math.sqrt(float(tol)) * (1.0 + float(c0.abs().max())) def residual(c: Tensor) -> Tensor: """Stack the projected rate equations and conservation constraints. @@ -233,15 +251,26 @@ def residual(c: Tensor) -> Tensor: step = torch.linalg.lstsq(jac, -f.unsqueeze(-1)).solution.squeeze(-1) alpha = 1.0 norm = float(f.norm()) - for _ in range(40): + for _ in range(_MAX_BACKTRACKS): candidate = (c + alpha * step).clamp(min=0) f_candidate = residual(candidate) - if float(f_candidate.norm()) < norm or norm == 0.0: + if float(f_candidate.norm()) < norm: break alpha *= 0.5 + else: + if norm <= residual_tol: + return c + msg = ( + f"Newton line search stalled for {system.name!r} " + f"(residual norm {norm:.3e})" + ) + raise RuntimeError(msg) moved = float((candidate - c).abs().max()) c, f = candidate, f_candidate - if moved <= tol * (1.0 + float(c.abs().max())): + if ( + moved <= tol * (1.0 + float(c.abs().max())) + and float(f.norm()) <= residual_tol + ): return c msg = ( f"Newton did not converge for {system.name!r} in {max_iter} iterations " @@ -270,7 +299,8 @@ def oracle_spectrum(system: ReactionSystem, state: Tensor) -> OracleSpectrum: ``J = N ∂r/∂c`` has rank at most ``rank(N)``, so at least ``num_species - rank(N)`` eigenvalues are exactly zero. These are the ones removed. Extra zeros (e.g. at a boundary state) stay among the - dynamic eigenvalues and show up as ``inf`` timescales. + dynamic eigenvalues and show up as ``inf`` timescales. Real parts within + ``1e-10 ‖J‖_F`` of zero count as zero. """ c = _state(system, state, "state") jac = mass_action_jacobian(system, c) @@ -286,6 +316,7 @@ def oracle_spectrum(system: ReactionSystem, state: Tensor) -> OracleSpectrum: dynamic_eigenvalues=dynamic, conservation_eigenvalues=conservation, num_conservation_laws=num_laws, + zero_tolerance=_ZERO_REAL_RTOL * float(torch.linalg.matrix_norm(jac)), ) @@ -307,33 +338,83 @@ def oracle_timescales(system: ReactionSystem, state: Tensor) -> Tensor: return oracle_spectrum(system, state).timescales -def _linear_equilibrium(system: ReactionSystem, c0: Tensor, w: Tensor) -> Tensor: - """Solve ``A c + b = 0`` with ``W c = W c0`` for a linear system. +def _stoichiometric_basis(system: ReactionSystem) -> tuple[int, Tensor, Tensor]: + """Return ``rank(N)`` and orthonormal bases of ``range(N)`` and its complement. Parameters ---------- system See signature. - c0 + + Returns + ------- + ``(rank, q, w)`` with ``q`` of shape ``(num_species, rank)`` and the + conservation laws ``w`` of shape ``(num_species - rank, num_species)``, + from one SVD.""" + _require_system(system) + u, _vh, rank = _svd_rank(system.stoichiometric_matrix) + return rank, u[:, :rank], u[:, rank:].T.contiguous() + + +def _svd_rank(matrix: Tensor) -> tuple[Tensor, Tensor, int]: + """Return the full singular vectors and numerical rank of ``matrix``. + + Parameters + ---------- + matrix See signature. - w + + Returns + ------- + ``(u, vh, rank)``, with the ``torch.linalg.matrix_rank`` default + tolerance ``max(m, n) eps σ_max``.""" + u, singular, vh = torch.linalg.svd(matrix, full_matrices=True) + eps = torch.finfo(matrix.dtype).eps + cutoff = max(matrix.shape) * eps * float(singular.max()) + return u, vh, int((singular > cutoff).sum()) + + +def _linear_equilibrium(system: ReactionSystem, c0: Tensor) -> Tensor: + """Return ``c_p + P0 (c0 - c_p)`` for the linear field ``A c + b``. + + Parameters + ---------- + system + See signature. + c0 See signature. Returns ------- See signature.""" a, b = linear_rate_equations(system) - lhs = torch.cat([a, w]) - rhs = torch.cat([-b, w @ c0]) - solution = torch.linalg.lstsq(lhs, rhs.unsqueeze(-1)).solution.squeeze(-1) - scale = 1.0 + float(rhs.abs().max()) - if float((lhs @ solution - rhs).abs().max()) > 1e-9 * scale: + particular = torch.linalg.lstsq(a, -b.unsqueeze(-1)).solution.squeeze(-1) + scale = 1.0 + max(float(b.abs().max()), float(c0.abs().max())) + if float((a @ particular + b).abs().max()) > 1e-9 * scale: msg = ( f"{system.name!r} has no steady state for this initial state " "(a source feeds a species with no outflow)" ) raise ValueError(msg) - return solution + # P0 = K (L K)^-1 L projects onto ker(A) along range(A), where K and L + # are orthonormal right and left kernel bases. L K is singular exactly + # when the zero eigenvalue is defective and the trajectory never settles. + u, vh, rank = _svd_rank(a) + solution = particular + if rank < system.num_species: + right = vh[rank:].T + left = u[:, rank:].T + coupling = left @ right + if float(torch.linalg.svdvals(coupling).min()) < 1e-8: + msg = ( + f"{system.name!r} has no limiting steady state for this " + "initial state (the zero eigenvalue of its rate matrix is " + "defective)" + ) + raise ValueError(msg) + projector = right @ torch.linalg.solve(coupling, left) + solution = particular + projector @ (c0 - particular) + return torch.where(solution.abs() <= 1e-12 * scale, 0.0, solution) def _state(system: ReactionSystem, value: Tensor, name: str) -> Tensor: diff --git a/tests/datasets/test_reaction_catalog.py b/tests/datasets/test_reaction_catalog.py index 9bfb435..594a80d 100644 --- a/tests/datasets/test_reaction_catalog.py +++ b/tests/datasets/test_reaction_catalog.py @@ -88,6 +88,20 @@ def test_brusselator_rejects_bad_params(a: float, b: float) -> None: brusselator(a, b) +@pytest.mark.parametrize( + ("factory", "kwargs"), + [ + (first_order_chain, {"rate_constants": (1.0, 0.0)}), + (reversible_dimerization, {"k_reverse": 0.0}), + (michaelis_menten, {"k_cat": math.nan}), + (lotka_volterra, {"gamma": -1.0}), + ], +) +def test_catalog_rejects_non_positive_rates(factory: object, kwargs: dict) -> None: + with pytest.raises(ValueError, match="must be finite and > 0"): + factory(**kwargs) # type: ignore[operator] + + def test_lotka_volterra_orbit_is_closed() -> None: """Conserved quantity ``V = βx - γ ln x + βy - α ln y`` along RK4 orbits.""" system = lotka_volterra() diff --git a/tests/datasets/test_reaction_oracles.py b/tests/datasets/test_reaction_oracles.py index 16a89ad..b4e682f 100644 --- a/tests/datasets/test_reaction_oracles.py +++ b/tests/datasets/test_reaction_oracles.py @@ -8,6 +8,7 @@ import torch from koopman_graph.datasets.reaction import ( + OracleSpectrum, ReactionSystem, brusselator, conservation_laws, @@ -42,6 +43,18 @@ def _fed_chain(*, outflow: bool) -> ReactionSystem: ) +def _system( + reactants: list[list[float]], products: list[list[float]], rates: list[float] +) -> ReactionSystem: + """Build a small system from nested lists.""" + return ReactionSystem( + species=tuple(f"S{i}" for i in range(len(reactants[0]))), + reactant_stoichiometry=torch.tensor(reactants, dtype=_F64), + product_stoichiometry=torch.tensor(products, dtype=_F64), + rate_constants=torch.tensor(rates, dtype=_F64), + ) + + def test_conservation_laws_span_left_nullspace() -> None: system = michaelis_menten() w = conservation_laws(system) @@ -92,6 +105,53 @@ def test_linear_open_equilibrium_and_missing_steady_state() -> None: equilibrium(_fed_chain(outflow=False), torch.zeros(2, dtype=_F64)) +def test_linear_equilibrium_with_non_unique_steady_state_is_the_limit() -> None: + # A -> B (1), A -> C (3): every (0, B, 1 - B) is steady; the limit + # splits A in the ratio of the rates. + system = _system([[1.0, 0.0, 0.0]] * 2, [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [1, 3]) + c0 = torch.tensor([1.0, 0.0, 0.0], dtype=_F64) + c_eq = equilibrium(system, c0) + assert c_eq[0] == 0.0 # rounding noise is snapped to zero, not negative + assert torch.allclose(c_eq, torch.tensor([0.0, 0.25, 0.75], dtype=_F64)) + late = integrate_rk4(system, c0, torch.tensor([0.0, 40.0]), substeps=4000)[-1] + assert torch.allclose(late, c_eq, atol=1e-10) + + +def test_linear_equilibrium_rejects_defective_zero_mode() -> None: + # S1 -> S0 + S1: S0 grows linearly while S1 > 0 and never settles. + system = _system([[0.0, 1.0]], [[1.0, 1.0]], [1.0]) + with pytest.raises(ValueError, match="no limiting steady state"): + equilibrium(system, torch.tensor([0.0, 1.0], dtype=_F64)) + + +def test_newton_stall_raises_instead_of_returning_non_root() -> None: + # -> X (1), 2 X -> 3 X (1): dX/dt = 1 + X^2 > 0 has no root. From X = 0 + # the Jacobian vanishes, Newton cannot move, and it must not report + # convergence. + system = _system([[0.0], [2.0]], [[1.0], [3.0]], [1.0, 1.0]) + with pytest.raises(RuntimeError, match="stalled"): + equilibrium(system, torch.zeros(1, dtype=_F64)) + + +def test_near_zero_real_parts_are_neutral() -> None: + eigenvalues = torch.tensor([1e-17 + 1j, -1e-17 - 1j], dtype=torch.complex128) + fields = { + "state": torch.zeros(2, dtype=_F64), + "jacobian": torch.zeros(2, 2, dtype=_F64), + "dynamic_eigenvalues": eigenvalues, + "conservation_eigenvalues": torch.zeros(0, dtype=torch.complex128), + "num_conservation_laws": 0, + } + exact = OracleSpectrum(**fields) + assert torch.isfinite(exact.timescales).all() + tolerant = OracleSpectrum(**fields, zero_tolerance=1e-12) + assert torch.isinf(tolerant.timescales).all() + assert (tolerant.timescales > 0).all() + assert not tolerant.is_stable + spectrum = oracle_spectrum(lotka_volterra(), torch.ones(2, dtype=_F64)) + assert spectrum.zero_tolerance > 0.0 + + def test_dimerization_equilibrium_matches_closed_form_and_rk4() -> None: system = reversible_dimerization() c0 = torch.tensor([1.0, 0.0], dtype=_F64)