OP System

op_system

op_system.

Domain-agnostic RHS specification + compilation utilities.

Public API (v1)

Primary user entrypoints: - compile_spec: Validate, normalize, and compile a RHS specification in one step. - normalize_rhs: Validate and normalize a YAML-friendly RHS specification. - compile_rhs: Compile a NormalizedRhs into an efficient callable RHS.

Core data structures: - NormalizedRhs - CompiledRhs - OperatorDescriptor

Design guarantees: - No dependency on provider/adapters (eg flepimop2). - Stable interface for downstream engines. - Forward-compatible with multiphysics extensions.

IdentifierString = Annotated[str, AfterValidator(_validate_identifier_string)] module-attribute

Custom pydantic type for validated identifier strings used in op_system.

Identifier strings are used for state names, dimension names, and other keys in the system. They must be non-empty, contain only alphanumeric characters, and start with a letter. Leading and trailing whitespace is stripped before validation.

Examples:

>>> from pydantic import BaseModel
>>> from op_system import IdentifierString
>>> class ExampleModel(BaseModel):
...     identifier: IdentifierString
...
>>> ExampleModel(identifier="S")
ExampleModel(identifier='S')
>>> ExampleModel(identifier="  Foobar  ")
ExampleModel(identifier='Foobar')
>>> ExampleModel(identifier="123abc")
Traceback (most recent call last):
    ...
pydantic_core._pydantic_core.ValidationError: 1 validation error for ExampleModel
identifier
Value error, IdentifierString must contain only alphanumerical characters and start with a letter. [...]
    For further information visit ...
>>> ExampleModel(identifier="")
Traceback (most recent call last):
    ...
pydantic_core._pydantic_core.ValidationError: 1 validation error for ExampleModel
identifier
Value error, IdentifierString must not be empty. [...]
    For further information visit ...

Array

Bases: Protocol

Structural Array-API protocol.

Any object whose runtime type implements shape, dtype, __array_namespace__ and item satisfies this protocol. NumPy

= 2.0 ndarrays and JAX arrays (concrete and traced) qualify directly.

Runtime evaluation is deliberately broader: it discovers namespaces via :func:array_api_compat.array_namespace, so native arrays such as torch.Tensor are accepted even though they do not structurally satisfy this protocol. No compile-time backend selector is needed.

AxisDeclaration

Bases: NamedTuple

Physical interpretation of a named axis.

Attributes:

Name	Type	Description
units	str | None	Coordinate units, independent of measure units.
measure	AxisMeasure | None	Explicit measure, or None when undeclared.

AxisMeasure

Bases: NamedTuple

Explicit physical measure in the axis coordinate order.

Attributes:

Name	Type	Description
kind	_MeasureKind	Counting, cell widths/volumes, or point quadrature weights.
weights	tuple[float, ...]	Finite strictly positive weights, one per coordinate.
units	str | None	Optional measure units, without dimensional verification.

from_mapping(value, *, size) classmethod

Validate a measure without inferring geometry from coordinates.

Parameters:

Name	Type	Description	Default
value	object	Mapping with kind, optional units, and weights for cell/quadrature.	required
size	int	Number of axis coordinates.	required

Returns:

Type	Description
AxisMeasure	Immutable measure in coordinate order.

Raises:

Type	Description
InvalidRhsSpecError	If the kind, units, weights or length are invalid.

Source code in src/op_system/_scientific_meta.py

@classmethod
def from_mapping(cls, value: object, *, size: int) -> AxisMeasure:
    """Validate a measure without inferring geometry from coordinates.

    Args:
        value: Mapping with kind, optional units, and weights for cell/quadrature.
        size: Number of axis coordinates.

    Returns:
        Immutable measure in coordinate order.

    Raises:
        InvalidRhsSpecError: If the kind, units, weights or length are invalid.
    """
    if not isinstance(size, int) or isinstance(size, bool) or size < 1:
        raise InvalidRhsSpecError(
            detail="axis measure size must be a positive integer"
        )
    if not isinstance(value, Mapping) or set(value) - {"kind", "weights", "units"}:
        raise InvalidRhsSpecError(
            detail="axis measure must contain kind, weights and/or units"
        )
    kind = value.get("kind")
    if not isinstance(kind, str) or kind not in {"counting", "cell", "quadrature"}:
        raise InvalidRhsSpecError(
            detail="axis measure kind must be counting, cell or quadrature"
        )
    if kind == "counting":
        if "weights" in value:
            raise InvalidRhsSpecError(
                detail="counting measure has unit weights; omit weights"
            )
        weights = (1.0,) * size
    else:
        raw_weights = value.get("weights")
        if not isinstance(raw_weights, (list, tuple)) or len(raw_weights) != size:
            raise InvalidRhsSpecError(
                detail="axis measure weights must match the coordinate count"
            )
        weights = tuple(_positive_weight(weight) for weight in raw_weights)
    return cls(cast("_MeasureKind", kind), weights, _optional_text(value, "units"))

BalanceCoverageGap

Bases: NamedTuple

An omission or unsupported contribution preventing issued laws.

Attributes:

Name	Type	Description
origin	str	Location in compiler metadata or reaction artifacts.
reason	str	Machine-readable coverage limitation.
name	str | None	Associated reaction/state name when known.
reaction_gap	ReactionGap | None	Original compiler omission, retaining its selectors/reason.

BlockAxisInfo(name, size, state_axis_pos, param_axis_pos) dataclass

Metadata for one block-diagonal (factorizable) axis.

A BlockAxisInfo is attached to :class:~op_system.compile.CompiledRhs for each axis declared in spec["factorize_axes"] that passes the IR separability check in :func:analyze_block_axes. Engines (e.g. the diffrax plugin) consume this to partition independent solves with jax.vmap over the block axis.

Attributes:

Name	Type	Description
name	str	Axis name string (e.g. "loc").
size	int	Number of elements along the axis.
state_axis_pos	dict[str, int]	Maps each state-template base name to the integer position of this axis within that template's shape tuple. Only templates that carry this axis appear in the dict.
param_axis_pos	dict[str, int | None]	Maps each shaped parameter name to the integer position of this axis within the actual runtime array that the engine passes to the eval function, or None if the parameter does not carry this axis (broadcast). For non-time-varying shaped parameters the position is the index in the parameter's axis tuple. For time-varying parameters the runtime array has time prepended at index 0, so the position equals the index of the block axis in the full (time, *spatial_axes) tuple. Parameters that are entirely scalar (not shaped) do not appear in this dict and are always broadcast.

Note

BlockAxisInfo uses dict fields and therefore cannot be used as a hash key. It is stored on :class:~op_system.compile.CompiledRhs with hash=False and is fully pickle-stable.

Examples:

>>> info = BlockAxisInfo(
...     name="loc",
...     size=3,
...     state_axis_pos={"S": 1, "I": 1, "R": 1},
...     param_axis_pos={"rho": 0, "beta": None},
... )
>>> info.name
'loc'
>>> info.size
3
>>> info.state_axis_pos["S"]
1
>>> info.param_axis_pos["rho"]
0
>>> info.param_axis_pos["beta"] is None
True

BodyEvalFn

Bases: Protocol

Callable that evaluates history signal bodies at (t, y, **params).

Returns a mapping from signal_id to the evaluated body array/value.

CompiledObservable

Bases: NamedTuple

Native evaluator and explicit fixed-layout scientific metadata.

Attributes:

Name	Type	Description
definition	ObservableDefinition	Backend-neutral expression and reduction description.
output_shape	tuple[int, ...]	Shape in definition.axes order; empty for a scalar.
coordinates	tuple[tuple[str, tuple[str, ...]], ...]	Coordinate labels for each output axis, in output order.
state_dependencies	tuple[str, ...]	Referenced state base names, in sorted order.
parameter_shapes	tuple[tuple[str, tuple[int, ...]], ...]	Required parameters and their complete input shapes. Scalars have empty shapes; these are separate from RHS requirements.
time_dependent	bool	Whether the expression explicitly references t.
measure_requirements	tuple[tuple[str, AxisMeasure], ...]	Physical measures applied exactly once per axis.
eval_fn	ObservableEvalFn	Evaluator for a complete flat state vector.
pytree_eval_fn	ObservablePytreeEvalFn	Evaluator for complete shaped state buffers.

CompiledReactant(state_base, state_axes, full_axes, pinned, order) dataclass

Molecular reactant metadata aligned to an expanded reaction channel.

The record is numerical-array neutral. Providers combine state_axes with the enclosing reaction's channel coordinates and pinned indices to locate a state cell, then place order in their reactant-order matrix. Catalysts appear here even when their net stoichiometric change is zero.

CompiledReaction(name, from_base, from_axes, full_axes, to_base, to_axes, sum_axes, pinned, from_pinned, propensity_fn, reactants=(), reactants_complete=False, offsets=(), to_full_axes=None, routed_axes=(), dependencies=(), propensity_order=None, dependencies_complete=False, resolved_propensity_fn=None, time_varying_params=(), independent_axes=(), uses_time=False) dataclass

Compiled propensity + axis bookkeeping for one named transition.

Produced from :class:op_system._reactions.ReactionArtifactIR at compile time. A firing event at from-cell index i (in the natural N-D shape given by from_axes) depletes 1 from from_base at that cell and adds 1 to to_base at the cell obtained by: keeping every axis in to_axes at i's value for that axis, and using the fixed coordinate index from pinned for every axis in pinned (this includes both axes that are wildcard on from_axes and pinned on to -- a "collapse to a fixed target" transition, where multiple source cells share one destination and the axis also appears in sum_axes -- and axes pinned on BOTH from and to, a "point-to-point" shift between two specific coordinates on the same axis, e.g. a vaccination-dose-progression transition, which do NOT appear in sum_axes since there's only ever one source value on that axis). Consumers that need a single combined scatter/reduce target (e.g. an engine applying many simultaneous firings) sum over sum_axes when depositing into to_base -- op_system does not do that summation itself, since whether/how to combine simultaneous firings across a summed axis is an execution-semantics decision for the consumer, not a compile-time one.

