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