From 9bb094f4b011bf9a0ee4705f75305e0d515acae2 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Thu, 15 Oct 2026 15:38:24 -0400 Subject: [PATCH] Add random reaction networks, Van de Vusse CSTR, and Gillespie SSA Refs #65. Fills in koopman_graph.datasets.reaction.random_network, .cstr, and .gillespie: - random_network: sample_random_network builds compositions first (bare elements, then unique unions of existing species), enumerates every element-balanced reaction with order <= max_order and <= 2 products, and draws reactions and log-uniform rates from one seeded generator. SampledReactionNetwork records the seed and parameters. - cstr: VanDeVusseCSTR with Klatt & Engell (1998) Arrhenius kinetics (the 2A -> D pre-exponential is halved so mass action matches the literature balance), temperature and dilution-rate inputs with piecewise-constant schedules, RK4 simulation, and a Newton steady state that reproduces the benchmark operating point (c_A ~ 2.14, c_B ~ 1.09 at 114.2 C). - gillespie: seeded direct-method SSA with mass-action propensities (falling factorials, V^(1-order)) that converge to the ODE as V grows, absorbing-state and max_events handling, and last-value-hold binning. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 17 + .../datasets/reaction/__init__.py | 30 ++ src/koopman_graph/datasets/reaction/cstr.py | 340 ++++++++++++++++++ .../datasets/reaction/gillespie.py | 330 +++++++++++++++++ .../datasets/reaction/random_network.py | 322 +++++++++++++++++ tests/datasets/test_reaction_cstr.py | 125 +++++++ tests/datasets/test_reaction_gillespie.py | 161 +++++++++ .../datasets/test_reaction_random_network.py | 102 ++++++ 8 files changed, 1427 insertions(+) create mode 100644 tests/datasets/test_reaction_cstr.py create mode 100644 tests/datasets/test_reaction_gillespie.py create mode 100644 tests/datasets/test_reaction_random_network.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 101f26f..859b7f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `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. +- `koopman_graph.datasets.reaction.random_network`: `sample_random_network` + builds seeded, element-balanced mass-action networks. Species + compositions are sampled first, so every candidate reaction conserves + elements; reaction order is capped at 2 (`max_order=1` gives linear + dissociation networks); rates are log-uniform. The returned + `SampledReactionNetwork` records the seed and every sampler parameter. +- `koopman_graph.datasets.reaction.cstr`: `van_de_vusse_cstr` / + `VanDeVusseCSTR`, a Van de Vusse CSTR with Klatt & Engell (1998) + Arrhenius kinetics. Reactor temperature and dilution rate are inputs, + either constant or piecewise-constant schedules. It provides the + controlled right-hand side, RK4 simulation, and a Newton steady state; the + nominal point (114.2 °C, `F/V = 14.19 1/h`) gives `c_A ≈ 2.14`, + `c_B ≈ 1.09`. +- `koopman_graph.datasets.reaction.gillespie`: seeded Gillespie direct-method + SSA (`gillespie_ssa`, `SSATrajectory`) with mass-action propensities + (falling factorials, volume scaling) that match the ODE rates, plus + last-value-hold binning onto a grid (`bin_to_grid`). - `[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 4f58cfb..591a47d 100644 --- a/src/koopman_graph/datasets/reaction/__init__.py +++ b/src/koopman_graph/datasets/reaction/__init__.py @@ -51,6 +51,20 @@ michaelis_menten, reversible_dimerization, ) +from koopman_graph.datasets.reaction.cstr import ( + VAN_DE_VUSSE_FEED_A, + VAN_DE_VUSSE_NOMINAL_DILUTION, + VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K, + VanDeVusseCSTR, + van_de_vusse_cstr, +) +from koopman_graph.datasets.reaction.gillespie import ( + DEFAULT_MAX_EVENTS, + SSATrajectory, + bin_to_grid, + gillespie_ssa, + propensities, +) from koopman_graph.datasets.reaction.graphs import ( PRODUCT_RELATION, REACTANT_RELATION, @@ -79,6 +93,10 @@ oracle_timescales, stoichiometric_rank, ) +from koopman_graph.datasets.reaction.random_network import ( + SampledReactionNetwork, + sample_random_network, +) from koopman_graph.datasets.reaction.system import ( GAS_CONSTANT_J_PER_MOL_K, ArrheniusParams, @@ -99,24 +117,33 @@ "DEFAULT_DIMERIZATION_RATES", "DEFAULT_LOG_FLOOR", "DEFAULT_LOTKA_VOLTERRA_RATES", + "DEFAULT_MAX_EVENTS", "DEFAULT_MICHAELIS_MENTEN_RATES", "GAS_CONSTANT_J_PER_MOL_K", "PRODUCT_RELATION", "REACTANT_RELATION", "REACTION_NODE", "SPECIES_NODE", + "VAN_DE_VUSSE_FEED_A", + "VAN_DE_VUSSE_NOMINAL_DILUTION", + "VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K", "ArrheniusParams", "ElementBalanceReport", "IntegrationMethod", "LogConcentrationTransform", "OracleSpectrum", "ReactionSystem", + "SSATrajectory", + "SampledReactionNetwork", + "VanDeVusseCSTR", + "bin_to_grid", "brusselator", "check_element_balance", "conservation_laws", "element_balance_violation", "equilibrium", "first_order_chain", + "gillespie_ssa", "integrate_expm", "integrate_reaction_system", "integrate_rk4", @@ -127,11 +154,14 @@ "michaelis_menten", "oracle_spectrum", "oracle_timescales", + "propensities", "reaction_directed_hyperedges", "reaction_hyperedge_index", "require_scipy", "reversible_dimerization", + "sample_random_network", "species_influence_graph", "species_reaction_bipartite", "stoichiometric_rank", + "van_de_vusse_cstr", ] diff --git a/src/koopman_graph/datasets/reaction/cstr.py b/src/koopman_graph/datasets/reaction/cstr.py index e69de29..546d318 100644 --- a/src/koopman_graph/datasets/reaction/cstr.py +++ b/src/koopman_graph/datasets/reaction/cstr.py @@ -0,0 +1,340 @@ +"""Van de Vusse continuous stirred-tank reactor with feed and temperature inputs. + +The Van de Vusse scheme ``A → B → C`` and ``2 A → D`` is a standard +benchmark for nonlinear process control: the intermediate ``B`` is the +product, and its yield depends non-monotonically on dilution rate and +temperature. This module uses the kinetics of Klatt and Engell (1998), +with time in hours and concentrations in mol/L: + +============= ================== =================== +Reaction ``k0`` ``E_a / R`` (K) +============= ================== =================== +``A → B`` 1.287e12 1/h 9758.3 +``B → C`` 1.287e12 1/h 9758.3 +``2 A → D`` 9.043e9 L/(mol h) 8560.0 +============= ================== =================== + +The literature writes the third rate as ``r3 = k3 c_A²`` with +``dc_A/dt ⊃ -k3 c_A²``. Mass action on ``2 A → D`` consumes two ``A`` per +event, so the mass-action pre-exponential here is ``k0 / 2``, which +reproduces the literature balance exactly. + +The reactor temperature is treated as a direct input (no energy balance). +The controlled field is + +.. math:: + + \\frac{dc}{dt} = N r(c; T) + D (c_\\text{feed} - c), + +with dilution rate ``D = F / V`` (1/h) and feed ``c_feed = (c_A0, 0, 0, +0)``. At ``D = 14.19``, ``T = 387.35 K`` (114.2 °C) and ``c_A0 = 5.10`` the +steady state is ``c_A ≈ 2.14``, ``c_B ≈ 1.09``, the nominal operating point +of the CSTR control benchmark. + +References: K.-U. Klatt and S. Engell, "Gain-scheduling trajectory control +of a continuous stirred tank reactor", *Computers & Chemical Engineering* +22(4-5), 491-502, 1998 (kinetics); H. Chen, A. Kremling and F. Allgöwer, +"Nonlinear predictive control of a benchmark CSTR", *Proc. 3rd European +Control Conference*, 1995 (operating point). + +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 mass_action_jacobian +from koopman_graph.datasets.reaction.system import ( + GAS_CONSTANT_J_PER_MOL_K, + ArrheniusParams, + ReactionSystem, +) + +VAN_DE_VUSSE_FEED_A = 5.10 +"""Nominal feed concentration of ``A`` in mol/L (Klatt and Engell, 1998).""" + +VAN_DE_VUSSE_NOMINAL_DILUTION = 14.19 +"""Nominal dilution rate ``F / V`` in 1/h.""" + +VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K = 387.35 +"""Nominal reactor temperature in K (114.2 °C).""" + +_K0 = (1.287e12, 1.287e12, 9.043e9 / 2.0) +_EA_OVER_R = (9758.3, 9758.3, 8560.0) + + +@dataclass(frozen=True, eq=False) +class VanDeVusseCSTR: + """Van de Vusse CSTR: reaction system plus feed composition. + + Attributes + ---------- + system : ReactionSystem + ``A → B``, ``B → C``, ``2 A → D`` over species ``(A, B, C, D)`` with + Arrhenius parameters; its ``rate_constants`` are those at the + nominal temperature. + feed : Tensor + ``float64`` feed concentrations ``(4,)`` in mol/L. + """ + + system: ReactionSystem + feed: Tensor + + def rhs( + self, + concentrations: Tensor, + *, + temperature_k: float, + dilution_rate: float, + ) -> Tensor: + """Evaluate the controlled field ``N r(c; T) + D (c_feed - c)``. + + Parameters + ---------- + concentrations : Tensor + Floating tensor ``(..., 4)`` in mol/L. + temperature_k : float + Reactor temperature in K (finite, > 0). + dilution_rate : float + ``F / V`` in 1/h (finite, >= 0). + + Returns + ------- + Tensor + ``dc/dt`` in mol/(L h) with the shape and dtype of + ``concentrations``. + """ + k = self.system.rate_constants_at(temperature_k) + d = _dilution(dilution_rate) + reaction = self.system.time_derivative(concentrations, rate_constants=k) + feed = self.feed.to(dtype=reaction.dtype, device=reaction.device) + return reaction + d * (feed - concentrations) + + def simulate( + self, + initial: Tensor, + times: Tensor, + *, + temperature_k: float | Tensor, + dilution_rate: float | Tensor, + substeps: int = 50, + ) -> Tensor: + """Integrate with RK4 under piecewise-constant inputs. + + Parameters + ---------- + initial : Tensor + Floating initial concentrations ``(..., 4)``. + times : Tensor + Strictly increasing 1-D output times in hours. + temperature_k : float or Tensor + Constant temperature, or one value per interval with shape + ``(len(times) - 1,)`` held over ``[times[i], times[i + 1])``. + dilution_rate : float or Tensor + Constant dilution rate, or one value per interval. + substeps : int, optional + RK4 steps per output interval. Default is ``50``. + + Returns + ------- + Tensor + Concentrations ``(len(times), ..., 4)``. + + Raises + ------ + ValueError + If schedules have the wrong length or inputs are invalid. + """ + 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) + if not isinstance(times, Tensor) or times.ndim != 1 or times.numel() < 1: + msg = "times must be a non-empty 1-D Tensor" + raise ValueError(msg) + t = times.to(torch.float64) + if times.numel() > 1 and not bool((t[1:] > t[:-1]).all()): + msg = "times must be strictly increasing" + raise ValueError(msg) + num_intervals = int(times.numel()) - 1 + temps = _schedule(temperature_k, num_intervals, "temperature_k") + dilutions = _schedule(dilution_rate, num_intervals, "dilution_rate") + c = initial + states = [c] + for i in range(num_intervals): + h = float(t[i + 1] - t[i]) / substeps + kwargs = {"temperature_k": temps[i], "dilution_rate": dilutions[i]} + for _ in range(substeps): + k1 = self.rhs(c, **kwargs) + k2 = self.rhs(c + 0.5 * h * k1, **kwargs) + k3 = self.rhs(c + 0.5 * h * k2, **kwargs) + k4 = self.rhs(c + h * k3, **kwargs) + c = c + (h / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) + states.append(c) + return torch.stack(states) + + def steady_state( + self, + *, + temperature_k: float, + dilution_rate: float, + guess: Tensor | None = None, + tol: float = 1e-12, + max_iter: int = 100, + ) -> Tensor: + """Solve ``rhs(c) = 0`` with damped Newton (non-negative iterates). + + Parameters + ---------- + temperature_k : float + Reactor temperature in K. + dilution_rate : float + ``F / V`` in 1/h (> 0 for a unique steady state). + guess : Tensor or None, optional + Starting point ``(4,)``. Default is the feed. + tol : float, optional + Relative step tolerance. Default is ``1e-12``. + max_iter : int, optional + Iteration cap. Default is ``100``. + + Returns + ------- + Tensor + ``float64`` steady-state concentrations ``(4,)``. + + Raises + ------ + RuntimeError + If Newton does not converge. + """ + d = _dilution(dilution_rate) + at_t = self.system.at_temperature(temperature_k) + c = (self.feed if guess is None else guess).to(torch.float64).clone() + eye = torch.eye(4, dtype=torch.float64) + + def residual(state: Tensor) -> Tensor: + """Controlled field at ``temperature_k`` / ``dilution_rate``. + + Parameters + ---------- + state + See signature. + + Returns + ------- + See signature.""" + return at_t.time_derivative(state) + d * (self.feed - state) + + f = residual(c) + for _ in range(max_iter): + jac = mass_action_jacobian(at_t, c) - d * eye + step = torch.linalg.lstsq(jac, -f.unsqueeze(-1)).solution.squeeze(-1) + alpha, norm = 1.0, 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"CSTR steady state did not converge in {max_iter} iterations" + raise RuntimeError(msg) + + +def van_de_vusse_cstr(feed_a: float = VAN_DE_VUSSE_FEED_A) -> VanDeVusseCSTR: + """Return the Van de Vusse CSTR with Klatt and Engell (1998) kinetics. + + Parameters + ---------- + feed_a : float, optional + Feed concentration of ``A`` in mol/L (finite, > 0). Default is + :data:`VAN_DE_VUSSE_FEED_A`. + + Returns + ------- + VanDeVusseCSTR + Reactor with feed ``(feed_a, 0, 0, 0)``. + + Raises + ------ + ValueError + If ``feed_a`` is not finite and positive. + """ + if not math.isfinite(float(feed_a)) or float(feed_a) <= 0.0: + msg = f"feed_a must be finite and > 0, got {feed_a!r}" + raise ValueError(msg) + arrhenius = tuple( + ArrheniusParams(k0, ea_over_r * GAS_CONSTANT_J_PER_MOL_K) + for k0, ea_over_r in zip(_K0, _EA_OVER_R, strict=True) + ) + nominal = tuple( + float(p.rate_constant(VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K)) for p in arrhenius + ) + system = ReactionSystem( + species=("A", "B", "C", "D"), + reactant_stoichiometry=torch.tensor( + [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [2.0, 0.0, 0.0, 0.0]], + dtype=torch.float64, + ), + product_stoichiometry=torch.tensor( + [[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]], + dtype=torch.float64, + ), + rate_constants=torch.tensor(nominal, dtype=torch.float64), + reaction_names=("A -> B", "B -> C", "2 A -> D"), + arrhenius=arrhenius, + name="van_de_vusse_cstr", + ) + feed = torch.tensor([float(feed_a), 0.0, 0.0, 0.0], dtype=torch.float64) + return VanDeVusseCSTR(system=system, feed=feed) + + +def _dilution(value: float) -> float: + """Validate a finite, non-negative dilution rate. + + Parameters + ---------- + value + See signature. + + Returns + ------- + See signature.""" + d = float(value) + if not math.isfinite(d) or d < 0.0: + msg = f"dilution_rate must be finite and >= 0, got {value!r}" + raise ValueError(msg) + return d + + +def _schedule(value: float | Tensor, num_intervals: int, name: str) -> list[float]: + """Expand a scalar or per-interval tensor into a list of floats. + + Parameters + ---------- + value + See signature. + num_intervals + See signature. + name + See signature. + + Returns + ------- + See signature.""" + if isinstance(value, Tensor): + if value.shape != (num_intervals,): + msg = ( + f"{name} schedule must have shape ({num_intervals},), got " + f"{tuple(value.shape)}" + ) + raise ValueError(msg) + return [float(v) for v in value.tolist()] + return [float(value)] * num_intervals diff --git a/src/koopman_graph/datasets/reaction/gillespie.py b/src/koopman_graph/datasets/reaction/gillespie.py index e69de29..da39ffa 100644 --- a/src/koopman_graph/datasets/reaction/gillespie.py +++ b/src/koopman_graph/datasets/reaction/gillespie.py @@ -0,0 +1,330 @@ +"""Seeded stochastic simulation (Gillespie direct method) and grid binning. + +Molecule counts ``n`` live in a volume ``V`` (concentration ``c = n / V``). +The propensity of reaction ``j`` matches the deterministic mass-action rate +``r_j = k_j ∏ c_s^{α_js}`` used by +:meth:`~koopman_graph.datasets.reaction.ReactionSystem.reaction_rates`: + +.. math:: + + a_j(n) = k_j V^{1 - o_j} \\prod_s n_s (n_s - 1) \\cdots (n_s - α_{js} + 1), + +where ``o_j = Σ_s α_js`` is the reaction order. The ``n(n-1)`` falling +factorial makes ``2 A → …`` need two distinct molecules. As ``V`` grows +with ``n / V`` fixed, ``n(t) / V`` approaches the ODE solution. + +:func:`gillespie_ssa` uses one seeded ``torch.Generator``, so a seed +reproduces the whole event path. :func:`bin_to_grid` samples the +piecewise-constant path on a regular grid, holding the last value. + +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.system import ReactionSystem + +DEFAULT_MAX_EVENTS = 1_000_000 +"""Default cap on SSA events per trajectory.""" + + +@dataclass(frozen=True) +class SSATrajectory: + """Event path of one stochastic simulation. + + Attributes + ---------- + event_times : Tensor + ``float64`` times ``(num_events + 1,)``; entry 0 is ``0.0``. + counts : Tensor + ``long`` molecule counts ``(num_events + 1, num_species)`` right + after each event (row 0 is the initial state). + t_end : float + Simulated horizon; the state after the last event holds to it. + volume : float + System volume ``V`` relating counts to concentrations. + seed : int + Generator seed. + truncated : bool + ``True`` when ``max_events`` stopped the run before ``t_end``. + """ + + event_times: Tensor + counts: Tensor + t_end: float + volume: float + seed: int + truncated: bool + + @property + def num_events(self) -> int: + """Return the number of reaction events. + + Returns + ------- + int + ``len(event_times) - 1``. + """ + return int(self.event_times.shape[0]) - 1 + + def counts_on_grid(self, grid: Tensor) -> Tensor: + """Return counts held at each grid time; see :func:`bin_to_grid`. + + Parameters + ---------- + grid : Tensor + Non-decreasing 1-D times in ``[0, t_end]``. + + Returns + ------- + Tensor + ``long`` counts ``(len(grid), num_species)``. + """ + last_event = float(self.event_times[-1]) + if self.truncated and grid.numel() and float(grid.max()) > last_event: + msg = "grid extends past the last event of a truncated trajectory" + raise ValueError(msg) + return bin_to_grid(self.event_times, self.counts, grid, t_end=self.t_end) + + def concentrations_on_grid(self, grid: Tensor) -> Tensor: + """Return concentrations ``counts / volume`` on ``grid``. + + Parameters + ---------- + grid : Tensor + Non-decreasing 1-D times in ``[0, t_end]``. + + Returns + ------- + Tensor + ``float64`` concentrations ``(len(grid), num_species)``. + """ + return self.counts_on_grid(grid).to(torch.float64) / self.volume + + +def propensities(system: ReactionSystem, counts: Tensor, volume: float) -> Tensor: + """Return mass-action SSA propensities ``a_j(n)``. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + counts : Tensor + Non-negative molecule counts ``(num_species,)``. + volume : float + System volume ``V`` (finite, > 0). + + Returns + ------- + Tensor + ``float64`` propensities ``(num_reactions,)``. Zero whenever a + reactant has fewer molecules than its coefficient. + """ + n = counts.to(torch.float64) + alpha = system.reactant_stoichiometry + falling = torch.ones_like(alpha) + for m in range(int(alpha.max().item()) if alpha.numel() else 0): + # Multiply in (n - m) wherever the coefficient exceeds m. + falling = falling * torch.where(alpha > m, (n - m).clamp(min=0), 1.0) + scale = volume ** (1.0 - system.reaction_orders) + return system.rate_constants * scale * falling.prod(dim=-1) + + +def gillespie_ssa( + system: ReactionSystem, + initial_counts: Tensor, + t_end: float, + *, + volume: float = 1.0, + seed: int = 0, + max_events: int = DEFAULT_MAX_EVENTS, +) -> SSATrajectory: + """Simulate one exact stochastic trajectory with Gillespie's direct method. + + Parameters + ---------- + system : ReactionSystem + Reaction network. + initial_counts : Tensor + Non-negative integer counts ``(num_species,)`` (any integer or + integer-valued floating dtype). + t_end : float + Horizon (finite, > 0). + volume : float, optional + System volume ``V``. Default is ``1.0``. + seed : int, optional + Generator seed. Default is ``0``. + max_events : int, optional + Event cap; the run stops early with ``truncated=True`` when hit. + Default is :data:`DEFAULT_MAX_EVENTS`. + + Returns + ------- + SSATrajectory + Event times and counts. The run ends at ``t_end``, when all + propensities vanish (absorbing state), or at ``max_events``. + + Raises + ------ + ValueError + If arguments are invalid. + """ + if not isinstance(system, ReactionSystem): + msg = f"system must be a ReactionSystem, got {type(system).__name__}" + raise TypeError(msg) + n = _initial_counts(system, initial_counts) + horizon = _positive("t_end", t_end) + v = _positive("volume", volume) + for name, value, low in (("seed", seed, 0), ("max_events", max_events, 1)): + if isinstance(value, bool) or not isinstance(value, int) or value < low: + msg = f"{name} must be an int >= {low}, got {value!r}" + raise ValueError(msg) + + generator = torch.Generator().manual_seed(seed) + net_change = system.stoichiometric_matrix.T.to(torch.long) # (R, S) + times = [0.0] + states = [n.clone()] + t = 0.0 + truncated = False + while True: + a = propensities(system, n, v) + total = float(a.sum()) + if total <= 0.0: + break + u = torch.rand(2, generator=generator, dtype=torch.float64) + t += -math.log(1.0 - float(u[0])) / total + if t > horizon: + break + if len(times) - 1 >= max_events: + truncated = True + break + # First reaction whose cumulative propensity exceeds u * a0; right=True + # never selects a zero-propensity reaction on a tie. + threshold = torch.tensor(float(u[1]) * total, dtype=torch.float64) + j = int(torch.searchsorted(torch.cumsum(a, 0), threshold, right=True)) + j = min(j, system.num_reactions - 1) + n = n + net_change[j] + times.append(t) + states.append(n.clone()) + return SSATrajectory( + event_times=torch.tensor(times, dtype=torch.float64), + counts=torch.stack(states), + t_end=horizon, + volume=v, + seed=seed, + truncated=truncated, + ) + + +def bin_to_grid( + event_times: Tensor, + counts: Tensor, + grid: Tensor, + *, + t_end: float | None = None, +) -> Tensor: + """Sample a piecewise-constant event path at grid times (last-value hold). + + Parameters + ---------- + event_times : Tensor + Non-decreasing event times ``(K,)`` starting at the path's origin. + counts : Tensor + State after each event ``(K, num_species)``. + grid : Tensor + Non-decreasing 1-D query times ``>= event_times[0]``. + t_end : float or None, optional + Upper bound for ``grid`` (the path's horizon). Default allows any + time after the last event. + + Returns + ------- + Tensor + ``counts[i]`` for the last event ``i`` with ``event_times[i] <= t``, + one row per grid time, in the dtype of ``counts``. + + Raises + ------ + ValueError + If shapes disagree or the grid leaves ``[event_times[0], t_end]``. + """ + if event_times.ndim != 1 or counts.ndim != 2: + msg = "event_times must be 1-D and counts 2-D" + raise ValueError(msg) + if int(event_times.shape[0]) != int(counts.shape[0]) or event_times.numel() == 0: + msg = ( + "event_times and counts must have the same non-zero length, got " + f"{int(event_times.shape[0])} and {int(counts.shape[0])}" + ) + raise ValueError(msg) + if grid.ndim != 1 or (grid.numel() > 1 and bool((grid[1:] < grid[:-1]).any())): + msg = "grid must be a non-decreasing 1-D Tensor" + raise ValueError(msg) + g = grid.to(torch.float64) + if g.numel() and float(g.min()) < float(event_times[0]): + msg = "grid starts before the first event time" + raise ValueError(msg) + if t_end is not None and g.numel() and float(g.max()) > float(t_end): + msg = f"grid extends past t_end={t_end}" + raise ValueError(msg) + index = torch.searchsorted(event_times.to(torch.float64), g, right=True) - 1 + return counts[index] + + +def _initial_counts(system: ReactionSystem, initial_counts: Tensor) -> Tensor: + """Validate non-negative integer counts and return them as ``long``. + + Parameters + ---------- + system + See signature. + initial_counts + See signature. + + Returns + ------- + See signature.""" + if not isinstance(initial_counts, Tensor): + msg = f"initial_counts must be a Tensor, got {type(initial_counts).__name__}" + raise TypeError(msg) + if initial_counts.shape != (system.num_species,): + msg = ( + f"initial_counts must have shape ({system.num_species},), got " + f"{tuple(initial_counts.shape)}" + ) + raise ValueError(msg) + as_float = initial_counts.to(torch.float64) + if ( + not bool(torch.isfinite(as_float).all()) + or bool((as_float < 0).any()) + or not bool((as_float == as_float.round()).all()) + ): + msg = "initial_counts must be non-negative integers" + raise ValueError(msg) + return as_float.to(torch.long) + + +def _positive(name: str, value: float) -> float: + """Return ``value`` as a finite positive float or raise ``ValueError``. + + Parameters + ---------- + name + See signature. + value + See signature. + + Returns + ------- + See signature.""" + x = float(value) + if not math.isfinite(x) or x <= 0.0: + msg = f"{name} must be finite and > 0, got {value!r}" + raise ValueError(msg) + return x diff --git a/src/koopman_graph/datasets/reaction/random_network.py b/src/koopman_graph/datasets/reaction/random_network.py index e69de29..dca5e58 100644 --- a/src/koopman_graph/datasets/reaction/random_network.py +++ b/src/koopman_graph/datasets/reaction/random_network.py @@ -0,0 +1,322 @@ +"""Seeded sampler for random, element-balanced mass-action networks. + +:func:`sample_random_network` builds the composition first, so every +sampled reaction conserves elements by construction: + +1. The first ``num_elements`` species are the elements themselves (unit + composition vectors). +2. Each further species is the union of two existing species. Its + composition is the sum of theirs, capped at ``max_atoms`` atoms and + unique, which guarantees at least one balanced association reaction. +3. Every reaction ``reactants → products`` with equal element totals + becomes a candidate, where the reactants are a multiset of at most + ``max_order`` species (the reaction order) and the products a multiset + of at most two. Compositions are unique, so ``1 → 1`` isomerizations + never balance, and ``max_order=1`` gives a linear network of + dissociations. ``num_reactions`` candidates are drawn without + replacement and given log-uniform rate constants. + +All randomness comes from one ``torch.Generator`` seeded with ``seed``, and +candidates are enumerated in a fixed order, so a seed reproduces the +network exactly. The seed and every sampler argument are returned with the +network. + +Pure PyTorch; importing this module does not import scipy or cantera. +""" + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor + +from koopman_graph.datasets.reaction.system import ReactionSystem + +_MAX_SPECIES = 32 +_MAX_PRODUCTS = 2 +_MAX_COMPOSITION_TRIES = 1000 + + +@dataclass(frozen=True) +class SampledReactionNetwork: + """A sampled network plus the seed and parameters that reproduce it. + + Attributes + ---------- + system : ReactionSystem + Element-balanced mass-action system (``elements`` are ``e0``, + ``e1``, ...). + seed : int + Seed passed to :func:`sample_random_network`. + parameters : dict + Every sampler argument (including ``seed``), JSON-compatible. + num_candidates : int + Number of balanced candidate reactions the draw was made from. + """ + + system: ReactionSystem + seed: int + parameters: dict[str, Any] + num_candidates: int + + def to_dict(self) -> dict[str, Any]: + """Serialize the network and its provenance. + + Returns + ------- + dict + ``{"system": ..., "seed": ..., "parameters": ..., + "num_candidates": ...}`` with JSON-compatible values. + """ + return { + "system": self.system.to_dict(), + "seed": self.seed, + "parameters": dict(self.parameters), + "num_candidates": self.num_candidates, + } + + +def sample_random_network( + *, + num_species: int = 6, + num_elements: int = 2, + num_reactions: int = 8, + max_atoms: int = 4, + max_order: int = 2, + rate_log10_range: tuple[float, float] = (-1.0, 1.0), + seed: int = 0, +) -> SampledReactionNetwork: + """Sample a random element-balanced mass-action network. + + Parameters + ---------- + num_species : int, optional + Total species count, greater than ``num_elements`` and at most 32. + Default is ``6``. + num_elements : int, optional + Number of elements; the first ``num_elements`` species are the bare + elements. Default is ``2``. + num_reactions : int, optional + Reactions to draw from the balanced candidates. Default is ``8``. + max_atoms : int, optional + Cap on the total atom count of any species (>= 2). Default is ``4``. + max_order : int, optional + Maximum reaction order (reactant species counted with + multiplicity), ``1`` or ``2``. Products always have at most two + species, so ``1`` yields a linear network. Default is ``2``. + rate_log10_range : tuple of float, optional + ``(low, high)`` bounds of ``log10 k``; rates are log-uniform. + Default is ``(-1.0, 1.0)``. + seed : int, optional + Seed of the ``torch.Generator`` that drives every draw. Default is + ``0``. + + Returns + ------- + SampledReactionNetwork + The system, its seed and parameters, and the candidate count. + + Raises + ------ + ValueError + If arguments are out of range, ``max_atoms`` leaves too few distinct + compositions, or there are fewer balanced candidates than + ``num_reactions``. + """ + _check_int("num_elements", num_elements, low=1, high=_MAX_SPECIES - 1) + _check_int("num_species", num_species, low=num_elements + 1, high=_MAX_SPECIES) + _check_int("num_reactions", num_reactions, low=1, high=None) + _check_int("max_atoms", max_atoms, low=2, high=None) + _check_int("max_order", max_order, low=1, high=2) + _check_int("seed", seed, low=0, high=None) + low, high = (float(v) for v in rate_log10_range) + if not (math.isfinite(low) and math.isfinite(high)) or low > high: + msg = ( + "rate_log10_range must be finite with low <= high, got " + f"{rate_log10_range!r}" + ) + raise ValueError(msg) + + generator = torch.Generator().manual_seed(seed) + compositions = _sample_compositions( + num_species=num_species, + num_elements=num_elements, + max_atoms=max_atoms, + generator=generator, + ) + candidates = _balanced_candidates(compositions, max_order=max_order) + if len(candidates) < num_reactions: + msg = ( + f"only {len(candidates)} balanced candidate reactions for " + f"num_reactions={num_reactions}; raise num_species / max_order or " + "lower num_reactions" + ) + raise ValueError(msg) + order = torch.randperm(len(candidates), generator=generator)[:num_reactions] + chosen = [candidates[int(i)] for i in torch.sort(order).values.tolist()] + exponents = low + (high - low) * torch.rand( + num_reactions, generator=generator, dtype=torch.float64 + ) + + reactants = torch.zeros(num_reactions, num_species, dtype=torch.float64) + products = torch.zeros(num_reactions, num_species, dtype=torch.float64) + for row, (lhs, rhs) in enumerate(chosen): + for s in lhs: + reactants[row, s] += 1.0 + for s in rhs: + products[row, s] += 1.0 + species = tuple(f"X{i}" for i in range(num_species)) + names = tuple( + f"{_side(lhs, species)} -> {_side(rhs, species)}" for lhs, rhs in chosen + ) + parameters: dict[str, Any] = { + "num_species": num_species, + "num_elements": num_elements, + "num_reactions": num_reactions, + "max_atoms": max_atoms, + "max_order": max_order, + "rate_log10_range": [low, high], + "seed": seed, + } + system = ReactionSystem( + species=species, + reactant_stoichiometry=reactants, + product_stoichiometry=products, + rate_constants=torch.pow(10.0, exponents), + element_matrix=compositions.T.contiguous(), + elements=tuple(f"e{i}" for i in range(num_elements)), + reaction_names=names, + name=f"random_network_seed{seed}", + ) + return SampledReactionNetwork( + system=system, + seed=seed, + parameters=parameters, + num_candidates=len(candidates), + ) + + +def _sample_compositions( + *, + num_species: int, + num_elements: int, + max_atoms: int, + generator: torch.Generator, +) -> Tensor: + """Return unique species compositions ``(num_species, num_elements)``. + + Parameters + ---------- + num_species + See signature. + num_elements + See signature. + max_atoms + See signature. + generator + See signature. + + Returns + ------- + See signature.""" + rows: list[tuple[int, ...]] = [ + tuple(int(i == e) for i in range(num_elements)) for e in range(num_elements) + ] + seen = set(rows) + tries = 0 + while len(rows) < num_species: + tries += 1 + if tries > _MAX_COMPOSITION_TRIES: + msg = ( + f"could not build {num_species} distinct species with " + f"max_atoms={max_atoms}; raise max_atoms or lower num_species" + ) + raise ValueError(msg) + a, b = torch.randint(len(rows), (2,), generator=generator).tolist() + combined = tuple(x + y for x, y in zip(rows[a], rows[b], strict=True)) + if sum(combined) > max_atoms or combined in seen: + continue + rows.append(combined) + seen.add(combined) + return torch.tensor(rows, dtype=torch.float64) + + +def _balanced_candidates( + compositions: Tensor, + *, + max_order: int, +) -> list[tuple[tuple[int, ...], tuple[int, ...]]]: + """Enumerate element-balanced ``(reactants, products)`` multiset pairs. + + Parameters + ---------- + compositions + See signature. + max_order + See signature. + + Returns + ------- + See signature.""" + num_species = int(compositions.shape[0]) + comp = [tuple(int(v) for v in row) for row in compositions.tolist()] + sides: list[tuple[int, ...]] = [] + for size in range(1, _MAX_PRODUCTS + 1): + sides.extend(itertools.combinations_with_replacement(range(num_species), size)) + totals = { + side: tuple(sum(comp[s][e] for s in side) for e in range(len(comp[0]))) + for side in sides + } + return [ + (lhs, rhs) + for lhs in sides + if len(lhs) <= max_order + for rhs in sides + if lhs != rhs and totals[lhs] == totals[rhs] + ] + + +def _side(side: tuple[int, ...], species: tuple[str, ...]) -> str: + """Format a multiset of species indices as ``"2 X0 + X3"``. + + Parameters + ---------- + side + See signature. + species + See signature. + + Returns + ------- + See signature.""" + parts = [] + for s in sorted(set(side)): + count = side.count(s) + parts.append(species[s] if count == 1 else f"{count} {species[s]}") + return " + ".join(parts) + + +def _check_int(name: str, value: Any, *, low: int, high: int | None) -> None: + """Raise ``ValueError`` unless ``value`` is an int in ``[low, high]``. + + Parameters + ---------- + name + See signature. + value + See signature. + low + See signature. + high + See signature.""" + bad = isinstance(value, bool) or not isinstance(value, int) or value < low + if not bad and high is not None and value > high: + bad = True + if bad: + bound = f">= {low}" if high is None else f"in [{low}, {high}]" + msg = f"{name} must be an int {bound}, got {value!r}" + raise ValueError(msg) diff --git a/tests/datasets/test_reaction_cstr.py b/tests/datasets/test_reaction_cstr.py new file mode 100644 index 0000000..85ef2ae --- /dev/null +++ b/tests/datasets/test_reaction_cstr.py @@ -0,0 +1,125 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.cstr`.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + VAN_DE_VUSSE_FEED_A, + VAN_DE_VUSSE_NOMINAL_DILUTION, + VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K, + VanDeVusseCSTR, + mass_action_jacobian, + van_de_vusse_cstr, +) + +_NOMINAL = { + "temperature_k": VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K, + "dilution_rate": VAN_DE_VUSSE_NOMINAL_DILUTION, +} + + +def test_nominal_steady_state_matches_benchmark_operating_point() -> None: + reactor = van_de_vusse_cstr() + steady = reactor.steady_state(**_NOMINAL) + # Benchmark operating point: c_A ≈ 2.14, c_B ≈ 1.09 mol/L. + assert math.isclose(float(steady[0]), 2.14, abs_tol=5e-3) + assert math.isclose(float(steady[1]), 1.09, abs_tol=5e-3) + assert float(reactor.rhs(steady, **_NOMINAL).abs().max()) < 1e-9 + + +def test_literature_balance_for_a_is_reproduced() -> None: + """``dc_A/dt = -k1 c_A - k3 c_A² + D (c_A0 - c_A)`` with literature ``k3``.""" + reactor = van_de_vusse_cstr() + temperature = VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K + k1 = 1.287e12 * math.exp(-9758.3 / temperature) + k3 = 9.043e9 * math.exp(-8560.0 / temperature) + c = torch.tensor([2.0, 1.0, 0.5, 0.25], dtype=torch.float64) + d = 10.0 + rhs = reactor.rhs(c, temperature_k=temperature, dilution_rate=d) + expected_a = -k1 * 2.0 - k3 * 4.0 + d * (VAN_DE_VUSSE_FEED_A - 2.0) + assert math.isclose(float(rhs[0]), expected_a, rel_tol=1e-6) + # dc_D/dt = k3 c_A² / 2 - D c_D in the literature. + assert math.isclose(float(rhs[3]), 0.5 * k3 * 4.0 - d * 0.25, rel_tol=1e-6) + + +def test_simulation_converges_to_steady_state() -> None: + reactor = van_de_vusse_cstr() + steady = reactor.steady_state(**_NOMINAL) + states = reactor.simulate( + torch.zeros(4, dtype=torch.float64), + torch.linspace(0.0, 1.0, 11, dtype=torch.float64), + **_NOMINAL, + ) + assert states.shape == (11, 4) + assert torch.allclose(states[-1], steady, atol=1e-6) + # Linearized controlled field is stable at the operating point. + jac = mass_action_jacobian( + reactor.system.at_temperature(VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K), steady + ) - VAN_DE_VUSSE_NOMINAL_DILUTION * torch.eye(4, dtype=torch.float64) + assert bool((torch.linalg.eigvals(jac).real < 0).all()) + + +def test_piecewise_constant_inputs_switch_operating_point() -> None: + reactor = van_de_vusse_cstr() + times = torch.linspace(0.0, 2.0, 5, dtype=torch.float64) + temps = torch.tensor([387.35, 387.35, 400.0, 400.0], dtype=torch.float64) + dilutions = torch.tensor([14.19, 14.19, 20.0, 20.0], dtype=torch.float64) + states = reactor.simulate( + torch.zeros(4, dtype=torch.float64), + times, + temperature_k=temps, + dilution_rate=dilutions, + substeps=200, + ) + first = reactor.steady_state(temperature_k=387.35, dilution_rate=14.19) + second = reactor.steady_state(temperature_k=400.0, dilution_rate=20.0) + assert torch.allclose(states[2], first, atol=1e-6) + assert torch.allclose(states[4], second, atol=1e-6) + assert not torch.allclose(first, second, atol=1e-2) + + +def test_batched_simulation_and_rhs_shape() -> None: + reactor = van_de_vusse_cstr(feed_a=3.0) + assert reactor.feed.tolist() == [3.0, 0.0, 0.0, 0.0] + c0 = torch.zeros(3, 4, dtype=torch.float64) + states = reactor.simulate(c0, torch.tensor([0.0, 0.1]), **_NOMINAL, substeps=10) + assert states.shape == (2, 3, 4) + assert torch.allclose(states[-1, 0], states[-1, 2]) + + +def test_reactor_type_and_arrhenius_data() -> None: + reactor = van_de_vusse_cstr() + assert isinstance(reactor, VanDeVusseCSTR) + assert reactor.system.arrhenius is not None + assert reactor.system.reaction_names == ("A -> B", "B -> C", "2 A -> D") + assert torch.allclose( + reactor.system.rate_constants, + reactor.system.rate_constants_at(VAN_DE_VUSSE_NOMINAL_TEMPERATURE_K), + ) + + +def test_cstr_validation() -> None: + reactor = van_de_vusse_cstr() + c = torch.zeros(4, dtype=torch.float64) + for feed in (0.0, -1.0, math.nan): + with pytest.raises(ValueError, match="feed_a must be"): + van_de_vusse_cstr(feed) + with pytest.raises(ValueError, match="dilution_rate must be"): + reactor.rhs(c, temperature_k=400.0, dilution_rate=-1.0) + with pytest.raises(ValueError, match="temperature_k must be"): + reactor.rhs(c, temperature_k=0.0, dilution_rate=1.0) + times = torch.tensor([0.0, 1.0, 2.0]) + with pytest.raises(ValueError, match=r"schedule must have shape \(2,\)"): + reactor.simulate(c, times, temperature_k=torch.ones(3) * 400, dilution_rate=1.0) + with pytest.raises(ValueError, match="strictly increasing"): + reactor.simulate(c, torch.tensor([0.0, 0.0]), **_NOMINAL) + with pytest.raises(ValueError, match="non-empty 1-D"): + reactor.simulate(c, torch.zeros(0), **_NOMINAL) + with pytest.raises(ValueError, match="substeps must be"): + reactor.simulate(c, times, **_NOMINAL, substeps=0) + with pytest.raises(RuntimeError, match="did not converge"): + reactor.steady_state(**_NOMINAL, max_iter=1) diff --git a/tests/datasets/test_reaction_gillespie.py b/tests/datasets/test_reaction_gillespie.py new file mode 100644 index 0000000..fae61b8 --- /dev/null +++ b/tests/datasets/test_reaction_gillespie.py @@ -0,0 +1,161 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.gillespie`.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + ReactionSystem, + bin_to_grid, + first_order_chain, + gillespie_ssa, + integrate_rk4, + propensities, + reversible_dimerization, +) + + +def test_propensities_use_falling_factorials_and_volume() -> None: + system = reversible_dimerization() # 2 A -> A2 (3.0), A2 -> 2 A (0.5) + a = propensities(system, torch.tensor([10, 4]), volume=2.0) + # 2 A: k V^{-1} n (n - 1); A2: k n. + assert torch.allclose( + a, torch.tensor([3.0 / 2.0 * 10 * 9, 0.5 * 4], dtype=torch.float64) + ) + one = propensities(system, torch.tensor([1, 0]), volume=1.0) + assert one.tolist() == [0.0, 0.0] # one A molecule cannot dimerize + feed = ReactionSystem( + species=("A",), + reactant_stoichiometry=torch.zeros(1, 1), + product_stoichiometry=torch.ones(1, 1), + rate_constants=torch.tensor([2.0]), + ) + # Zeroth order scales with V. + assert propensities(feed, torch.tensor([0]), volume=5.0).tolist() == [10.0] + + +def test_seed_reproduces_path_exactly() -> None: + system = reversible_dimerization() + kwargs = {"volume": 100.0, "seed": 11} + a = gillespie_ssa(system, torch.tensor([100, 0]), 3.0, **kwargs) + b = gillespie_ssa(system, torch.tensor([100, 0]), 3.0, **kwargs) + assert torch.equal(a.event_times, b.event_times) + assert torch.equal(a.counts, b.counts) + c = gillespie_ssa(system, torch.tensor([100, 0]), 3.0, volume=100.0, seed=12) + assert not torch.equal(a.counts, c.counts) or not torch.equal( + a.event_times, c.event_times + ) + assert a.seed == 11 + assert a.counts.dtype == torch.long + assert float(a.event_times[0]) == 0.0 + assert bool((a.event_times[1:] > a.event_times[:-1]).all()) + assert float(a.event_times[-1]) <= 3.0 + assert not a.truncated + + +def test_path_conserves_monomer_units() -> None: + system = reversible_dimerization() + path = gillespie_ssa(system, torch.tensor([60, 20]), 5.0, volume=50.0, seed=0) + units = path.counts[:, 0] + 2 * path.counts[:, 1] + assert bool((units == 100).all()) + assert bool((path.counts >= 0).all()) + + +def test_large_volume_approaches_ode() -> None: + system = reversible_dimerization() + grid = torch.linspace(0.0, 4.0, 9, dtype=torch.float64) + ode = integrate_rk4( + system, torch.tensor([1.0, 0.0], dtype=torch.float64), grid, substeps=50 + ) + + def mean_error(volume: int) -> float: + """Mean max-abs SSA-vs-ODE error over three fixed seeds. + + Parameters + ---------- + volume + See signature. + + Returns + ------- + See signature.""" + errors = [] + for seed in range(3): + path = gillespie_ssa( + system, torch.tensor([volume, 0]), 4.0, volume=float(volume), seed=seed + ) + errors.append(float((path.concentrations_on_grid(grid) - ode).abs().max())) + return sum(errors) / len(errors) + + small, large = mean_error(40), mean_error(2000) + assert large < 0.05 + assert large < small / 2 + + +def test_absorbing_state_stops_early() -> None: + system = first_order_chain(2) # S0 -> S1 only + path = gillespie_ssa(system, torch.tensor([5, 0]), 1.0e6, seed=0) + assert path.num_events == 5 + assert path.counts[-1].tolist() == [0, 5] + held = path.counts_on_grid(torch.tensor([0.0, 1.0e6])) + assert held.tolist() == [[5, 0], [0, 5]] + + +def test_max_events_truncates() -> None: + system = reversible_dimerization() + path = gillespie_ssa( + system, torch.tensor([200, 0]), 100.0, volume=200.0, seed=0, max_events=10 + ) + assert path.truncated + assert path.num_events == 10 + with pytest.raises(ValueError, match="truncated"): + path.counts_on_grid(torch.tensor([0.0, 100.0])) + assert path.counts_on_grid(torch.tensor([0.0])).tolist() == [[200, 0]] + + +def test_bin_to_grid_holds_last_value() -> None: + times = torch.tensor([0.0, 1.0, 2.5], dtype=torch.float64) + counts = torch.tensor([[3, 0], [2, 1], [1, 2]]) + grid = torch.tensor([0.0, 0.5, 1.0, 2.0, 2.5, 3.0]) + held = bin_to_grid(times, counts, grid) + assert held.tolist() == [[3, 0], [3, 0], [2, 1], [2, 1], [1, 2], [1, 2]] + assert bin_to_grid(times, counts, torch.zeros(0)).shape == (0, 2) + with pytest.raises(ValueError, match="past t_end"): + bin_to_grid(times, counts, grid, t_end=2.9) + with pytest.raises(ValueError, match="before the first event"): + bin_to_grid(times, counts, torch.tensor([-1.0])) + with pytest.raises(ValueError, match="non-decreasing"): + bin_to_grid(times, counts, torch.tensor([1.0, 0.5])) + with pytest.raises(ValueError, match="1-D and counts 2-D"): + bin_to_grid(times, counts[:, 0], grid) + with pytest.raises(ValueError, match="same non-zero length"): + bin_to_grid(times[:2], counts, grid) + + +@pytest.mark.parametrize( + ("initial", "kwargs", "error", "match"), + [ + ([5, 0], {}, TypeError, "must be a Tensor"), + (torch.tensor([5]), {}, ValueError, r"shape \(2,\)"), + (torch.tensor([5.5, 0.0]), {}, ValueError, "non-negative integers"), + (torch.tensor([-1, 0]), {}, ValueError, "non-negative integers"), + (torch.tensor([5, 0]), {"t_end": 0.0}, ValueError, "t_end must be"), + (torch.tensor([5, 0]), {"volume": math.inf}, ValueError, "volume must be"), + (torch.tensor([5, 0]), {"seed": -1}, ValueError, "seed must be"), + (torch.tensor([5, 0]), {"max_events": 0}, ValueError, "max_events must be"), + ], +) +def test_ssa_validation( + initial: object, kwargs: dict[str, object], error: type[Exception], match: str +) -> None: + t_end = kwargs.pop("t_end", 1.0) + with pytest.raises(error, match=match): + gillespie_ssa(reversible_dimerization(), initial, t_end, **kwargs) # type: ignore[arg-type] + + +def test_ssa_requires_reaction_system() -> None: + with pytest.raises(TypeError, match="must be a ReactionSystem"): + gillespie_ssa("dimer", torch.tensor([1, 0]), 1.0) # type: ignore[arg-type] diff --git a/tests/datasets/test_reaction_random_network.py b/tests/datasets/test_reaction_random_network.py new file mode 100644 index 0000000..f790da9 --- /dev/null +++ b/tests/datasets/test_reaction_random_network.py @@ -0,0 +1,102 @@ +"""Tests for :mod:`koopman_graph.datasets.reaction.random_network`.""" + +from __future__ import annotations + +import json +import math + +import pytest +import torch + +from koopman_graph.datasets.reaction import ( + ReactionSystem, + check_element_balance, + integrate_rk4, + sample_random_network, +) + + +@pytest.mark.parametrize("seed", range(20)) +def test_sampled_networks_are_element_balanced(seed: int) -> None: + network = sample_random_network(seed=seed) + system = network.system + assert check_element_balance(system).balanced + assert system.num_species == 6 + assert system.num_reactions == 8 + assert int(system.reaction_orders.max()) <= 2 + assert int(system.product_stoichiometry.sum(dim=1).max()) <= 2 + # First species are the bare elements; compositions are unique. + assert system.element_matrix is not None + composition = system.element_matrix.T + assert torch.equal(composition[:2], torch.eye(2, dtype=torch.float64)) + assert len({tuple(row) for row in composition.tolist()}) == 6 + assert int(composition.sum(dim=1).max()) <= 4 + + +def test_seed_reproduces_network_and_provenance() -> None: + a = sample_random_network(seed=7) + b = sample_random_network(seed=7) + assert a.system.to_dict() == b.system.to_dict() + assert a.seed == 7 + assert a.parameters["seed"] == 7 + assert a.system.name == "random_network_seed7" + payload = json.loads(json.dumps(a.to_dict())) + assert ReactionSystem.from_dict(payload["system"]).to_dict() == a.system.to_dict() + assert payload["parameters"]["rate_log10_range"] == [-1.0, 1.0] + assert payload["num_candidates"] >= 8 + assert sample_random_network(seed=8).system.to_dict() != a.system.to_dict() + + +def test_rates_are_log_uniform_in_range() -> None: + network = sample_random_network( + num_species=8, num_reactions=20, rate_log10_range=(-2.0, 0.5), seed=3 + ) + log_k = torch.log10(network.system.rate_constants) + assert float(log_k.min()) >= -2.0 + assert float(log_k.max()) <= 0.5 + fixed = sample_random_network(rate_log10_range=(0.0, 0.0), seed=1) + assert torch.allclose( + fixed.system.rate_constants, torch.ones(8, dtype=torch.float64) + ) + + +def test_max_order_one_gives_linear_dissociation_network() -> None: + system = sample_random_network(num_reactions=4, max_order=1, seed=2).system + assert system.is_linear + assert check_element_balance(system).balanced + assert bool((system.reaction_orders == 1).all()) + # Unique compositions rule out 1 -> 1 isomerizations, so every reaction + # dissociates one species into two. + assert bool((system.product_stoichiometry.sum(dim=1) == 2).all()) + + +def test_sampled_network_integrates_and_conserves_elements() -> None: + system = sample_random_network(seed=4).system + c0 = torch.full((6,), 0.5, dtype=torch.float64) + states = integrate_rk4(system, c0, torch.linspace(0.0, 2.0, 11), substeps=200) + assert system.element_matrix is not None + drift = (states - c0) @ system.element_matrix.T + assert float(drift.abs().max()) < 1e-10 + assert bool(torch.isfinite(states).all()) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_elements": 0}, "num_elements must be an int"), + ({"num_species": 2, "num_elements": 2}, "num_species must be an int"), + ({"num_species": 40}, "num_species must be an int"), + ({"num_reactions": 0}, "num_reactions must be an int"), + ({"max_atoms": 1}, "max_atoms must be an int"), + ({"max_order": 3}, "max_order must be an int"), + ({"seed": -1}, "seed must be an int"), + ({"seed": True}, "seed must be an int"), + ({"rate_log10_range": (1.0, 0.0)}, "rate_log10_range"), + ({"rate_log10_range": (0.0, math.inf)}, "rate_log10_range"), + ({"num_species": 10, "max_atoms": 2}, "could not build 10 distinct species"), + ({"num_reactions": 500}, "balanced candidate reactions"), + ], +) +def test_sampler_validation(kwargs: dict[str, object], match: str) -> None: + with pytest.raises(ValueError, match=match): + sample_random_network(**kwargs) # type: ignore[arg-type]