from_base is None for a SOURCE-ONLY (from: null) reaction -- an exogenous hazard with no compartment to deplete (e.g. cross-district case importation). A consumer must special-case this: skip every from-side depletion/clamp step and apply only the to_base scatter, using from_axes/full_axes (which are the to-side template's own axes in this case, not a nonexistent from-side one -- see ReactionArtifactIR).

An axis-wide coord_shift publishes one reaction per state template with a non-empty offsets: the destination index on each offset axis is the source index plus its step (see offsets). The destination cell is then built from to_axes (copied), pinned (fixed), and offsets (shifted), which together cover full_axes.

Examples:

>>> import numpy as np
>>> from op_system import compile_spec
>>> compiled = compile_spec({
...     "kind": "transitions",
...     "state": ["S", "I"],
...     "transitions": [
...         {"name": "infect", "from": "S", "to": "I", "rate": "k"},
...     ],
... })
>>> reaction = compiled.reactions[0]
>>> float(
...     reaction.propensity_fn(
...         0.0, {"S": np.array(10.0), "I": np.array(0.0)}, k=0.2
...     )
... )
2.0
>>> reaction.time_varying_params, reaction.uses_time
((), False)

CompiledRhs(state_names, param_names, eval_fn, meta=(lambda: MappingProxyType({}))(), operators=tuple(), factorize_axes=tuple(), block_axes=tuple(), pytree_eval_fn=None, template_shapes=None, block_pytree_eval_fn=None, block_template_shapes=None, history_requirements=tuple(), history_eval_fn=None, body_eval_fn=None, block_history_eval_fn=None, block_body_eval_fn=None, reactions=tuple(), reaction_gaps=tuple(), _rhs=None, param_meta=(lambda: MappingProxyType({}))(), state_meta=(lambda: MappingProxyType({}))(), axis_meta=(lambda: MappingProxyType({}))(), observables=(lambda: MappingProxyType({}))(), gauges=None, gauge_fix=None) dataclass

Container for a compiled RHS evaluation function.

Instances produced by :func:compile_rhs retain a private reference to their source :class:NormalizedRhs so the container can be pickled and re-hydrated by re-running the compile pipeline on load. eval_fn itself is a closure (and on the vectorized path captures compiled code objects), so it is dropped from the pickle and rebuilt by :func:compile_rhs in :meth:__setstate__. Round-tripping a CompiledRhs therefore costs one compile on load and yields a functionally equivalent instance whose eval_fn produces identical outputs for identical inputs.

__getstate__()

Return picklable state.

The compiled eval_fn is a closure (and on the vectorized path captures compiled :class:types.CodeType objects), which is not portably picklable. Instead we serialize just the source :class:NormalizedRhs and let :meth:__setstate__ recompile.

Raises:

Type	Description
TypeError	If the source NormalizedRhs was not retained (i.e. the instance was constructed directly rather than via :func:compile_rhs).

Source code in src/op_system/compile.py

def __getstate__(self) -> dict[str, Any]:
    """Return picklable state.

    The compiled ``eval_fn`` is a closure (and on the vectorized path
    captures compiled :class:`types.CodeType` objects), which is not
    portably picklable. Instead we serialize just the source
    :class:`NormalizedRhs` and let :meth:`__setstate__` recompile.

    Raises:
        TypeError: If the source ``NormalizedRhs`` was not retained
            (i.e. the instance was constructed directly rather than
            via :func:`compile_rhs`).
    """
    if self._rhs is None:
        msg = (
            "CompiledRhs is not picklable: the source NormalizedRhs was "
            "not retained. Construct via compile_rhs() to produce a "
            "picklable CompiledRhs."
        )
        raise TypeError(msg)
    return {"_rhs": self._rhs}

__setstate__(state)

Restore by recompiling from the pickled :class:NormalizedRhs.

Source code in src/op_system/compile.py

def __setstate__(self, state: Mapping[str, Any]) -> None:
    """Restore by recompiling from the pickled :class:`NormalizedRhs`."""
    rhs = state["_rhs"]
    rebuilt = compile_rhs(rhs)
    # frozen+slots dataclass: bypass __setattr__ via object.__setattr__
    object.__setattr__(self, "state_names", rebuilt.state_names)
    object.__setattr__(self, "param_names", rebuilt.param_names)
    object.__setattr__(self, "param_meta", rebuilt.param_meta)
    object.__setattr__(self, "state_meta", rebuilt.state_meta)
    object.__setattr__(self, "axis_meta", rebuilt.axis_meta)
    object.__setattr__(self, "observables", rebuilt.observables)
    object.__setattr__(self, "gauges", rebuilt.gauges)
    object.__setattr__(self, "gauge_fix", rebuilt.gauge_fix)
    object.__setattr__(self, "eval_fn", rebuilt.eval_fn)
    object.__setattr__(self, "meta", rebuilt.meta)
    object.__setattr__(self, "operators", rebuilt.operators)
    object.__setattr__(self, "factorize_axes", rebuilt.factorize_axes)
    object.__setattr__(self, "block_axes", rebuilt.block_axes)
    object.__setattr__(self, "pytree_eval_fn", rebuilt.pytree_eval_fn)
    object.__setattr__(self, "template_shapes", rebuilt.template_shapes)
    object.__setattr__(self, "block_pytree_eval_fn", rebuilt.block_pytree_eval_fn)
    object.__setattr__(self, "block_template_shapes", rebuilt.block_template_shapes)
    object.__setattr__(self, "history_requirements", rebuilt.history_requirements)
    object.__setattr__(self, "history_eval_fn", rebuilt.history_eval_fn)
    object.__setattr__(self, "body_eval_fn", rebuilt.body_eval_fn)
    object.__setattr__(self, "block_history_eval_fn", rebuilt.block_history_eval_fn)
    object.__setattr__(self, "block_body_eval_fn", rebuilt.block_body_eval_fn)
    object.__setattr__(self, "reactions", rebuilt.reactions)
    object.__setattr__(self, "reaction_gaps", rebuilt.reaction_gaps)
    object.__setattr__(self, "_rhs", rhs)

bind(params)

Bind parameter values and return a 2-arg RHS: rhs(t, y) -> dydt.

Parameters:

Name	Type	Description	Default
params	Mapping[str, object]	Mapping of parameter names to values.	required

Returns:

Type	Description
Callable[[object, object], Float64Array]	A callable rhs(t, y) that evaluates the RHS with params fixed.

Source code in src/op_system/compile.py

def bind(
    self, params: Mapping[str, object]
) -> Callable[[object, object], Float64Array]:
    """Bind parameter values and return a 2-arg RHS: rhs(t, y) -> dydt.

    Args:
        params: Mapping of parameter names to values.

    Returns:
        A callable `rhs(t, y)` that evaluates the RHS with `params` fixed.
    """
    params_dict = dict(params)

    def rhs(t: object, y: object) -> Float64Array:
        return self.eval_fn(t, y, **params_dict)

    return rhs

coordinate_view()

Build deterministic evaluators in the declared state coordinates.

Returns:

Type	Description
StateCoordinateView	An explicit view with initial-state encoders and natural-scale
StateCoordinateView	observations. The original evaluator/reactions remain natural.

Source code in src/op_system/compile.py

def coordinate_view(self) -> StateCoordinateView:
    """Build deterministic evaluators in the declared state coordinates.

    Returns:
        An explicit view with initial-state encoders and natural-scale
        observations. The original evaluator/reactions remain natural.
    """
    from .coordinates import _build_state_coordinates  # ruff: ignore[import-outside-top-level]

    return _build_state_coordinates(self)

CoordinateResult

Bases: NamedTuple, Generic[_StateT]

Transformed state and a namespace-native scalar domain flag.

Attributes:

Name	Type	Description
state	_StateT	Array or PyTree after transformation; no floor or clamp is applied.
valid	Array	Finite state with strictly positive selected physical components. For decoding this also flags exp underflow to zero and overflow.

EvalFn

Bases: Protocol

Callable RHS evaluator supporting runtime parameter kwargs.

Accepts a flat (n_state,) state array and returns a flat (n_state,) derivative array in the same array namespace.

ExprRhs(state_names, equations, aliases, param_names, all_symbols, meta, state_templates=(), shaped_params=(), time_varying_params=(), aliases_ir=dict(), equations_ir=(), aliases_ir_reduce=dict(), equations_ir_reduce=(), alias_templates=()) dataclass

Bases: _RhsBase

Normalized RHS for kind="expr" specs (explicit d(state)/dt equations).

Produced by :func:normalize_expr_rhs. Use :data:NormalizedRhs as the union type when you need to accept both kinds.

ExpressionString(source) dataclass

Validated expression wrapper with cached AST and symbol names.

as_ir()

Return the typed IR representation for this expression.

Returns:

Type	Description
Expr	Parsed typed IR tree.

Source code in src/op_system/_symbols.py

def as_ir(self) -> Expr:
    """Return the typed IR representation for this expression.

    Returns:
        Parsed typed IR tree.
    """
    return to_ir(self.ast)

as_lowered_ir()

Return helper-lowered typed IR for this expression.

Returns:

Type	Description
Expr	Typed IR tree with helper calls lowered to Reduce nodes.

Source code in src/op_system/_symbols.py

def as_lowered_ir(self) -> Expr:
    """Return helper-lowered typed IR for this expression.

    Returns:
        Typed IR tree with helper calls lowered to ``Reduce`` nodes.
    """
    return parse_expr_to_ir(self.source, lower_helpers=True)

GaugeCoordinates

Bases: NamedTuple

A representative together with the group coordinates it discarded.

Attributes:

Name	Type	Description
state	Array	Flat state in the fixed representative's units.
parameters	dict[str, object]	Remaining parameters; anchor values are omitted.
log_scales	Array	Group coordinates needed to invert this particular map.
valid	Array	Native scalar flag for positive anchors and finite coordinates. Invalid domains produce nonfinite arithmetic; no floor is applied.

GaugeFixMap

Bases: NamedTuple

Positive scalar anchors for a full-rank shared scaling action.

Attributes:

Name	Type	Description
fixed_values	tuple[tuple[str, float], ...]	Scalar parameter anchors and target values.
generators	tuple[ScalingGauge, ...]	Independent group basis used by the map.
state_weights	tuple[tuple[float, ...], ...]	One generator-weight row per expanded state component.
inverse_anchor_matrix	tuple[tuple[float, ...], ...]	Inverse anchor-weight matrix, stored immutably.

decode(coordinates)

Recover original values using retained group coordinates.

Parameters:

Name	Type	Description	Default
coordinates	GaugeCoordinates	Encoded representative with its original log scales.	required

Returns:

Type	Description
Array	Original flat state and complete parameter mapping. A quotient alone
dict[str, object]	has no inverse; these retained scales select the original member.

Source code in src/op_system/gauges.py

def decode(self, coordinates: GaugeCoordinates) -> tuple[Array, dict[str, object]]:
    """Recover original values using retained group coordinates.

    Args:
        coordinates: Encoded representative with its original log scales.

    Returns:
        Original flat state and complete parameter mapping. A quotient alone
        has no inverse; these retained scales select the original member.
    """
    state, remaining, logs, _ = coordinates
    xp = array_namespace(state)
    factors = xp.exp(xp.asarray(self.state_weights, dtype=state.dtype) @ logs)
    factors = xp.reshape(
        factors, (len(self.state_weights),) + (1,) * (len(state.shape) - 1)
    )
    params = {**remaining, **dict(self.fixed_values)}
    original = {
        name: xp.asarray(value, dtype=state.dtype)
        * xp.exp(
            -xp.sum(
                xp.asarray(
                    [
                        dict(generator.parameters).get(name, 0.0)
                        for generator in self.generators
                    ],
                    dtype=state.dtype,
                )
                * logs
            )
        )
        for name, value in params.items()
    }
    return state / factors, original

encode(state, parameters)

Map natural state/parameters to the anchored representative.

Parameters:

Name	Type	Description	Default
state	Array	Real floating flat state, with optional trailing batch axes.	required
parameters	Mapping[str, object]	Original parameter values, including positive anchors.	required

Returns:

Type	Description
GaugeCoordinates	Representative, original group coordinates and native domain flag.
GaugeCoordinates	Check valid before using a result in an inference pipeline.

Raises:

Type	Description
ValueError	If state axes or anchor shapes do not match this map.
TypeError	If states are not real floating arrays.

Source code in src/op_system/gauges.py

def encode(
    self, state: Array, parameters: Mapping[str, object]
) -> GaugeCoordinates:
    """Map natural state/parameters to the anchored representative.

    Args:
        state: Real floating flat state, with optional trailing batch axes.
        parameters: Original parameter values, including positive anchors.

    Returns:
        Representative, original group coordinates and native domain flag.
        Check ``valid`` before using a result in an inference pipeline.

    Raises:
        ValueError: If state axes or anchor shapes do not match this map.
        TypeError: If states are not real floating arrays.
    """
    if not state.shape or state.shape[0] != len(self.state_weights):
        msg = "Gauge map state axis does not match the compiled model"
        raise ValueError(msg)
    xp = array_namespace(state)
    if not xp.isdtype(state.dtype, "real floating"):
        msg = "Gauge coordinates require real floating state arrays"
        raise TypeError(msg)
    values = xp.stack([
        xp.asarray(parameters[name], dtype=state.dtype)
        for name, _ in self.fixed_values
    ])
    if len(values.shape) != 1:
        msg = "Gauge anchors must be scalar; use vmap for separate parameter draws"
        raise ValueError(msg)
    targets = xp.asarray(
        [value for _, value in self.fixed_values], dtype=state.dtype
    )
    matrix = xp.asarray(self.inverse_anchor_matrix, dtype=state.dtype)
    logs = matrix @ (xp.log(targets) - xp.log(values))
    valid = xp.all(xp.logical_and(xp.isfinite(values), values > 0))
    valid = xp.logical_and(valid, xp.all(xp.isfinite(logs)))
    weights = xp.asarray(self.state_weights, dtype=state.dtype)
    factors = xp.exp(weights @ logs)
    valid = xp.logical_and(
        valid, xp.all(xp.logical_and(xp.isfinite(factors), factors > 0))
    )
    factors = xp.reshape(
        factors, (len(self.state_weights),) + (1,) * (len(state.shape) - 1)
    )
    transformed = {
        name: xp.asarray(value, dtype=state.dtype)
        * xp.exp(
            xp.sum(
                xp.asarray(
                    [
                        dict(generator.parameters).get(name, 0.0)
                        for generator in self.generators
                    ],
                    dtype=state.dtype,
                )
                * logs
            )
        )
        for name, value in parameters.items()
        if name not in dict(self.fixed_values)
    }
    mapped_state = state * factors
    valid = xp.logical_and(valid, xp.all(xp.isfinite(mapped_state)))
    valid = xp.logical_and(
        valid,
        xp.all(xp.logical_or(xp.not_equal(mapped_state, 0), xp.equal(state, 0))),
    )
    for name, value in transformed.items():
        original = xp.asarray(parameters[name], dtype=state.dtype)
        finite = xp.all(xp.isfinite(value))
        retained = xp.all(xp.logical_or(value != 0, original == 0))
        valid = xp.logical_and(valid, xp.logical_and(finite, retained))
    return GaugeCoordinates(mapped_state, transformed, logs, valid)

GaugeInfo

Bases: NamedTuple

Declared/detected shared-template scalings and their rational ranks.

Attributes:

Name	Type	Description
declared	tuple[ScalingGauge, ...]	User-supplied generators verified at compile time.
detected	tuple[ScalingGauge, ...]	Conservative structural nullspace basis, if detection requested.
generators	tuple[ScalingGauge, ...]	Independent basis spanning the generators being reported.
rank	int	Rank on states and parameters together.
parameter_rank	int	Rank of the parameter projection of that action.
parameter_names	tuple[str, ...]	Scalar/shaped parameter bases in coordinate order.
invariant_coordinates	tuple[InvariantCoordinate, ...]	Integer nullspace monomials for parameters. This count is not an identifiable dimension estimate.
verification	GaugeVerification | None	Finite-probe record for declared generators, or None.
detection_requested	bool	Distinguishes no detection from an empty nullspace.

GaugeVerification

Bases: NamedTuple

Reproducible finite probes, rather than a global numerical proof.

Attributes:

Name	Type	Description
seed	int	Local NumPy random seed; no global random state is changed.
n_probes	int	Points tested per generator, at both reported scales.
scales	tuple[float, ...]	Positive group elements tested independently per generator.
rtol	float	Relative tolerance used for every compared component.
atol	float	Absolute tolerance used for every compared component.
max_absolute_error	float	Largest absolute covariance/invariance residual.
observables	tuple[str, ...]	All declared output names checked for invariance.
structural	bool	Whether supported homogeneity rules also verify declarations.

InvariantCheck

Bases: NamedTuple

A residual-to-zero balance check with explicit applicability conditions.

InvariantCoordinate

Bases: NamedTuple

A monomial invariant of the reported parameter exponent matrix.

Attributes:

Name	Type	Description
exponents	tuple[tuple[str, int], ...]	Parameter base names and nonzero integer powers.

evaluate(parameters)

Evaluate a monomial using supplied values, without domain correction.

Parameters:

Name	Type	Description	Default
parameters	Mapping[str, object]	Values with native array broadcasting semantics.	required

Returns:

Type	Description
object	Product of the named powers. Negative powers require nonzero values.

Examples:

>>> invariant = InvariantCoordinate((("K", 1), ("c", 1)))
>>> invariant.evaluate({"K": 4.0, "c": 3.0})
12.0

Source code in src/op_system/gauges.py

def evaluate(self, parameters: Mapping[str, object]) -> object:
    """Evaluate a monomial using supplied values, without domain correction.

    Args:
        parameters: Values with native array broadcasting semantics.

    Returns:
        Product of the named powers. Negative powers require nonzero values.

    Examples:
        >>> invariant = InvariantCoordinate((("K", 1), ("c", 1)))
        >>> invariant.evaluate({"K": 4.0, "c": 3.0})
        12.0
    """
    result: object = 1.0
    for name, exponent in self.exponents:
        # Arithmetic is namespace-native for concrete values and tracers.
        result *= parameters[name] ** exponent  # type: ignore[operator]
    return result

ObservableDefinition

Bases: NamedTuple

Backend-neutral definition of one instantaneous latent observable.

Attributes:

Name	Type	Description
name	str	Unique output name within a compiled model.
expression	str	Restricted pointwise state/parameter/time expression.
axes	tuple[str, ...]	Output axes, after any explicit reduction.
reduction	ObservableReduction | None	Optional unweighted sum or physical integral.
units	str | None	Optional declared output units, without dimensional verification.
description	str | None	Optional human-readable interpretation.

from_mapping(name, value) classmethod

Validate the bounded observable schema.

Parameters:

Name	Type	Description	Default
name	str	Output identifier.	required
value	object	Mapping with expr, optional axes/reduction/units/description.	required

Returns:

Type	Description
ObservableDefinition	Detached immutable definition.

Raises:

Type	Description
InvalidRhsSpecError	If identifiers, fields or reductions are invalid.

Source code in src/op_system/_observables.py

@classmethod
def from_mapping(cls, name: str, value: object) -> ObservableDefinition:
    """Validate the bounded observable schema.

    Args:
        name: Output identifier.
        value: Mapping with expr, optional axes/reduction/units/description.

    Returns:
        Detached immutable definition.

    Raises:
        InvalidRhsSpecError: If identifiers, fields or reductions are invalid.
    """
    if not _identifier(name):
        raise InvalidRhsSpecError(detail="observable names must be identifiers")
    if not isinstance(value, Mapping) or set(value) - {
        "expr",
        "axes",
        "reduction",
        "units",
        "description",
    }:
        raise InvalidRhsSpecError(detail=f"observable {name!r} has invalid fields")
    expression = _text(value, "expr", required=True)
    axes = _axes(value.get("axes", ()))
    reduction = None
    if "reduction" in value:
        raw = value["reduction"]
        if not isinstance(raw, Mapping) or set(raw) != {"kind", "axes"}:
            raise InvalidRhsSpecError(
                detail="observable reduction needs kind and axes"
            )
        kind = raw["kind"]
        reduction_axes = _axes(raw["axes"])
        if (
            not isinstance(kind, str)
            or kind not in {"sum", "integral"}
            or not reduction_axes
        ):
            raise InvalidRhsSpecError(
                detail="observable reduction must be a nonempty sum or integral"
            )
        reduction = ObservableReduction(
            cast('Literal["sum", "integral"]', kind), reduction_axes
        )
        if set(axes) & set(reduction_axes):
            raise InvalidRhsSpecError(
                detail="observable output and reduction axes must be disjoint"
            )
    return cls(
        name,
        cast("str", expression),
        axes,
        reduction,
        _text(value, "units"),
        _text(value, "description"),
    )

ObservableEvalFn

Bases: Protocol

Evaluate an instantaneous observable from a native flat state array.

ObservablePytreeEvalFn

Bases: Protocol

Evaluate an instantaneous observable from complete shaped state buffers.

ObservableReduction

Bases: NamedTuple

Explicit reduction independent of legacy equation quadrature.

Attributes:

Name	Type	Description
kind	Literal['sum', 'integral']	Unweighted sum or integral with declared physical measures.
axes	tuple[str, ...]	Named axes to remove, in the declared reduction order.

OperatorDescriptor(axis, kind=None, bc=None, velocity=None, rate=None, kernel=None, name=None, apply_to=None, direction=None) dataclass

Typed description of a spatial operator declared in an op_system RHS spec.

Captures the model-level description of an operator that is known at compile time: which axis it acts on, its kind (advection, diffusion, etc.), optional boundary condition, and the names of any runtime parameters (velocity, rate) it consumes. Grid geometry, CN matrix construction, and solver staging remain downstream concerns handled by the engine.

Attributes:

Name	Type	Description
axis	str	Name of the axis the operator acts on (e.g. "loc").
kind	str | None	Operator type string, e.g. "advection", "diffusion", "transport", "jump_integral", or None if unspecified.
bc	str | None	Boundary condition, e.g. "absorbing", "periodic", "neumann", "reflecting", or None if unspecified.
velocity	str | float | None	Parameter name or numeric constant for an advection velocity, or None.
rate	str | float | None	Parameter name or numeric constant for a rate/coefficient, or None.
kernel	Mapping[str, Any] | None	Mixing-kernel sub-specification, or None.
name	str | None	Optional operator name from the specification.
apply_to	tuple[str, ...] | None	Concrete state names selected by the operator, or None when the operator applies to every compatible state.
direction	str | None	Optional normalized direction. For advection and transport, "increasing" preserves the resolved velocity sign and "decreasing" reverses it; when omitted, velocity is used as a signed coefficient. For jump integrals, "up", "down", or "both" masks source-to-target movement in coordinate order.

Examples:

>>> od = OperatorDescriptor(axis="loc")
>>> od.axis
'loc'
>>> od.kind is None
True
>>> od.bc is None
True
>>> od.velocity is None
True
>>> od2 = OperatorDescriptor(
...     axis="loc",
...     kind="advection",
...     bc="absorbing",
...     velocity=0.25,
...     name="drift",
...     apply_to=("S",),
... )
>>> od2.kind
'advection'
>>> od2.bc
'absorbing'
>>> od2.velocity
0.25
>>> od2.name
'drift'
>>> od2.apply_to
('S',)

ParameterDeclaration

Bases: NamedTuple

Optional scientific metadata for an inferred runtime parameter.

Attributes:

Name	Type	Description
support	ParameterSupport | None	Declared domain, or None when it was not specified.
units	str | None	Unit annotation, without dimensional-analysis verification.
description	str | None	Optional human-readable parameter description.
fixed	bool	Whether inference should treat the parameter as fixed. This does not bind a value or remove a runtime parameter.

from_mapping(value) classmethod

Validate and copy one parameter declaration.

Parameters:

Name	Type	Description	Default
value	object	Mapping with optional support, units, description and fixed.	required

Returns:

Type	Description
ParameterDeclaration	Immutable parameter declaration.

Raises:

Type	Description
InvalidRhsSpecError	If an annotation is malformed or unknown.

Source code in src/op_system/_parameter_meta.py

@classmethod
def from_mapping(cls, value: object) -> ParameterDeclaration:
    """Validate and copy one parameter declaration.

    Args:
        value: Mapping with optional support, units, description and fixed.

    Returns:
        Immutable parameter declaration.

    Raises:
        InvalidRhsSpecError: If an annotation is malformed or unknown.
    """
    if not isinstance(value, Mapping) or any(
        key not in {"support", "units", "description", "fixed"} for key in value
    ):
        raise InvalidRhsSpecError(
            detail=(
                "parameter declaration must contain only support, units, "
                "description and fixed"
            )
        )
    fixed = value.get("fixed", False)
    if not isinstance(fixed, bool):
        raise InvalidRhsSpecError(detail="parameter fixed must be a boolean")
    support = (
        ParameterSupport.from_mapping(value["support"])
        if "support" in value
        else None
    )
    return cls(
        support,
        _optional_text(value, "units"),
        _optional_text(value, "description"),
        fixed,
    )

ParameterSupport

Bases: NamedTuple

Declared domain; bounds are inclusive and positivity is strict.

Attributes:

Name	Type	Description
kind	_SupportKind	Domain identifier; unit_interval includes zero and one.
lower	float | None	Finite lower bound for bounded support, otherwise None.
upper	float | None	Finite upper bound for bounded support, otherwise None.

from_mapping(value) classmethod

Parse a domain name or bounded-domain mapping.

Parameters:

Name	Type	Description	Default
value	object	Domain string or mapping with kind, lower and upper.	required

Returns:

Type	Description
ParameterSupport	Immutable support declaration.

Raises:

Type	Description
InvalidRhsSpecError	If the domain or bounds are invalid.

Source code in src/op_system/_parameter_meta.py

@classmethod
def from_mapping(cls, value: object) -> ParameterSupport:
    """Parse a domain name or bounded-domain mapping.

    Args:
        value: Domain string or mapping with kind, lower and upper.

    Returns:
        Immutable support declaration.

    Raises:
        InvalidRhsSpecError: If the domain or bounds are invalid.
    """
    if isinstance(value, str) and value in {"real", "positive", "unit_interval"}:
        return cls(cast("_SupportKind", value))
    if not isinstance(value, Mapping) or set(value) != {"kind", "lower", "upper"}:
        raise InvalidRhsSpecError(
            detail=(
                "parameter support must be real, positive, unit_interval, "
                "or a bounded mapping with kind, lower and upper"
            )
        )
    lower = _finite_bound(value["lower"])
    upper = _finite_bound(value["upper"])
    if value["kind"] != "bounded" or lower >= upper:
        raise InvalidRhsSpecError(
            detail="bounded parameter support requires finite numeric lower < upper"
        )
    return cls("bounded", float(lower), float(upper))

PytreeEvalFn

Bases: Protocol

Callable RHS evaluator operating on shaped PyTree state dicts.

Accepts y as a StateDict (mapping from state-template base name to a shaped array with the template's natural N-D shape) and returns a StateDict of the same structure containing the derivative. Enables the engine to skip the flatten/unflatten step entirely and expose the full tensor structure to JAX/XLA.

QualificationConfig(rtol=1e-07, atol=1e-09, derivative_step=1e-05, derivative_rtol=0.0001, derivative_atol=1e-06) dataclass

Reference and derivative tolerances for eager qualification.

Defaults assume double precision. Choose wider tolerances and a larger finite-difference step for float32 fixtures.

__post_init__()

Validate controls.

Raises:

Type	Description
ValueError	If a tolerance or finite-difference step is invalid.

Source code in src/op_system/qualification.py

def __post_init__(self) -> None:
    """Validate controls.

    Raises:
        ValueError: If a tolerance or finite-difference step is invalid.
    """
    for name in (
        "rtol",
        "atol",
        "derivative_rtol",
        "derivative_atol",
        "derivative_step",
    ):
        value = getattr(self, name)
        numeric = isinstance(value, Real) and not isinstance(value, bool)
        if (
            not numeric
            or not math.isfinite(value)
            or value < 0
            or (name == "derivative_step" and value == 0)
        ):
            bound = "positive" if name == "derivative_step" else "non-negative"
            msg = f"{name} must be finite and {bound}"
            raise ValueError(msg)

QualificationFixture

Bases: NamedTuple

Named structured state/parameter fixture and optional state direction.

QualificationReport(fixtures) dataclass

Inspectable results for all requested fixture/evaluator combinations.

ok property

Whether all fixtures and requested evaluators qualified.

to_dict()

Return JSON-compatible fields, including per-fixture pass/fail status.

Returns:

Type	Description
dict[str, object]	A mapping suitable for json.dumps(..., allow_nan=False).

Source code in src/op_system/qualification.py

def to_dict(self) -> dict[str, object]:
    """Return JSON-compatible fields, including per-fixture pass/fail status.

    Returns:
        A mapping suitable for ``json.dumps(..., allow_nan=False)``.
    """
    return {
        "ok": self.ok,
        "fixtures": [
            {
                **asdict(item),
                "ok": item.ok,
                "arrays": [value._asdict() for value in item.arrays],
                "invariants": [value._asdict() for value in item.invariants],
                "derivative": item.derivative._asdict()
                if item.derivative
                else None,
            }
            for item in self.fixtures
        ],
    }

ReactionBalanceLaw

Bases: NamedTuple

A conditional identity d(coefficients @ state)/dt = source fluxes.

Attributes:

Name	Type	Description
name	str	Component label stable within this report's compiled state order.
coefficients	tuple[int, ...]	Exact integer weights in the report's state order.
reaction_names	tuple[str, ...]	All reaction identities covered by the analysis.
sources	tuple[ReactionBalanceSource, ...]	Nonzero external terms, resolved by reaction name.
assumptions	tuple[str, ...]	Conditions needed to apply this structural identity.
origin	Literal['compiled_unit_transfers']	Compiler unit-transfer derivation, rather than a user assertion.

closed property

Whether this identity has no represented external source term.

ReactionBalanceReport

Bases: NamedTuple

Immutable scalar balance analysis with explicit scope and coverage.

Attributes:

Name	Type	Description
state_names	tuple[str, ...]	Coefficient-vector coordinate order.
laws	tuple[ReactionBalanceLaw, ...]	Issued identities; empty if any coverage gap exists.
coverage_gaps	tuple[BalanceCoverageGap, ...]	Unsupported/omitted dynamics or inconsistent metadata.
state_meta	tuple[tuple[str, StateDeclaration], ...]	State declarations in compiled state order; unknowns stay absent.
axis_meta	tuple[tuple[str, AxisDeclaration], ...]	Retained scientific axis declarations; no weights are inferred.
assumptions	tuple[str, ...]	Shared applicability conditions, including unknown physical units.
scope	Literal['scalar_unit_transfers']	Bounded analysis capability, excluding general stoichiometry/meshes.

fully_covered property

Whether this bounded analysis covers the supplied structure.

ReactionBalanceSource

Bases: NamedTuple

One external flux term in a component's balance equation.

Attributes:

Name	Type	Description
reaction_name	str	Source-only reaction whose propensity supplies the flux.
coefficient	int	Exact multiplier of that propensity in the balance.
target_state	str	Compiled scalar state receiving the source.

ReactionGap(origin, name, source, target, reason) dataclass

Model dynamics that publish no compiled reaction artifact.

A consumer that executes only CompiledRhs.reactions (for example a pure stochastic simulation) silently loses these dynamics. An empty CompiledRhs.reaction_gaps means every transition has an artifact.

Attributes:

Name	Type	Description
origin	str	Where the dynamics come from in the spec, such as transitions[2], chain[0].forward[1], or equations for an expr spec.
name	str | None	The transition's name, or None when it has none.
source	str | None	The transition's from selector, or None for a source-only transition or an expr spec.
target	str | None	The transition's to selector, or None for an expr spec.
reason	str	Why no artifact exists: unnamed: the transition has no name. target_axis_not_on_source: to has a wildcard axis that from lacks. rate_axis_out_of_scope: the rate has a free axis outside the reaction's channel axes. unsupported_layout: the compiled RHS has no vectorized state layout to index propensities against. unresolved_alias: the rate still references an alias after inlining: one on a reference cycle, a templated alias referenced by bare name, or a declaration out of scope for inlining. compile_failed: the propensity could not be lowered or its coordinates could not be resolved. gauge_fix_requires_jump_contract: equation/output covariance does not verify a jump process's distribution. Gauge-fixed models omit reaction exports until a separate jump contract is supplied. expr_spec: an expr spec, whose equations are not decomposed into transitions.

ReactionPropensityFn

Bases: Protocol

Callable propensity evaluator for one named reaction/transition.

Accepts y as a StateDict and returns an array shaped like the reaction's from_axes (one independent rate per source cell) -- see :class:CompiledReaction.

ScalingGauge

Bases: NamedTuple

One positive scaling generator, shared over each named template.

For scale s > 0, an exponent w acts as value * s**w. Declaration does not establish covariance or output invariance; compiled analysis records those checks separately.

Attributes:

Name	Type	Description
name	str	Unique generator label.
states	tuple[tuple[str, float], ...]	State base names and finite scaling exponents.
parameters	tuple[tuple[str, float], ...]	Parameter base names and finite scaling exponents.

ShapeGroup(cells, example_cell, example_expression) dataclass

Cells of one state template that share an expression shape.

StateCoordinateView(natural, log_indices, log_bases, eval_fn, pytree_eval_fn, block_pytree_eval_fn, observables) dataclass

Deterministic evaluators in explicitly selected solver coordinates.

Natural model state, reactions, scientific declarations and observations remain available on natural. This view does not export jump artifacts, balance laws or original derivative-affinity claims in solver coordinates.

Attributes:

Name	Type	Description
natural	CompiledRhs	Source model, in natural coordinates (possibly gauge-fixed).
log_indices	tuple[int, ...]	Selected expanded flat state components.
log_bases	tuple[str, ...]	State bases using log coordinates, shared over their templates.
eval_fn	EvalFn	Coordinate derivative, with the source flat evaluator contract.
pytree_eval_fn	PytreeEvalFn | None	Optional coordinate derivative for shaped state buffers.
block_pytree_eval_fn	PytreeEvalFn | None	Optional coordinate derivative for stripped blocks.
observables	Mapping[str, CompiledObservable]	Natural-scale observables evaluated from encoded states.

Examples:

>>> import numpy as np
>>> from op_system import compile_spec
>>> model = compile_spec({
...     "kind": "expr",
...     "state": ["x"],
...     "state_metadata": {"x": {"scale": "log"}},
...     "equations": {"x": "-rate*x"},
... })
>>> view = model.coordinate_view()
>>> z0 = view.encode(np.array([2.0]))
>>> np.allclose(view.eval_fn(0, z0, rate=0.5), [-0.5])
True
>>> np.allclose(view.decode(z0), [2.0])
True

param_names property

Source equation parameter requirements, omitting gauge anchors.

state_names property

Source component order; selected values now store log states.

__post_init__()

Cache component selection once, independently of execution namespace.

Source code in src/op_system/coordinates.py

def __post_init__(self) -> None:
    """Cache component selection once, independently of execution namespace."""
    selected = frozenset(self.log_indices)
    object.__setattr__(
        self,
        "_mask",
        tuple(index in selected for index in range(len(self.natural.state_names))),
    )

__reduce__()

Rebuild evaluators from the retained natural source when unpickling.

Returns:

Type	Description
tuple[object, tuple[CompiledRhs]]	Module factory and source model, without serialized closures.

Source code in src/op_system/coordinates.py

def __reduce__(self) -> tuple[object, tuple[CompiledRhs]]:
    """Rebuild evaluators from the retained natural source when unpickling.

    Returns:
        Module factory and source model, without serialized closures.
    """
    return _build_state_coordinates, (self.natural,)

bind(params)

Bind source parameters to the coordinate derivative.

Returns:

Type	Description
Callable[[object, object], Float64Array]	A two-argument deterministic RHS, with solver-coordinate state.

Source code in src/op_system/coordinates.py

def bind(
    self, params: Mapping[str, object]
) -> Callable[[object, object], Float64Array]:
    """Bind source parameters to the coordinate derivative.

    Returns:
        A two-argument deterministic RHS, with solver-coordinate state.
    """
    values = dict(params)

    def evaluate(t: object, state: object) -> Float64Array:
        return self.eval_fn(t, state, **values)

    return evaluate

decode(state, *, state_axis=0)

Return native decoded values without an eager domain check.

Parameters:

Name	Type	Description	Default
state	Array	Floating solver coordinates.	required
state_axis	int	Component axis; use 1 for output trajectories.	0

Returns:

Type	Description
Array	Natural-scale array, for transformed computations and observations.
Array	Use decode_native or decode_checked when checking representation.

Source code in src/op_system/coordinates.py

def decode(self, state: Array, *, state_axis: int = 0) -> Array:
    """Return native decoded values without an eager domain check.

    Args:
        state: Floating solver coordinates.
        state_axis: Component axis; use 1 for output trajectories.

    Returns:
        Natural-scale array, for transformed computations and observations.
        Use decode_native or decode_checked when checking representation.
    """
    return self.decode_native(state, state_axis=state_axis).state

decode_checked(state, *, state_axis=0)

Decode and eagerly reject underflow/overflow or invalid values.

Parameters:

Name	Type	Description	Default
state	Array	Floating solver coordinates.	required
state_axis	int	Component axis in source order.	0

Returns:

Type	Description
Array	A finite natural state, positive in selected components.

Source code in src/op_system/coordinates.py

def decode_checked(self, state: Array, *, state_axis: int = 0) -> Array:
    """Decode and eagerly reject underflow/overflow or invalid values.

    Args:
        state: Floating solver coordinates.
        state_axis: Component axis in source order.

    Returns:
        A finite natural state, positive in selected components.
    """
    result = self.decode_native(state, state_axis=state_axis)
    _require_valid(result.valid, "decode")
    return result.state

decode_native(state, *, state_axis=0)

Decode natural values and flag finite-precision representation failures.

Parameters:

Name	Type	Description	Default
state	Array	Floating solver coordinates.	required
state_axis	int	Component axis; use 1 for output trajectories.	0

Returns:

Type	Description
CoordinateResult[Array]	Natural values and finite/positive-domain flag. exp can underflow or
CoordinateResult[Array]	overflow; the flag is not an integration-accuracy assessment.

Source code in src/op_system/coordinates.py

def decode_native(
    self, state: Array, *, state_axis: int = 0
) -> CoordinateResult[Array]:
    """Decode natural values and flag finite-precision representation failures.

    Args:
        state: Floating solver coordinates.
        state_axis: Component axis; use 1 for output trajectories.

    Returns:
        Natural values and finite/positive-domain flag. exp can underflow or
        overflow; the flag is not an integration-accuracy assessment.
    """
    return _native_flat(state, self._mask, state_axis, encode=False)

decode_pytree_native(state)

Decode complete buffers and return native representation flags.

Returns:

Type	Description
CoordinateResult[dict[str, Array]]	Natural buffers with a native finite/positive flag.

Source code in src/op_system/coordinates.py

def decode_pytree_native(
    self, state: Mapping[str, Array]
) -> CoordinateResult[dict[str, Array]]:
    """Decode complete buffers and return native representation flags.

    Returns:
        Natural buffers with a native finite/positive flag.
    """
    return self._tree(state, encode=False)

encode(state, *, state_axis=0)

Check a physical initial state eagerly before taking logarithms.

Parameters:

Name	Type	Description	Default
state	Array	Finite floating natural state, positive in selected components.	required
state_axis	int	Component axis, in source order.	0

Returns:

Type	Description
Array	Initial state in solver coordinates. Use encode_native under JIT/AD.

Source code in src/op_system/coordinates.py

def encode(self, state: Array, *, state_axis: int = 0) -> Array:
    """Check a physical initial state eagerly before taking logarithms.

    Args:
        state: Finite floating natural state, positive in selected components.
        state_axis: Component axis, in source order.

    Returns:
        Initial state in solver coordinates. Use encode_native under JIT/AD.
    """
    xp = array_namespace(state)
    values, mask = _array_and_mask(state, self._mask, state_axis)
    valid = xp.logical_and(
        xp.all(xp.isfinite(values)),
        xp.all(xp.logical_or(xp.logical_not(mask), xp.greater(values, 0))),
    )
    _require_valid(valid, "encode")
    return self.encode_native(state, state_axis=state_axis).state

encode_native(state, *, state_axis=0)

Map physical arrays to selected log coordinates, with native flags.

Parameters:

Name	Type	Description	Default
state	Array	Real floating natural state, with component axis in source order.	required
state_axis	int	Component axis; use 1 for output trajectories.	0

Returns:

Type	Description
CoordinateResult[Array]	Encoded state and scalar flag without host extraction. Invalid
CoordinateResult[Array]	selected inputs produce nonfinite logarithms; no floor is applied.

Source code in src/op_system/coordinates.py

def encode_native(
    self, state: Array, *, state_axis: int = 0
) -> CoordinateResult[Array]:
    """Map physical arrays to selected log coordinates, with native flags.

    Args:
        state: Real floating natural state, with component axis in source order.
        state_axis: Component axis; use 1 for output trajectories.

    Returns:
        Encoded state and scalar flag without host extraction. Invalid
        selected inputs produce nonfinite logarithms; no floor is applied.
    """
    return _native_flat(state, self._mask, state_axis, encode=True)

encode_pytree_native(state)

Encode complete buffers and return native finite/positive flags.

Returns:

Type	Description
CoordinateResult[dict[str, Array]]	Solver-coordinate buffers with a native domain flag.

Source code in src/op_system/coordinates.py

def encode_pytree_native(
    self, state: Mapping[str, Array]
) -> CoordinateResult[dict[str, Array]]:
    """Encode complete buffers and return native finite/positive flags.

    Returns:
        Solver-coordinate buffers with a native domain flag.
    """
    return self._tree(state, encode=True)

StateDeclaration

Bases: NamedTuple

Optional representation and physical reduction requirements of a state.

Attributes:

Name	Type	Description
representation	_Representation | None	Extensive counts/mass/counter or intensive density/ concentration; None when unknown. Cumulative means an extensive accumulated counter, without asserting monotonicity of its equations.
units	str | None	State value units, without dimensional verification.
description	str | None	Human-readable annotation.
measure_axes	tuple[str, ...] | None	Axes whose measures are required for an intensive integral. None means undeclared; an empty tuple explicitly needs none.
scale	Literal['identity', 'log']	Requested identity/log coordinate for an explicit coordinate_view. Primary compiled equations and reactions retain natural coordinates.

from_mapping(value) classmethod

Validate a state annotation without changing its equations.

Parameters:

Name	Type	Description	Default
value	object	Mapping of representation, units, description and measure_axes.	required

Returns:

Type	Description
StateDeclaration	Immutable declaration; omitted information remains unknown.

Raises:

Type	Description
InvalidRhsSpecError	If fields or their values are invalid.

Source code in src/op_system/_scientific_meta.py

@classmethod
def from_mapping(cls, value: object) -> StateDeclaration:
    """Validate a state annotation without changing its equations.

    Args:
        value: Mapping of representation, units, description and measure_axes.

    Returns:
        Immutable declaration; omitted information remains unknown.

    Raises:
        InvalidRhsSpecError: If fields or their values are invalid.
    """
    if not isinstance(value, Mapping) or set(value) - {
        "representation",
        "units",
        "description",
        "measure_axes",
        "scale",
    }:
        raise InvalidRhsSpecError(
            detail="state metadata contains unknown fields or is not a mapping"
        )
    scale = value.get("scale", "identity")
    if not isinstance(scale, str) or scale not in {"identity", "log"}:
        raise InvalidRhsSpecError(detail="state scale must be identity or log")
    representation = value.get("representation")
    if "representation" in value and (
        not isinstance(representation, str)
        or representation
        not in {"counts", "cell_mass", "density", "concentration", "cumulative"}
    ):
        raise InvalidRhsSpecError(
            detail=(
                "state representation must be counts, cell_mass, density, "
                "concentration or cumulative"
            )
        )
    measure_axes = None
    if "measure_axes" in value:
        axes = value["measure_axes"]
        if (
            not isinstance(axes, (list, tuple))
            or any(not isinstance(axis, str) or not axis.strip() for axis in axes)
            or len(set(axes)) != len(axes)
        ):
            raise InvalidRhsSpecError(
                detail="state measure_axes must list unique nonempty axis names"
            )
        measure_axes = tuple(axes)
        if representation in {"counts", "cell_mass", "cumulative"} and measure_axes:
            raise InvalidRhsSpecError(
                detail="extensive states do not require measure weights"
            )
    return cls(
        cast("_Representation | None", representation),
        _optional_text(value, "units"),
        _optional_text(value, "description"),
        measure_axes,
        cast('Literal["identity", "log"]', scale),
    )

StateString

Bases: BaseModel

Structured representation of a state string.

A state string is either a bare state name like "S" or a state name followed immediately by bracketed dimensions like "R[age,vax]".

Examples:

>>> StateString.model_validate("S")
StateString(name='S', dims=())
>>> recovery = StateString.model_validate("R[age,vax]")
>>> recovery
StateString(name='R', dims=('age', 'vax'))
>>> print(recovery)
R[age,vax]
>>> recovery.model_dump()
'R[age,vax]'
>>> StateString.model_validate("Foobar[ age , vax ]")
StateString(name='Foobar', dims=('age', 'vax'))

__str__()

Return the compact string form.

Returns:

Type	Description
str	Compact state string.

Examples:

>>> str(StateString(name="S", dims=()))
'S'
>>> str(StateString(name="R", dims=("age",)))
'R[age]'
>>> str(StateString(name="R", dims=("age", "vax")))
'R[age,vax]'
>>> str(StateString(name="lambda", dims=("age", "vax", "state")))
'lambda[age,vax,state]'

Source code in src/op_system/_state_string.py

def __str__(self) -> str:
    """
    Return the compact string form.

    Returns:
        Compact state string.

    Examples:
        >>> str(StateString(name="S", dims=()))
        'S'
        >>> str(StateString(name="R", dims=("age",)))
        'R[age]'
        >>> str(StateString(name="R", dims=("age", "vax")))
        'R[age,vax]'
        >>> str(StateString(name="lambda", dims=("age", "vax", "state")))
        'lambda[age,vax,state]'
    """
    if not self.dims:
        return self.name
    dims = ",".join(self.dims)
    return f"{self.name}[{dims}]"

TransitionsRhs(state_names, equations, aliases, param_names, all_symbols, meta, state_templates=(), shaped_params=(), time_varying_params=(), aliases_ir=dict(), equations_ir=(), aliases_ir_reduce=dict(), equations_ir_reduce=(), alias_templates=(), reactions_ir=(), reaction_gaps_ir=()) dataclass

Bases: _RhsBase

Normalized RHS for kind="transitions" specs (per-capita hazard diagram).

Produced by :func:normalize_transitions_rhs. Use :data:NormalizedRhs as the union type when you need to accept both kinds.

ValidationReport(stages, errors=list(), cost=dict(), parameters=dict(), shape_groups=dict(), gauges=None) dataclass

Result of validating one op_system specification.

Attributes:

Name	Type	Description
stages	dict[str, str]	"passed", "failed", or "skipped" for normalize, compile, and vectorize.
errors	list[str]	Error messages from the failed stage.
cost	dict[str, int]	Compile-cost drivers (expanded states, templates, transitions, coordinate-pinned transitions, operators).
parameters	dict[str, tuple[str, ...]]	Parameters the spec consumes, mapped to their declared axes (empty tuple for scalars); operator parameters use kernel.param_axes when declared.
gauges	dict[str, object] | None	Optional reported scaling counts and finite-probe scope. Does not establish identifiable dimensions.
shape_groups	dict[str, list[ShapeGroup]]	When vectorization fails, the templates whose cells have more than one expression shape, with each distinct shape's count and an example cell. Empty otherwise, because comparing every cell's expanded equation is slow on large specs.

ok property

Whether every stage passed.

array_namespace(value)

Return an Array-API-compatible namespace for value.

array-api-compat recognizes both arrays that implement __array_namespace__ and native arrays such as torch.Tensor that need a compatibility namespace. Optional backends are imported only when one of their arrays is supplied.

Parameters:

Name	Type	Description	Default
value	object	Numerical array whose namespace should be selected.	required

Returns:

Type	Description
Any	The Array-API-compatible namespace for value.

Raises:

Type	Description
TypeError	If value is not a supported array object.

Source code in src/op_system/_array.py

def array_namespace(value: object) -> Any:  # ruff: ignore[any-type]
    """Return an Array-API-compatible namespace for ``value``.

    ``array-api-compat`` recognizes both arrays that implement
    ``__array_namespace__`` and native arrays such as ``torch.Tensor`` that
    need a compatibility namespace. Optional backends are imported only when
    one of their arrays is supplied.

    Args:
        value: Numerical array whose namespace should be selected.

    Returns:
        The Array-API-compatible namespace for ``value``.

    Raises:
        TypeError: If ``value`` is not a supported array object.
    """
    try:
        return _array_namespace(value)
    except TypeError as error:
        msg = (
            "op_system numerical inputs must be supported array objects; "
            f"got {type(value).__name__}."
        )
        raise TypeError(msg) from error

axis_kernel_generator_rhs(state, generator, *, axis, velocity=1.0)

Return velocity * state @ generator applied along axis.

Parameters:

Name	Type	Description	Default
state	Any	Array with the kernel axis at position axis.	required
generator	Any	Square matrix over the kernel axis; rows are sources.	required
axis	int	Position of the kernel axis in state.	required
velocity	Any	Scalar multiplier (for example a waning rate).	1.0

Returns:

Type	Description
Any	The derivative contribution, with the same shape as state.

Source code in src/op_system/_axis_kernel.py

def axis_kernel_generator_rhs(
    state: Any,  # ruff: ignore[any-type]
    generator: Any,  # ruff: ignore[any-type]
    *,
    axis: int,
    velocity: Any = 1.0,  # ruff: ignore[any-type]
) -> Any:  # ruff: ignore[any-type]
    """Return ``velocity * state @ generator`` applied along ``axis``.

    Args:
        state: Array with the kernel axis at position ``axis``.
        generator: Square matrix over the kernel axis; rows are sources.
        axis: Position of the kernel axis in ``state``.
        velocity: Scalar multiplier (for example a waning rate).

    Returns:
        The derivative contribution, with the same shape as ``state``.
    """
    xp = array_namespace(state)
    moved = xp.moveaxis(state, axis, -1)
    return xp.moveaxis(velocity * (moved @ generator), -1, axis)

axis_kernel_redistribute(flux, kernel, *, axis=-1)

Return flux @ kernel - flux along axis for a transferred flux.

Parameters:

Name	Type	Description	Default
flux	Any	Per-coordinate flux already moved by an ordinary transfer.	required
kernel	Any	Row-stochastic matrix over the kernel axis.	required
axis	int	Position of the kernel axis in flux.	-1

Returns:

Type	Description
Any	The redistribution to add to the target slice.

Source code in src/op_system/_axis_kernel.py

def axis_kernel_redistribute(
    flux: Any,  # ruff: ignore[any-type]
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis: int = -1,
) -> Any:  # ruff: ignore[any-type]
    """Return ``flux @ kernel - flux`` along ``axis`` for a transferred flux.

    Args:
        flux: Per-coordinate flux already moved by an ordinary transfer.
        kernel: Row-stochastic matrix over the kernel axis.
        axis: Position of the kernel axis in ``flux``.

    Returns:
        The redistribution to add to the target slice.
    """
    xp = array_namespace(flux)
    moved = xp.moveaxis(flux, axis, -1)
    return xp.moveaxis(moved @ kernel - moved, -1, axis)

compile_rhs(rhs, *, xp=None, detect_gauges=False, gauge_fix=None)

Compile a normalized RHS into a runnable evaluation function.

Always uses the vectorized eval path that operates on shaped buffers (one tensor expression per state template) for specs that declare axes. Specs without axes (genuinely scalar models) fall back to the scalar path. Raising :class:UnsupportedFeatureError if an axis-indexed spec cannot be vectorized, rather than silently falling back to the catastrophically slow scalar path.

The returned eval_fn is namespace-polymorphic: it infers its array namespace from the input y at call time (through array_api_compat.array_namespace), and returns arrays in that same namespace. Calling it with JAX arrays (or tracers) yields a JAX-native computation suitable for jax.jit / jax.vmap without any correctness wrapping.

Parameters:

Name	Type	Description	Default
rhs	NormalizedRhs	Normalized RHS produced by op_system.specs.normalize_rhs.	required
detect_gauges	bool	Opt in to conservative structural scaling detection.	False
gauge_fix	Mapping[str, float] | None	Positive scalar anchors for the reported scaling action.	None
xp	object | None	Deprecated. Formerly the compile-time array namespace. Now ignored — the namespace is resolved per call from the input y. Will be removed in a future release.	None

Returns:

Type	Description
CompiledRhs	A CompiledRhs containing an eval_fn(t, y, **params) -> dydt.
CompiledRhs	For axis-indexed specs the returned object also carries
CompiledRhs	pytree_eval_fn and template_shapes. If the spec declares
CompiledRhs	axes but the vectorizer cannot build a plan an
CompiledRhs	UnsupportedFeatureError is raised (see bail reason in the detail
CompiledRhs	message) rather than silently degrading to the scalar path.

Source code in src/op_system/compile.py

def compile_rhs(  # ruff: ignore[too-many-locals]
    rhs: NormalizedRhs,
    *,
    xp: object | None = None,
    detect_gauges: bool = False,
    gauge_fix: Mapping[str, float] | None = None,
) -> CompiledRhs:
    """Compile a normalized RHS into a runnable evaluation function.

    Always uses the vectorized eval path that operates on shaped buffers
    (one tensor expression per state template) for specs that declare axes.
    Specs without axes (genuinely scalar models) fall back to the scalar path.
    Raising :class:`UnsupportedFeatureError` if an axis-indexed spec cannot be
    vectorized, rather than silently falling back to the catastrophically slow
    scalar path.

    The returned ``eval_fn`` is **namespace-polymorphic**: it infers its
    array namespace from the input ``y`` at call time
    (through ``array_api_compat.array_namespace``), and returns arrays in that same
    namespace. Calling it with JAX arrays (or tracers) yields a JAX-native
    computation suitable for ``jax.jit`` / ``jax.vmap`` without any
    correctness wrapping.

    Args:
        rhs: Normalized RHS produced by `op_system.specs.normalize_rhs`.
        detect_gauges: Opt in to conservative structural scaling detection.
        gauge_fix: Positive scalar anchors for the reported scaling action.
        xp: **Deprecated.** Formerly the compile-time array namespace.
            Now ignored — the namespace is resolved per call
            from the input ``y``. Will be removed in a future release.

    Returns:
        A `CompiledRhs` containing an `eval_fn(t, y, **params) -> dydt`.
        For axis-indexed specs the returned object also carries
        ``pytree_eval_fn`` and ``template_shapes``.  If the spec declares
        axes but the vectorizer cannot build a plan an
        ``UnsupportedFeatureError`` is raised (see bail reason in the detail
        message) rather than silently degrading to the scalar path.
    """
    _warn_on_deprecated_xp(xp)
    _validate_rhs_type(rhs)

    _validate_history_kinds(
        _history_requirements_from_ir(
            aliases_ir=rhs.aliases_ir,
            equations_ir=rhs.equations_ir,
        )
        if _reduce_ir_has_history(rhs)
        else ()
    )

    vec, plan, eval_fn, pytree_eval_fn, template_shapes = _build_primary_eval_artifacts(
        rhs
    )

    # Apply both time-varying and synth-const wrappers to ``pytree_eval_fn``
    # BEFORE ``_build_history_artifacts`` consumes it.  history_eval_fn /
    # body_eval_fn capture the pytree_eval_fn reference at construction
    # time, so any later re-wrapping would not propagate into them; that
    # would leave runtime history-body calls missing time-varying param
    # slicing and synthesized constants (e.g. __op_system_mask__* one-hot
    # arrays for pinned transition selectors).
    synth_consts: Mapping[str, object] | None = None
    if eval_fn is not None:
        eval_fn, pytree_eval_fn = _wrap_time_varying_artifacts(
            rhs=rhs,
            eval_fn=eval_fn,
            pytree_eval_fn=pytree_eval_fn,
        )
        eval_fn, pytree_eval_fn, synth_consts = _apply_synth_const_wrappers(
            rhs=rhs,
            eval_fn=eval_fn,
            pytree_eval_fn=pytree_eval_fn,
        )

    history_requirements, history_eval_fn, body_eval_fn = _build_history_artifacts(
        rhs=rhs,
        plan=plan,
        pytree_eval_fn=pytree_eval_fn,
    )

    eval_fn = _resolve_eval_fn(
        rhs=rhs,
        eval_fn=eval_fn,
        history_requirements=history_requirements,
    )

    block_axes = analyze_block_axes(rhs)

    # ------------------------------------------------------------------
    # Block-stripped compile: produce a per-block-coord pytree_eval_fn by
    # stripping the first factorize axis from the RHS and re-running the
    # vectorizer.  Engines can jax.vmap this over the block axis instead
    # of baking literal axis indices that break under vmap.
    # ------------------------------------------------------------------
    block_pytree_eval_fn, block_template_shapes = _build_block_pytree_artifacts(
        rhs=rhs,
        vec=vec,
        pytree_eval_fn=pytree_eval_fn,
        block_axes=block_axes,
        synth_consts=synth_consts,
    )

    # Build per-block history / body eval fns when both the block compile and
    # history path succeeded.  These wrap ``block_pytree_eval_fn`` exactly as
    # ``history_eval_fn`` / ``body_eval_fn`` wrap ``pytree_eval_fn``.
    block_history_eval_fn, block_body_eval_fn = _build_block_history_artifacts(
        block_pytree_eval_fn=block_pytree_eval_fn,
        history_requirements=history_requirements,
    )

    reaction_artifacts = _build_reaction_artifacts(rhs=rhs, plan=plan, vec=vec)
    compiled = CompiledRhs(
        state_names=rhs.state_names,
        param_names=tuple(rhs.param_names),
        eval_fn=eval_fn,
        meta=rhs.meta,
        param_meta=MappingProxyType(
            dict(
                cast(
                    "Mapping[str, ParameterDeclaration]",
                    rhs.meta.get("parameter_declarations", {}),
                )
            )
        ),
        state_meta=MappingProxyType(
            dict(
                cast(
                    "Mapping[str, StateDeclaration]",
                    rhs.meta.get("state_declarations", {}),
                )
            )
        ),
        axis_meta=MappingProxyType(
            _axis_declarations(
                cast("Sequence[Mapping[str, object]]", rhs.meta.get("axes", ()))
            )
        ),
        operators=_parse_operator_descriptors(rhs.meta),
        observables=MappingProxyType(_compile_observables(rhs)),
        factorize_axes=_parse_factorize_axes(rhs.meta),
        block_axes=block_axes,
        pytree_eval_fn=pytree_eval_fn,
        template_shapes=template_shapes,
        block_pytree_eval_fn=block_pytree_eval_fn,
        block_template_shapes=block_template_shapes,
        history_requirements=history_requirements,
        history_eval_fn=history_eval_fn,
        body_eval_fn=body_eval_fn,
        block_history_eval_fn=block_history_eval_fn,
        block_body_eval_fn=block_body_eval_fn,
        reactions=reaction_artifacts.reactions,
        reaction_gaps=reaction_artifacts.gaps,
        _rhs=rhs,
    )

    return _gauge_options(compiled, detect=detect_gauges, anchors=gauge_fix)

compile_spec(spec, *, xp=None, backend=DEFAULT_ARRAY_BACKEND, detect_gauges=False, gauge_fix=None)

Validate, normalize, and compile a RHS specification in one call.

This is the recommended public entrypoint for most users and adapters.

The compiled eval_fn is namespace-polymorphic: it infers its array namespace from the input y at call time (through array_api_compat.array_namespace), so one callable handles NumPy, JAX (concrete and traced), and other qualified array namespaces natively. This selects array operations; a downstream numerical engine separately selects integration methods. No compile-time namespace selection is required.

Parameters:

Name	Type	Description	Default
spec	dict[str, object]	Raw RHS specification mapping (YAML/JSON friendly).	required
detect_gauges	bool	Opt in to conservative structural scaling detection.	False
gauge_fix	Mapping[str, float] | None	Positive scalar anchors for the reported scaling action.	None
xp	object | None	Deprecated. Formerly the compile-time array namespace. Now ignored — see compile_rhs for details. Will be removed in a future release.	None
backend	Literal['numpy', 'jax']	Deprecated. Formerly selected the compile-time array namespace ("numpy" or "jax"). Now ignored. Will be removed in a future release.	DEFAULT_ARRAY_BACKEND

Returns:

Name	Type	Description
CompiledRhs	CompiledRhs	Runnable RHS callable container.

Examples:

>>> import numpy as np
>>> compiled = compile_spec({
...     "kind": "expr",
...     "state": ["x"],
...     "equations": {"x": "-rate*x"},
... })
>>> compiled.template_shapes
{'x': ()}
>>> compiled.eval_fn(0.0, np.array([4.0]), rate=0.5).tolist()
[-2.0]

Source code in src/op_system/__init__.py

def compile_spec(  # ruff: ignore[non-empty-init-module]
    spec: dict[str, object],
    *,
    xp: object | None = None,
    backend: Literal["numpy", "jax"] = DEFAULT_ARRAY_BACKEND,
    detect_gauges: bool = False,
    gauge_fix: Mapping[str, float] | None = None,
) -> CompiledRhs:
    """
    Validate, normalize, and compile a RHS specification in one call.

    This is the recommended public entrypoint for most users and adapters.

    The compiled ``eval_fn`` is **namespace-polymorphic**: it infers its
    array namespace from the input ``y`` at call time
    (through ``array_api_compat.array_namespace``), so one callable handles
    NumPy, JAX (concrete and traced), and other qualified array namespaces
    natively. This selects array operations; a downstream numerical engine
    separately selects integration methods. No compile-time namespace
    selection is required.

    Args:
        spec: Raw RHS specification mapping (YAML/JSON friendly).
        detect_gauges: Opt in to conservative structural scaling detection.
        gauge_fix: Positive scalar anchors for the reported scaling action.
        xp: **Deprecated.** Formerly the compile-time array namespace.
            Now ignored — see ``compile_rhs`` for details.
            Will be removed in a future release.
        backend: **Deprecated.** Formerly selected the compile-time
            array namespace (``"numpy"`` or ``"jax"``). Now ignored. Will be
            removed in a future release.

    Returns:
        CompiledRhs: Runnable RHS callable container.

    Examples:
        >>> import numpy as np
        >>> compiled = compile_spec({
        ...     "kind": "expr",
        ...     "state": ["x"],
        ...     "equations": {"x": "-rate*x"},
        ... })
        >>> compiled.template_shapes
        {'x': ()}
        >>> compiled.eval_fn(0.0, np.array([4.0]), rate=0.5).tolist()
        [-2.0]
    """
    if xp is not None or backend != DEFAULT_ARRAY_BACKEND:
        warnings.warn(
            "compile_spec(xp=..., backend=...) is deprecated and ignored. "
            "The compiled eval_fn now infers its array namespace from the "
            "input `y` at call time; pass JAX arrays for a JAX-native call, "
            "NumPy arrays for a NumPy call, or Torch tensors for a Torch "
            "call. These kwargs will be removed in a future release.",
            DeprecationWarning,
            stacklevel=2,
        )

    rhs = normalize_rhs(spec)
    return compile_rhs(rhs, detect_gauges=detect_gauges, gauge_fix=gauge_fix)

jump_integral_generator(kernel, *, axis_type, direction='both', boundary='reflecting', quadrature_weights=None)

Build the conservative row-source generator for a jump kernel.

kernel[i, j] is the non-negative rate density from source coordinate i to target coordinate j. Its diagonal is ignored. up keeps only j > i; down keeps only j < i; both keeps every off-diagonal entry. For a continuous axis each target column j is multiplied by quadrature_weights[j] before the diagonal loss is set.

Parameters:

Name	Type	Description	Default
kernel	Any	Square source-by-target rate-density matrix.	required
axis_type	str	categorical, ordinal, or continuous.	required
direction	str	Permitted movement in declared coordinate order.	'both'
boundary	str	Boundary contract. Only reflecting is currently defined.	'reflecting'
quadrature_weights	Any | None	Positive target weights required for continuous axes and forbidden for categorical or ordinal axes.	None

Returns:

Type	Description
Any	A generator in the kernel's Array-API namespace. Rows sum to zero.

Raises:

Type	Description
ValueError	If static choices, shapes, or quadrature usage are invalid.

Notes

Call :func:validate_jump_integral_kernel with concrete parameter values at an orchestration boundary to check finiteness and signs. This builder deliberately keeps array values dynamic under transforms.

Source code in src/op_system/_jump_integral.py

def jump_integral_generator(
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis_type: str,
    direction: str = "both",
    boundary: str = "reflecting",
    quadrature_weights: Any | None = None,  # ruff: ignore[any-type]
) -> Any:  # ruff: ignore[any-type]
    """Build the conservative row-source generator for a jump kernel.

    ``kernel[i, j]`` is the non-negative rate density from source coordinate
    ``i`` to target coordinate ``j``. Its diagonal is ignored. ``up`` keeps
    only ``j > i``; ``down`` keeps only ``j < i``; ``both`` keeps every
    off-diagonal entry. For a continuous axis each target column ``j`` is
    multiplied by ``quadrature_weights[j]`` before the diagonal loss is set.

    Args:
        kernel: Square source-by-target rate-density matrix.
        axis_type: ``categorical``, ``ordinal``, or ``continuous``.
        direction: Permitted movement in declared coordinate order.
        boundary: Boundary contract. Only ``reflecting`` is currently defined.
        quadrature_weights: Positive target weights required for continuous
            axes and forbidden for categorical or ordinal axes.

    Returns:
        A generator in the kernel's Array-API namespace. Rows sum to zero.

    Raises:
        ValueError: If static choices, shapes, or quadrature usage are invalid.

    Notes:
        Call :func:`validate_jump_integral_kernel` with concrete parameter
        values at an orchestration boundary to check finiteness and signs.
        This builder deliberately keeps array values dynamic under transforms.
    """
    _validate_static_contract(
        axis_type=axis_type,
        direction=direction,
        boundary=boundary,
    )
    if kernel.ndim != 2 or kernel.shape[0] != kernel.shape[1]:
        msg = "jump_integral kernel must be square"
        raise ValueError(msg)

    xp = array_namespace(kernel)
    upper = xp.triu(xp.ones_like(kernel, dtype=xp.bool), k=1)
    lower = xp.tril(xp.ones_like(kernel, dtype=xp.bool), k=-1)
    if direction == "up":
        allowed = upper
    elif direction == "down":
        allowed = lower
    else:
        allowed = xp.logical_or(upper, lower)
    rates = xp.where(allowed, kernel, xp.zeros_like(kernel))

    if axis_type == "continuous":
        if quadrature_weights is None:
            msg = "continuous jump_integral axes require quadrature_weights"
            raise ValueError(msg)
        weights = xp.asarray(
            quadrature_weights,
            dtype=kernel.dtype,
            **_device_kwargs(kernel),
        )
        if weights.ndim != 1 or weights.shape[0] != kernel.shape[0]:
            msg = "quadrature_weights must have one entry per kernel target"
            raise ValueError(msg)
        rates *= xp.reshape(weights, (1, weights.shape[0]))
    elif quadrature_weights is not None:
        msg = "quadrature_weights are only valid for continuous jump_integral axes"
        raise ValueError(msg)

    row_sums = xp.sum(rates, axis=1)
    identity = xp.eye(
        kernel.shape[0],
        dtype=kernel.dtype,
        **_device_kwargs(kernel),
    )
    return rates - identity * xp.reshape(row_sums, (row_sums.shape[0], 1))

jump_integral_rhs(state, kernel, *, axis, axis_type, rate=1.0, direction='both', boundary='reflecting', quadrature_weights=None)

Return a conservative jump-integral derivative contribution.

Parameters:

Name	Type	Description	Default
state	Any	State array containing the jump axis.	required
kernel	Any	Source-by-target rate-density matrix.	required
axis	int	Position of the jump axis in state.	required
axis_type	str	categorical, ordinal, or continuous.	required
rate	Any	Scalar multiplier with inverse-time units after any continuous target quadrature has been applied.	1.0
direction	str	up, down, or both.	'both'
boundary	str	Boundary contract; currently only reflecting.	'reflecting'
quadrature_weights	Any | None	Target quadrature weights for continuous axes.	None

Returns:

Type	Description
Any	Derivative contribution with the same shape as state.

Source code in src/op_system/_jump_integral.py

def jump_integral_rhs(  # ruff: ignore[too-many-arguments]
    state: Any,  # ruff: ignore[any-type]
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis: int,
    axis_type: str,
    rate: Any = 1.0,  # ruff: ignore[any-type]
    direction: str = "both",
    boundary: str = "reflecting",
    quadrature_weights: Any | None = None,  # ruff: ignore[any-type]
) -> Any:  # ruff: ignore[any-type]
    """Return a conservative jump-integral derivative contribution.

    Args:
        state: State array containing the jump axis.
        kernel: Source-by-target rate-density matrix.
        axis: Position of the jump axis in ``state``.
        axis_type: ``categorical``, ``ordinal``, or ``continuous``.
        rate: Scalar multiplier with inverse-time units after any continuous
            target quadrature has been applied.
        direction: ``up``, ``down``, or ``both``.
        boundary: Boundary contract; currently only ``reflecting``.
        quadrature_weights: Target quadrature weights for continuous axes.

    Returns:
        Derivative contribution with the same shape as ``state``.
    """
    generator = jump_integral_generator(
        kernel,
        axis_type=axis_type,
        direction=direction,
        boundary=boundary,
        quadrature_weights=quadrature_weights,
    )
    return axis_kernel_generator_rhs(state, generator, axis=axis, velocity=rate)

normalize_expr_rhs(spec)

Normalize an expression-based RHS specification.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any]	Raw RHS specification mapping.	required

Returns:

Type	Description
ExprRhs	Backend-facing normalized RHS representation.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_expr_rhs(spec: Mapping[str, Any]) -> ExprRhs:  # ruff: ignore[complex-structure, too-many-locals, too-many-statements]
    """Normalize an expression-based RHS specification.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.

    Raises:
        InvalidRhsSpecError: If validation fails.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        raise InvalidRhsSpecError(detail="state contains duplicate names")

    equations_map = spec.get("equations")
    if not isinstance(equations_map, dict):
        raise InvalidRhsSpecError(detail="equations must be a mapping of state->expr")

    equations_map = {_normalize_bracket_key(k): v for k, v in equations_map.items()}

    time_axis_name = _resolve_time_axis_name(spec)
    axes_meta = _normalize_axes(spec.get("axes"), time_axis_name=time_axis_name)
    meta_parts = _normalize_common_meta(
        spec,
        axis_names={"subgroup"} | {ax["name"] for ax in axes_meta},
        state_set=set(state_raw),
        operator_state_set=set(state_raw),
        axes=axes_meta,
    )

    meta: dict[str, Any] = {
        "axes": axes_meta,
        "state_axes": meta_parts[1],
        "kernels": meta_parts[2],
        "operators": meta_parts[3],
    }
    meta["op_system_may_have_history"] = _may_have_history(spec)
    for reserved_key in ("sources", "couplings", "constraints"):
        if reserved_key in spec:
            meta[reserved_key] = spec.get(reserved_key)

    state_expanded, state_template_map = _expand_state_templates(
        state_raw, axes=axes_meta
    )
    if len(state_expanded) != len(set(state_expanded)):
        raise InvalidRhsSpecError(detail="expanded state contains duplicates")

    _validate_bare_template_references(
        state_raw,
        state_expanded,
        chain(
            (
                (f"aliases[{key!r}]", value)
                for key, value in (meta_parts[0] or {}).items()
            ),
            ((f"equations[{key!r}]", value) for key, value in equations_map.items()),
        ),
    )

    # Pre-scan raw aliases + equations for shaped-parameter references before
    # alias/template expansion mangles bracketed bases into per-coord names.
    axis_lookup_dict: dict[str, list[str]] = build_axis_lookup(axes_meta)
    aliases_raw_map = meta_parts[0] or {}
    name_blocklist = (
        {parse_selector(s)[0] for s in state_raw}
        | set(state_expanded)
        | {parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map}
        | set(aliases_raw_map.keys())
    )
    raw_expressions: list[str] = [
        v for v in aliases_raw_map.values() if isinstance(v, str)
    ]
    raw_expressions.extend(v for v in equations_map.values() if isinstance(v, str))
    shaped_params = _scan_shaped_param_refs(
        raw_expressions,
        name_blocklist=name_blocklist,
        axis_lookup=axis_lookup_dict,
    )
    _reject_legacy_time_varying_field(spec)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    meta.update(
        _normalize_time_metadata(
            spec,
            time_axis_name=time_axis_name,
            axes_meta=axes_meta,
            time_varying_params=time_varying_full,
        )
    )
    if time_varying_full:
        _strip_time_axis_in_mapping(
            equations_map,
            tv_full_axes=time_varying_full,
            time_axis_name=time_axis_name,
        )
        if isinstance(meta_parts[0], dict):
            _strip_time_axis_in_mapping(
                meta_parts[0],
                tv_full_axes=time_varying_full,
                time_axis_name=time_axis_name,
            )

    aliases_ir_map, aliases_ir_reduce_map, alias_template_map = (
        _build_aliases_ir_from_raw(
            meta_parts[0],
            axes=axes_meta,
            shaped_params=shaped_params,
            axis_lookup=axis_lookup_dict,
        )
    )
    template_map_all = {**state_template_map, **alias_template_map}

    chain_block = spec.get("chain")
    if chain_block:
        if not isinstance(chain_block, list):
            raise InvalidRhsSpecError(detail="chain must be a list if provided")
        _apply_expr_chains(
            chains=chain_block,
            state_expanded=state_expanded,
            equations_map=equations_map,
        )

    unknown_keys = [
        k
        for k in equations_map
        if k not in state_expanded and k not in template_map_all
    ]
    if unknown_keys:
        raise InvalidRhsSpecError(
            detail=f"unknown equation key(s): {sorted(unknown_keys)}"
        )

    equations_ir_built, equations_ir_reduce, all_syms = _build_equations_ir_from_raw(
        state_expanded=state_expanded,
        equations_map=equations_map,
        template_map=template_map_all,
        axes=axes_meta,
        shaped_params=shaped_params,
        axis_lookup=axis_lookup_dict,
        aliases_ir=aliases_ir_map,
    )
    # Collect free symbols from alias bodies (alias bodies may reference
    # params not appearing in any equation directly). Walk the *reduce*
    # map (Reduce nodes still folded) rather than the fully expanded
    # map: alias inlining is identical between the two, but the reduce
    # form is orders of magnitude smaller for continuum specs. Share a
    # single id-keyed memo across the per-cell entries so common
    # subtrees are visited at most once (issue #145).
    fs_memo: dict[int, frozenset[str]] = {}
    for alias_ir_val in aliases_ir_reduce_map.values():
        all_syms |= free_symbols(alias_ir_val, memo=fs_memo)

    _maybe_attach_initial_state(
        meta,
        spec.get("initial_state"),
        axes=axes_meta,
        template_map=template_map_all,
    )

    meta["shaped_params"] = tuple(sorted(shaped_params.items()))
    meta["time_axis"] = time_axis_name
    meta["time_varying_params"] = tuple(sorted(time_varying_full.items()))

    factorize_axes_raw = spec.get("factorize_axes")
    if factorize_axes_raw:
        known_axes = {ax["name"] for ax in axes_meta}
        meta["factorize_axes"] = [
            a for a in factorize_axes_raw if isinstance(a, str) and a in known_axes
        ]

    state_set = set(state_expanded)
    shaped_set = set(shaped_params)
    time_varying_set = set(time_varying_full)
    axis_name_set = set(axis_lookup_dict)
    template_base_set = {parse_selector(k)[0] for k in template_map_all}
    # Share an identity-keyed unparse memo between the equation and
    # alias rendering passes so subexpressions that recur across alias
    # bodies and equations (e.g. coord-pinned copies of the same
    # alias template body) are rendered only once (issue #145).
    unparse_memo: dict[tuple[int, int, bool], str] = {}
    equations_strings = _derive_equation_strings_lazy(
        equations_ir_built, _unparse_memo=unparse_memo
    )
    aliases_strings = _derive_alias_strings(
        aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo
    )

    param_names = _sorted_unique(
        sym
        for sym in all_syms
        if sym not in state_set
        and sym not in aliases_ir_map
        and sym not in shaped_set
        and sym not in time_varying_set
        and sym not in axis_name_set
        and sym not in template_base_set
        and sym not in _SHAPED_PARAM_BUILTIN_NAMES
    )

    state_templates = _build_state_templates(
        state_raw,
        axes=axes_meta,
        state_template_map=state_template_map,
        state_expanded=state_expanded,
    )
    meta.update(
        _state_declaration_metadata(
            spec.get("state_metadata"),
            templates=state_templates,
            axes=axes_meta,
            state_axes=meta_parts[1],
        )
    )
    normalized = ExprRhs(
        state_names=tuple(state_expanded),
        equations=equations_strings,
        aliases=aliases_strings,
        param_names=param_names,
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta=meta,
        state_templates=state_templates,
        shaped_params=tuple(sorted(shaped_params.items())),
        time_varying_params=tuple(sorted(time_varying_full.items())),
        aliases_ir=aliases_ir_map,
        equations_ir=equations_ir_built,
        aliases_ir_reduce=aliases_ir_reduce_map,
        equations_ir_reduce=equations_ir_reduce,
        alias_templates=_build_alias_templates(
            aliases_raw_map,
            axes=axes_meta,
            alias_template_map=alias_template_map,
        ),
    )
    return _finish_observables(normalized, spec)

normalize_rhs(spec)

Normalize a RHS specification dict into a backend-facing representation.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any] | None	Raw RHS specification mapping.	required

Returns:

Type	Description
NormalizedRhs	Backend-facing normalized RHS representation.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.
UnsupportedFeatureError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_rhs(spec: Mapping[str, Any] | None) -> NormalizedRhs:
    """Normalize a RHS specification dict into a backend-facing representation.

    Args:
        spec: Raw RHS specification mapping.

    Returns:
        Backend-facing normalized RHS representation.

    Raises:
        InvalidRhsSpecError: If validation fails.
        UnsupportedFeatureError: If validation fails.
    """
    if spec is None:
        raise InvalidRhsSpecError(detail="rhs specification is required")

    kind = str(spec.get("kind", "expr")).strip().lower()

    if kind == "expr":
        return normalize_expr_rhs(spec)

    if kind == "transitions":
        return normalize_transitions_rhs(spec)

    raise UnsupportedFeatureError(
        feature=f"rhs.kind={kind}",
        detail="Only 'expr' and 'transitions' are supported in v1.",
    )

normalize_transitions_rhs(spec)

Normalize a transition-based RHS specification (diagram/hazard semantics).

Returns:

Type	Description
TransitionsRhs	Backend-facing normalized RHS representation for the transitions kind.

Raises:

Type	Description
InvalidRhsSpecError	If validation fails.

Source code in src/op_system/_normalize.py

def normalize_transitions_rhs(  # ruff: ignore[complex-structure, too-many-branches, too-many-locals, too-many-statements]
    spec: Mapping[str, Any],
) -> TransitionsRhs:
    """Normalize a transition-based RHS specification (diagram/hazard semantics).

    Returns:
        Backend-facing normalized RHS representation for the transitions kind.

    Raises:
        InvalidRhsSpecError: If validation fails.
    """
    state_raw = _ensure_str_list(spec.get("state"), name="state")
    if len(state_raw) != len(set(state_raw)):
        raise InvalidRhsSpecError(detail="state contains duplicate names")

    transitions_raw = spec.get("transitions")
    if transitions_raw is None:
        transitions_raw = []
    elif isinstance(transitions_raw, list):
        # Copy each entry: time-axis stripping rewrites rates in place.
        # Stamp each with its spec position for reaction coverage records.
        transitions_raw = [
            {**tr, ORIGIN_KEY: f"transitions[{idx}]"}
            if isinstance(tr, _MappingABC)
            else tr
            for idx, tr in enumerate(transitions_raw)
        ]
        _hoist_coord_shift_rates(transitions_raw)
    else:
        raise InvalidRhsSpecError(detail="transitions must be a list")

    time_axis_name = _resolve_time_axis_name(spec)
    axes_meta = _normalize_axes(spec.get("axes"), time_axis_name=time_axis_name)

    meta_parts = _normalize_common_meta(
        spec,
        axis_names={"subgroup"} | {ax["name"] for ax in axes_meta},
        state_set=None,
        operator_state_set=set(state_raw),
        axes=axes_meta,
    )

    meta: dict[str, Any] = {
        "transitions": transitions_raw,
        "axes": axes_meta,
        "kernels": meta_parts[2],
        "operators": meta_parts[3],
    }
    meta["op_system_may_have_history"] = _may_have_history(spec)
    meta.update({
        k: spec[k] for k in ("sources", "couplings", "constraints") if k in spec
    })

    chain_block = spec.get("chain")
    if chain_block:
        if not isinstance(chain_block, list):
            raise InvalidRhsSpecError(detail="chain must be a list if provided")
        _apply_transition_chains(
            chains=chain_block,
            state_raw=state_raw,
            transitions_raw=transitions_raw,
            state_set=set(state_raw),
        )

    state_expanded, state_template_map = _expand_state_templates(
        state_raw, axes=axes_meta
    )
    if len(state_expanded) != len(set(state_expanded)):
        raise InvalidRhsSpecError(detail="expanded state contains duplicates")

    _validate_bare_template_references(
        state_raw,
        state_expanded,
        chain(
            (
                (f"aliases[{key!r}]", value)
                for key, value in (meta_parts[0] or {}).items()
            ),
            (
                (f"{tr.get(ORIGIN_KEY, f'transitions[{index}]')}.rate", tr.get("rate"))
                for index, tr in enumerate(transitions_raw)
                if isinstance(tr, _MappingABC)
            ),
        ),
    )

    # Pre-scan raw aliases + transition rates for shaped-parameter references.
    axis_lookup_dict: dict[str, list[str]] = build_axis_lookup(axes_meta)
    aliases_raw_map = meta_parts[0] or {}
    name_blocklist = (
        {parse_selector(s)[0] for s in state_raw}
        | set(state_expanded)
        | {parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map}
        | set(aliases_raw_map.keys())
    )
    raw_expressions: list[str] = [
        v for v in aliases_raw_map.values() if isinstance(v, str)
    ]
    for tr in transitions_raw:
        if isinstance(tr, _MappingABC):
            r = tr.get("rate")
            if isinstance(r, str):
                raw_expressions.append(r)
            n = tr.get("name")
            if isinstance(n, str):
                raw_expressions.append(n)
    shaped_params = _scan_shaped_param_refs(
        raw_expressions,
        name_blocklist=name_blocklist,
        axis_lookup=axis_lookup_dict,
    )
    _reject_legacy_time_varying_field(spec)
    shaped_params, time_varying_full = _partition_time_varying_shaped(
        shaped_params,
        time_axis_name=time_axis_name,
        axis_lookup=axis_lookup_dict,
    )
    meta.update(
        _normalize_time_metadata(
            spec,
            time_axis_name=time_axis_name,
            axes_meta=axes_meta,
            time_varying_params=time_varying_full,
        )
    )
    if time_varying_full:
        if isinstance(meta_parts[0], dict):
            _strip_time_axis_in_mapping(
                meta_parts[0],
                tv_full_axes=time_varying_full,
                time_axis_name=time_axis_name,
            )
        for tr in transitions_raw:
            if not isinstance(tr, _MappingABC):
                continue
            r = tr.get("rate")
            if isinstance(r, str):
                tr["rate"] = _strip_time_axis_in_expr(  # type: ignore[index]
                    r,
                    tv_full_axes=time_varying_full,
                    time_axis_name=time_axis_name,
                )
            n = tr.get("name")
            if isinstance(n, str):
                tr["name"] = _strip_time_axis_in_expr(  # type: ignore[index]
                    n,
                    tv_full_axes=time_varying_full,
                    time_axis_name=time_axis_name,
                )

    aliases_ir_map, aliases_ir_reduce_map, alias_template_map = (
        _build_aliases_ir_from_raw(
            aliases_raw_map,
            axes=axes_meta,
            shaped_params=shaped_params,
            axis_lookup=axis_lookup_dict,
        )
    )
    template_map_all = {**state_template_map, **alias_template_map}

    state_set = set(state_expanded)
    # Collect alias symbols from IR. Walk the Reduce-folded map (same
    # inlined symbol set as the full-expansion map but vastly smaller
    # on continuum specs) and share an id-keyed memo across per-cell
    # entries that share subtrees by identity (issue #145).
    all_syms: set[str] = set()
    fs_memo: dict[int, frozenset[str]] = {}
    for alias_expr in aliases_ir_reduce_map.values():
        all_syms |= free_symbols(alias_expr, memo=fs_memo)

    _apply_coord_shifts(
        transitions_raw=transitions_raw,
        state_expanded=state_expanded,
        axes=axes_meta,
        state_template_map=state_template_map,
    )

    if not transitions_raw:
        raise InvalidRhsSpecError(
            detail="transitions must be non-empty after applying chain expansion"
        )

    # Build equations and collect expanded transitions in one IR-native pass
    pinned_mask_names, pinned_mask_values = _discover_pinned_token_masks(
        transitions_raw, axis_lookup=axis_lookup_dict
    )
    shift_axes, shift_values = _discover_coord_shift_constants(
        transitions_raw, axis_lookup=axis_lookup_dict
    )
    # Routing reactions mask their no-op diagonal (see op_system._reactions).
    for tr in transitions_raw:
        offdiag_axis = routing_offdiag_axis(tr) if isinstance(tr, dict) else None
        if offdiag_axis is None:
            continue
        n_coords = len(axis_lookup_dict[offdiag_axis])
        offdiag_name = offdiag_constant_name(offdiag_axis)
        shift_axes[offdiag_name] = (offdiag_axis, offdiag_axis)
        shift_values[offdiag_name] = tuple(
            tuple(0.0 if i == j else 1.0 for j in range(n_coords))
            for i in range(n_coords)
        )
    if pinned_mask_values or shift_values:
        # Register one-hot masks and shift matrices as shaped params so the
        # vectorizer's extra-param-buffers plumbing assembles them at eval
        # time; stash the actual values under ``meta`` so ``compile_rhs``
        # can inject them into the eval_fn's ``params`` automatically.
        for (axis, _coord), mask_name in pinned_mask_names.items():
            shaped_params[mask_name] = (axis,)
        shaped_params.update(shift_axes)
        meta["op_system_synth_constants"] = {**pinned_mask_values, **shift_values}

    equations_ir_pre_inline, equations_ir_reduce, transitions_expanded, rate_syms = (
        _build_transition_equations_ir(
            transitions_raw,
            state_set=state_set,
            state_expanded=state_expanded,
            axes=axes_meta,
            axis_lookup=axis_lookup_dict,
            template_map=template_map_all,
            shaped_params=shaped_params,
            mask_names=pinned_mask_names,
            time_axis_name=time_axis_name,
            alias_bases={
                parse_selector(_normalize_bracket_key(k))[0] for k in aliases_raw_map
            },
        )
    )
    all_syms |= rate_syms

    reactions_ir, reaction_gaps_ir = build_reaction_artifacts_ir(
        transitions_raw,
        axes=axes_meta,
        axis_lookup=axis_lookup_dict,
        shaped_params=shaped_params,
        time_axis_name=time_axis_name,
        aliases_raw=aliases_raw_map,
        state_axes={
            base: tuple(tok.axis for tok in tokens)
            for state_s in state_raw
            for base, tokens in (parse_selector(state_s),)
        },
    )

    _maybe_attach_initial_state(
        meta,
        spec.get("initial_state"),
        axes=axes_meta,
        template_map=template_map_all,
    )

    meta["shaped_params"] = tuple(sorted(shaped_params.items()))
    meta["time_axis"] = time_axis_name
    meta["time_varying_params"] = tuple(sorted(time_varying_full.items()))

    factorize_axes_raw = spec.get("factorize_axes")
    if factorize_axes_raw:
        known_axes = {ax["name"] for ax in axes_meta}
        meta["factorize_axes"] = [
            a for a in factorize_axes_raw if isinstance(a, str) and a in known_axes
        ]

    shaped_set = set(shaped_params)
    time_varying_set = set(time_varying_full)
    axis_name_set = set(axis_lookup_dict)
    template_base_set = {parse_selector(k)[0] for k in template_map_all}
    # Share an identity-keyed unparse memo between the equation and
    # alias rendering passes so subexpressions that recur across alias
    # bodies and equations are rendered only once (issue #145).
    unparse_memo_final: dict[tuple[int, int, bool], str] = {}
    eqs_tuple = _derive_equation_strings(
        equations_ir_pre_inline, _unparse_memo=unparse_memo_final
    )
    # Inline aliases directly into the per-cell IR we already built in
    # ``_build_transition_equations_ir`` rather than re-parsing the
    # round-tripped equation strings. Avoids 73k x ``parse_expr_to_ir``
    # plus a redundant IR rebuild on large continuum specs (issue #145).
    # ``aliases_ir_map`` is already fully alias-inlined inside
    # ``_build_aliases_ir_from_raw`` so cycle detection is redundant here.
    #
    # The dominant cost on large specs is the per-cell ``inline_aliases``
    # call. Because synthesized transitions install the SAME ``synth_to`` /
    # ``synth_neg`` IR object into every cell of a template, the *terms*
    # of the per-cell sum are shared across many cells even though the
    # outer ``Apply(op="+", ...)`` wrapper is unique per cell. Inlining
    # term-by-term with a shared ``result_memo`` keyed on ``id(term)``
    # collapses the alias-substitution work from O(n_state) to
    # O(n_unique_terms) (issue #145).
    # Both memos key ``id()`` of terms owned by ``equations_ir_pre_inline``,
    # which outlives them -- the invariant that makes an identity key sound
    # (issue #200). Passing a freshly-built expression to either would
    # silently return another expression's inlined result.
    alias_inline_memo = _InlineMemo()
    alias_inline_result_memo: dict[int, Expr] = {}

    def _inline_one(expr: Expr) -> Expr:
        return inline_aliases(
            expr,
            aliases_ir_map,
            memo=alias_inline_memo,
            skip_cycle_check=True,
            result_memo=alias_inline_result_memo,
        )

    # Dedup the post-inline outer Apply by tuple of arg ids so cells
    # whose inlined-term tuple is identical share one wrapper, letting
    # downstream identity-keyed memos (e.g. unparse_ir) cache the
    # rendered string once per unique equation. Issue #145.
    apply_plus_dedup: dict[tuple[int, ...], Expr] = {}
    # The upstream ``sum_dedup_full`` collapses ``equations_ir_pre_inline``
    # to a handful of unique outer Apply objects shared across many
    # cells (e.g. 7 unique exprs across 72,828 cells on the COVID19_USA
    # continuum spec). Cache the per-expr inlined result by ``id(expr)``
    # so we skip the per-term ``_inline_one`` calls entirely on cache
    # hits -- this collapses ~1.9M memo-hit calls to ~7 real inlines
    # plus 72k dict lookups (issue #145).
    outer_inline_memo: dict[int, Expr] = {}

    def _inline_outer(expr: Expr) -> Expr:
        """Run alias-inlining for one outer equation expression.

        Returns:
            The inlined expression (possibly the original ``expr`` when no
            inlining was needed).
        """
        if isinstance(expr, Apply) and expr.op == "+":
            new_args = tuple(_inline_one(a) for a in expr.args)
            if all(n is o for n, o in zip(new_args, expr.args, strict=True)):
                return expr
            dedup_key = tuple(id(a) for a in new_args)
            cached_apply = apply_plus_dedup.get(dedup_key)
            if cached_apply is None:
                cached_apply = Apply(op="+", args=new_args)
                apply_plus_dedup[dedup_key] = cached_apply
            return cached_apply
        return _inline_one(expr)

    equations_ir_built_list: list[Expr | None] = []
    for expr in equations_ir_pre_inline:
        if expr is None:
            equations_ir_built_list.append(None)
            continue
        if not aliases_ir_map:
            equations_ir_built_list.append(expr)
            continue
        cached_outer = outer_inline_memo.get(id(expr))
        if cached_outer is not None:
            equations_ir_built_list.append(cached_outer)
            continue
        try:
            result_expr = _inline_outer(expr)
        except (ValueError, RecursionError):
            result_expr = expr
        outer_inline_memo[id(expr)] = result_expr
        equations_ir_built_list.append(result_expr)
    equations_ir_built = tuple(equations_ir_built_list)
    param_names = _sorted_unique(
        sym
        for sym in all_syms
        if sym not in state_set
        and sym not in aliases_ir_map
        and sym not in shaped_set
        and sym not in time_varying_set
        and sym not in axis_name_set
        and sym not in template_base_set
        and sym not in _SHAPED_PARAM_BUILTIN_NAMES
    )

    state_templates = _build_state_templates(
        state_raw,
        axes=axes_meta,
        state_template_map=state_template_map,
        state_expanded=state_expanded,
    )
    meta.update(
        _state_declaration_metadata(
            spec.get("state_metadata"),
            templates=state_templates,
            axes=axes_meta,
            state_axes=meta_parts[1],
        )
    )
    normalized = TransitionsRhs(
        reactions_ir=reactions_ir,
        reaction_gaps_ir=reaction_gaps_ir,
        state_names=tuple(state_expanded),
        equations=eqs_tuple,
        aliases=_derive_alias_strings(
            aliases_ir_map, aliases_ir_map, _unparse_memo=unparse_memo_final
        ),
        param_names=param_names,
        all_symbols=frozenset(all_syms | set(aliases_ir_map.keys())),
        meta={**meta, "transitions": transitions_expanded},
        state_templates=state_templates,
        shaped_params=tuple(sorted(shaped_params.items())),
        time_varying_params=tuple(sorted(time_varying_full.items())),
        aliases_ir=aliases_ir_map,
        equations_ir=equations_ir_built,
        aliases_ir_reduce=aliases_ir_reduce_map,
        equations_ir_reduce=equations_ir_reduce,
        alias_templates=_build_alias_templates(
            aliases_raw_map,
            axes=axes_meta,
            alias_template_map=alias_template_map,
        ),
    )
    return _finish_observables(normalized, spec)

qualify_rhs(compiled, fixtures, reference, *, evaluators=('flat', 'structured'), config=None, invariants=())

Compare compiled layouts against independently supplied equations.

Parameters:

Name	Type	Description	Default
compiled	CompiledRhs	Compiled RHS to inspect; structural validation is separate.	required
fixtures	Sequence[QualificationFixture]	Named structured state/parameter fixtures. Supply a JAX state direction to request an additional derivative check.	required
reference	ReferenceRhs	Independent reference(t, state, params) returning derivatives keyed by the same state templates.	required
evaluators	Sequence[Evaluator]	Requested layouts. Unavailable layouts produce failed records rather than silently falling back to another evaluator.	('flat', 'structured')
config	QualificationConfig | None	Reference and derivative tolerances. Defaults assume double precision; adjust the tolerances and difference step for float32.	None
invariants	Sequence[InvariantCheck]	Residual-to-zero callbacks with stated conditions and individual absolute tolerances, such as closed-system balance.	()

Returns:

Type	Description
QualificationReport	Eager host diagnostics for every requested combination. Fixture-based
QualificationReport	agreement is not a proof of scientific validity or solution accuracy.

Raises:

Type	Description
ValueError	If fixtures/evaluators are empty, fixture names/time are invalid, or an invariant lacks valid conditions/tolerances.

Examples:

>>> import numpy as np
>>> from op_system import compile_spec
>>> compiled = compile_spec({
...     "kind": "expr",
...     "state": ["x"],
...     "equations": {"x": "-rate*x"},
... })
>>> fixture = QualificationFixture(
...     "decay", {"x": np.array(4.0)}, {"rate": 0.5}
... )
>>> def reference(time, state, params):
...     return {"x": -params["rate"] * state["x"]}
>>> report = qualify_rhs(compiled, [fixture], reference)
>>> report.ok
True
>>> len(report.fixtures)
2

Source code in src/op_system/qualification.py

def qualify_rhs(  # ruff: ignore[too-many-arguments]
    compiled: CompiledRhs,
    fixtures: Sequence[QualificationFixture],
    reference: ReferenceRhs,
    *,
    evaluators: Sequence[Evaluator] = ("flat", "structured"),
    config: QualificationConfig | None = None,
    invariants: Sequence[InvariantCheck] = (),
) -> QualificationReport:
    """Compare compiled layouts against independently supplied equations.

    Args:
        compiled: Compiled RHS to inspect; structural validation is separate.
        fixtures: Named structured state/parameter fixtures. Supply a JAX
            state ``direction`` to request an additional derivative check.
        reference: Independent ``reference(t, state, params)`` returning
            derivatives keyed by the same state templates.
        evaluators: Requested layouts. Unavailable layouts produce failed
            records rather than silently falling back to another evaluator.
        config: Reference and derivative tolerances. Defaults assume double
            precision; adjust the tolerances and difference step for float32.
        invariants: Residual-to-zero callbacks with stated conditions and
            individual absolute tolerances, such as closed-system balance.

    Returns:
        Eager host diagnostics for every requested combination. Fixture-based
        agreement is not a proof of scientific validity or solution accuracy.

    Raises:
        ValueError: If fixtures/evaluators are empty, fixture names/time are
            invalid, or an invariant lacks valid conditions/tolerances.

    Examples:
        >>> import numpy as np
        >>> from op_system import compile_spec
        >>> compiled = compile_spec({
        ...     "kind": "expr",
        ...     "state": ["x"],
        ...     "equations": {"x": "-rate*x"},
        ... })
        >>> fixture = QualificationFixture(
        ...     "decay", {"x": np.array(4.0)}, {"rate": 0.5}
        ... )
        >>> def reference(time, state, params):
        ...     return {"x": -params["rate"] * state["x"]}
        >>> report = qualify_rhs(compiled, [fixture], reference)
        >>> report.ok
        True
        >>> len(report.fixtures)
        2
    """
    controls = config or QualificationConfig()
    if not fixtures or not evaluators:
        msg = "at least one fixture and evaluator are required"
        raise ValueError(msg)
    names = [item.name for item in fixtures]
    if len(set(names)) != len(names) or any(not name.strip() for name in names):
        msg = "fixture names must be unique and non-empty"
        raise ValueError(msg)
    if any(not math.isfinite(item.time) for item in fixtures):
        msg = "fixture times must be finite"
        raise ValueError(msg)
    if any(name not in {"flat", "structured", "block"} for name in evaluators):
        msg = "unknown evaluator"
        raise ValueError(msg)
    if any(
        not item.condition.strip() or not math.isfinite(item.atol) or item.atol < 0
        for item in invariants
    ):
        msg = "invariant checks need conditions and finite non-negative tolerances"
        raise ValueError(msg)
    return QualificationReport(
        tuple(
            _qualify_fixture(
                compiled, item, reference, evaluator, controls, invariants=invariants
            )
            for item in fixtures
            for evaluator in evaluators
        )
    )

reaction_balances(compiled)

Derive conditional scalar totals from fully represented unit transfers.

Each internal incidence column has one consumed and one produced unit. Component-indicator vectors therefore give an exact integer left-nullspace basis. Source-only channels appear explicitly on the right-hand side of their component balances. Any unsupported contribution prevents all laws; omitted terms are never assigned zero flux. Propensities are not evaluated.

Parameters:

Name	Type	Description	Default
compiled	CompiledRhs	Compiled RHS with complete named scalar reaction artifacts.	required

Returns:

Type	Description
ReactionBalanceReport	Immutable identities, coverage records, declarations and assumptions.
ReactionBalanceReport	Physical interpretation and numerical/runtime realization require
ReactionBalanceReport	independent qualification; fully_covered is a structural result.

Raises:

Type	Description
TypeError	If compiled is not a CompiledRhs container.

Examples:

>>> from op_system import compile_spec
>>> compiled = compile_spec({
...     "kind": "transitions",
...     "state": ["S", "I"],
...     "transitions": [
...         {"name": "move", "from": "S", "to": "I", "rate": "k"},
...     ],
... })
>>> report = reaction_balances(compiled)
>>> report.fully_covered, len(report.laws)
(True, 1)
>>> report.laws[0].coefficients, report.laws[0].closed
((1, 1), True)

Source code in src/op_system/balances.py

def reaction_balances(compiled: CompiledRhs) -> ReactionBalanceReport:
    """Derive conditional scalar totals from fully represented unit transfers.

    Each internal incidence column has one consumed and one produced unit.
    Component-indicator vectors therefore give an exact integer left-nullspace
    basis. Source-only channels appear explicitly on the right-hand side of
    their component balances. Any unsupported contribution prevents all laws;
    omitted terms are never assigned zero flux. Propensities are not evaluated.

    Args:
        compiled: Compiled RHS with complete named scalar reaction artifacts.

    Returns:
        Immutable identities, coverage records, declarations and assumptions.
        Physical interpretation and numerical/runtime realization require
        independent qualification; ``fully_covered`` is a structural result.

    Raises:
        TypeError: If ``compiled`` is not a CompiledRhs container.

    Examples:
        >>> from op_system import compile_spec
        >>> compiled = compile_spec({
        ...     "kind": "transitions",
        ...     "state": ["S", "I"],
        ...     "transitions": [
        ...         {"name": "move", "from": "S", "to": "I", "rate": "k"},
        ...     ],
        ... })
        >>> report = reaction_balances(compiled)
        >>> report.fully_covered, len(report.laws)
        (True, 1)
        >>> report.laws[0].coefficients, report.laws[0].closed
        ((1, 1), True)
    """
    if not isinstance(compiled, CompiledRhs):
        msg = "compiled must be a CompiledRhs"
        raise TypeError(msg)
    states = tuple(compiled.state_names)
    gaps = _coverage(compiled)
    components = () if gaps else _components(states, compiled.reactions)
    gaps.extend(_scientific_coverage(compiled, components))
    names = tuple(reaction.name for reaction in compiled.reactions)
    laws = (
        ()
        if gaps
        else tuple(
            ReactionBalanceLaw(
                name=f"component:{index}",
                coefficients=tuple(int(name in component) for name in states),
                reaction_names=names,
                sources=tuple(
                    ReactionBalanceSource(reaction.name, 1, reaction.to_base)
                    for reaction in compiled.reactions
                    if reaction.from_base is None and reaction.to_base in component
                ),
                assumptions=_ASSUMPTIONS,
            )
            for index, component in enumerate(components)
        )
    )
    return ReactionBalanceReport(
        states,
        laws,
        tuple(gaps),
        tuple(
            (name, compiled.state_meta[name])
            for name in states
            if name in compiled.state_meta
        ),
        tuple(compiled.axis_meta.items()),
        _ASSUMPTIONS,
    )

validate_axis_kernel_matrix(matrix, *, form, tolerance=1e-09)

Return problems with a numeric kernel matrix for form (empty if valid).

Parameters:

Name	Type	Description	Default
matrix	Any	Candidate square matrix.	required
form	str	"generator" or "stochastic".	required
tolerance	float	Absolute tolerance for sign and row-sum checks.	1e-09

Returns:

Type	Description
list[str]	Human-readable problems; an empty list means the matrix is valid.

Raises:

Type	Description
ValueError	If form is not a supported axis-kernel form.

Source code in src/op_system/_axis_kernel.py

def validate_axis_kernel_matrix(
    matrix: Any,  # ruff: ignore[any-type]
    *,
    form: str,
    tolerance: float = 1e-9,
) -> list[str]:
    """Return problems with a numeric kernel matrix for ``form`` (empty if valid).

    Args:
        matrix: Candidate square matrix.
        form: ``"generator"`` or ``"stochastic"``.
        tolerance: Absolute tolerance for sign and row-sum checks.

    Returns:
        Human-readable problems; an empty list means the matrix is valid.

    Raises:
        ValueError: If ``form`` is not a supported axis-kernel form.
    """
    if form not in AXIS_KERNEL_FORMS:
        msg = f"unknown axis_kernel form {form!r}"
        raise ValueError(msg)
    xp = array_namespace(matrix)
    problems: list[str] = []
    if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
        return ["matrix must be square"]
    n = matrix.shape[0]
    off_diagonal = matrix * (1.0 - xp.eye(n, dtype=matrix.dtype))
    if bool(xp.any(off_diagonal < -tolerance)):
        problems.append("off-diagonal entries must be nonnegative")
    row_sums = xp.sum(matrix, axis=1)
    if form == "generator":
        if bool(xp.any(xp.abs(row_sums) > tolerance)):
            problems.append("generator rows must sum to zero")
    else:
        if bool(xp.any(xp.linalg.diagonal(matrix) < -tolerance)):
            problems.append("stochastic diagonal entries must be nonnegative")
        if bool(xp.any(xp.abs(row_sums - 1.0) > tolerance)):
            problems.append("stochastic rows must sum to one")
    return problems

validate_jump_integral_kernel(kernel, *, axis_type, direction='both', boundary='reflecting', quadrature_weights=None, tolerance=1e-09)

Return problems with concrete jump-kernel values (empty if valid).

This eager validation belongs at parameter resolution or another host orchestration boundary, not inside a traced numerical solve.

Parameters:

Name	Type	Description	Default
kernel	Any	Candidate source-by-target rate-density matrix.	required
axis_type	str	categorical, ordinal, or continuous.	required
direction	str	up, down, or both.	'both'
boundary	str	Boundary contract; currently only reflecting.	'reflecting'
quadrature_weights	Any | None	Target quadrature weights for continuous axes.	None
tolerance	float	Absolute tolerance for sign and diagonal checks.	1e-09

Returns:

Type	Description
list[str]	Human-readable problems; an empty list means the values conform.

Raises:

Type	Description
ValueError	If a static contract choice is unsupported.

Source code in src/op_system/_jump_integral.py

def validate_jump_integral_kernel(  # ruff: ignore[complex-structure, too-many-arguments, too-many-branches]
    kernel: Any,  # ruff: ignore[any-type]
    *,
    axis_type: str,
    direction: str = "both",
    boundary: str = "reflecting",
    quadrature_weights: Any | None = None,  # ruff: ignore[any-type]
    tolerance: float = 1e-9,
) -> list[str]:
    """Return problems with concrete jump-kernel values (empty if valid).

    This eager validation belongs at parameter resolution or another host
    orchestration boundary, not inside a traced numerical solve.

    Args:
        kernel: Candidate source-by-target rate-density matrix.
        axis_type: ``categorical``, ``ordinal``, or ``continuous``.
        direction: ``up``, ``down``, or ``both``.
        boundary: Boundary contract; currently only ``reflecting``.
        quadrature_weights: Target quadrature weights for continuous axes.
        tolerance: Absolute tolerance for sign and diagonal checks.

    Returns:
        Human-readable problems; an empty list means the values conform.

    Raises:
        ValueError: If a static contract choice is unsupported.
    """
    try:
        _validate_static_contract(
            axis_type=axis_type,
            direction=direction,
            boundary=boundary,
        )
    except ValueError as exc:
        raise ValueError(str(exc)) from exc
    if kernel.ndim != 2 or kernel.shape[0] != kernel.shape[1]:
        return ["kernel must be square"]
    xp = array_namespace(kernel)
    problems: list[str] = []
    if bool(xp.any(xp.logical_not(xp.isfinite(kernel)))):
        problems.append("kernel entries must be finite")
    if bool(xp.any(kernel < -tolerance)):
        problems.append("kernel entries must be nonnegative")
    if bool(xp.any(xp.abs(xp.linalg.diagonal(kernel)) > tolerance)):
        problems.append("kernel diagonal entries must be zero")

    if axis_type == "continuous":
        if quadrature_weights is None:
            problems.append("continuous axes require quadrature_weights")
        else:
            weights = xp.asarray(
                quadrature_weights,
                dtype=kernel.dtype,
                **_device_kwargs(kernel),
            )
            if weights.ndim != 1 or weights.shape[0] != kernel.shape[0]:
                problems.append("quadrature_weights must match kernel size")
            else:
                if bool(xp.any(xp.logical_not(xp.isfinite(weights)))):
                    problems.append("quadrature_weights must be finite")
                if bool(xp.any(weights <= 0.0)):
                    problems.append("quadrature_weights must be positive")
    elif quadrature_weights is not None:
        problems.append("quadrature_weights require a continuous axis")
    return problems

validate_spec(spec, *, detect_gauges=False, gauge_fix=None)

Validate an op_system specification without raising.

Parameters:

Name	Type	Description	Default
spec	Mapping[str, Any]	A raw RHS specification mapping.	required
detect_gauges	bool	Opt in to conservative structural scaling detection.	False
gauge_fix	Mapping[str, float] | None	Optional positive scalar anchors for the reported action.	None

Returns:

Type	Description
ValidationReport	A report with per-stage status, error messages, compile-cost drivers,
ValidationReport	consumed parameters, and expression-shape groups for templates whose
ValidationReport	cells differ.

Source code in src/op_system/_validate.py

def validate_spec(
    spec: Mapping[str, Any],
    *,
    detect_gauges: bool = False,
    gauge_fix: Mapping[str, float] | None = None,
) -> ValidationReport:
    """Validate an op_system specification without raising.

    Args:
        spec: A raw RHS specification mapping.
        detect_gauges: Opt in to conservative structural scaling detection.
        gauge_fix: Optional positive scalar anchors for the reported action.

    Returns:
        A report with per-stage status, error messages, compile-cost drivers,
        consumed parameters, and expression-shape groups for templates whose
        cells differ.
    """
    from op_system import compile_rhs, normalize_rhs  # ruff: ignore[import-outside-top-level]

    stages = dict.fromkeys(_STAGES, "skipped")
    try:
        rhs = normalize_rhs(spec)
    except Exception as exc:  # ruff: ignore[blind-except] - reported, not raised
        stages["normalize"] = "failed"
        return ValidationReport(
            stages=stages, errors=[str(exc)], cost=_cost(spec, None)
        )
    stages["normalize"] = "passed"
    report_kwargs: dict[str, Any] = {
        "cost": _cost(spec, rhs),
        "parameters": _parameters(rhs),
    }
    try:
        compiled = compile_rhs(rhs, detect_gauges=detect_gauges, gauge_fix=gauge_fix)
    except UnsupportedFeatureError as exc:
        vectorize = "vectorized eval path" in str(exc)
        stages["compile"] = "passed" if vectorize else "failed"
        stages["vectorize"] = "failed" if vectorize else "skipped"
        if vectorize:
            report_kwargs["shape_groups"] = _shape_groups(rhs)
        return ValidationReport(stages=stages, errors=[str(exc)], **report_kwargs)
    except Exception as exc:  # ruff: ignore[blind-except]
        stages["compile"] = "failed"
        return ValidationReport(stages=stages, errors=[str(exc)], **report_kwargs)
    stages["compile"] = "passed"
    stages["vectorize"] = "passed"
    report_kwargs["gauges"] = _gauge_summary(compiled)
    if compiled.gauge_fix is not None:
        for name, _ in compiled.gauge_fix.fixed_values:
            report_kwargs["parameters"].pop(name, None)
    return ValidationReport(stages=stages, **report_kwargs)