OP Engine

op_engine

Numerical methods for structured deterministic and stochastic dynamics.

RandomStream = np.random.Generator | np.random.SeedSequence module-attribute

An advancing NumPy generator or a replayable seed-sequence initialization.

Scalar = float | Array module-attribute

A Python float or a backend-native rank-zero numerical array.

ActionOperator(shape, dtype, namespace, matvec, transpose_matvec=None) dataclass

Caller-owned, real linear action at a fixed linearization/context.

matvec maps (shape[1],) to (shape[0],) without changing namespace, dtype or device. Matrix application applies that callback column by column. transpose_matvec represents the Euclidean transpose of the same fixed map, not an inverse or a weighted adjoint. Linearity and the correctness of the transpose are the provider's obligations; neither is inferred from a numerical probe. Only real float32/float64 actions are supported.

Attributes:

Name	Type	Description
shape	tuple[int, int]	Output and input dimensions.
dtype	object	Native real floating dtype shared by inputs and outputs.
namespace	Any	Array-API namespace returned by :func:array_namespace.
matvec	Callable[[Array], Array]	Linear vector application.
transpose_matvec	Callable[[Array], Array] | None	Optional Euclidean transpose vector application.

T property

Return the explicitly supplied real Euclidean transpose.

Raises:

Type	Description
NotImplementedError	If no transpose callback was supplied.

__matmul__(other)

Apply the action without densifying or coercing the input.

Returns:

Type	Description
Array	Vector or matrix of independent output columns.

Source code in src/op_engine/linear_operator.py

def __matmul__(self, other: Array, /) -> Array:
    """Apply the action without densifying or coercing the input.

    Returns:
        Vector or matrix of independent output columns.
    """
    self._validate(other, size=self.shape[1])
    if len(other.shape) == 1:
        return self._vector(other)
    xp = self.namespace
    if other.shape[1] == 0:
        return cast(
            "Array",
            xp.broadcast_to(
                xp.sum(other, axis=0, keepdims=True), (self.shape[0], 0)
            ),
        )
    # Snapshot each output before the next call can reuse its storage.
    # A neutral native operation retains AD graphs in all supported namespaces.
    columns = [
        xp.multiply(self._vector(other[:, index]), 1)  # type: ignore[index]
        for index in range(other.shape[1])
    ]
    return cast("Array", xp.stack(columns, axis=1))

__post_init__()

Validate static metadata and callback availability.

Raises:

Type	Description
ValueError	If dimensions are invalid.
TypeError	If dtype or callbacks are unsupported.

Source code in src/op_engine/linear_operator.py

def __post_init__(self) -> None:
    """Validate static metadata and callback availability.

    Raises:
        ValueError: If dimensions are invalid.
        TypeError: If dtype or callbacks are unsupported.
    """
    if (
        not isinstance(self.shape, tuple)
        or len(self.shape) != 2
        or any(
            not isinstance(size, Integral) or isinstance(size, bool) or size < 1
            for size in self.shape
        )
    ):
        msg = "action shape must contain two positive integer dimensions"
        raise ValueError(msg)
    if self.dtype not in (self.namespace.float32, self.namespace.float64):  # noqa: PLR6201 - equal NumPy dtype objects can have different hashes
        msg = "actions require native real float32 or float64 dtype"
        raise TypeError(msg)
    if not callable(self.matvec) or (
        self.transpose_matvec is not None and not callable(self.transpose_matvec)
    ):
        msg = "action callbacks must be callable"
        raise TypeError(msg)

AdaptiveScheduleProvenance

Bases: NamedTuple

Static compatibility fields and descriptive adaptive-discovery metadata.

Backend, dtype, discovery controls, and reference regime are descriptive. Replay checks method, state layout, relevant operator/kernel settings, and the caller's structural model identity. Callback semantics are not inferred.

AdaptiveTauLeapingConfig(leap_tolerance=0.03, critical_threshold=10, exact_fallback_multiplier=10.0, max_steps=1000000, max_retries=20, forcing_breakpoints=()) dataclass

Controls for bounded adaptive tau-leaping.

Attributes:

Name	Type	Description
leap_tolerance	float	Relative leap-condition tolerance.
critical_threshold	int	A reaction is critical when fewer than this many firings would exhaust one of its reactants. Zero disables critical classification.
exact_fallback_multiplier	float	Use exact SSA when the proposed leap is less than this multiple of the mean time to the next event. Zero disables this fallback.
max_steps	int	Maximum accepted leaps or exact events per output interval.
max_retries	int	Maximum post-leap rejections before failing.
forcing_breakpoints	tuple[float, ...]	Strictly increasing finite forcing-change times. Leaps end at these boundaries; exact events at or beyond a boundary are discarded. Rates must be right-continuous and constant in time between boundaries for exact fallback.

__post_init__()

Validate adaptive controls.

Raises:

Type	Description
ValueError	If a control is outside its valid range.

Source code in src/op_engine/adaptive_tau.py

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

    Raises:
        ValueError: If a control is outside its valid range.
    """
    if (
        not isinstance(self.leap_tolerance, Real)
        or isinstance(self.leap_tolerance, bool)
        or not np.isfinite(self.leap_tolerance)
        or not 0.0 < self.leap_tolerance < 1.0
    ):
        msg = "leap_tolerance must be finite and strictly between zero and one"
        raise ValueError(msg)
    if (
        not isinstance(self.critical_threshold, Integral)
        or isinstance(self.critical_threshold, bool)
        or self.critical_threshold < 0
    ):
        msg = "critical_threshold must be a non-negative integer"
        raise ValueError(msg)
    if (
        not isinstance(self.exact_fallback_multiplier, Real)
        or isinstance(self.exact_fallback_multiplier, bool)
        or not np.isfinite(self.exact_fallback_multiplier)
        or self.exact_fallback_multiplier < 0.0
    ):
        msg = "exact_fallback_multiplier must be finite and non-negative"
        raise ValueError(msg)
    if (
        not isinstance(self.max_steps, Integral)
        or isinstance(self.max_steps, bool)
        or self.max_steps < 1
    ):
        msg = "max_steps must be a positive integer"
        raise ValueError(msg)
    if (
        not isinstance(self.max_retries, Integral)
        or isinstance(self.max_retries, bool)
        or self.max_retries < 0
    ):
        msg = "max_retries must be a non-negative integer"
        raise ValueError(msg)
    schedule = _ForcingSchedule(self.forcing_breakpoints)
    object.__setattr__(self, "forcing_breakpoints", schedule.breakpoints)

AdaptiveTauLeapingSolver(core, stoichiometry, reactant_stoichiometry, *, reaction_axis='state', dependency_incidence=None, propensity_orders=None)

Adaptive non-negative Poisson tau-leaping with exact critical events.

The implementation follows the Cao-Gillespie-Petzold species-based pre-leap selector. Reactions near exhaustion are bounded to at most one collective exact event. A negative post-leap proposal is rejected and retried at half the attempted tau; populations and firing counts are never clipped.

Initialize the reaction and reactant structures.

Parameters:

Name	Type	Description	Default
core	ModelCore	Model state and output-time container.	required
stoichiometry	Array	Net integer state changes by reaction.	required
reactant_stoichiometry	Array	Non-negative integer reactant counts by species and reaction. This cannot be inferred safely from net stoichiometry for catalytic reactions. For a reaction with a positive propensity_orders entry, list only the species it consumes; they bound its critical firings.	required
reaction_axis	str | int	State axis changed by reaction firings.	'state'
dependency_incidence	Array | None	Optional (n_species, n_reactions) non-negative integers; a positive entry means the reaction's propensity reads that species. Use it for propensities that are not mass action, such as frequency-dependent infection.	None
propensity_orders	Array | None	Optional (n_reactions,) non-negative integers. A positive entry E_j bounds the reaction's total elasticity sum_i |d log a_j / d log x_i| and declares it through dependency_incidence; zero keeps the reactant-order rules. Orders above three are allowed.	None

Source code in src/op_engine/adaptive_tau.py

def __init__(  # noqa: PLR0913
    self,
    core: ModelCore,
    stoichiometry: Array,
    reactant_stoichiometry: Array,
    *,
    reaction_axis: str | int = "state",
    dependency_incidence: Array | None = None,
    propensity_orders: Array | None = None,
) -> None:
    """Initialize the reaction and reactant structures.

    Args:
        core: Model state and output-time container.
        stoichiometry: Net integer state changes by reaction.
        reactant_stoichiometry: Non-negative integer reactant counts by
            species and reaction. This cannot be inferred safely from net
            stoichiometry for catalytic reactions. For a reaction with a
            positive ``propensity_orders`` entry, list only the species it
            consumes; they bound its critical firings.
        reaction_axis: State axis changed by reaction firings.
        dependency_incidence: Optional ``(n_species, n_reactions)``
            non-negative integers; a positive entry means the reaction's
            propensity reads that species. Use it for propensities that
            are not mass action, such as frequency-dependent infection.
        propensity_orders: Optional ``(n_reactions,)`` non-negative
            integers. A positive entry ``E_j`` bounds the reaction's
            total elasticity ``sum_i |d log a_j / d log x_i|`` and
            declares it through ``dependency_incidence``; zero keeps the
            reactant-order rules. Orders above three are allowed.
    """
    self.core = core
    self._network = _ReactionNetwork(core, stoichiometry, reaction_axis)
    self._reactants = _ReactantStructure(
        reactant_stoichiometry,
        self._network,
        dependencies=dependency_incidence,
        propensity_orders=propensity_orders,
    )

n_reactions property

Return the number of reaction channels.

run(propensity_func, poisson_sampler=None, ssa_sampler=None, *, config=None, rng=None)

Advance a bounded adaptive trajectory through the output grid.

Parameters:

Name	Type	Description	Default
propensity_func	PropensityFunction	Reaction-channel propensities.	required
poisson_sampler	PoissonSampler | None	Backend-specific noncritical Poisson sampler.	None
ssa_sampler	SSASampler | None	Backend-specific exact-event sampler.	None
config	AdaptiveTauLeapingConfig | None	Optional adaptive, retry, and forcing controls.	None
rng	RandomStream | None	One NumPy stream shared by call-local Poisson/exact samplers; supply instead of both explicit samplers.	None

Raises:

Type	Description
ValueError	If either sampler is absent or rng conflicts with a sampler.
TypeError	If rng is unsupported or used with non-NumPy state.
RuntimeError	If a step guard, retry guard, or non-negativity invariant fails.

Source code in src/op_engine/adaptive_tau.py

def run(  # noqa: C901, PLR0912, PLR0914, PLR0915
    self,
    propensity_func: PropensityFunction,
    poisson_sampler: PoissonSampler | None = None,
    ssa_sampler: SSASampler | None = None,
    *,
    config: AdaptiveTauLeapingConfig | None = None,
    rng: RandomStream | None = None,
) -> None:
    """Advance a bounded adaptive trajectory through the output grid.

    Args:
        propensity_func: Reaction-channel propensities.
        poisson_sampler: Backend-specific noncritical Poisson sampler.
        ssa_sampler: Backend-specific exact-event sampler.
        config: Optional adaptive, retry, and forcing controls.
        rng: One NumPy stream shared by call-local Poisson/exact samplers;
            supply instead of both explicit samplers.

    Raises:
        ValueError: If either sampler is absent or rng conflicts with a sampler.
        TypeError: If rng is unsupported or used with non-NumPy state.
        RuntimeError: If a step guard, retry guard, or non-negativity
            invariant fails.
    """  # noqa: DOC502 - stream type validation delegates to the shared helper
    if rng is not None:
        if poisson_sampler is not None or ssa_sampler is not None:
            msg = "rng and explicit adaptive tau samplers are mutually exclusive"
            raise ValueError(msg)
        _require_numpy_state(self.core.get_current_state())
        generator = _generator(rng)
        poisson_sampler = NumpyPoissonSampler(rng=generator)
        ssa_sampler = NumpySSASampler(rng=generator)
    if poisson_sampler is None or ssa_sampler is None:
        msg = "provide both adaptive tau samplers or rng"
        raise ValueError(msg)
    cfg = config or AdaptiveTauLeapingConfig()
    forcing = _ForcingSchedule(cfg.forcing_breakpoints)
    state = self.core.get_current_state()
    self._network.validate_finite_nonnegative(
        state,
        message="Initial state is invalid",
    )
    time_grid = np.asarray(self.core.time_grid, dtype=float)
    t = float(time_grid[0])
    poisson_draw_index = 0
    ssa_draw_index = 0
    pending_exact_time: float | None = None
    pending_exact_index: int | None = None
    absorbing = False

    for output_index in range(int(self.core.n_timesteps) - 1):
        target = float(time_grid[output_index + 1])
        interval_steps = 0
        while not absorbing and t < target:
            boundary = forcing.next_after(t)
            if pending_exact_time is not None:
                if boundary <= target and boundary <= pending_exact_time:
                    t = boundary
                    pending_exact_time = None
                    pending_exact_index = None
                    continue
                if pending_exact_time > target:
                    break
                if interval_steps >= cfg.max_steps:
                    raise RuntimeError(_MAX_STEPS)
                if pending_exact_index is None:
                    msg = "Adaptive tau exact-event state is inconsistent"
                    raise RuntimeError(msg)
                state = self._apply_exact_event(state, pending_exact_index)
                t = pending_exact_time
                pending_exact_time = None
                pending_exact_index = None
                interval_steps += 1
                continue

            if interval_steps >= cfg.max_steps:
                raise RuntimeError(_MAX_STEPS)
            propensity = self._network.evaluate_propensity(
                propensity_func,
                t=t,
                state=state,
            )
            xp = _namespace_of(state)
            total_rate = cast(
                "Array",
                xp.asarray(xp.sum(propensity), dtype=state.dtype),
            )
            total_rate_value = float(total_rate.item())
            if total_rate_value == 0.0:
                if np.isfinite(boundary):
                    pending_exact_time = boundary
                    continue
                absorbing = True
                break

            critical = self._reactants.critical_mask(
                state,
                propensity,
                self._network,
                threshold=cfg.critical_threshold,
            )
            noncritical_active = _array_any(
                cast(
                    "Array",
                    xp.logical_and(
                        xp.logical_not(critical),
                        xp.greater(propensity, 0),
                    ),
                )
            )
            tau_candidate = self._reactants.select_tau(
                state,
                propensity,
                critical,
                self._network,
                leap_tolerance=cfg.leap_tolerance,
            )
            fallback_threshold = cfg.exact_fallback_multiplier / total_rate_value
            use_exact = not noncritical_active or tau_candidate < fallback_threshold
            if use_exact:
                event = _draw_ssa_event(
                    propensity,
                    ssa_sampler,
                    draw_index=ssa_draw_index,
                )
                if event is None:
                    absorbing = True
                    break
                waiting_time, pending_exact_index = event
                pending_exact_time = t + waiting_time
                if pending_exact_time <= t:
                    raise RuntimeError(_STEP_UNDERFLOW)
                ssa_draw_index += 1
                continue

            limit = min(target, boundary)
            remaining = limit - t
            result = self._attempt_leap(
                state,
                propensity,
                critical,
                poisson_sampler,
                ssa_sampler,
                tau_candidate=tau_candidate,
                remaining=remaining,
                forcing_boundary=boundary <= target,
                poisson_draw_index=poisson_draw_index,
                ssa_draw_index=ssa_draw_index,
                config=cfg,
            )
            state = result.state
            if t + result.dt == t:
                raise RuntimeError(_STEP_UNDERFLOW)
            t = limit if result.dt == remaining else t + result.dt
            poisson_draw_index = result.poisson_draw_index
            ssa_draw_index = result.ssa_draw_index
            interval_steps += 1

        self.core.advance_timestep(state)

ApplicableOperator

Bases: Protocol

Shape-bearing operator with vector and column-matrix application.

dtype property

Return the operator's native real floating dtype.

shape property

Return the two-dimensional operator shape.

__matmul__(other)

Apply the operator to a vector or matrix of independent columns.

Source code in src/op_engine/linear_operator.py

def __matmul__(self, other: Array, /) -> Array:
    """Apply the operator to a vector or matrix of independent columns."""
    ...

Array

Bases: Protocol

Minimal Array-API duck type used at public numerical boundaries.

The contract intentionally matches :class:op_system.Array and :class:flepimop2.typing.Array. It captures only the surface inspected at package boundaries: shape, dtype, and scalar extraction. Runtime namespace discovery is delegated to :func:op_engine.array_namespace, which accepts standard-compliant arrays and native arrays such as torch.Tensor through array-api-compat. No backend module is stored in solver configuration.

dtype property

Backend-defined array dtype.

shape property

Array shape.

item()

Return a scalar value from a zero-dimensional array.

Source code in src/op_engine/_typing.py

def item(self) -> Any:  # noqa: ANN401
    """Return a scalar value from a zero-dimensional array."""

BandedOperator(bands, lower, upper) dataclass

Real NumPy operator in SciPy's diagonal-ordered band storage.

bands[upper + i - j, j] stores matrix entry (i, j). Construction copies the input and application touches only valid band entries. The retained storage is read-only; construct a new operator when coefficients, mesh or stage change. No dense matrix or factorization cache is created.

T property

Return the real Euclidean transpose without dense assembly.

dtype property

Return the native storage dtype.

shape property

Return square operator dimensions.

__matmul__(other)

Apply valid bands to a vector or independent RHS columns.

Returns:

Type	Description
NDArray[floating]	NumPy array with the same shape and dtype as the RHS.

Raises:

Type	Description
TypeError	If namespace or dtype differs.
ValueError	If RHS dimensions disagree.

Source code in src/op_engine/scipy_operator.py

def __matmul__(self, other: Array, /) -> NDArray[np.floating]:
    """Apply valid bands to a vector or independent RHS columns.

    Returns:
        NumPy array with the same shape and dtype as the RHS.

    Raises:
        TypeError: If namespace or dtype differs.
        ValueError: If RHS dimensions disagree.
    """
    if not isinstance(other, np.ndarray) or other.dtype != self.dtype:
        msg = "banded application requires the same NumPy dtype"
        raise TypeError(msg)
    n = self.shape[0]
    if other.ndim not in {1, 2} or other.shape[0] != n:
        msg = "banded application requires vector/matrix RHS with matching rows"
        raise ValueError(msg)
    result = np.zeros_like(other)
    for offset in range(-self.upper, self.lower + 1):
        start, stop = max(0, -offset), min(n, n - offset)
        weights = self.bands[self.upper + offset, start:stop]
        if other.ndim == 2:
            weights = weights[:, None]
        result[start + offset : stop + offset] += weights * other[start:stop]
    return result

__post_init__()

Copy and validate storage and static bandwidths.

Raises:

Type	Description
ValueError	If dimensions, bandwidths or valid band entries are invalid.
TypeError	If storage is not native NumPy float32/float64.

Source code in src/op_engine/scipy_operator.py

def __post_init__(self) -> None:
    """Copy and validate storage and static bandwidths.

    Raises:
        ValueError: If dimensions, bandwidths or valid band entries are invalid.
        TypeError: If storage is not native NumPy float32/float64.
    """
    real_dtypes = (np.float32, np.float64)
    if (
        not isinstance(self.bands, np.ndarray)
        or self.bands.dtype not in real_dtypes
    ):
        msg = "band storage must be a NumPy float32/float64 array"
        raise TypeError(msg)
    if self.bands.ndim != 2 or self.bands.shape[1] < 1:
        msg = "band storage must have shape (lower + upper + 1, n), n >= 1"
        raise ValueError(msg)
    n = self.bands.shape[1]
    if (
        any(
            not isinstance(v, Integral) or isinstance(v, bool) or not 0 <= v < n
            for v in (self.lower, self.upper)
        )
        or self.bands.shape[0] != self.lower + self.upper + 1
    ):
        msg = "bandwidths must be integers in [0, n) matching storage rows"
        raise ValueError(msg)
    data = self.bands.copy()
    # Canonicalize unused padding; it is not part of the represented matrix.
    for offset in range(-self.upper, self.lower + 1):
        start, stop = max(0, -offset), min(n, n - offset)
        row = self.upper + offset
        data[row, :start] = 0
        data[row, stop:] = 0
    if not np.isfinite(data).all():
        msg = "valid band entries must be finite"
        raise ValueError(msg)
    data.flags.writeable = False
    object.__setattr__(self, "bands", data)

BatchedDirectSSASolver(core, stoichiometry)

Advance independent replicates with local clocks and one local network.

A callback taking (times, state) returns (channel, replicate) hazards. It must use each replicate's own time, preserve the state, and remain constant between jumps and declared forcing boundaries. Generic callbacks assert that independence themselves; the compiled-model adapter checks it structurally.

Examples:

>>> import numpy as np
>>> from op_engine import FiringRecorder, ModelCore
>>> core = ModelCore(2, 3, np.array([0.0, 1.0, 2.0]))
>>> core.set_initial_state(np.array([[10.0] * 3, [0.0] * 3]))
>>> counts = FiringRecorder((0,))
>>> solver = BatchedDirectSSASolver(core, np.array([[-1.0], [1.0]]))
>>> def hazards(times, state):
...     return 0.5 * state[:1]
>>> solver.run(hazards, rng=np.random.SeedSequence(7), recorder=counts)
>>> core.state_array.shape, counts.counts.shape
((3, 2, 3), (2, 1, 3))
>>> bool(np.all(core.state_array.sum(axis=1) == 10.0))
True

Validate a NumPy local stoichiometric matrix against the bank state.

Raises:

Type	Description
ValueError	The local network or core layout is invalid.
TypeError	State or stoichiometry is not a NumPy array.

Source code in src/op_engine/batched_ssa.py

def __init__(self, core: ModelCore, stoichiometry: FloatBank) -> None:
    """Validate a NumPy local stoichiometric matrix against the bank state.

    Raises:
        ValueError: The local network or core layout is invalid.
        TypeError: State or stoichiometry is not a NumPy array.
    """
    state = core.get_current_state()
    if not isinstance(state, np.ndarray) or not isinstance(
        stoichiometry, np.ndarray
    ):
        msg = "Batched direct SSA currently requires NumPy state and stoichiometry"
        raise TypeError(msg)
    if (
        state.ndim != 2
        or stoichiometry.ndim != 2
        or stoichiometry.shape[0] != state.shape[0]
        or not state.shape[1]
        or not stoichiometry.shape[1]
    ):
        msg = (
            "Expected state (local species, replicate) and "
            "stoichiometry (local species, channel)"
        )
        raise ValueError(msg)
    if not np.all(np.isfinite(stoichiometry)) or np.any(
        stoichiometry != np.round(stoichiometry)
    ):
        msg = "Stoichiometry must contain finite integers"
        raise ValueError(msg)
    self.core = core
    self.stoichiometry = stoichiometry.copy()
    self.sweeps = 0
    self.events = 0

run(propensity_func, *, rng, replicate_ids=None, config=None, recorder=None)

Store bank trajectories and optional counts at the core's output times.

A root SeedSequence is framed by stable unsigned 64-bit replicate IDs using scheme OPB1 and native PCG64DXSM. Explicit generators advance in place. Replay is bank/order/shard invariant for identical inputs and schedules; changing an actual forcing declaration can change paths.

Raises:

Type	Description
ValueError	IDs, state, schedules or callback hazards are invalid.
RuntimeError	A population becomes invalid, time stagnates or a per-replicate interval event limit is exceeded.

Source code in src/op_engine/batched_ssa.py

def run(  # noqa: C901, PLR0912, PLR0914, PLR0915 - one independent-clock event loop
    self,
    propensity_func: BatchPropensity,
    *,
    rng: BankStreams,
    replicate_ids: Sequence[int] | None = None,
    config: BatchedSSAConfig | None = None,
    recorder: FiringRecorder | None = None,
) -> None:
    """Store bank trajectories and optional counts at the core's output times.

    A root SeedSequence is framed by stable unsigned 64-bit replicate IDs
    using scheme OPB1 and native PCG64DXSM. Explicit generators advance in
    place. Replay is bank/order/shard invariant for identical inputs and
    schedules; changing an actual forcing declaration can change paths.

    Raises:
        ValueError: IDs, state, schedules or callback hazards are invalid.
        RuntimeError: A population becomes invalid, time stagnates or a
            per-replicate interval event limit is exceeded.
    """
    state = np.asarray(self.core.get_current_state()).copy()
    _n_species, n_replicates = state.shape
    n_channels = self.stoichiometry.shape[1]
    ids = (
        tuple(range(n_replicates))
        if replicate_ids is None
        else tuple(replicate_ids)
    )
    if (
        len(ids) != n_replicates
        or len(set(ids)) != len(ids)
        or any(
            isinstance(i, bool) or not isinstance(i, Integral) or not 0 <= i < 2**64
            for i in ids
        )
    ):
        msg = (
            "replicate_ids must be distinct unsigned 64-bit integers, "
            "one per replicate"
        )
        raise ValueError(msg)
    streams = _replicate_streams(rng, tuple(int(i) for i in ids))
    cfg = config or BatchedSSAConfig()
    if cfg.replicate_breakpoints and len(cfg.replicate_breakpoints) != n_replicates:
        msg = "replicate_breakpoints requires one schedule per replicate"
        raise ValueError(msg)
    if (
        np.any(~np.isfinite(state))
        or np.any(state < 0)
        or np.any(state != np.round(state))
    ):
        msg = "Initial counts must be finite non-negative integers"
        raise ValueError(msg)
    local = cfg.replicate_breakpoints or ((),) * n_replicates
    schedules = [
        np.asarray(sorted({*cfg.forcing_breakpoints, *row}), dtype=float)
        for row in local
    ]
    times = np.full(n_replicates, self.core.get_time_at(0), dtype=float)
    next_forcing = np.array([
        row[np.searchsorted(row, times[i], side="right")]
        if np.searchsorted(row, times[i], side="right") < len(row)
        else np.inf
        for i, row in enumerate(schedules)
    ])
    pending = np.full(n_replicates, np.nan)
    channels = np.full(n_replicates, -1, dtype=np.int64)
    n_intervals = int(self.core.n_timesteps) - 1
    if recorder is not None:
        recorder.start(n_intervals, (n_channels, n_replicates), 0)
    self.sweeps = self.events = 0
    for interval in range(n_intervals):
        target = self.core.get_time_at(interval + 1)
        interval_events = np.zeros(n_replicates, dtype=np.int64)
        while np.any(times < target):
            self.sweeps += 1
            active = times < target
            need = active & np.isnan(pending)
            if np.any(need):
                hazards = propensity_func(times.copy(), state)
                if not isinstance(hazards, np.ndarray) or hazards.shape != (
                    n_channels,
                    n_replicates,
                ):
                    msg = (
                        "Batch propensity must return NumPy shape "
                        f"{(n_channels, n_replicates)}"
                    )
                    raise ValueError(msg)
                if np.any(~np.isfinite(hazards)) or np.any(hazards < 0):
                    msg = "Reaction propensities must be finite and non-negative"
                    raise ValueError(msg)
                cumulative = np.cumsum(hazards, axis=0)
                total = cumulative[-1]
                if not np.all(np.isfinite(total)):
                    msg = "Total reaction propensities must remain finite"
                    raise ValueError(msg)
                pending[need] = np.inf
                for column in np.flatnonzero(need & (total > 0)):
                    wait = streams[column].exponential() / total[column]
                    event_time = times[column] + wait
                    if not np.isfinite(event_time) or event_time <= times[column]:
                        msg = "Direct SSA event time failed to advance"
                        raise RuntimeError(msg)
                    threshold = min(
                        streams[column].random() * total[column],
                        np.nextafter(total[column], 0.0),
                    )
                    channels[column] = np.searchsorted(
                        cumulative[:, column], threshold, side="right"
                    )
                    pending[column] = event_time
            force = active & (next_forcing <= target) & (next_forcing <= pending)
            fire = active & ~force & (pending <= target)
            columns = np.flatnonzero(fire)
            if np.any(interval_events[columns] >= cfg.max_events):
                msg = (
                    "Exceeded max_events for a replicate while "
                    "advancing to an output time"
                )
                raise RuntimeError(msg)
            if columns.size:
                proposed = (
                    state[:, columns] + self.stoichiometry[:, channels[columns]]
                )
                if np.any(~np.isfinite(proposed)) or np.any(proposed < 0):
                    msg = "Direct SSA produced a negative or non-finite population"
                    raise RuntimeError(msg)
                state[:, columns] = proposed
                times[columns] = pending[columns]
                pending[columns] = np.nan
                interval_events[columns] += 1
                self.events += int(columns.size)
                if recorder is not None:
                    for column in columns:
                        recorder.record(
                            interval,
                            int(channels[column]) * n_replicates + int(column),
                        )
            times[force] = next_forcing[force]
            pending[force] = np.nan
            for column in np.flatnonzero(force):
                row = schedules[column]
                position = np.searchsorted(row, times[column], side="right")
                next_forcing[column] = (
                    row[position] if position < len(row) else np.inf
                )
            times[active & ~force & ~fire] = target
        self.core.advance_timestep(state.copy())

BatchedSSAConfig(max_events=1000000, forcing_breakpoints=(), replicate_breakpoints=()) dataclass

Exact event limits and shared/per-replicate forcing declarations.

max_events applies to each replicate in each output interval. Local schedules follow the input replicate order and are unioned with the shared schedule. Output times do not change the forcing or resample pending events.

__post_init__()

Validate and snapshot all declarations before executing callbacks.

Source code in src/op_engine/batched_ssa.py

def __post_init__(self) -> None:
    """Validate and snapshot all declarations before executing callbacks."""
    shared = DirectSSAConfig(self.max_events, self.forcing_breakpoints)
    object.__setattr__(self, "forcing_breakpoints", shared.forcing_breakpoints)
    local = tuple(
        _ForcingSchedule(row).breakpoints for row in self.replicate_breakpoints
    )
    object.__setattr__(self, "replicate_breakpoints", local)

CompiledObservableLike

Bases: Protocol

Read-only structural compiler surface; core imports no compiler or provider.

coordinates property

Output axes and their coordinate labels.

definition property

Scientific name, units and description.

eval_fn property

Native flat-state evaluator.

measure_requirements property

Physical measures already applied by the evaluator.

output_shape property

Output dimensions excluding time.

parameter_shapes property

Direct parameters and complete native buffer shapes.

state_dependencies property

Direct state base dependencies.

time_dependent property

Whether the expression explicitly references time.

CompiledReactionNetwork(stoichiometry, reactant_stoichiometry, channel_reactions, channel_names, _blocks, _reactions, _event_shapes, _params, reactants_complete=False, incomplete_reactions=(), dependency_incidence=None, propensity_orders=None) dataclass

Flat stochastic network compiled from typed reaction artifacts.

channel_reactions retains the parent reaction name for every expanded source-cell channel. channel_names adds integer coordinates and is intended for diagnostics only.

n_channels property

Return the number of expanded reaction channels.

n_state property

Return the number of flattened state cells.

mean_drift(time, state)

Return deterministic mean drift from the compiled channels.

Returns:

Type	Description
Array	stoichiometry @ propensity in the state's namespace.

Source code in src/op_engine/reactions.py

def mean_drift(self, time: float, state: Array) -> Array:
    """Return deterministic mean drift from the compiled channels.

    Returns:
        ``stoichiometry @ propensity`` in the state's namespace.
    """
    xp = _namespace_of(state)
    propensity = self.propensity(time, state)
    stoichiometry = cast(
        "Array",
        xp.asarray(self.stoichiometry, dtype=state.dtype),
    )
    return cast("Array", xp.matmul(stoichiometry, propensity))

propensity(time, state)

Evaluate every reaction and concatenate its source-cell rates.

The stochastic core stores the state as (n_state, 1). A flat state is accepted as well so the same callback can construct hybrid mean drift. Returned shape mirrors the input convention.

Returns:

Type	Description
Array	Propensities in the evolving state's array namespace.

Source code in src/op_engine/reactions.py

def propensity(self, time: float, state: Array) -> Array:
    """Evaluate every reaction and concatenate its source-cell rates.

    The stochastic core stores the state as ``(n_state, 1)``. A flat
    state is accepted as well so the same callback can construct hybrid
    mean drift. Returned shape mirrors the input convention.

    Returns:
        Propensities in the evolving state's array namespace.
    """
    xp = _namespace_of(state)
    batched = state.shape == (self.n_state, 1)
    if not batched and state.shape != (self.n_state,):
        msg = (
            f"Reaction propensity received state shape {state.shape}; expected "
            f"{(self.n_state,)} or {(self.n_state, 1)}."
        )
        raise ValueError(msg)
    flat_state = cast("Array", xp.reshape(state, (self.n_state,)))
    indexable_state = cast("_IndexableArray", flat_state)
    state_dict = {
        block.base: cast(
            "Array",
            xp.reshape(
                indexable_state[block.offset : block.offset + block.size],
                block.shape,
            ),
        )
        for block in self._blocks
    }

    values: list[Array] = []
    for reaction, expected_shape in zip(
        self._reactions,
        self._event_shapes,
        strict=True,
    ):
        raw = reaction.propensity_fn(time, state_dict, **self._params)
        raw_namespace = getattr(raw, "__array_namespace__", None)
        if raw_namespace is not None and raw_namespace() is not xp:
            msg = f"Reaction {reaction.name!r} changed the state array namespace."
            raise TypeError(msg)
        value = cast("Array", xp.asarray(raw, dtype=state.dtype))
        if value.shape != expected_shape:
            msg = (
                f"Reaction {reaction.name!r} returned propensity shape "
                f"{value.shape}; expected {expected_shape}."
            )
            raise ValueError(msg)
        values.append(cast("Array", xp.reshape(value, (-1,))))

    combined = cast("Array", xp.concat(tuple(values), axis=0))
    if batched:
        return cast("Array", xp.reshape(combined, (self.n_channels, 1)))
    return combined

CoreSolver(core, operators=None, *, operator_axis='state')

Semi-implicit solver operating on a ModelCore time/state grid.

Initialize CoreSolver.

Parameters:

Name	Type	Description	Default
core	ModelCore	ModelCore instance to solve.	required
operators	CoreOperators | StageOperatorFactory | None	Default operator spec (tuple or factory) for implicit stages.	None
operator_axis	str | int	Axis along which operators act (name or index).	'state'

Source code in src/op_engine/core_solver.py

def __init__(
    self,
    core: ModelCore,
    operators: CoreOperators | StageOperatorFactory | None = None,
    *,
    operator_axis: str | int = "state",
) -> None:
    """Initialize CoreSolver.

    Args:
        core: ModelCore instance to solve.
        operators: Default operator spec (tuple or factory) for implicit stages.
        operator_axis: Axis along which operators act (name or index).
    """
    self.core = core
    self.dtype = core.dtype
    self.state_shape = core.state_shape
    self.state_ndim = len(self.state_shape)

    # Operator axis resolution
    self._op_axis = operator_axis
    self._op_axis_idx: int | None = None
    self._op_axis_len: int | None = None

    # Default operator spec (tuple or factory or None)
    self._default_operator_spec: CoreOperators | StageOperatorFactory | None = (
        operators
    )

    # Preallocate buffers (full tensor shape)
    self._rhs_buffer: NDArray[np.floating] = np.zeros(
        self.state_shape,
        dtype=self.dtype,
    )
    self._next_state_buffer: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Shared stepping buffers
    self._f_n: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._f_pred: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._state_pred: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Adaptive buffers
    self._y_full: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_half: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_two_half: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_low: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._err: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # Working state buffers (avoid allocating per substep)
    self._y_curr: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._y_try: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)

    # TR-BDF2 additional buffers
    self._y_stage1: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._f_stage1: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._f_extrap: NDArray[np.floating] = np.zeros_like(self._rhs_buffer)
    self._sparse_identity: csr_matrix | None = None

    self._last_adaptive_schedule: AdaptiveStepSchedule | None = None
    self._last_nonlinear_diagnostics: NonlinearIntegrationDiagnostics | None = None
    self._state_monitor: _OutputStateMonitor | None = None

    # Validate operator sizes if default spec is a static tuple
    if operators is not None and not callable(operators):
        _predictor, left_op, right_op = self._normalize_ops_tuple(operators)
        if left_op is not None and right_op is not None:
            self._resolve_operator_axis()
            self._validate_operator_sizes(left_op, right_op)

last_adaptive_schedule property

Return the schedule recorded or replayed by the latest adaptive run.

last_nonlinear_diagnostics property

Return nonlinear diagnostics from the latest SDIRK run or replay.

last_state_diagnostics property

Return output-state checks from the latest run/replay, if enabled.

adaptive_explicit_step(rhs_func, *, method, t, dt, y, first_stage=None, forcing_breakpoints=())

Attempt one explicit step and return its local-error estimate.

This functional boundary does not mutate :class:ModelCore history and does not extract host scalars. Provider integrations can compose it with backend-native adaptive controller loops.

Parameters:

Name	Type	Description	Default
rhs_func	RHSFunction	Function computing the explicit RHS F(t, y).	required
method	str	Explicit solver method name.	required
t	Scalar	Step start time as a Python float or backend-native scalar.	required
dt	Scalar	Step size as a Python float or backend-native scalar.	required
y	Array	State at the step start.	required
first_stage	Array | None	Optional cached derivative at (t, y).	None
forcing_breakpoints	tuple[float, ...]	Static known jumps. The caller must clip the step at the next jump. RHS times use native one-sided neighbors; declared forcing disables first-stage reuse.	()

Returns:

Type	Description
ExplicitStepResult	Candidate state, local error estimate, controller order, and
ExplicitStepResult	reusable derivative stages.

Raises:

Type	Description
ValueError	If method is not an explicit method.

Source code in src/op_engine/core_solver.py

def adaptive_explicit_step(  # noqa: PLR0913
    self,
    rhs_func: RHSFunction,
    *,
    method: str,
    t: Scalar,
    dt: Scalar,
    y: Array,
    first_stage: Array | None = None,
    forcing_breakpoints: tuple[float, ...] = (),
) -> ExplicitStepResult:
    """Attempt one explicit step and return its local-error estimate.

    This functional boundary does not mutate :class:`ModelCore` history
    and does not extract host scalars. Provider integrations can compose
    it with backend-native adaptive controller loops.

    Args:
        rhs_func: Function computing the explicit RHS F(t, y).
        method: Explicit solver method name.
        t: Step start time as a Python float or backend-native scalar.
        dt: Step size as a Python float or backend-native scalar.
        y: State at the step start.
        first_stage: Optional cached derivative at ``(t, y)``.
        forcing_breakpoints: Static known jumps. The caller must clip the
            step at the next jump. RHS times use native one-sided neighbors;
            declared forcing disables first-stage reuse.

    Returns:
        Candidate state, local error estimate, controller order, and
        reusable derivative stages.

    Raises:
        ValueError: If ``method`` is not an explicit method.
    """
    normalized_method = _normalize_method(method)
    if normalized_method not in _EXPLICIT_METHODS:
        msg = f"Method '{method}' is not an explicit solver method"
        raise ValueError(msg)
    return self._attempt_explicit_step(
        rhs_func,
        method=normalized_method,
        t=cast("float", t),
        dt=cast("float", dt),
        y=y,
        first_stage=first_stage,
        forcing_breakpoints=forcing_breakpoints,
    )

fixed_explicit_step(rhs_func, *, method, t, dt, y, first_stage=None, forcing_breakpoints=())

Take one functional fixed step with an explicit method.

This boundary does not mutate :class:ModelCore history, so callers can compose it with backend-native loop primitives and apply the completed trajectory once. It is the supported public step kernel for external drivers such as jax.lax.scan; the method and state shape are static, while t, dt, y, and RHS parameters may be traced. RHS results must preserve the configured state shape and the namespace of y.

Parameters:

Name	Type	Description	Default
rhs_func	RHSFunction	Function computing the explicit RHS F(t, y).	required
method	str	Explicit solver method name.	required
t	Scalar	Step start time as a Python float or backend-native scalar.	required
dt	Scalar	Step size as a Python float or backend-native scalar.	required
y	Array	State at the step start.	required
first_stage	Array | None	Optional derivative at (t, y). Dormand--Prince returns this cache for the next step; reuse it only when time, state, and RHS parameters are unchanged at that start.	None
forcing_breakpoints	tuple[float, ...]	Static known jumps. The caller must clip the step at the next jump. RHS times use native one-sided neighbors; declared forcing disables first-stage reuse.	()

Returns:

Type	Description
Array	Next state and an FSAL derivative for Dormand--Prince, or
Array | None	None for other explicit methods. Seed a scan's FSAL carry
tuple[Array, Array | None]	with one step outside the scan so its structure stays fixed.

Raises:

Type	Description
ValueError	If method is not an explicit method.

Source code in src/op_engine/core_solver.py

def fixed_explicit_step(  # noqa: PLR0913
    self,
    rhs_func: RHSFunction,
    *,
    method: str,
    t: Scalar,
    dt: Scalar,
    y: Array,
    first_stage: Array | None = None,
    forcing_breakpoints: tuple[float, ...] = (),
) -> tuple[Array, Array | None]:
    """Take one functional fixed step with an explicit method.

    This boundary does not mutate :class:`ModelCore` history, so callers
    can compose it with backend-native loop primitives and apply the
    completed trajectory once. It is the supported public step kernel
    for external drivers such as ``jax.lax.scan``; the method and state
    shape are static, while ``t``, ``dt``, ``y``, and RHS parameters may
    be traced. RHS results must preserve the configured state shape and
    the namespace of ``y``.

    Args:
        rhs_func: Function computing the explicit RHS F(t, y).
        method: Explicit solver method name.
        t: Step start time as a Python float or backend-native scalar.
        dt: Step size as a Python float or backend-native scalar.
        y: State at the step start.
        first_stage: Optional derivative at ``(t, y)``. Dormand--Prince
            returns this cache for the next step; reuse it only when
            time, state, and RHS parameters are unchanged at that start.
        forcing_breakpoints: Static known jumps. The caller must clip the
            step at the next jump. RHS times use native one-sided neighbors;
            declared forcing disables first-stage reuse.

    Returns:
        Next state and an FSAL derivative for Dormand--Prince, or
        ``None`` for other explicit methods. Seed a scan's FSAL carry
        with one step outside the scan so its structure stays fixed.

    Raises:
        ValueError: If ``method`` is not an explicit method.
    """
    normalized_method = _normalize_method(method)
    if normalized_method not in _EXPLICIT_METHODS:
        msg = f"Method '{method}' is not an explicit solver method"
        raise ValueError(msg)
    return self._step_explicit_fixed(
        rhs_func,
        method=normalized_method,
        t=t,
        dt=dt,
        y=y,
        first_stage=first_stage,
        forcing_breakpoints=forcing_breakpoints,
    )

replay_adaptive_schedule(rhs_func, schedule, *, config)

Replay a recorded adaptive mesh through the configured method.

Replay bypasses error norms and accept/reject decisions while invoking the same high-order step kernels used by a live adaptive solve. With a JAX state, the static Python schedule can therefore be traced by jax.jit and differentiated with respect to array-valued model inputs.

Parameters:

Name	Type	Description	Default
rhs_func	RHSFunction	Function computing the explicit RHS F(t, y).	required
schedule	AdaptiveStepSchedule	Accepted step mesh recorded by an adaptive run.	required
config	RunConfig	Matching adaptive run configuration.	required

Returns:

Type	Description
NonlinearIntegrationDiagnostics | None	Array-valued nonlinear diagnostics for SDIRK2, otherwise None.
NonlinearIntegrationDiagnostics | None	Opt-in output-state checks are available through
NonlinearIntegrationDiagnostics | None	attr:last_state_diagnostics, independently of this return value.

Raises:

Type	Description
TypeError	If schedule has the wrong type or the state changes array ecosystems during replay.
ValueError	If config is not adaptive, the output grid differs, required provenance is absent, or structural provenance differs.

Source code in src/op_engine/core_solver.py

def replay_adaptive_schedule(
    self,
    rhs_func: RHSFunction,
    schedule: AdaptiveStepSchedule,
    *,
    config: RunConfig,
) -> NonlinearIntegrationDiagnostics | None:
    """Replay a recorded adaptive mesh through the configured method.

    Replay bypasses error norms and accept/reject decisions while invoking
    the same high-order step kernels used by a live adaptive solve. With a
    JAX state, the static Python schedule can therefore be traced by
    ``jax.jit`` and differentiated with respect to array-valued model inputs.

    Args:
        rhs_func: Function computing the explicit RHS F(t, y).
        schedule: Accepted step mesh recorded by an adaptive run.
        config: Matching adaptive run configuration.

    Returns:
        Array-valued nonlinear diagnostics for SDIRK2, otherwise ``None``.
        Opt-in output-state checks are available through
        :attr:`last_state_diagnostics`, independently of this return value.

    Raises:
        TypeError: If schedule has the wrong type or the state changes array
            ecosystems during replay.
        ValueError: If config is not adaptive, the output grid differs,
            required provenance is absent, or structural provenance differs.
    """
    if not isinstance(schedule, AdaptiveStepSchedule):
        msg = "schedule must be an AdaptiveStepSchedule"
        raise TypeError(msg)
    if not config.adaptive:
        msg = "Schedule replay requires config.adaptive=True"
        raise ValueError(msg)

    self._last_adaptive_schedule = None
    self._last_nonlinear_diagnostics = None
    self._state_monitor = None
    self._validate_schedule_time_grid(schedule)
    plan = self._resolve_run_plan(config)
    self._validate_schedule_provenance(schedule, config, plan)
    _ForcingSchedule(config.forcing_breakpoints).validate_mesh(
        schedule.output_times, schedule.step_sizes
    )
    self._start_state_diagnostics(config)

    adapter = None
    if config.replay_loop != "unroll":
        adapter = get_loop_adapter(_namespace_of(self.core.get_current_state()))
        if plan.method == "sdirk2":
            adapter = None
        if config.replay_loop == "scan" and adapter is None:
            msg = (
                "Scan replay requires a loop adapter and a method other than SDIRK2"
            )
            raise ValueError(msg)
    if adapter is not None:
        self._replay_scan_schedule(
            rhs_func,
            plan=plan,
            schedule=schedule,
            adapter=adapter,
            checkpoint=config.replay_checkpoint,
            forcing_breakpoints=config.forcing_breakpoints,
        )
        self._last_adaptive_schedule = schedule
        return None

    if plan.method == "sdirk2":
        diagnostics = self._replay_sdirk_schedule(
            rhs_func,
            plan=plan,
            schedule=schedule,
        )
        self._last_adaptive_schedule = schedule
        return diagnostics
    if plan.method in _EXPLICIT_METHODS:
        self._replay_explicit_schedule(
            rhs_func,
            plan=plan,
            schedule=schedule,
            forcing_breakpoints=config.forcing_breakpoints,
        )
    else:
        current_state = self.core.get_current_state()
        if plan.method == "imex-ark3" or not isinstance(current_state, np.ndarray):
            self._replay_array_implicit_schedule(
                rhs_func,
                plan=plan,
                schedule=schedule,
            )
        else:
            self._replay_numpy_implicit_schedule(
                rhs_func,
                plan=plan,
                schedule=schedule,
            )

    self._last_adaptive_schedule = schedule
    return None

run(rhs_func, *, config=None)

Advance the ModelCore state through its time grid.

Parameters:

Name	Type	Description	Default
rhs_func	RHSFunction	Function computing the explicit RHS F(t, y).	required
config	RunConfig | None	Optional run configuration. If None, defaults are used.	None

Returns:

Type	Description
NonlinearIntegrationDiagnostics | None	Array-valued nonlinear diagnostics for SDIRK2, otherwise None.
NonlinearIntegrationDiagnostics | None	Opt-in output-state checks are available through
NonlinearIntegrationDiagnostics | None	attr:last_state_diagnostics, independently of this return value.

Raises:

Type	Description
TypeError	If the state changes array ecosystems during a run.
ValueError	If invalid parameters are provided.

Examples:

>>> import numpy as np
>>> from op_engine import ModelCore
>>> from op_engine.core_solver import RunConfig
>>> core = ModelCore(1, 1, np.linspace(0.0, 1.0, 11))
>>> core.set_initial_state(np.array([[1.0]]))
>>> def decay(time, state):
...     return -state
>>> CoreSolver(core).run(decay, config=RunConfig(method="rk4"))
>>> np.allclose(core.state_array[-1, 0, 0], np.exp(-1.0), rtol=1e-5)
True

Source code in src/op_engine/core_solver.py

def run(
    self,
    rhs_func: RHSFunction,
    *,
    config: RunConfig | None = None,
) -> NonlinearIntegrationDiagnostics | None:
    """Advance the ModelCore state through its time grid.

    Args:
        rhs_func: Function computing the explicit RHS F(t, y).
        config: Optional run configuration. If None, defaults are used.

    Returns:
        Array-valued nonlinear diagnostics for SDIRK2, otherwise ``None``.
        Opt-in output-state checks are available through
        :attr:`last_state_diagnostics`, independently of this return value.

    Raises:
        TypeError: If the state changes array ecosystems during a run.
        ValueError: If invalid parameters are provided.

    Examples:
        >>> import numpy as np
        >>> from op_engine import ModelCore
        >>> from op_engine.core_solver import RunConfig
        >>> core = ModelCore(1, 1, np.linspace(0.0, 1.0, 11))
        >>> core.set_initial_state(np.array([[1.0]]))
        >>> def decay(time, state):
        ...     return -state
        >>> CoreSolver(core).run(decay, config=RunConfig(method="rk4"))
        >>> np.allclose(core.state_array[-1, 0, 0], np.exp(-1.0), rtol=1e-5)
        True
    """
    self._last_adaptive_schedule = None
    self._last_nonlinear_diagnostics = None
    self._state_monitor = None
    cfg = config or RunConfig()
    plan = self._resolve_run_plan(cfg)
    self._start_state_diagnostics(cfg)

    if plan.method == "sdirk2":
        return self._run_sdirk(rhs_func, plan=plan, config=cfg)

    if plan.method in _EXPLICIT_METHODS:
        self._run_explicit(rhs_func, plan=plan, config=cfg)
        return None

    current_state = self.core.get_current_state()
    if plan.method == "imex-ark3" or not isinstance(current_state, np.ndarray):
        self._run_array_implicit(rhs_func, plan=plan, config=cfg)
        return None

    time_grid = np.asarray(self.core.time_grid, dtype=float)
    n_steps = int(self.core.n_timesteps)
    history = MultistepHistory[NDArray[np.floating]](
        capacity=BDF2.stored_state_history if plan.method == "bdf2" else 0,
    )
    schedule_steps: list[tuple[float, ...]] = []

    for idx in range(n_steps - 1):
        t0 = float(time_grid[idx])
        t1 = float(time_grid[idx + 1])
        if t1 <= t0:
            raise ValueError(_TIME_GRID_INCREASING_ERROR_MSG)
        dt_out = t1 - t0

        state = self.core.get_current_state()
        if not isinstance(state, np.ndarray):
            raise TypeError(
                _NUMPY_ARRAY_ERROR_MSG.format(type_name=type(state).__name__)
            )
        np.copyto(self._y_curr, state)

        if not cfg.adaptive:
            y_next = self._advance_nonadaptive_to_time(
                rhs_func,
                plan=plan,
                t0=t0,
                dt_out=dt_out,
                y0=self._y_curr,
                history=history.states,
            )
            history = history.push(self._y_curr)
            self._advance_output_state(y_next)
            continue

        accepted_steps: list[float] = []
        y_end = self._advance_adaptive_to_time(
            rhs_func,
            AdaptiveAdvanceParams(
                plan=plan,
                t0=t0,
                t1=t1,
                y0=self._y_curr,
                adaptive_cfg=cfg.adaptive_cfg,
                dt_ctrl=cfg.dt_controller,
            ),
            accepted_steps,
        )
        schedule_steps.append(tuple(accepted_steps))
        history = history.push(self._y_curr)
        self._advance_output_state(y_end)

    if cfg.adaptive:
        self._last_adaptive_schedule = AdaptiveStepSchedule(
            output_times=tuple(float(value) for value in time_grid),
            step_sizes=tuple(schedule_steps),
            provenance=self._schedule_provenance(cfg, plan),
        )
    return None

DenseNewtonSolver(config=NewtonConfig()) dataclass

Portable dense Newton solver with statically unrolled iterations.

The solver has no hidden warm-start state. Every call begins from the explicit initial_guess supplied by the numerical method. Automatic differentiation follows the performed Newton iterations; this class does not install an implicit-function or custom derivative rule.

Attributes:

Name	Type	Description
config	NewtonConfig	Static iteration, tolerance, and damping configuration.

solve(problem, initial_guess)

Run a fixed number of dense Newton updates.

Parameters:

Name	Type	Description	Default
problem	NonlinearProblem	Residual and dense Jacobian callbacks.	required
initial_guess	Array	Explicit warm start in the desired array namespace.	required

Returns:

Type	Description
NonlinearSolveResult	Candidate root with array-valued convergence diagnostics.

Raises:

Type	Description
TypeError	If callbacks change namespace or dtype.
ValueError	If no dense Jacobian is supplied or shapes are invalid.

Source code in src/op_engine/nonlinear_solver.py

def solve(  # noqa: PLR0914
    self,
    problem: NonlinearProblem,
    initial_guess: Array,
    /,
) -> NonlinearSolveResult:
    """Run a fixed number of dense Newton updates.

    Args:
        problem: Residual and dense Jacobian callbacks.
        initial_guess: Explicit warm start in the desired array namespace.

    Returns:
        Candidate root with array-valued convergence diagnostics.

    Raises:
        TypeError: If callbacks change namespace or dtype.
        ValueError: If no dense Jacobian is supplied or shapes are invalid.
    """
    jacobian = problem.jacobian
    if jacobian is None:
        msg = "DenseNewtonSolver requires NonlinearProblem.jacobian"
        raise ValueError(msg)

    xp = _namespace_of(initial_guess)
    n_unknowns = math.prod(initial_guess.shape)
    if n_unknowns < 1:
        msg = "initial_guess must contain at least one value"
        raise ValueError(msg)
    if not (
        xp.isdtype(initial_guess.dtype, "real floating")
        or xp.isdtype(initial_guess.dtype, "complex floating")
    ):
        msg = "initial_guess must have a floating-point dtype"
        raise TypeError(msg)

    iterate = cast(
        "Array",
        xp.asarray(initial_guess, dtype=initial_guess.dtype),
    )
    residual = problem.residual(iterate)
    _validate_like(residual, iterate, label="residual")
    initial_residual_norm = _rms_norm(residual)
    applied_step = cast("Array", xp.zeros_like(iterate))

    for _ in range(self.config.max_iterations):
        jacobian_value = jacobian(iterate)
        if _namespace_of(jacobian_value) is not xp:
            msg = "jacobian must preserve the iterate array namespace"
            raise TypeError(msg)
        expected_jacobian_shape = (n_unknowns, n_unknowns)
        if jacobian_value.shape != expected_jacobian_shape:
            msg = (
                f"jacobian shape {jacobian_value.shape} does not match "
                f"flattened system shape {expected_jacobian_shape}"
            )
            raise ValueError(msg)
        if jacobian_value.dtype != iterate.dtype:
            msg = "jacobian must preserve the iterate dtype"
            raise TypeError(msg)

        residual_flat = xp.reshape(residual, (n_unknowns,))
        delta_flat = xp.linalg.solve(
            jacobian_value,
            xp.negative(residual_flat),
        )
        delta = xp.reshape(delta_flat, iterate.shape)
        applied_step = cast(
            "Array",
            xp.multiply(delta, self.config.damping),
        )
        iterate = cast("Array", xp.add(iterate, applied_step))
        residual = problem.residual(iterate)
        _validate_like(residual, iterate, label="residual")

    residual_norm = _rms_norm(residual)
    step_norm = _rms_norm(applied_step)
    tolerance = xp.add(
        xp.asarray(self.config.atol, dtype=residual_norm.dtype),
        xp.multiply(initial_residual_norm, self.config.rtol),
    )
    finite = xp.all(xp.isfinite(residual))
    converged = cast(
        "Array",
        xp.logical_and(finite, xp.less_equal(residual_norm, tolerance)),
    )
    diagnostics = NonlinearSolveDiagnostics(
        converged=converged,
        iterations=self.config.max_iterations,
        residual_evaluations=self.config.max_iterations + 1,
        jacobian_evaluations=self.config.max_iterations,
        initial_residual_norm=initial_residual_norm,
        residual_norm=residual_norm,
        step_norm=step_norm,
    )
    return NonlinearSolveResult(iterate, residual, diagnostics)

DerivativeArtifact

Bases: Protocol

Structural compiler boundary, with no runtime compiler dependency.

Matrices have shapes (N,N), (N,Q), (M,N), (M,Q). The producer declares supported operations and must represent the same smooth RHS/measurements.

model property

Identify the derivative producer's source model.

observable_names property

Give the ordered scalar measurement coordinates.

parameter_names property

Give RHS and measurement parameter coordinates.

state_names property

Give the ordered expanded state coordinates.

evaluate(t, y, **parameters)

Evaluate native derivatives without changing coordinates or graphs.

Source code in src/op_engine/sensitivities.py

def evaluate(
    self, t: object, y: object, **parameters: object
) -> DerivativeMatrices:
    """Evaluate native derivatives without changing coordinates or graphs."""
    ...

DerivativeMatrices

Bases: Protocol

Four native matrices in declared physical state/parameter coordinates.

observable_parameters property

Return dg/dtheta with shape (M,Q).

observable_state property

Return dg/dy with shape (M,N).

parameters property

Return df/dtheta with shape (N,Q).

state property

Return df/dy with shape (N,N).

DiffusionConfig(coeff, dtype=np.float64, bc='neumann') dataclass

Configuration for diffusion-like linear operators.

Attributes:

Name	Type	Description
coeff	float	Physical diffusion coefficient D (units length^2 / time).
dtype	DTypeLike	Floating dtype (e.g. np.float64).
bc	str	Boundary condition; either "neumann" or "absorbing".

__post_init__()

Validate context-free diffusion parameters.

Raises:

Type	Description
ValueError	If the coefficient or boundary condition is invalid.

Source code in src/op_engine/matrix_ops.py

def __post_init__(self) -> None:
    """Validate context-free diffusion parameters.

    Raises:
        ValueError: If the coefficient or boundary condition is invalid.
    """
    if not np.isfinite(self.coeff) or self.coeff < 0.0:
        msg = "coeff must be finite and non-negative"
        raise ValueError(msg)
    if self.bc not in {"neumann", "absorbing"}:
        raise ValueError(_UNKNOWN_BC_ERROR.format(bc=self.bc))
    np.dtype(self.dtype)

DirectSSAConfig(max_events=1000000, forcing_breakpoints=()) dataclass

Configuration for Gillespie's exact direct SSA.

Attributes:

Name	Type	Description
max_events	int	Maximum number of events applied in one output interval. This guard detects explosive or otherwise pathological processes.
forcing_breakpoints	tuple[float, ...]	Strictly increasing times at which external forcing may change. Propensities must be constant in time between these boundaries while the state is unchanged. Forcing is right-continuous; an event at a boundary is discarded and redrawn.

__post_init__()

Validate the event-count guard and snapshot forcing boundaries.

Raises:

Type	Description
ValueError	If the event-count guard or forcing schedule is invalid.

Source code in src/op_engine/stochastic_solver.py

def __post_init__(self) -> None:
    """Validate the event-count guard and snapshot forcing boundaries.

    Raises:
        ValueError: If the event-count guard or forcing schedule is invalid.
    """
    if (
        not isinstance(self.max_events, Integral)
        or isinstance(self.max_events, bool)
        or self.max_events < 1
    ):
        msg = "max_events must be a positive integer"
        raise ValueError(msg)
    schedule = _ForcingSchedule(self.forcing_breakpoints)
    object.__setattr__(self, "forcing_breakpoints", schedule.breakpoints)

DirectSSASolver(core, stoichiometry, *, reaction_axis='state')

Gillespie direct SSA for an exact continuous-time reaction process.

All reaction channels and batch cells form one flattened categorical event space. This is the superposition of the independent batched processes, so each accepted event changes exactly one batch cell. Propensities must be constant in time between events and declared forcing boundaries. Requested output times only observe the path and never cause an event to be resampled.

Initialize a reaction network.

Parameters:

Name	Type	Description	Default
core	ModelCore	Model state and output-time container.	required
stoichiometry	Array	Integer-valued matrix shaped (n_species, n_reactions).	required
reaction_axis	str | int	State axis changed by reaction firings.	'state'

Source code in src/op_engine/stochastic_solver.py

def __init__(
    self,
    core: ModelCore,
    stoichiometry: Array,
    *,
    reaction_axis: str | int = "state",
) -> None:
    """Initialize a reaction network.

    Args:
        core: Model state and output-time container.
        stoichiometry: Integer-valued matrix shaped
            (n_species, n_reactions).
        reaction_axis: State axis changed by reaction firings.

    """
    self.core = core
    self._network = _ReactionNetwork(core, stoichiometry, reaction_axis)

n_reactions property

Return the number of reaction channels.

run(propensity_func, ssa_sampler=None, *, config=None, rng=None, recorder=None)

Advance an exact reaction trajectory through the core output grid.

An event drawn beyond an output boundary is retained for the following interval. A pending event at or beyond a forcing boundary is discarded. Zero total propensity waits for the next forcing boundary without sampling; with no future boundary it is absorbing.

Parameters:

Name	Type	Description	Default
propensity_func	PropensityFunction	Reaction-channel propensities, constant in time between events and configured forcing boundaries.	required
ssa_sampler	SSASampler | None	Backend-specific exponential/categorical sampler.	None
config	DirectSSAConfig | None	Optional event-count guard and forcing schedule.	None
rng	RandomStream | None	Explicit NumPy stream; supply instead of ssa_sampler.	None
recorder	FiringRecorder | None	Optional accepted-firing diagnostics, reset for this run.	None

Raises:

Type	Description
ValueError	If a sampler/stream is absent or both are supplied.
TypeError	If rng is unsupported or used with non-NumPy state.
RuntimeError	If an interval exceeds its event limit or a reaction produces a negative population.

Source code in src/op_engine/stochastic_solver.py

def run(  # noqa: C901 - pending events and optional recording share one event loop
    self,
    propensity_func: PropensityFunction,
    ssa_sampler: SSASampler | None = None,
    *,
    config: DirectSSAConfig | None = None,
    rng: RandomStream | None = None,
    recorder: FiringRecorder | None = None,
) -> None:
    """Advance an exact reaction trajectory through the core output grid.

    An event drawn beyond an output boundary is retained for the following
    interval. A pending event at or beyond a forcing boundary is discarded.
    Zero total propensity waits for the next forcing boundary without
    sampling; with no future boundary it is absorbing.

    Args:
        propensity_func: Reaction-channel propensities, constant in time
            between events and configured forcing boundaries.
        ssa_sampler: Backend-specific exponential/categorical sampler.
        config: Optional event-count guard and forcing schedule.
        rng: Explicit NumPy stream; supply instead of ssa_sampler.
        recorder: Optional accepted-firing diagnostics, reset for this run.

    Raises:
        ValueError: If a sampler/stream is absent or both are supplied.
        TypeError: If rng is unsupported or used with non-NumPy state.
        RuntimeError: If an interval exceeds its event limit or a reaction
            produces a negative population.
    """  # noqa: DOC502 - random-source validation is shared
    ssa_sampler = _select_sampler(
        ssa_sampler,
        rng,
        state=self.core.get_current_state(),
        make_sampler=lambda generator: NumpySSASampler(rng=generator),
    )
    cfg = config or DirectSSAConfig()
    forcing = _ForcingSchedule(cfg.forcing_breakpoints)
    state = self.core.get_current_state()
    self._network.validate_finite_nonnegative(
        state,
        message="Initial state is invalid",
    )

    if recorder is not None:
        recorder.start(
            int(self.core.n_timesteps) - 1,
            self._network.reaction_shape,
            self._network.reaction_axis,
        )

    t = self.core.get_time_at(0)
    draw_index = 0
    pending_time: float | None = None
    pending_event_index: int | None = None
    absorbing = False

    for output_index in range(int(self.core.n_timesteps) - 1):
        target = self.core.get_time_at(output_index + 1)
        interval_events = 0
        while not absorbing and t < target:
            boundary = forcing.next_after(t)
            if pending_time is None:
                propensity = self._network.evaluate_propensity(
                    propensity_func,
                    t=t,
                    state=state,
                )
                event = _draw_ssa_event(
                    propensity,
                    ssa_sampler,
                    draw_index=draw_index,
                )
                if event is None:
                    if not np.isfinite(boundary):
                        absorbing = True
                        break
                    pending_time = boundary
                else:
                    pending_time = _ssa_event_time(t, event[0])
                    pending_event_index = event[1]
                    draw_index += 1

            if boundary <= target and boundary <= pending_time:
                t = boundary
                pending_time = None
                pending_event_index = None
                continue

            if pending_time > target:
                break
            if interval_events >= cfg.max_events:
                raise RuntimeError(_MAX_EVENTS)
            if pending_event_index is None:
                msg = "Direct SSA event state is inconsistent"
                raise RuntimeError(msg)

            state = self._apply_event(pending_event_index, state)
            if recorder is not None:
                recorder.record(output_index, pending_event_index)
            t = pending_time
            pending_time = None
            pending_event_index = None
            interval_events += 1

        self.core.advance_timestep(state)

EquilibriumDomain

Bases: StrEnum

Whether declared one-sided state constraints affect interpretation.

FiringRecorder(channels=None)

Record selected channels in output intervals, without adding state rows.

Counts are NumPy int64 diagnostics, even when the evolving state uses another namespace. Intervals are (t[i], t[i+1]]. A forcing change wins a tie with a candidate event, so that discarded event is not counted. Reusing a recorder starts a fresh record. Recording never consumes random numbers.

Examples:

>>> import numpy as np
>>> from op_engine import DirectSSASolver, ModelCore
>>> core = ModelCore(2, 1, np.array([0.0, 1.0]))
>>> core.set_initial_state(np.array([[10.0], [0.0]]))
>>> counts = FiringRecorder((0,))
>>> solver = DirectSSASolver(core, np.array([[-1.0], [1.0]]))
>>> def hazards(time, state):
...     return state[:1]
>>> solver.run(hazards, rng=np.random.default_rng(4), recorder=counts)
>>> counts.counts.shape, counts.cumulative.shape
((1, 1, 1), (2, 1, 1))
>>> int(counts.counts.sum()) == int(core.state_array[-1, 1, 0])
True

Select local reaction indices; None records every channel.

Raises:

Type	Description
ValueError	Channel selections contain invalid or duplicate indices.

Source code in src/op_engine/firing.py

def __init__(self, channels: tuple[int, ...] | None = None) -> None:
    """Select local reaction indices; None records every channel.

    Raises:
        ValueError: Channel selections contain invalid or duplicate indices.
    """
    if channels is not None and (
        any(
            isinstance(i, bool) or not isinstance(i, Integral) or i < 0
            for i in channels
        )
        or len(set(channels)) != len(channels)
    ):
        msg = "channels must contain distinct non-negative integer indices"
        raise ValueError(msg)
    self._selection = None if channels is None else tuple(int(i) for i in channels)
    self.channel_indices: tuple[int, ...] = ()
    self._rows: dict[int, int] = {}
    self._counts: NDArray[np.int64] | None = None
    self._reaction_shape: tuple[int, ...] = ()
    self._reaction_axis = 0

counts property

Return read-only (interval, selected channel, batch...) counts.

Raises:

Type	Description
RuntimeError	The recorder has not been attached to a run.

cumulative property

Return cumulative counts at output times, including an initial zero.

record(interval, flat_event_index)

Count one accepted event in its channel and batch cell.

Raises:

Type	Description
RuntimeError	Recording has not started or an interval count overflows.

Source code in src/op_engine/firing.py

def record(self, interval: int, flat_event_index: int) -> None:
    """Count one accepted event in its channel and batch cell.

    Raises:
        RuntimeError: Recording has not started or an interval count overflows.
    """
    if self._counts is None:
        msg = "FiringRecorder has not started a run"
        raise RuntimeError(msg)
    coordinate = np.unravel_index(flat_event_index, self._reaction_shape)
    row = self._rows.get(int(coordinate[self._reaction_axis]))
    if row is None:
        return
    batch = tuple(
        int(v) for i, v in enumerate(coordinate) if i != self._reaction_axis
    )
    index = (interval, row, *batch)
    if self._counts[index] == np.iinfo(np.int64).max:
        msg = "Accepted firing count exceeds int64 capacity"
        raise RuntimeError(msg)
    self._counts[index] += 1

start(n_intervals, reaction_shape, reaction_axis)

Allocate a fresh record for a solver's validated reaction layout.

Raises:

Type	Description
ValueError	A selected channel is outside the reaction layout.

Source code in src/op_engine/firing.py

def start(
    self, n_intervals: int, reaction_shape: tuple[int, ...], reaction_axis: int
) -> None:
    """Allocate a fresh record for a solver's validated reaction layout.

    Raises:
        ValueError: A selected channel is outside the reaction layout.
    """
    n_channels = reaction_shape[reaction_axis]
    selection = (
        self._selection if self._selection is not None else tuple(range(n_channels))
    )
    if any(i >= n_channels for i in selection):
        msg = "Selected firing channel is outside the reaction layout"
        raise ValueError(msg)
    self.channel_indices = selection
    self._rows = {channel: row for row, channel in enumerate(selection)}
    batch_shape = tuple(
        n for i, n in enumerate(reaction_shape) if i != reaction_axis
    )
    self._counts = np.zeros(
        (n_intervals, len(selection), *batch_shape), dtype=np.int64
    )
    self._reaction_shape = reaction_shape
    self._reaction_axis = reaction_axis

GridGeometry(n, dx) dataclass

Geometry of a 1D spatial grid.

Attributes:

Name	Type	Description
n	int	Number of grid points.
dx	float	Grid spacing.

__post_init__()

Validate context-free grid geometry.

Raises:

Type	Description
ValueError	If the grid size or spacing is invalid.

Source code in src/op_engine/matrix_ops.py

def __post_init__(self) -> None:
    """Validate context-free grid geometry.

    Raises:
        ValueError: If the grid size or spacing is invalid.
    """
    if not isinstance(self.n, Integral) or isinstance(self.n, bool) or self.n < 1:
        msg = "n must be a positive integer"
        raise ValueError(msg)
    if not np.isfinite(self.dx) or self.dx <= 0.0:
        msg = "dx must be finite and positive"
        raise ValueError(msg)

LinearSolveCandidate(value, success, iterations, message='') dataclass

Callback output before the engine verifies the original residual.

success reports the solver's own termination status. A false status cannot be overridden by a small residual. Iterations use the concrete callback's declared unit and must fit its requested budget.

LinearSolveDiagnostics(converged, residual_norm, tolerance, iterations, message) dataclass

Eager convergence evidence, one entry per RHS column (one for vectors).

LinearSolveError(result)

Bases: RuntimeError

A checked eager solve failed; result retains per-column evidence.

Retain the failed result and identify failed RHS columns.

Source code in src/op_engine/linear_operator.py

def __init__(self, result: LinearSolveResult) -> None:
    """Retain the failed result and identify failed RHS columns."""
    self.result = result
    failed = np.flatnonzero(~result.diagnostics.converged).tolist()
    super().__init__(
        f"linear solve failed for RHS columns {failed}: "
        f"{result.diagnostics.message}"
    )

LinearSolveOptions(rtol=1e-08, atol=1e-10, max_iterations=100) dataclass

Per-column true-residual tolerance and callback iteration budget.

A column passes when ||b - A x||_2 <= max(atol, rtol * ||b||_2). max_iterations counts the callback's documented iteration unit; a concrete adapter must declare whether this means inner steps or cycles.

__post_init__()

Validate static tolerances and iteration budget.

Raises:

Type	Description
ValueError	If a tolerance or iteration budget is invalid.

Source code in src/op_engine/linear_operator.py

def __post_init__(self) -> None:
    """Validate static tolerances and iteration budget.

    Raises:
        ValueError: If a tolerance or iteration budget is invalid.
    """
    if any(not np.isfinite(v) or v < 0 for v in (self.rtol, self.atol)):
        msg = "linear solve tolerances must be finite and non-negative"
        raise ValueError(msg)
    if (
        not isinstance(self.max_iterations, Integral)
        or isinstance(self.max_iterations, bool)
        or self.max_iterations < 1
    ):
        msg = "max_iterations must be a positive integer"
        raise ValueError(msg)

LinearSolveRequest(operator, rhs, options=LinearSolveOptions(), preconditioner=None, context=None) dataclass

Explicit eager solve inputs; assembly and cache ownership stay external.

preconditioner is an optional approximate inverse action with the same square shape, namespace and dtype as the operator. The concrete solver must document its side and use, or reject it. context is caller-owned metadata identifying mesh, coefficients, stage size/scale and preconditioner version. A callback that caches preparation must invalidate whenever any of these change. The engine creates no cache for this callback path.

LinearSolveResult(value, diagnostics) dataclass

Candidate solution with independently recomputed residual diagnostics.

LocalSpectrum

Bases: StrEnum

Strict classification of an assessed reduced linear spectrum.

LocalStabilityConfig(residual_tolerance, spectral_absolute_tolerance, spectral_relative_tolerance, boundary_tolerance, finite_difference_step=None) dataclass

Explicit numerical thresholds for :func:local_stability.

residual_tolerance and boundary_tolerance may be scalars or one value per state coordinate. Spectral tolerances have inverse-time units; their effective band is absolute + relative * max(abs(eigenvalue)).

__post_init__()

Reject invalid scalar controls before evaluating a model.

Raises:

Type	Description
ValueError	A spectral tolerance is negative, non-finite or not scalar.

Source code in src/op_engine/equilibrium_stability.py

def __post_init__(self) -> None:
    """Reject invalid scalar controls before evaluating a model.

    Raises:
        ValueError: A spectral tolerance is negative, non-finite or not scalar.
    """
    for name in ("spectral_absolute_tolerance", "spectral_relative_tolerance"):
        value = getattr(self, name)
        if (
            not isinstance(value, (int, float))
            or isinstance(value, bool)
            or not math.isfinite(value)
            or value < 0
        ):
            msg = f"{name} must be finite and non-negative"
            raise ValueError(msg)

LocalStabilityResult(status, classification, message, state, time, coordinates, coordinate_system, time_unit, jacobian_source, jacobian_reference, model_digest, solver_converged, residual_tolerance, boundary_tolerance, fixed_mask, nonnegative_mask, domain=EquilibriumDomain.UNSPECIFIED, raw_residual=None, active_residual=None, residual_scaled_rms=None, boundary_coordinates=(), finite_difference_steps=None, jacobian=None, full_eigenvalues=None, active_jacobian=None, active_eigenvalues=None, invariant_rows=None, invariant_rank=0, constraint_rank=0, constraint_svd_tolerance=None, tangent_basis=None, tangent_jacobian=None, tangent_eigenvalues=None, conservation_error=None, conservation_tolerance=None, invariant_flow_error=None, invariant_flow_tolerance=None, expected_conservation_neutral_modes=0, observed_active_neutral_modes=None, effective_spectral_tolerance=None, spectral_abscissa=None, leading_eigenvalue=None, leading_timescale=None, timescale_kind=None) dataclass

Evidence and interpretation from one local linearization.

LocalStabilityStatus

Bases: StrEnum

Numerical status, independent of any spectral classification.

ModelCore(n_states, n_subgroups, time_grid, *, options=None)

Core state and time manager for time-evolving models.

Examples:

>>> import numpy as np
>>> core = ModelCore(1, 1, np.array([0.0, 0.5, 1.0]))
>>> core.set_initial_state(np.array([[2.0]]))
>>> core.get_current_state().shape
(1, 1)
>>> core.get_time_at(1)
0.5

Initialize ModelCore.

Parameters:

Name	Type	Description	Default
n_states	int	Number of state variables.	required
n_subgroups	int	Number of subgroups (e.g., age groups).	required
time_grid	ndarray	1D array of times, shape (n_timesteps,).	required
options	ModelCoreOptions | None	Optional ModelCoreOptions for additional configuration.	None

Raises:

Type	Description
ValueError	if time_grid is invalid or any shapes mismatch.

Source code in src/op_engine/model_core.py

def __init__(
    self,
    n_states: int,
    n_subgroups: int,
    time_grid: np.ndarray,
    *,
    options: ModelCoreOptions | None = None,
) -> None:
    """
    Initialize ModelCore.

    Args:
        n_states: Number of state variables.
        n_subgroups: Number of subgroups (e.g., age groups).
        time_grid: 1D array of times, shape (n_timesteps,).
        options: Optional ModelCoreOptions for additional configuration.

    Raises:
        ValueError: if time_grid is invalid or any shapes mismatch.
    """
    opts = options or ModelCoreOptions()

    self.dtype = np.dtype(opts.dtype)

    self.time_grid = np.asarray(time_grid, dtype=self.dtype)
    if self.time_grid.ndim != 1:
        raise ValueError(_TIMEGRID_1D_ERROR)

    self.n_timesteps = int(self.time_grid.size)
    if self.n_timesteps < 1:
        raise ValueError(_TIMEGRID_MIN_POINTS_ERROR)

    if self.n_timesteps > 1:
        dt_arr = np.diff(self.time_grid)
        if np.any(dt_arr <= 0):
            raise ValueError(_TIMEGRID_MONOTONE_ERROR)
        self.dt_grid = np.asarray(dt_arr, dtype=self.dtype)
        self.dt = float(self.dt_grid.mean())
    else:
        self.dt_grid = np.asarray([], dtype=self.dtype)
        self.dt = 0.0

    self.n_states = int(n_states)
    self.n_subgroups = int(n_subgroups)
    self._other_axes = tuple(int(x) for x in opts.other_axes)

    self.state_shape = (self.n_states, self.n_subgroups, *self._other_axes)
    self.store_history = bool(opts.store_history)

    if opts.axis_names is None:
        base = ["state", "subgroup"]
        extra = [f"axis{i}" for i in range(2, len(self.state_shape))]
        self.axis_names = tuple(base + extra)
    else:
        if len(opts.axis_names) != len(self.state_shape):
            raise ValueError(
                _AXIS_NAMES_LEN_ERROR.format(
                    actual=len(opts.axis_names),
                    expected=len(self.state_shape),
                )
            )
        self.axis_names = tuple(opts.axis_names)

    self.axis_coords: dict[str, np.ndarray] = {}
    if opts.axis_coords is not None:
        self.axis_coords = {
            str(k): np.asarray(v, dtype=self.dtype)
            for k, v in opts.axis_coords.items()
        }

    self.current_step = 0

    # Per-timestep working state (contiguous).
    self.current_state: Array = np.zeros(self.state_shape, dtype=self.dtype)

    # Optional full history: (n_timesteps, *state_shape)
    self.state_array: Array | None
    if self.store_history:
        self.state_array = cast(
            "Array",
            np.zeros((self.n_timesteps, *self.state_shape), dtype=self.dtype),
        )
    else:
        self.state_array = None

current_time property

Current simulation time t = time_grid[current_step].

Returns:

Type	Description
float	Current time as a float.

state_ndim property

Number of state tensor dimensions (rank).

Returns:

Type	Description
int	State tensor rank as an integer.

advance_timestep(next_state)

Alias for apply_next_state, for solver-friendly naming.

Source code in src/op_engine/model_core.py

def advance_timestep(self, next_state: Array) -> None:
    """Alias for apply_next_state, for solver-friendly naming."""
    self.apply_next_state(next_state)

apply_deltas(deltas)

Apply state deltas, advancing the timestep.

Supports alternate solver implementations (e.g., splitting updates, additive increments) without forcing allocation of y_next.

Parameters:

Name	Type	Description	Default
deltas	Array	State deltas to apply, shape state_shape.	required

Raises:

Type	Description
ValueError	if deltas has incorrect shape.

Source code in src/op_engine/model_core.py

def apply_deltas(self, deltas: Array) -> None:
    """
    Apply state deltas, advancing the timestep.

    Supports alternate solver implementations (e.g., splitting updates,
    additive increments) without forcing allocation of y_next.

    Args:
        deltas: State deltas to apply, shape state_shape.

    Raises:
        ValueError: if deltas has incorrect shape.
    """
    xp = _namespace_of(self.current_state)
    deltas_arr = _coerce_state(
        deltas, self.current_state.dtype, reference=self.current_state
    )
    if deltas_arr.shape != self.state_shape:
        raise ValueError(
            _DELTAS_SHAPE_ERROR.format(
                actual=deltas_arr.shape, expected=self.state_shape
            )
        )

    self._check_can_advance()

    if isinstance(self.current_state, np.ndarray) and isinstance(
        deltas_arr, np.ndarray
    ):
        self.current_state += deltas_arr
    else:
        self.current_state = cast("Array", xp.add(self.current_state, deltas_arr))

    self.current_step += 1
    self._store_history_value(self.current_step, self.current_state)

apply_next_state(next_state)

Set the next state directly, advancing the timestep.

Parameters:

Name	Type	Description	Default
next_state	Array	State at the next timestep, shape state_shape.	required

Raises:

Type	Description
ValueError	if next_state has incorrect shape.

Source code in src/op_engine/model_core.py

def apply_next_state(self, next_state: Array) -> None:
    """
    Set the next state directly, advancing the timestep.

    Args:
        next_state: State at the next timestep, shape state_shape.

    Raises:
        ValueError: if next_state has incorrect shape.
    """
    next_state_arr = _coerce_state(
        next_state, self.current_state.dtype, reference=self.current_state
    )
    if next_state_arr.shape != self.state_shape:
        raise ValueError(
            _NEXT_STATE_SHAPE_ERROR.format(
                actual=next_state_arr.shape, expected=self.state_shape
            )
        )

    self._check_can_advance()

    if isinstance(self.current_state, np.ndarray) and isinstance(
        next_state_arr, np.ndarray
    ):
        np.copyto(self.current_state, next_state_arr)
    else:
        self.current_state = next_state_arr

    self.current_step += 1
    self._store_history_value(self.current_step, self.current_state)

apply_trajectory(trajectory)

Adopt a complete trajectory produced by a whole-grid solver.

This provides one functional update for strategies such as Diffrax that return every requested output at once. It avoids replaying immutable history through one concatenation per time point.

Parameters:

Name	Type	Description	Default
trajectory	Array	Array shaped (n_timesteps, *state_shape).	required

Raises:

Type	Description
TypeError	If the trajectory changes the active array namespace.
ValueError	If the trajectory shape is invalid.

Source code in src/op_engine/model_core.py

def apply_trajectory(self, trajectory: Array) -> None:
    """Adopt a complete trajectory produced by a whole-grid solver.

    This provides one functional update for strategies such as Diffrax that
    return every requested output at once. It avoids replaying immutable
    history through one concatenation per time point.

    Args:
        trajectory: Array shaped ``(n_timesteps, *state_shape)``.

    Raises:
        TypeError: If the trajectory changes the active array namespace.
        ValueError: If the trajectory shape is invalid.
    """
    state_namespace = _namespace_of(self.current_state)
    trajectory_namespace = _namespace_of(trajectory)
    if trajectory_namespace is not state_namespace:
        msg = "Trajectory must preserve the current state array namespace."
        raise TypeError(msg)

    adopted = _coerce_state(
        trajectory, self.current_state.dtype, reference=self.current_state
    )
    expected_shape = (self.n_timesteps, *self.state_shape)
    if adopted.shape != expected_shape:
        msg = (
            f"Trajectory shape {adopted.shape} does not match expected "
            f"{expected_shape}."
        )
        raise ValueError(msg)

    self.current_state = _array_slice(adopted, -1)
    self.state_array = adopted if self.store_history else None
    self.current_step = self.n_timesteps - 1

axis_index(axis)

Resolve an axis name or integer into an axis index.

Parameters:

Name	Type	Description	Default
axis	str | int	Axis name or index.	required

Returns:

Type	Description
int	Axis index as an integer.

Raises:

Type	Description
IndexError	if axis index is out of bounds.
ValueError	if axis name is unknown.

Source code in src/op_engine/model_core.py

def axis_index(self, axis: str | int) -> int:
    """
    Resolve an axis name or integer into an axis index.

    Args:
        axis: Axis name or index.

    Returns:
        Axis index as an integer.

    Raises:
        IndexError: if axis index is out of bounds.
        ValueError: if axis name is unknown.
    """
    if isinstance(axis, int):
        if not (0 <= axis < self.state_ndim):
            raise IndexError(_AXIS_INDEX_OOB_ERROR.format(axis=axis))
        return axis

    try:
        return self.axis_names.index(axis)
    except ValueError as exc:
        raise ValueError(_AXIS_UNKNOWN_ERROR.format(axis=axis)) from exc

get_axis_coords(axis)

Return coordinate array for a given axis, or None if not set.

Parameters:

Name	Type	Description	Default
axis	str | int	Axis name or index.	required

Returns:

Type	Description
ndarray | None	Coordinate array for the axis, or None if not set.

Source code in src/op_engine/model_core.py

def get_axis_coords(self, axis: str | int) -> np.ndarray | None:
    """
    Return coordinate array for a given axis, or None if not set.

    Args:
        axis: Axis name or index.

    Returns:
        Coordinate array for the axis, or None if not set.
    """
    idx = self.axis_index(axis)
    name = self.axis_names[idx]
    return self.axis_coords.get(name)

get_current_state()

Return the current state.

Returns:

Type	Description
Array	Current state, shape state_shape.

Source code in src/op_engine/model_core.py

def get_current_state(self) -> Array:
    """
    Return the current state.

    Returns:
        Current state, shape state_shape.
    """
    return self.current_state

get_dt(step_idx)

Return dt for the step [t_step_idx, t_step_idx+1].

Parameters:

Name	Type	Description	Default
step_idx	int	Timestep index in [0, n_timesteps - 1].	required

Returns:

Type	Description
float	dt as a float.

Raises:

Type	Description
IndexError	if step_idx is out of bounds.

Source code in src/op_engine/model_core.py

def get_dt(self, step_idx: int) -> float:
    """
    Return dt for the step [t_step_idx, t_step_idx+1].

    Args:
        step_idx: Timestep index in [0, n_timesteps - 1].

    Returns:
        dt as a float.

    Raises:
        IndexError: if step_idx is out of bounds.
    """
    if self.n_timesteps <= 1:
        return 0.0
    if not (0 <= step_idx < self.n_timesteps - 1):
        raise IndexError(_DT_INDEX_OOB_ERROR.format(idx=step_idx))
    return float(self.dt_grid[step_idx])

get_state_at(step)

Return the state at a given timestep from history.

Parameters:

Name	Type	Description	Default
step	int	Timestep index in [0, n_timesteps).	required

Returns:

Type	Description
Array	State at the given timestep, shape state_shape.

Raises:

Type	Description
RuntimeError	if history is not stored.
IndexError	if step is out of bounds.

Source code in src/op_engine/model_core.py

def get_state_at(self, step: int) -> Array:
    """
    Return the state at a given timestep from history.

    Args:
        step: Timestep index in [0, n_timesteps).

    Returns:
        State at the given timestep, shape state_shape.

    Raises:
        RuntimeError: if history is not stored.
        IndexError: if step is out of bounds.
    """
    if not self.store_history or self.state_array is None:
        raise RuntimeError(_HISTORY_NOT_STORED_ERROR)

    if not (0 <= step < self.n_timesteps):
        raise IndexError(_STEP_OOB_ERROR)

    return _array_slice(self.state_array, step)

get_time_at(step_idx)

Return time at a given step index.

Parameters:

Name	Type	Description	Default
step_idx	int	Timestep index in [0, n_timesteps).	required

Returns:

Type	Description
float	Time as a float.

Raises:

Type	Description
IndexError	if step_idx is out of bounds.

Source code in src/op_engine/model_core.py

def get_time_at(self, step_idx: int) -> float:
    """
    Return time at a given step index.

    Args:
        step_idx: Timestep index in [0, n_timesteps).

    Returns:
        Time as a float.

    Raises:
        IndexError: if step_idx is out of bounds.
    """
    if not (0 <= step_idx < self.n_timesteps):
        raise IndexError(_TIME_INDEX_OOB_ERROR.format(idx=step_idx))
    return float(self.time_grid[step_idx])

reshape_for_axis_solve(x, axis)

Reshape a state-like tensor into 2D for an axis-local operator solve.

Parameters:

Name	Type	Description	Default
x	ndarray	State-like tensor of shape state_shape.	required
axis	str | int	Axis name or index along which to reshape.	required

Returns:

Type	Description
tuple[ndarray, tuple[int, ...], int]	(x2d, original_shape, axis_index)

Raises:

Type	Description
ValueError	if x does not have shape state_shape.

Source code in src/op_engine/model_core.py

def reshape_for_axis_solve(
    self,
    x: np.ndarray,
    axis: str | int,
) -> tuple[np.ndarray, tuple[int, ...], int]:
    """Reshape a state-like tensor into 2D for an axis-local operator solve.

    Args:
        x: State-like tensor of shape state_shape.
        axis: Axis name or index along which to reshape.

    Returns:
        (x2d, original_shape, axis_index)

    Raises:
        ValueError: if x does not have shape state_shape.
    """
    x_arr = np.asarray(x, dtype=self.dtype)
    if x_arr.shape != self.state_shape:
        raise ValueError(
            _NEXT_STATE_SHAPE_ERROR.format(
                actual=x_arr.shape, expected=self.state_shape
            )
        )

    original_shape = x_arr.shape
    axis_idx = self.axis_index(axis)
    axis_len = int(original_shape[axis_idx])

    moved = np.moveaxis(x_arr, axis_idx, 0)
    x2d = moved.reshape(axis_len, -1)
    return x2d, original_shape, axis_idx

set_initial_state(initial_state)

Set the initial state at time_grid[0].

Parameters:

Name	Type	Description	Default
initial_state	Array	Initial state, shape state_shape.	required

Raises:

Type	Description
ValueError	if initial_state has incorrect shape.

Source code in src/op_engine/model_core.py

def set_initial_state(self, initial_state: Array) -> None:
    """
    Set the initial state at time_grid[0].

    Args:
        initial_state: Initial state, shape state_shape.

    Raises:
        ValueError: if initial_state has incorrect shape.
    """
    xp = _namespace_of(initial_state)
    initial_state_arr = _coerce_state(initial_state, getattr(xp, self.dtype.name))
    if initial_state_arr.shape != self.state_shape:
        raise ValueError(
            _INITIAL_STATE_SHAPE_ERROR.format(
                actual=initial_state_arr.shape,
                expected=self.state_shape,
            )
        )

    if isinstance(self.current_state, np.ndarray) and isinstance(
        initial_state_arr, np.ndarray
    ):
        np.copyto(self.current_state, initial_state_arr)
    else:
        self.current_state = initial_state_arr
    if self.store_history:
        self.state_array = cast(
            "Array",
            xp.zeros(
                (self.n_timesteps, *self.state_shape),
                dtype=initial_state_arr.dtype,
                device=_device(initial_state_arr),
            ),
        )
        self._store_history_value(0, self.current_state)

    self.current_step = 0

unreshape_from_axis_solve(x2d, original_shape, axis)

Inverse of reshape_for_axis_solve.

Parameters:

Name	Type	Description	Default
x2d	ndarray	2D array of shape (axis_len, batch).	required
original_shape	tuple[int, ...]	Original full shape before reshape.	required
axis	str | int	Axis name or index along which the reshape was done.	required

Returns:

Type	Description
ndarray	Reconstructed array of shape original_shape.

Source code in src/op_engine/model_core.py

def unreshape_from_axis_solve(
    self,
    x2d: np.ndarray,
    original_shape: tuple[int, ...],
    axis: str | int,
) -> np.ndarray:
    """
    Inverse of reshape_for_axis_solve.

    Args:
        x2d: 2D array of shape (axis_len, batch).
        original_shape: Original full shape before reshape.
        axis: Axis name or index along which the reshape was done.

    Returns:
        Reconstructed array of shape original_shape.
    """
    axis_idx = self.axis_index(axis)
    axis_len = int(original_shape[axis_idx])

    # Rebuild shape with axis leading, then move it back.
    trailing = tuple(d for i, d in enumerate(original_shape) if i != axis_idx)
    arr = np.asarray(x2d, dtype=self.dtype).reshape((axis_len, *trailing))
    return np.moveaxis(arr, 0, axis_idx)

validate_state_shape(arr, *, msg=None)

Validate that arr has state_shape.

This is intentionally small and solver-friendly; higher layers can reuse it to avoid duplicating shape checks for intermediate/stage states.

Parameters:

Name	Type	Description	Default
arr	Array	Array to validate.	required
msg	str | None	Optional custom error message prefix.	None

Raises:

Type	Description
ValueError	if arr does not have shape state_shape.

Source code in src/op_engine/model_core.py

def validate_state_shape(self, arr: Array, *, msg: str | None = None) -> None:
    """
    Validate that arr has state_shape.

    This is intentionally small and solver-friendly; higher layers can reuse it
    to avoid duplicating shape checks for intermediate/stage states.

    Args:
        arr: Array to validate.
        msg: Optional custom error message prefix.

    Raises:
        ValueError: if arr does not have shape state_shape.
    """
    arr_shape = arr.shape
    if arr_shape != self.state_shape:
        raise ValueError(
            (msg or _NEXT_STATE_SHAPE_ERROR).format(
                actual=arr_shape, expected=self.state_shape
            )
        )

NewtonConfig(max_iterations=8, rtol=1e-08, atol=1e-10, damping=1.0) dataclass

Static configuration for the portable dense Newton solver.

Attributes:

Name	Type	Description
max_iterations	int	Exact number of Newton updates to unroll.
rtol	float	Relative convergence tolerance against the initial residual norm.
atol	float	Absolute convergence tolerance.
damping	float	Fixed multiplier applied to every Newton update.

__post_init__()

Validate static solver parameters.

Raises:

Type	Description
ValueError	If an iteration count or tolerance is invalid.

Source code in src/op_engine/nonlinear_solver.py

def __post_init__(self) -> None:
    """Validate static solver parameters.

    Raises:
        ValueError: If an iteration count or tolerance is invalid.
    """
    if (
        not isinstance(self.max_iterations, Integral)
        or isinstance(self.max_iterations, bool)
        or self.max_iterations < 1
    ):
        msg = "max_iterations must be a positive integer"
        raise ValueError(msg)
    if not np.isfinite(self.rtol) or self.rtol < 0.0:
        msg = "rtol must be finite and non-negative"
        raise ValueError(msg)
    if not np.isfinite(self.atol) or self.atol < 0.0:
        msg = "atol must be finite and non-negative"
        raise ValueError(msg)
    if not np.isfinite(self.damping) or not (0.0 < self.damping <= 1.0):
        msg = "damping must be finite and in (0, 1]"
        raise ValueError(msg)

NewtonKrylovConfig(max_iterations=8, rtol=1e-08, atol=1e-10, damping=1.0, linear_options=(lambda: LinearSolveOptions(rtol=1e-06, atol=0.0))(), max_backtracks=12, backtrack_factor=0.5, armijo=0.0001) dataclass

Bases: NewtonConfig

Outer RMS tolerances, inner true-residual gate and Armijo backtracking.

Inherited max_iterations is a maximum number of applied Newton updates, not an exact unroll count. The outer gate is atol + rtol * initial_rms. Inner options independently use the checked linear solver's L2 gate max(atol, rtol * ||rhs||) and its documented iteration budget.

Attributes:

Name	Type	Description
linear_options	LinearSolveOptions	Explicit inner tolerance and iteration budget.
max_backtracks	int	Maximum reductions after the initial damped trial.
backtrack_factor	float	Multiplicative step-length reduction, strictly in (0,1).
armijo	float	Required fractional RMS decrease per unit step, strictly in (0,1).

__post_init__()

Validate outer, inner and globalization settings.

Raises:

Type	Description
TypeError	If inner options do not use the existing checked contract.
ValueError	If backtracking settings are invalid.

Source code in src/op_engine/newton_krylov.py

def __post_init__(self) -> None:
    """Validate outer, inner and globalization settings.

    Raises:
        TypeError: If inner options do not use the existing checked contract.
        ValueError: If backtracking settings are invalid.
    """
    NewtonConfig.__post_init__(self)
    if not isinstance(self.linear_options, LinearSolveOptions):
        msg = "linear_options must be LinearSolveOptions"
        raise TypeError(msg)
    if (
        not isinstance(self.max_backtracks, Integral)
        or isinstance(self.max_backtracks, bool)
        or self.max_backtracks < 0
    ):
        msg = "max_backtracks must be a nonnegative integer"
        raise ValueError(msg)
    if any(
        not np.isfinite(value) or not 0.0 < value < 1.0
        for value in (self.backtrack_factor, self.armijo)
    ):
        msg = "backtrack_factor and armijo must be finite and strictly in (0,1)"
        raise ValueError(msg)

NewtonKrylovSolver(config=NewtonKrylovConfig(), linear_solver=ScipyGmresSolver(), preconditioner_factory=None) dataclass

Stateless eager NumPy Newton-GMRES with a per-linearization inverse factory.

The default inner solver is :class:ScipyGmresSolver (left inverse preconditioning). A failed checked inner solve returns nonlinear nonconvergence without applying its candidate; malformed callbacks and arbitrary linear-solver exceptions remain fatal. The integration boundary applies its existing fixed/adaptive/replay failure policy.

Attributes:

Name	Type	Description
config	NewtonKrylovConfig	Outer/inner tolerances and bounded line search.
linear_solver	LinearSolver	Existing checked eager linear-solver callback.
preconditioner_factory	InversePreconditionerFactory | None	Optional factory receiving the exact current request, operator and :class:NewtonLinearization context.

__post_init__()

Validate explicit configuration and callbacks.

Raises:

Type	Description
TypeError	If configuration or callbacks are unsupported.

Source code in src/op_engine/newton_krylov.py

def __post_init__(self) -> None:
    """Validate explicit configuration and callbacks.

    Raises:
        TypeError: If configuration or callbacks are unsupported.
    """
    if not isinstance(self.config, NewtonKrylovConfig):
        msg = "config must be NewtonKrylovConfig"
        raise TypeError(msg)
    if not callable(self.linear_solver) or (
        self.preconditioner_factory is not None
        and not callable(self.preconditioner_factory)
    ):
        msg = "linear solver and preconditioner factory must be callable"
        raise TypeError(msg)

solve(problem, initial_guess)

Return a root candidate and explicit eager termination evidence.

Returns:

Type	Description
NonlinearSolveResult	Owned candidate and residual with checked inner/outer diagnostics.

Raises:

Type	Description
TypeError	If input or callback arrays are not native real NumPy float32/float64 with matching shape/dtype/namespace.
ValueError	If the state is empty or no residual JVP is supplied.

Source code in src/op_engine/newton_krylov.py

def solve(
    self, problem: NonlinearProblem, initial_guess: Array, /
) -> NonlinearSolveResult:
    """Return a root candidate and explicit eager termination evidence.

    Returns:
        Owned candidate and residual with checked inner/outer diagnostics.

    Raises:
        TypeError: If input or callback arrays are not native real NumPy
            float32/float64 with matching shape/dtype/namespace.
        ValueError: If the state is empty or no residual JVP is supplied.
    """
    real_dtypes = (np.float32, np.float64)
    if (
        not isinstance(initial_guess, np.ndarray)
        or initial_guess.dtype not in real_dtypes
    ):
        msg = "NewtonKrylovSolver requires eager NumPy float32/float64 arrays"
        raise TypeError(msg)
    if initial_guess.size < 1 or problem.jvp is None:
        msg = (
            "NewtonKrylovSolver requires a nonempty state and NonlinearProblem.jvp"
        )
        raise ValueError(msg)
    if not np.isfinite(initial_guess).all():
        msg = "initial_guess must be finite"
        raise ValueError(msg)
    iterate = initial_guess.copy()
    iterate.flags.writeable = False
    work = _Work()
    residual = self._residual(problem, iterate, work)
    initial_norm = _rms(residual)
    tolerance = self.config.atol + self.config.rtol * initial_norm
    step_norm = 0.0
    termination = NewtonKrylovTermination.MAX_ITERATIONS
    for _ in range(self.config.max_iterations):
        norm = _rms(residual)
        if not np.isfinite(norm):
            termination = NewtonKrylovTermination.NONFINITE_RESIDUAL
            break
        if norm <= tolerance:
            termination = NewtonKrylovTermination.CONVERGED
            break
        direction = self._direction(problem, iterate, residual, work)
        if direction is None:
            termination = NewtonKrylovTermination.LINEAR_FAILURE
            break
        trial = self._trial(
            problem, iterate, direction, norm=norm, tolerance=tolerance, work=work
        )
        if trial is None:
            termination = NewtonKrylovTermination.LINE_SEARCH_FAILURE
            break
        candidate, residual = trial
        step_norm = _rms(candidate - iterate)
        iterate = candidate
        work.iterations += 1
    norm = _rms(residual)
    converged = np.isfinite(norm) and norm <= tolerance
    if converged:
        termination = NewtonKrylovTermination.CONVERGED
    diagnostics = NonlinearSolveDiagnostics(
        converged=np.asarray(converged),
        iterations=work.iterations,
        residual_evaluations=work.residuals,
        jacobian_evaluations=0,
        initial_residual_norm=np.asarray(initial_norm, dtype=iterate.dtype),
        residual_norm=np.asarray(norm, dtype=iterate.dtype),
        step_norm=np.asarray(step_norm, dtype=iterate.dtype),
        linear_solves=tuple(work.linear),
        termination=int(termination),
        jvp_evaluations=work.actions,
        backtracks=work.backtracks,
    )
    return NonlinearSolveResult(iterate, residual, diagnostics)

NewtonKrylovTermination

Bases: IntEnum

Eager termination codes, retained in stage/integration diagnostics.

NewtonLinearization(iteration, iterate, problem_context) dataclass

Unique per-linearization preconditioner context, with an owned iterate.

iterate is a read-only snapshot. Its full shape, dtype and coefficients, problem_context (including SDIRK scale/time when applicable) and this record's identity define the preparation lifetime. A factory is invoked afresh for every linearization; a caller retaining a cache owns invalidation. No iterate history is retained in solve diagnostics.

NonlinearConvergenceError(result)

Bases: RuntimeError

Raised by an eager boundary when a nonlinear solve did not converge.

Store the failed result and its diagnostics.

Source code in src/op_engine/nonlinear_solver.py

def __init__(self, result: NonlinearSolveResult) -> None:
    """Store the failed result and its diagnostics."""
    self.result = result
    diagnostics = result.diagnostics
    super().__init__(
        "Nonlinear solve did not converge after "
        f"{diagnostics.iterations} iterations; "
        f"residual RMS={diagnostics.residual_norm.item()!r}"
    )

NonlinearIntegrationConvergenceError(diagnostics)

Bases: RuntimeError

Raised after validating failed array-valued integration diagnostics.

Store diagnostics for the invalid integration or frozen mesh.

Source code in src/op_engine/core_solver.py

def __init__(self, diagnostics: NonlinearIntegrationDiagnostics) -> None:
    """Store diagnostics for the invalid integration or frozen mesh."""
    self.diagnostics = diagnostics
    failed = np.logical_and(
        np.asarray(diagnostics.step_accepted),
        np.logical_not(np.asarray(diagnostics.step_converged)),
    )
    super().__init__(
        f"{int(np.count_nonzero(failed))} accepted nonlinear step(s) "
        "did not converge"
    )

NonlinearIntegrationDiagnostics

Bases: NamedTuple

Array-valued nonlinear diagnostics for one integration or replay.

Every array is in the state array's namespace, so the record can cross a compiled JAX boundary. stages_per_step maps flattened stage fields back to attempted steps. Rejected adaptive attempts remain present with step_accepted=False.

require_converged()

Return valid diagnostics or invalidate a failed compiled replay.

Returns:

Type	Description
NonlinearIntegrationDiagnostics	This unchanged diagnostic record.

Raises:

Type	Description
NonlinearIntegrationConvergenceError	If an accepted or replayed step contains a failed nonlinear stage.

Source code in src/op_engine/core_solver.py

def require_converged(self) -> NonlinearIntegrationDiagnostics:
    """Return valid diagnostics or invalidate a failed compiled replay.

    Returns:
        This unchanged diagnostic record.

    Raises:
        NonlinearIntegrationConvergenceError: If an accepted or replayed
            step contains a failed nonlinear stage.
    """
    if not bool(self.converged.item()):
        raise NonlinearIntegrationConvergenceError(self)
    return self

NonlinearMethodConfig(rhs_jacobian=None, solver=DenseNewtonSolver(), rhs_jvp=None) dataclass

Configuration for fully nonlinear integration methods.

Derivatives act on the entire RHS state, including cross-axis coupling. They are deliberately separate from :attr:RunConfig.jacobian, whose operators act only along CoreSolver.operator_axis for linearly implicit methods. Supply at least one derivative form and a solver that supports it. The default dense solver requires rhs_jacobian.

Attributes:

Name	Type	Description
rhs_jacobian	FullRhsJacobianFunction | None	Dense full-system Jacobian with shape (state.size, state.size).
solver	NonlinearSolver	Backend-neutral nonlinear solver implementation.
rhs_jvp	FullRhsJvpFunction | None	Optional full RHS derivative action (t, state, direction) returning the state shape, dtype and namespace. No dense Jacobian is inferred from this callback.

__post_init__()

Validate the method boundary without selecting an array backend.

Raises:

Type	Description
TypeError	If a callback or nonlinear solver is invalid, or neither derivative form is supplied.

Source code in src/op_engine/core_solver.py

def __post_init__(self) -> None:
    """Validate the method boundary without selecting an array backend.

    Raises:
        TypeError: If a callback or nonlinear solver is invalid, or neither
            derivative form is supplied.
    """
    for name, callback in (
        ("rhs_jacobian", self.rhs_jacobian),
        ("rhs_jvp", self.rhs_jvp),
    ):
        if callback is not None and not callable(callback):
            msg = f"{name} must be callable"
            raise TypeError(msg)
    if self.rhs_jacobian is None and self.rhs_jvp is None:
        msg = "Supply at least one of rhs_jacobian or rhs_jvp"
        raise TypeError(msg)
    if not isinstance(self.solver, NonlinearSolver):
        msg = "solver must implement the NonlinearSolver protocol"
        raise TypeError(msg)

NonlinearProblem(residual, jacobian=None, jvp=None, context=None) dataclass

Residual and optional derivative actions for a nonlinear equation.

residual(x) must return an array with the same shape, dtype, and array namespace as x. jacobian(x) returns a dense square array acting on the flattened residual. jvp(x, vector) returns J(x) @ vector with the same structure as x. Solvers may require either derivative form; :class:DenseNewtonSolver specifically requires jacobian.

Attributes:

Name	Type	Description
residual	NonlinearResidual	Residual function whose root is sought.
jacobian	NonlinearJacobian | None	Optional dense Jacobian function.
jvp	JacobianVectorProduct | None	Optional matrix-free Jacobian-vector product.
context	object | None	Caller-owned stage/linearization metadata; no cache is created.

__post_init__()

Validate callback presence.

Raises:

Type	Description
TypeError	If a supplied callback is not callable.

Source code in src/op_engine/nonlinear_solver.py

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

    Raises:
        TypeError: If a supplied callback is not callable.
    """
    for name, callback in (
        ("residual", self.residual),
        ("jacobian", self.jacobian),
        ("jvp", self.jvp),
    ):
        if callback is not None and not callable(callback):
            msg = f"Nonlinear problem {name} must be callable"
            raise TypeError(msg)

NonlinearSolveDiagnostics(converged, iterations, residual_evaluations, jacobian_evaluations, initial_residual_norm, residual_norm, step_norm, linear_solves=(), termination=0, jvp_evaluations=0, backtracks=0) dataclass

Array-safe diagnostics from a nonlinear solve.

converged and the norm fields are zero-dimensional arrays in the initial guess's namespace. Keeping them on-device avoids an accidental host synchronization or tracer conversion in compiled code.

Attributes:

Name	Type	Description
converged	Array	Boolean array indicating finite residual convergence.
iterations	int	Number of nonlinear updates performed.
residual_evaluations	int	Number of residual evaluations performed.
jacobian_evaluations	int	Number of dense Jacobian evaluations performed.
initial_residual_norm	Array	RMS norm before the first update.
residual_norm	Array	RMS norm after the final update.
step_norm	Array	RMS norm of the final damped update.
linear_solves	tuple[LinearSolveDiagnostics, ...]	Checked eager inner solves in execution order.
termination	int	Optional eager solver termination code (zero is unspecified).
jvp_evaluations	int	Number of residual derivative actions evaluated.
backtracks	int	Number of step-length reductions attempted.

NonlinearSolveResult(value, residual, diagnostics) dataclass

Candidate root, final residual, and convergence diagnostics.

NonlinearSolver

Bases: Protocol

Backend-neutral interface implemented by nonlinear solvers.

solve(problem, initial_guess)

Return a candidate root and explicit convergence diagnostics.

Source code in src/op_engine/nonlinear_solver.py

def solve(
    self,
    problem: NonlinearProblem,
    initial_guess: Array,
    /,
) -> NonlinearSolveResult:
    """Return a candidate root and explicit convergence diagnostics."""

NonlinearStageContext(time, step_size, diagonal_step, stage_index) dataclass

Explicit SDIRK stage metadata, freshly bound on every stage and retry.

NonlinearWorkDiagnostics

Bases: NamedTuple

Ragged eager inner-solve evidence using only namespace-preserving arrays.

solves_per_stage maps linear solves to flattened stages; columns_per_solve maps residual/threshold/convergence entries to those solves. Iteration counts are per solve. Termination, JVP and backtracking counts are per stage. Zero termination means unspecified (the dense default). Tolerances and L2 residuals use the integration diagnostic dtype; the checked inner convergence flag was computed before this conversion.

NumpyPoissonSampler(seed=None, *, rng=None)

Seeded stateful Poisson sampler for NumPy tau-leaping runs.

Create a NumPy generator.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed passed to :func:numpy.random.default_rng.	None
rng	RandomStream | None	Advancing generator or replayable SeedSequence; excludes seed.	None

Raises:

Type	Description
ValueError	If both seed and rng are supplied.
TypeError	If rng has an unsupported type.

Source code in src/op_engine/stochastic_solver.py

def __init__(
    self, seed: int | None = None, *, rng: RandomStream | None = None
) -> None:
    """Create a NumPy generator.

    Args:
        seed: Optional seed passed to :func:`numpy.random.default_rng`.
        rng: Advancing generator or replayable SeedSequence; excludes seed.

    Raises:
        ValueError: If both seed and rng are supplied.
        TypeError: If rng has an unsupported type.
    """  # noqa: DOC502 - documented errors originate in the shared helper
    self._rng = _generator(rng, seed=seed)

__call__(mean, step_index)

Draw Poisson firing counts in the NumPy namespace.

Parameters:

Name	Type	Description	Default
mean	Array	Non-negative Poisson means.	required
step_index	int	Accepted-step index; unused by the stateful generator.	required

Returns:

Type	Description
Array	NumPy integer firing-count array.

Raises:

Type	Description
TypeError	If mean is not a NumPy array.

Source code in src/op_engine/stochastic_solver.py

def __call__(self, mean: Array, step_index: int, /) -> Array:
    """Draw Poisson firing counts in the NumPy namespace.

    Args:
        mean: Non-negative Poisson means.
        step_index: Accepted-step index; unused by the stateful generator.

    Returns:
        NumPy integer firing-count array.

    Raises:
        TypeError: If ``mean`` is not a NumPy array.
    """
    del step_index
    if not isinstance(mean, np.ndarray):
        msg = "NumpyPoissonSampler requires a NumPy mean array"
        raise TypeError(msg)
    return cast("Array", self._rng.poisson(mean))

NumpySSASampler(seed=None, *, rng=None)

Seeded stateful direct-SSA sampler for NumPy runs.

Create a NumPy generator.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed passed to :func:numpy.random.default_rng.	None
rng	RandomStream | None	Advancing generator or replayable SeedSequence; excludes seed.	None

Raises:

Type	Description
ValueError	If both seed and rng are supplied.
TypeError	If rng has an unsupported type.

Source code in src/op_engine/stochastic_solver.py

def __init__(
    self, seed: int | None = None, *, rng: RandomStream | None = None
) -> None:
    """Create a NumPy generator.

    Args:
        seed: Optional seed passed to :func:`numpy.random.default_rng`.
        rng: Advancing generator or replayable SeedSequence; excludes seed.

    Raises:
        ValueError: If both seed and rng are supplied.
        TypeError: If rng has an unsupported type.
    """  # noqa: DOC502 - documented errors originate in the shared helper
    self._rng = _generator(rng, seed=seed)

__call__(total_rate, probabilities, draw_index)

Draw an exponential wait and flattened categorical event.

Parameters:

Name	Type	Description	Default
total_rate	Array	Positive scalar sum of all propensities.	required
probabilities	Array	Normalized batched reaction probabilities.	required
draw_index	int	Draw index; unused by the stateful generator.	required

Returns:

Type	Description
SSASample	A NumPy scalar waiting time and integer event index.

Raises:

Type	Description
TypeError	If either input is not a NumPy array.

Source code in src/op_engine/stochastic_solver.py

def __call__(
    self,
    total_rate: Array,
    probabilities: Array,
    draw_index: int,
    /,
) -> SSASample:
    """Draw an exponential wait and flattened categorical event.

    Args:
        total_rate: Positive scalar sum of all propensities.
        probabilities: Normalized batched reaction probabilities.
        draw_index: Draw index; unused by the stateful generator.

    Returns:
        A NumPy scalar waiting time and integer event index.

    Raises:
        TypeError: If either input is not a NumPy array.
    """
    del draw_index
    if not isinstance(total_rate, np.ndarray) or not isinstance(
        probabilities, np.ndarray
    ):
        msg = "NumpySSASampler requires NumPy rate and probability arrays"
        raise TypeError(msg)
    rate = float(total_rate.item())
    probability_vector = np.reshape(probabilities, (-1,))
    waiting_time = np.asarray(
        self._rng.exponential(scale=1.0 / rate),
        dtype=probabilities.dtype,
    )
    flat_event_index = np.asarray(
        self._rng.choice(probability_vector.size, p=probability_vector),
        dtype=np.int64,
    )
    return SSASample(
        cast("Array", waiting_time),
        cast("Array", flat_event_index),
    )

NumpyThinningSampler(seed=None, *, rng=None)

Seeded stateful thinning sampler for NumPy reaction networks.

Create or reuse a NumPy generator.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional legacy integer seed.	None
rng	RandomStream | None	Advancing generator or replayable SeedSequence; excludes seed.	None

Raises:

Type	Description
ValueError	If both seed and rng are supplied.
TypeError	If rng has an unsupported type.

Source code in src/op_engine/thinning_ssa.py

def __init__(
    self, seed: int | None = None, *, rng: RandomStream | None = None
) -> None:
    """Create or reuse a NumPy generator.

    Args:
        seed: Optional legacy integer seed.
        rng: Advancing generator or replayable SeedSequence; excludes seed.

    Raises:
        ValueError: If both seed and rng are supplied.
        TypeError: If rng has an unsupported type.
    """  # noqa: DOC502 - documented errors originate in the shared helper
    self._rng = _generator(rng, seed=seed)

__call__(bound_rate, draw_index)

Draw a candidate wait and uniform in the input dtype.

Returns:

Type	Description
ThinningSample	Scalar NumPy sampling arrays.

Raises:

Type	Description
TypeError	If the bound rate is not a NumPy array.

Source code in src/op_engine/thinning_ssa.py

def __call__(self, bound_rate: Array, draw_index: int, /) -> ThinningSample:
    """Draw a candidate wait and uniform in the input dtype.

    Returns:
        Scalar NumPy sampling arrays.

    Raises:
        TypeError: If the bound rate is not a NumPy array.
    """
    del draw_index
    if not isinstance(bound_rate, np.ndarray):
        msg = "NumpyThinningSampler requires a NumPy bound-rate array"
        raise TypeError(msg)
    wait = np.asarray(
        self._rng.exponential(scale=1.0 / float(bound_rate.item())),
        dtype=bound_rate.dtype,
    )
    uniform = np.asarray(self._rng.random(size=(), dtype=bound_rate.dtype))
    return ThinningSample(cast("Array", wait), cast("Array", uniform))

ObservationDiagnostics

Bases: NamedTuple

Describe evaluation coverage without inventing a numerical error estimate.

Accuracy inherits both the supplied trajectory and evaluator. No method has a qualified continuous extension through this interface yet. Spatial evaluation is owned by the callback; the compiler adapter selects declared nodes and applies declared measures without interpolation.

ObservationEvaluator

Bases: Protocol

Evaluate one native output from an instantaneous state and parameters.

__call__(t, y, **params)

Return one native output with the declared shape and state dtype.

Source code in src/op_engine/observations.py

def __call__(self, t: object, y: object, **params: object) -> object:
    """Return one native output with the declared shape and state dtype."""
    ...

ObservationMeasure

Bases: NamedTuple

Describe weights already applied by an evaluator, without applying them again.

ObservationMetadata

Bases: NamedTuple

Static scientific output metadata, separate from solver state layout.

Attributes:

Name	Type	Description
output_shape	tuple[int, ...]	Shape after evaluator reductions, excluding time.
coordinates	tuple[tuple[str, tuple[str, ...]], ...]	Ordered output axis names and coordinate labels.
units	str | None	Declared output units; no dimensional verification is implied.
state_dependencies	tuple[str, ...]	Direct state base dependencies.
parameter_shapes	tuple[tuple[str, tuple[int, ...]], ...]	Direct parameters and complete input shapes.
measure_requirements	tuple[tuple[str, ObservationMeasure], ...]	Named measures already applied by the evaluator.
time_dependent	bool	Whether the expression explicitly references time.
description	str | None	Optional scientific interpretation.

ObservationResult

Bases: NamedTuple

Static times, read-only named series mapping, and evaluation diagnostics.

ObservationSeries

Bases: NamedTuple

Native (n_output_times, *output_shape) values with their specification.

ObservationSpec

Bases: NamedTuple

A named evaluator and its immutable output description.

Custom evaluators receive each state in the trajectory's natural shape. Use :func:observation_from_compiled to flatten each state for an op_system compiled observable, preserving that compiler's state ordering.

ObservationSummary

Bases: NamedTuple

A sampled window and bounded-size named reductions, without a trajectory.

The window is exactly the first/last supplied time. There is no off-grid clipping, extrapolation or qualified method-specific dense output. The diagnostics report callback count without an independent error estimate.

ObservationSummaryValue

Bases: NamedTuple

Native summary arrays with the original scientific output specification.

Attributes:

Name	Type	Description
spec	ObservationSpec	Instantaneous output shape, coordinates, units and interpretation.
initial	Array	Output at the first supplied sample.
final	Array	Output at the last supplied sample.
change	Array	Final minus initial; appropriate for a cumulative counter.
integral	Array	Composite-trapezoid time integral over the supplied samples. Its units are output units multiplied by the caller's time unit.

OutputStateCheckConfig(nonnegative_indices=None, monotone_indices=(), negative_tolerance=0.0, monotone_tolerance=0.0) dataclass

Components and tolerances for checks at requested output times.

Attributes:

Name	Type	Description
nonnegative_indices	tuple[int, ...] | None	State components expected to be nonnegative; None selects every component. Selectors address the state axis.
monotone_indices	tuple[int, ...]	Components expected to be nondecreasing, such as cumulative counters. Empty disables increment checks.
negative_tolerance	float	A value below minus this absolute tolerance flags an output. This changes the diagnostic threshold, not the state.
monotone_tolerance	float	An increment below minus this absolute tolerance flags an output interval. It is not scaled by interval duration.

__post_init__()

Detach selectors and validate finite nonnegative tolerances.

Raises:

Type	Description
TypeError	If a selector or tolerance has the wrong type.
ValueError	If a selector or tolerance is invalid.

Source code in src/op_engine/state_diagnostics.py

def __post_init__(self) -> None:
    """Detach selectors and validate finite nonnegative tolerances.

    Raises:
        TypeError: If a selector or tolerance has the wrong type.
        ValueError: If a selector or tolerance is invalid.
    """
    if self.nonnegative_indices is not None:
        object.__setattr__(
            self,
            "nonnegative_indices",
            _indices(
                self.nonnegative_indices,
                name="nonnegative_indices",
                allow_empty=False,
            ),
        )
    object.__setattr__(
        self,
        "monotone_indices",
        _indices(self.monotone_indices, name="monotone_indices", allow_empty=True),
    )
    for name in ("negative_tolerance", "monotone_tolerance"):
        value = getattr(self, name)
        if isinstance(value, bool) or not isinstance(value, Real):
            msg = f"{name} must be a finite nonnegative number"
            raise TypeError(msg)
        normalized = float(value)
        if not np.isfinite(normalized) or normalized < 0:
            msg = f"{name} must be a finite nonnegative number"
            raise ValueError(msg)
        object.__setattr__(self, name, normalized)

OutputStateDiagnostics

Bases: NamedTuple

Per-subgroup/field-cell checks over the requested output grid.

Arrays retain every state axis except the state-component axis. Counts count outputs or intervals with at least one flagged component, not the number of components. Values and flags stay in the input namespace.

Attributes:

Name	Type	Description
n_outputs	int	Number of inspected outputs, including the initial state.
complete	bool	Whether every requested output was inspected.
min_state	Array	Minimum selected nonnegative-state value across outputs. NaNs propagate; inspect nonfinite counts too.
n_negative_outputs	Array	Outputs with any selected state below tolerance.
n_nonfinite_outputs	Array	Outputs with any nonfinite state component, including components outside nonnegative_indices.
min_increment	Array | None	Minimum increment of the selected monotone components, or None if disabled. Positive infinity when no interval exists.
n_decreasing_intervals	Array | None	Output intervals with any selected decreasing component, or None if monotonicity checks are disabled.

valid property

Return native flags for complete, finite, unflagged output records.

PoissonSampler

Bases: Protocol

Backend-specific Poisson sampler injected into tau-leaping.

step_index is a stable, zero-based accepted-step index. Functional PRNG ecosystems can derive a distinct key from it without mutable state.

__call__(mean, step_index)

Draw independent Poisson counts with the shape of mean.

Source code in src/op_engine/stochastic_solver.py

def __call__(self, mean: Array, step_index: int, /) -> Array:
    """Draw independent Poisson counts with the shape of ``mean``."""

RateBoundFunction

Bases: Protocol

Bound total propensity at a fixed state on [t, valid_until).

The callback must return t < valid_until <= limit. limit is the next forcing change or solve endpoint, independent of observation times.

__call__(t, state, limit)

Return a certified bound and its exclusive expiry time.

Source code in src/op_engine/thinning_ssa.py

def __call__(self, t: float, state: Array, limit: float, /) -> TotalRateBound:
    """Return a certified bound and its exclusive expiry time."""

ReplicateReactionNetwork(stoichiometry, channel_names, channel_reactions, n_replicates, forcing_breakpoints, _batch_axis, _states, _reactions, _params, _time_axis, _time_coordinates, _table_axes, _axis_sizes) dataclass

One local stoichiometry plus a vectorized, independent propensity bank.

Local state/channel ordering follows the existing flat adapter after removing the declared replicate axis. Full input parameter arrays retain declared axis order. Parameters are read on every call; changing their values is allowed between runs and no interpolation result is cached.

flatten_history(history)

Restore compiler flat cell order from a local bank history.

Returns:

Type	Description
FloatBank	History shaped (output time, full compiler state).

Raises:

Type	Description
ValueError	The bank history has an incompatible layout.

Source code in src/op_engine/replicate_network.py

def flatten_history(self, history: FloatBank) -> FloatBank:
    """Restore compiler flat cell order from a local bank history.

    Returns:
        History shaped (output time, full compiler state).

    Raises:
        ValueError: The bank history has an incompatible layout.
    """
    if history.ndim != 3 or history.shape[1:] != (
        self.stoichiometry.shape[0],
        self.n_replicates,
    ):
        msg = "Expected history (output time, local species, replicate)"
        raise ValueError(msg)
    pieces = []
    offset = 0
    for _name, shape, position in self._states:
        local_shape = tuple(n for i, n in enumerate(shape) if i != position)
        size = int(np.prod(local_shape))
        values = history[:, offset : offset + size].reshape((
            len(history),
            *local_shape,
            self.n_replicates,
        ))
        pieces.append(
            np.moveaxis(values, -1, position + 1).reshape((len(history), -1))
        )
        offset += size
    return np.concatenate(pieces, axis=1)

pack_state(flat_state)

Convert compiler flat state order into (local species, replicate).

Returns:

Type	Description
FloatBank	A local state bank preserving each template's declared axis order.

Raises:

Type	Description
ValueError	The flat state has an incompatible size.

Source code in src/op_engine/replicate_network.py

def pack_state(self, flat_state: FloatBank) -> FloatBank:
    """Convert compiler flat state order into (local species, replicate).

    Returns:
        A local state bank preserving each template's declared axis order.

    Raises:
        ValueError: The flat state has an incompatible size.
    """
    expected = self.stoichiometry.shape[0] * self.n_replicates
    if flat_state.shape != (expected,):
        msg = f"Expected flat compiler state {(expected,)}"
        raise ValueError(msg)
    pieces = []
    offset = 0
    for _name, shape, position in self._states:
        size = int(np.prod(shape))
        values = flat_state[offset : offset + size].reshape(shape)
        pieces.append(
            np.moveaxis(values, position, -1).reshape((-1, self.n_replicates))
        )
        offset += size
    return np.concatenate(pieces, axis=0)

propensity(times, state)

Evaluate one hazard bank with every table resolved at local clocks.

Returns:

Type	Description
FloatBank	NumPy propensities shaped (local channel, replicate).

Raises:

Type	Description
ValueError	State, clocks, parameter tables or reaction outputs have incompatible shapes.

Source code in src/op_engine/replicate_network.py

def propensity(self, times: FloatBank, state: FloatBank) -> FloatBank:
    """Evaluate one hazard bank with every table resolved at local clocks.

    Returns:
        NumPy propensities shaped (local channel, replicate).

    Raises:
        ValueError: State, clocks, parameter tables or reaction outputs have
            incompatible shapes.
    """
    expected = (self.stoichiometry.shape[0], self.n_replicates)
    if state.shape != expected or times.shape != (self.n_replicates,):
        msg = f"Expected bank state {expected} and one clock per replicate"
        raise ValueError(msg)
    states = {}
    offset = 0
    for name, shape, position in self._states:
        local_shape = tuple(n for i, n in enumerate(shape) if i != position)
        size = int(np.prod(local_shape))
        state_values = state[offset : offset + size].reshape((
            *local_shape,
            self.n_replicates,
        ))
        states[name] = np.moveaxis(state_values, -1, position)
        offset += size
    params = self._resolve_tables(times)
    values = []
    for reaction in self._reactions:
        fn = reaction.resolved_propensity_fn
        raw = np.asarray(fn(times, states, **params))
        axes = reaction.from_axes + tuple(getattr(reaction, "routed_axes", ()))
        expected_shape = tuple(self._axis_sizes[a] for a in axes)
        if raw.shape != expected_shape:
            msg = (
                f"Reaction {reaction.name!r} expected propensity {expected_shape}, "
                f"received {raw.shape}"
            )
            raise ValueError(msg)
        values.append(
            np.moveaxis(raw, axes.index(self._batch_axis), -1).reshape((
                -1,
                self.n_replicates,
            ))
        )
    return np.concatenate(values, axis=0)

SSASample

Bases: NamedTuple

One direct-SSA random draw in the active array namespace.

Attributes:

Name	Type	Description
waiting_time	Array	Scalar exponential waiting time.
flat_event_index	Array	Scalar categorical index into the flattened batched propensity array.

SSASampler

Bases: Protocol

Backend-specific direct-SSA sampler.

draw_index is a stable zero-based index for each random draw. Functional PRNG ecosystems can derive a distinct key from it without mutable state. Probabilities have the full batched reaction shape; the returned event index addresses their flattened representation.

__call__(total_rate, probabilities, draw_index)

Draw an exponential wait and one categorical event index.

Source code in src/op_engine/stochastic_solver.py

def __call__(
    self,
    total_rate: Array,
    probabilities: Array,
    draw_index: int,
    /,
) -> SSASample:
    """Draw an exponential wait and one categorical event index."""

ScipyGmresSolver(restart=20) dataclass

Uncached eager GMRES with explicit left inverse preconditioning.

The requested iteration budget is inner iterations per RHS column. iterations in the candidate/checked diagnostics is the maximum used by any column, not the total across columns. Each column starts at zero. There is no warm-start, basis, factorization or preconditioner cache. SciPy minimizes a preconditioned residual; solve_linear independently gates the original residual before stage adoption. No differentiation is supplied.

__call__(request)

Solve independent columns without forming the represented matrix.

Returns:

Type	Description
LinearSolveCandidate	Candidate and solver termination evidence for original-residual checks.

Source code in src/op_engine/scipy_operator.py

def __call__(self, request: LinearSolveRequest) -> LinearSolveCandidate:
    """Solve independent columns without forming the represented matrix.

    Returns:
        Candidate and solver termination evidence for original-residual checks.
    """
    n = request.rhs.shape[0]
    operator = LinearOperator(
        (n, n),
        dtype=request.rhs.dtype,
        matvec=lambda v: _apply_numpy(request.operator, v),
    )
    inverse = request.preconditioner
    preconditioner = (
        None
        if inverse is None
        else LinearOperator(
            (n, n),
            dtype=request.rhs.dtype,
            matvec=lambda v: _apply_numpy(inverse, v),
        )
    )
    columns = request.rhs[:, None] if request.rhs.ndim == 1 else request.rhs
    result = np.empty_like(columns)
    counts: list[int] = []
    statuses: list[int] = []
    # The legacy callback mode fixes maxiter to inner iterations on both
    # old (tol) and current (rtol) SciPy; explicit atol avoids legacy defaults.
    tolerance_name = "rtol" if "rtol" in signature(gmres).parameters else "tol"
    for column in range(columns.shape[1]):
        count = 0

        def update(_residual: object) -> None:
            nonlocal count
            count += 1

        value, status = cast("Any", gmres)(
            operator,
            columns[:, column],
            **{tolerance_name: request.options.rtol},
            atol=request.options.atol,
            restart=self.restart,
            maxiter=request.options.max_iterations,
            M=preconditioner,
            callback=update,
            callback_type="legacy",
        )
        result[:, column] = value
        counts.append(count)
        statuses.append(int(status))
    output = result[:, 0] if request.rhs.ndim == 1 else result
    return LinearSolveCandidate(
        output,
        success=all(status == 0 for status in statuses),
        iterations=max(counts, default=0),
        message=f"GMRES inner iterations per column={counts}; info={statuses}",
    )

__post_init__()

Validate the static restart length.

Raises:

Type	Description
ValueError	If restart is not a positive integer.

Source code in src/op_engine/scipy_operator.py

def __post_init__(self) -> None:
    """Validate the static restart length.

    Raises:
        ValueError: If restart is not a positive integer.
    """
    if (
        not isinstance(self.restart, Integral)
        or isinstance(self.restart, bool)
        or self.restart < 1
    ):
        msg = "GMRES restart must be a positive integer"
        raise ValueError(msg)

SensitivityInitial(state, parameter_derivatives) dataclass

Native state (N,) and its required parameter derivative (N,P).

Pass an explicit zero matrix for parameter-independent initial states. This avoids silently dropping initial-condition derivatives.

SensitivityOptions(solver=(lambda: RunConfig(method='rk4'))(), store_history=False, max_augmented_states=16384, max_output_bytes=64 * 1024 * 1024) dataclass

Bound fixed explicit integration and retained numerical output.

Only the final output is retained by default. History adds requested output times, not internal stages. Byte limits cover returned numeric arrays, not temporary work, AD graphs, compiler memory or the whole process.

__post_init__()

Reject unsupported numerical paths instead of silently downshifting.

Raises:

Type	Description
ValueError	Solver mode or storage bounds are unsupported.
TypeError	History is not boolean.

Source code in src/op_engine/sensitivities.py

def __post_init__(self) -> None:
    """Reject unsupported numerical paths instead of silently downshifting.

    Raises:
        ValueError: Solver mode or storage bounds are unsupported.
        TypeError: History is not boolean.
    """
    cfg = self.solver
    if not isinstance(cfg, RunConfig):
        msg = "Sensitivity solver options require a RunConfig"
        raise TypeError(msg)
    if (
        cfg.adaptive
        or not cfg.strict
        or cfg.method not in {"euler", "heun", "rk4", "dopri5"}
        or cfg.forcing_breakpoints
        or any(
            value is not None
            for value in (
                cfg.state_checks,
                cfg.nonlinear,
                cfg.jacobian,
                cfg.operators.default,
                cfg.operators.tr,
                cfg.operators.bdf2,
            )
        )
    ):
        msg = "Sensitivities require a smooth fixed explicit ODE without operators"
        raise ValueError(msg)
    if type(self.store_history) is not bool:
        msg = "Sensitivity store_history must be boolean"
        raise TypeError(msg)
    if any(
        type(bound) is not int or bound < 1
        for bound in (self.max_augmented_states, self.max_output_bytes)
    ):
        msg = "Sensitivity storage bounds must be positive integers"
        raise ValueError(msg)

SensitivitySolution(source_model, state_names, parameter_names, observable_names, times, mesh, states, parameter_derivatives, observable_derivatives, augmented_dimension, options) dataclass

Native output states and parameter derivatives with static provenance.

Array shapes are (T,N), (T,N,P) and (T,M,P). Measurement sensitivities are dg/dy @ S + dg/dtheta in the same parameter order. They do not include observation noise, prior curvature or nuisance elimination. times contains only retained outputs; mesh is the requested solver mesh.

SensitivitySystem(rhs, derivatives, initial_parameters=()) dataclass

Bind one RHS to derivatives and explicitly named initial-only parameters.

rhs accepts every derivative-artifact parameter, including measurement- only parameters. Initial-only names are appended to that ordered layout and excluded from RHS/derivative calls; their effects enter through the required initial derivative matrix. The caller is responsible for binding the same RHS as the derivative producer and for declaring a smooth deterministic ODE.

parameter_names property

Return derivative parameters followed by initial-only parameters.

__post_init__()

Reject ambiguous coordinate layouts before numerical work.

Raises:

Type	Description
ValueError	Names are empty, duplicate or overlap.
TypeError	The RHS is not callable.

Source code in src/op_engine/sensitivities.py

def __post_init__(self) -> None:
    """Reject ambiguous coordinate layouts before numerical work.

    Raises:
        ValueError: Names are empty, duplicate or overlap.
        TypeError: The RHS is not callable.
    """
    if not callable(self.rhs):
        msg = "Sensitivity RHS must be callable"
        raise TypeError(msg)
    if not isinstance(self.initial_parameters, tuple):
        msg = "Initial-only parameter names must be a tuple"
        raise TypeError(msg)
    groups = (
        self.derivatives.state_names,
        self.parameter_names,
        self.derivatives.observable_names,
    )
    if (
        not self.derivatives.state_names
        or not self.parameter_names
        or any(
            not isinstance(names, tuple)
            or any(not isinstance(name, str) or not name.strip() for name in names)
            or len(set(names)) != len(names)
            for names in groups
        )
    ):
        msg = "Sensitivity layouts require nonempty unique state/parameter names"
        raise ValueError(msg)

SparseAdapter(ecosystem_id, issparse, factorize=None, solve=None, cache_key=None, linear_solver=None, solve_options=LinearSolveOptions(), preconditioner=None, context=None, accepts_right=None) dataclass

Function bundle in the existing sparse/operator registry.

Sparse acceleration is optional: dense operators always fall back to the input array's Array-API linalg.solve implementation. Adapters are selected structurally with their own issparse predicate. Existing factorize/solve/cache_key bundles retain their historical identity cache. Alternatively, linear_solver opts into a checked, uncached eager NumPy solve. The two routes are mutually exclusive. accepts_right optionally recognizes mixed right operators; otherwise issparse must accept both. New callbacks own preparation, context invalidation and preconditioner use.

__post_init__()

Reject incomplete or conflicting dispatch routes.

Raises:

Type	Description
ValueError	If the identifier or route configuration is invalid.
TypeError	If a predicate or callback is not callable.

Source code in src/op_engine/matrix_ops.py

def __post_init__(self) -> None:
    """Reject incomplete or conflicting dispatch routes.

    Raises:
        ValueError: If the identifier or route configuration is invalid.
        TypeError: If a predicate or callback is not callable.
    """
    if (
        not isinstance(self.ecosystem_id, str)
        or not self.ecosystem_id
        or self.ecosystem_id.strip() != self.ecosystem_id
    ):
        msg = "adapter ecosystem_id must be a nonempty stripped string"
        raise ValueError(msg)
    if not callable(self.issparse):
        msg = "adapter issparse predicate must be callable"
        raise TypeError(msg)
    legacy = (self.factorize, self.solve, self.cache_key)
    if (self.linear_solver is None and any(v is None for v in legacy)) or (
        self.linear_solver is not None and any(v is not None for v in legacy)
    ):
        msg = "supply either factorize/solve/cache_key or linear_solver"
        raise ValueError(msg)
    for callback in (
        self.issparse,
        *legacy,
        self.linear_solver,
        self.accepts_right,
    ):
        if callback is not None and not callable(callback):
            msg = "adapter callbacks and predicates must be callable"
            raise TypeError(msg)
    if self.linear_solver is None and (
        self.preconditioner is not None
        or self.context is not None
        or self.solve_options != LinearSolveOptions()
    ):
        msg = "solve options, preconditioner and context require linear_solver"
        raise ValueError(msg)

StabilityJacobianSource

Bases: StrEnum

Provenance of the dense full-state Jacobian.

StabilityTimescaleKind

Bases: StrEnum

Interpretation of an e-folding timescale.

SteadyStateConfig(dt0=1.0, dt_max=1000000000000.0, min_growth=2.0, max_growth=10.0, max_iterations=100, residual_tol=1e-09, step_tol=1e-09, atol=1e-08, early_exit=True) dataclass

Controls for :func:steady_state.

Attributes:

Name	Type	Description
dt0	float	Initial pseudo-time step, in the RHS's time units.
dt_max	float	Largest pseudo-time step.
min_growth	float	Smallest factor by which dt grows when the unscaled residual did not increase. Without it, a slowly decaying mode barely shrinks the residual and dt stalls.
max_growth	float	Largest factor by which dt grows in one iteration. While the unscaled residual decreases, dt follows its ratio between the growth bounds (switched evolution relaxation). Residual growth is allowed through a transient: dt is held on accepted steps. Candidates whose scaled or unscaled residual exceeds the current one by this factor are rejected, as are non-finite or infeasible candidates. Rejection holds the state and halves dt for the next attempt.
max_iterations	int	Iteration budget; the loop is unrolled this many times when early_exit is false.
residual_tol	float	Bound on the RMS of rhs / (atol + |y|) over the unknowns, in inverse time units.
step_tol	float	Bound on the RMS of the Newton correction (I / dt_max - J)^{-1} rhs / (atol + |y|), a relative state-error estimate.
atol	float	Absolute floor of the per-state scale atol + |y|.
early_exit	bool	With the default Python loop, stop as soon as the solve converges. This converts the convergence flag to a host boolean, so under jax.jit or jax.vmap either set it to False or, better, pass loop=jax.lax.fori_loop.

__post_init__()

Validate the controls.

Raises:

Type	Description
ValueError	If a control is outside its valid range.

Source code in src/op_engine/steady_state.py

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

    Raises:
        ValueError: If a control is outside its valid range.
    """
    for name in ("dt0", "dt_max", "residual_tol", "step_tol", "atol"):
        value = getattr(self, name)
        if (
            not isinstance(value, Real)
            or isinstance(value, bool)
            or not math.isfinite(value)
            or value <= 0.0
        ):
            msg = f"{name} must be finite and positive"
            raise ValueError(msg)
    if self.dt_max < self.dt0:
        msg = "dt_max must be at least dt0"
        raise ValueError(msg)
    for name in ("min_growth", "max_growth"):
        value = getattr(self, name)
        if (
            not isinstance(value, Real)
            or isinstance(value, bool)
            or not math.isfinite(value)
            or value < 1.0
        ):
            msg = f"{name} must be finite and at least one"
            raise ValueError(msg)
    if self.max_growth < self.min_growth:
        msg = "max_growth must be at least min_growth"
        raise ValueError(msg)
    if (
        not isinstance(self.max_iterations, Integral)
        or isinstance(self.max_iterations, bool)
        or self.max_iterations < 1
    ):
        msg = "max_iterations must be a positive integer"
        raise ValueError(msg)

SteadyStateConvergenceError(result)

Bases: RuntimeError

Raised at an eager boundary when a steady-state solve did not converge.

Store the failed result.

Source code in src/op_engine/steady_state.py

def __init__(self, result: SteadyStateResult) -> None:
    """Store the failed result."""
    self.result = result
    super().__init__(
        "Steady-state solve did not converge after "
        f"{int(result.iterations.item())} iterations; scaled residual "
        f"{float(result.residual_norm.item())!r}, Newton correction "
        f"{float(result.step_norm.item())!r}"
    )

SteadyStateResult(state, residual, converged, iterations, residual_norm, step_norm, dt, invariant_drift) dataclass

Candidate steady state and array-valued diagnostics.

Every diagnostic is a zero-dimensional array in the state's namespace, so a traced solve returns them without a host conversion; check them afterwards, or call :func:require_steady_state at an eager boundary.

Attributes:

Name	Type	Description
state	Array	The candidate steady state, including fixed states.
residual	Array	rhs(t, state), zero at fixed states.
converged	Array	Whether both the residual and the Newton correction met their tolerances.
iterations	Array	PTC attempts performed before convergence, including rejected steps.
residual_norm	Array	Scaled residual RMS at state.
step_norm	Array	Scaled RMS of the last Newton correction.
dt	Array	Final pseudo-time step.
invariant_drift	Array	Largest |W @ state - W @ y0| relative to 1 + |W @ y0|; zero without invariants.

SymbolicDerivativeArtifact

Bases: Protocol

Structural surface consumed from op_system.SymbolicDerivatives.

evaluate(t, y, **parameters)

Evaluate derivative matrices in the artifact's declared order.

Source code in src/op_engine/equilibrium_stability.py

def evaluate(self, t: object, y: object, **parameters: object) -> _DerivativeValues:
    """Evaluate derivative matrices in the artifact's declared order."""
    ...

TauLeapingConfig(max_step=None, max_steps=1000000, forcing_breakpoints=()) dataclass

Configuration for fixed-step explicit tau-leaping.

Attributes:

Name	Type	Description
max_step	float | None	Maximum internal tau. None takes one leap per output interval, split at forcing boundaries.
max_steps	int	Maximum number of internal leaps per output interval.
forcing_breakpoints	tuple[float, ...]	Strictly increasing finite forcing-change times. Leaps end at these boundaries and reevaluate right-continuous rates.

__post_init__()

Validate fixed tau controls.

Raises:

Type	Description
ValueError	If a control is outside its valid range.

Source code in src/op_engine/stochastic_solver.py

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

    Raises:
        ValueError: If a control is outside its valid range.
    """
    if self.max_step is not None and (
        not np.isfinite(self.max_step) or self.max_step <= 0.0
    ):
        msg = "max_step must be finite and positive when provided"
        raise ValueError(msg)
    if (
        not isinstance(self.max_steps, Integral)
        or isinstance(self.max_steps, bool)
        or self.max_steps < 1
    ):
        msg = "max_steps must be a positive integer"
        raise ValueError(msg)
    schedule = _ForcingSchedule(self.forcing_breakpoints)
    object.__setattr__(self, "forcing_breakpoints", schedule.breakpoints)

TauLeapingSolver(core, stoichiometry, *, reaction_axis='state')

Fixed-step explicit Poisson tau-leaping on a ModelCore state.

The stoichiometric matrix has shape (n_species, n_reactions). The configured reaction axis identifies the species dimension in the state; every other dimension is treated as an independent batch. A propensity function returns the same batched shape with n_reactions replacing n_species on that axis.

Initialize a reaction network.

Parameters:

Name	Type	Description	Default
core	ModelCore	Model state and output-time container.	required
stoichiometry	Array	Integer-valued matrix shaped (n_species, n_reactions).	required
reaction_axis	str | int	State axis changed by reaction firings.	'state'

Source code in src/op_engine/stochastic_solver.py

def __init__(
    self,
    core: ModelCore,
    stoichiometry: Array,
    *,
    reaction_axis: str | int = "state",
) -> None:
    """Initialize a reaction network.

    Args:
        core: Model state and output-time container.
        stoichiometry: Integer-valued matrix shaped
            ``(n_species, n_reactions)``.
        reaction_axis: State axis changed by reaction firings.

    """
    self.core = core
    self._network = _ReactionNetwork(core, stoichiometry, reaction_axis)

n_reactions property

Return the number of reaction channels.

run(propensity_func, poisson_sampler=None, *, config=None, rng=None)

Advance the reaction network through the core output grid.

Parameters:

Name	Type	Description	Default
propensity_func	PropensityFunction	Reaction-channel propensity function.	required
poisson_sampler	PoissonSampler | None	Backend-specific Poisson sampler.	None
config	TauLeapingConfig | None	Optional fixed-tau controls and forcing schedule.	None
rng	RandomStream | None	Explicit NumPy stream; supply instead of poisson_sampler.	None

Raises:

Type	Description
ValueError	If a sampler/stream is absent or both are supplied.
TypeError	If rng is unsupported or used with non-NumPy state.
RuntimeError	If an interval exceeds its step limit or a leap produces a negative population.

Source code in src/op_engine/stochastic_solver.py

def run(
    self,
    propensity_func: PropensityFunction,
    poisson_sampler: PoissonSampler | None = None,
    *,
    config: TauLeapingConfig | None = None,
    rng: RandomStream | None = None,
) -> None:
    """Advance the reaction network through the core output grid.

    Args:
        propensity_func: Reaction-channel propensity function.
        poisson_sampler: Backend-specific Poisson sampler.
        config: Optional fixed-tau controls and forcing schedule.
        rng: Explicit NumPy stream; supply instead of poisson_sampler.

    Raises:
        ValueError: If a sampler/stream is absent or both are supplied.
        TypeError: If rng is unsupported or used with non-NumPy state.
        RuntimeError: If an interval exceeds its step limit or a leap
            produces a negative population.
    """  # noqa: DOC502 - random-source validation is shared
    poisson_sampler = _select_sampler(
        poisson_sampler,
        rng,
        state=self.core.get_current_state(),
        make_sampler=lambda generator: NumpyPoissonSampler(rng=generator),
    )
    cfg = config or TauLeapingConfig()
    forcing = _ForcingSchedule(cfg.forcing_breakpoints)
    state = self.core.get_current_state()
    self._network.validate_finite_nonnegative(
        state,
        message="Initial state is invalid",
    )

    time_grid = np.asarray(self.core.time_grid, dtype=float)
    step_index = 0
    for output_index in range(int(self.core.n_timesteps) - 1):
        t = float(time_grid[output_index])
        target = float(time_grid[output_index + 1])
        state = self.core.get_current_state()
        interval_steps = 0
        while t < target:
            if interval_steps >= cfg.max_steps:
                raise RuntimeError(_MAX_STEPS)
            limit = min(target, forcing.next_after(t))
            remaining = limit - t
            dt = remaining if cfg.max_step is None else min(cfg.max_step, remaining)
            if t + dt <= t:
                msg = "Tau-leaping step size underflowed"
                raise RuntimeError(msg)
            state = self._step(
                propensity_func,
                poisson_sampler,
                t=t,
                dt=dt,
                state=state,
                step_index=step_index,
            )
            t = limit if dt == remaining else t + dt
            step_index += 1
            interval_steps += 1
        self.core.advance_timestep(state)

ThinningSSAConfig(max_candidates=1000000, forcing_breakpoints=()) dataclass

Controls for exact bounded-thinning SSA.

Attributes:

Name	Type	Description
max_candidates	int	Maximum candidate draws across the entire run, including rejected and expired candidates.
forcing_breakpoints	tuple[float, ...]	Strictly increasing finite forcing changes. Forcing and bound expiry take precedence over tied candidates.

__post_init__()

Validate the candidate guard and snapshot the forcing schedule.

Raises:

Type	Description
ValueError	If either control is invalid.

Source code in src/op_engine/thinning_ssa.py

def __post_init__(self) -> None:
    """Validate the candidate guard and snapshot the forcing schedule.

    Raises:
        ValueError: If either control is invalid.
    """
    if (
        not isinstance(self.max_candidates, Integral)
        or isinstance(self.max_candidates, bool)
        or self.max_candidates < 1
    ):
        msg = "max_candidates must be a positive integer"
        raise ValueError(msg)
    schedule = _ForcingSchedule(self.forcing_breakpoints)
    object.__setattr__(self, "forcing_breakpoints", schedule.breakpoints)

ThinningSSASolver(core, stoichiometry, *, reaction_axis='state')

Exact thinning SSA with user-certified total-rate bounds.

Reaction and batch channels share one candidate clock. Rates may vary smoothly while the state is unchanged; a valid bound must cover their total over the entire declared interval. Observation times only observe the trajectory and never refresh bounds or discard candidates.

Initialize the shared validated reaction network.

Parameters:

Name	Type	Description	Default
core	ModelCore	Model state and observation-time container.	required
stoichiometry	Array	Integer-valued species-by-reaction matrix.	required
reaction_axis	str | int	State axis changed by reaction firings.	'state'

Source code in src/op_engine/thinning_ssa.py

def __init__(
    self,
    core: ModelCore,
    stoichiometry: Array,
    *,
    reaction_axis: str | int = "state",
) -> None:
    """Initialize the shared validated reaction network.

    Args:
        core: Model state and observation-time container.
        stoichiometry: Integer-valued species-by-reaction matrix.
        reaction_axis: State axis changed by reaction firings.
    """
    self.core = core
    self._network = _ReactionNetwork(core, stoichiometry, reaction_axis)

n_reactions property

Return the number of reaction channels.

run(propensity_func, thinning_sampler=None, *, rate_bound, config=None, rng=None)

Advance an exact trajectory with bounded candidate thinning.

Parameters:

Name	Type	Description	Default
propensity_func	PropensityFunction	Time-dependent batched reaction propensities.	required
thinning_sampler	ThinningSampler | None	Namespace-preserving candidate sampler.	None
rate_bound	float | RateBoundFunction	Constant bound or callback returning a certified bound on [t, valid_until) at the supplied unchanged state.	required
config	ThinningSSAConfig | None	Candidate guard and forcing schedule.	None
rng	RandomStream | None	Explicit NumPy stream; supply instead of thinning_sampler.	None

Raises:

Type	Description
ValueError	If a sampler/stream is absent or both are supplied.
TypeError	If rng is unsupported or used with non-NumPy state.
RuntimeError	If the candidate limit is exceeded, time cannot advance, or a reaction produces invalid populations.

Source code in src/op_engine/thinning_ssa.py

def run(
    self,
    propensity_func: PropensityFunction,
    thinning_sampler: ThinningSampler | None = None,
    *,
    rate_bound: float | RateBoundFunction,
    config: ThinningSSAConfig | None = None,
    rng: RandomStream | None = None,
) -> None:
    """Advance an exact trajectory with bounded candidate thinning.

    Args:
        propensity_func: Time-dependent batched reaction propensities.
        thinning_sampler: Namespace-preserving candidate sampler.
        rate_bound: Constant bound or callback returning a certified bound
            on ``[t, valid_until)`` at the supplied unchanged state.
        config: Candidate guard and forcing schedule.
        rng: Explicit NumPy stream; supply instead of thinning_sampler.

    Raises:
        ValueError: If a sampler/stream is absent or both are supplied.
        TypeError: If rng is unsupported or used with non-NumPy state.
        RuntimeError: If the candidate limit is exceeded, time cannot
            advance, or a reaction produces invalid populations.
    """  # noqa: DOC502 - random-source validation is shared
    thinning_sampler = _select_sampler(
        thinning_sampler,
        rng,
        state=self.core.get_current_state(),
        make_sampler=lambda generator: NumpyThinningSampler(rng=generator),
    )
    cfg = config or ThinningSSAConfig()
    forcing = _ForcingSchedule(cfg.forcing_breakpoints)
    state = self._initial_state(rate_bound)
    xp = _namespace_of(state)
    times = np.asarray(self.core.time_grid, dtype=float)
    t, end = float(times[0]), float(times[-1])
    bound: TotalRateBound | None = None
    pending: tuple[float, float] | None = None
    draw_index = 0
    for target in times[1:]:
        while t < target:
            if bound is None:
                bound = _resolve_bound(
                    rate_bound, t, state, min(forcing.next_after(t), end)
                )
                self._rates(propensity_func, t, state, bound)
            if pending is None:
                if bound.rate == 0:
                    pending = (bound.valid_until, 0.0)
                else:
                    if draw_index >= cfg.max_candidates:
                        msg = "Exceeded max_candidates during thinning SSA run"
                        raise RuntimeError(msg)
                    pending = _candidate(
                        t,
                        cast("Array", xp.asarray(bound.rate, dtype=state.dtype)),
                        thinning_sampler,
                        draw_index,
                    )
                    draw_index += 1
            if bound.valid_until <= target and bound.valid_until <= pending[0]:
                t = bound.valid_until
                bound, pending = None, None
                continue
            if pending[0] > target:
                break
            t, uniform = pending
            pending = None
            state, accepted = self._accept(
                propensity_func, t, uniform, state, bound
            )
            if accepted:
                bound = None
        self.core.advance_timestep(state)

ThinningSample

Bases: NamedTuple

Independent exponential wait and uniform in the active array namespace.

Both values are real floating scalar arrays. waiting_time is finite and positive; uniform is finite and belongs to [0, 1).

ThinningSampler

Bases: Protocol

Sample candidates with globally increasing zero-based draw indices.

__call__(bound_rate, draw_index)

Draw a wait at bound_rate and an independent uniform.

Source code in src/op_engine/thinning_ssa.py

def __call__(self, bound_rate: Array, draw_index: int, /) -> ThinningSample:
    """Draw a wait at ``bound_rate`` and an independent uniform."""

TotalRateBound(rate, valid_until) dataclass

Upper bound on all reaction/batch rates until an exclusive endpoint.

Attributes:

Name	Type	Description
rate	float	Finite non-negative total-rate bound at the unchanged state.
valid_until	float	Finite exclusive endpoint, checked against the current time and forcing/solve limit when the callback returns it.

__post_init__()

Normalize and validate the scalar controls.

Raises:

Type	Description
ValueError	If either control is invalid.

Source code in src/op_engine/thinning_ssa.py

def __post_init__(self) -> None:
    """Normalize and validate the scalar controls.

    Raises:
        ValueError: If either control is invalid.
    """
    object.__setattr__(
        self, "rate", _nonnegative_real(self.rate, name="bound rate")
    )
    try:
        endpoint = float(self.valid_until)
    except (TypeError, ValueError, OverflowError):
        endpoint = math.nan
    if (
        not isinstance(self.valid_until, Real)
        or isinstance(self.valid_until, bool)
        or not math.isfinite(endpoint)
    ):
        msg = "valid_until must be a finite real scalar"
        raise ValueError(msg)
    object.__setattr__(self, "valid_until", endpoint)

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.

Raises:

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

Source code in src/op_engine/_array.py

def array_namespace(value: object) -> Any:  # noqa: ANN401
    """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.

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

build_advection_matrix(n, dx, velocity, *, bc='absorbing', reference=None)

Build a dense first-order upwind finite-volume operator.

The returned matrix A acts on a column state as dy = A @ y. Positive velocity transports values toward increasing coordinate indices; negative velocity transports toward decreasing indices. The velocity may be an Array-API scalar, so its value stays dynamic under transformations such as JAX jit and grad.

Boundary modes are:

• absorbing: zero inflow at the upstream boundary and free outflow at the downstream boundary;
• reflecting: zero flux through the downstream boundary;
• periodic: downstream outflow wraps to the upstream boundary.

The namespace comes from reference when provided, then from velocity; plain numeric velocities use NumPy. Grid geometry and the boundary mode are static structural inputs.

Parameters:

Name	Type	Description	Default
n	int	Number of uniformly spaced finite-volume cells.	required
dx	float	Positive cell width.	required
velocity	float | Array	Signed scalar transport velocity.	required
bc	str	Boundary mode: "absorbing", "reflecting", or "periodic".	'absorbing'
reference	Array | None	Optional array whose namespace and floating dtype control the result.	None

Returns:

Type	Description
Array	Dense (n, n) operator in the selected Array-API namespace.

Raises:

Type	Description
ValueError	If the grid, velocity shape/value, or boundary mode is invalid.

Source code in src/op_engine/matrix_ops.py

def build_advection_matrix(
    n: int,
    dx: float,
    velocity: float | Array,
    *,
    bc: str = "absorbing",
    reference: Array | None = None,
) -> Array:
    """Build a dense first-order upwind finite-volume operator.

    The returned matrix ``A`` acts on a column state as ``dy = A @ y``.
    Positive velocity transports values toward increasing coordinate indices;
    negative velocity transports toward decreasing indices. The velocity may be
    an Array-API scalar, so its value stays dynamic under transformations such
    as JAX ``jit`` and ``grad``.

    Boundary modes are:

    - ``absorbing``: zero inflow at the upstream boundary and free outflow at
      the downstream boundary;
    - ``reflecting``: zero flux through the downstream boundary;
    - ``periodic``: downstream outflow wraps to the upstream boundary.

    The namespace comes from ``reference`` when provided, then from ``velocity``;
    plain numeric velocities use NumPy. Grid geometry and the boundary mode are
    static structural inputs.

    Args:
        n: Number of uniformly spaced finite-volume cells.
        dx: Positive cell width.
        velocity: Signed scalar transport velocity.
        bc: Boundary mode: ``"absorbing"``, ``"reflecting"``, or ``"periodic"``.
        reference: Optional array whose namespace and floating dtype control the
            result.

    Returns:
        Dense ``(n, n)`` operator in the selected Array-API namespace.

    Raises:
        ValueError: If the grid, velocity shape/value, or boundary mode is
            invalid.
    """
    if n < 2:
        msg = f"advection grid size must be at least 2; got {n}."
        raise ValueError(msg)
    if not np.isfinite(dx) or dx <= 0.0:
        msg = f"advection dx must be finite and positive; got {dx}."
        raise ValueError(msg)

    bc_normalized = str(bc).strip().lower()
    if bc_normalized not in {"absorbing", "reflecting", "periodic"}:
        msg = (
            f"Unknown advection bc: {bc!r}; expected absorbing, reflecting, "
            "or periodic."
        )
        raise ValueError(msg)

    namespace_source: object = reference if reference is not None else velocity
    try:
        _namespace_of(namespace_source)
    except TypeError:
        namespace_source = np.asarray(namespace_source)
    xp = _namespace_of(namespace_source)
    source_dtype = cast("Any", namespace_source).dtype
    dtype = (
        source_dtype
        if xp.isdtype(source_dtype, "real floating")
        else xp.result_type(xp.asarray(0.0).dtype, source_dtype)
    )
    velocity_array = cast("Array", xp.asarray(velocity, dtype=dtype))
    if velocity_array.shape != ():
        msg = f"advection velocity must be scalar; got shape {velocity_array.shape}."
        raise ValueError(msg)
    if isinstance(velocity_array, np.ndarray) and not np.isfinite(velocity_array).all():
        msg = "advection velocity must be finite."
        raise ValueError(msg)

    positive = -np.eye(n, dtype=np.float64)
    negative = -np.eye(n, dtype=np.float64)
    indices = np.arange(n - 1)
    positive[indices + 1, indices] = 1.0
    negative[indices, indices + 1] = 1.0

    if bc_normalized == "reflecting":
        positive[-1, -1] = 0.0
        negative[0, 0] = 0.0
    elif bc_normalized == "periodic":
        positive[0, -1] = 1.0
        negative[-1, 0] = 1.0

    positive_array = xp.asarray(positive, dtype=dtype)
    negative_array = xp.asarray(negative, dtype=dtype)
    zero = xp.asarray(0.0, dtype=dtype)
    stencil = xp.where(
        xp.greater_equal(velocity_array, zero),
        positive_array,
        negative_array,
    )
    speed = xp.divide(xp.abs(velocity_array), dx)
    return cast("Array", xp.multiply(stencil, speed))

build_crank_nicolson_operator(geom, cfg, dt)

Build Crank-Nicolson operators with dense/sparse autodispatch.

Parameters:

Name	Type	Description	Default
geom	GridGeometry	Grid geometry.	required
cfg	DiffusionConfig	Diffusion configuration.	required
dt	float	Time step.	required

Returns:

Type	Description
tuple[Operator, Operator]	Tuple of (L, R) operators for Crank-Nicolson scheme.

Source code in src/op_engine/matrix_ops.py

def build_crank_nicolson_operator(
    geom: GridGeometry,
    cfg: DiffusionConfig,
    dt: float,
) -> tuple[Operator, Operator]:
    """
    Build Crank-Nicolson operators with dense/sparse autodispatch.

    Args:
        geom: Grid geometry.
        cfg: Diffusion configuration.
        dt: Time step.

    Returns:
        Tuple of (L, R) operators for Crank-Nicolson scheme.
    """
    if geom.n < _DISPATCH_THRESHOLD:
        return _build_crank_nicolson_dense(geom, cfg, dt)
    return _build_crank_nicolson_sparse(geom, cfg, dt)

build_diffusion_matrix(n, dx, coefficient, *, grid=None, bc='neumann', reference=None)

Build a dense second-order 1D finite-volume diffusion operator.

The returned matrix A acts on a column state as dy = A @ y. Supply either a uniform cell width with dx or strictly increasing cell-center coordinates with grid. On a non-uniform grid, interior flux differences are divided by the local Voronoi-cell width. Boundary faces are inferred one half-spacing beyond the first and last centers.

The coefficient may be an Array-API scalar, so it stays dynamic under transformations such as JAX jit and grad. Static geometry is assembled with NumPy and transferred once to the namespace selected by reference or coefficient.

Boundary modes are:

• neumann or reflecting: zero flux at both boundaries;
• absorbing: zero-valued exterior ghost cells;
• periodic: the two grid ends are adjacent. A grid passed explicitly must be uniform because its coordinates do not determine wrap spacing.

Parameters:

Name	Type	Description	Default
n	int	Number of finite-volume cells.	required
dx	float | None	Positive uniform cell width, or None when grid is supplied.	required
coefficient	float | Array	Non-negative scalar diffusion coefficient.	required
grid	ArrayLike | None	Optional strictly increasing cell-center coordinates of shape (n,). Exactly one of dx and grid must be supplied.	None
bc	str	Boundary mode: neumann, reflecting, absorbing, or periodic.	'neumann'
reference	Array | None	Optional array whose namespace and floating dtype control the result.	None

Returns:

Type	Description
Array	Dense (n, n) operator in the selected Array-API namespace.

Raises:

Type	Description
ValueError	If the geometry, coefficient, or boundary mode is invalid.

Source code in src/op_engine/matrix_ops.py

def build_diffusion_matrix(  # noqa: PLR0913
    n: int,
    dx: float | None,
    coefficient: float | Array,
    *,
    grid: ArrayLike | None = None,
    bc: str = "neumann",
    reference: Array | None = None,
) -> Array:
    """Build a dense second-order 1D finite-volume diffusion operator.

    The returned matrix ``A`` acts on a column state as ``dy = A @ y``.
    Supply either a uniform cell width with ``dx`` or strictly increasing
    cell-center coordinates with ``grid``. On a non-uniform grid, interior
    flux differences are divided by the local Voronoi-cell width. Boundary
    faces are inferred one half-spacing beyond the first and last centers.

    The coefficient may be an Array-API scalar, so it stays dynamic under
    transformations such as JAX ``jit`` and ``grad``. Static geometry is
    assembled with NumPy and transferred once to the namespace selected by
    ``reference`` or ``coefficient``.

    Boundary modes are:

    - ``neumann`` or ``reflecting``: zero flux at both boundaries;
    - ``absorbing``: zero-valued exterior ghost cells;
    - ``periodic``: the two grid ends are adjacent. A grid passed explicitly
      must be uniform because its coordinates do not determine wrap spacing.

    Args:
        n: Number of finite-volume cells.
        dx: Positive uniform cell width, or ``None`` when ``grid`` is supplied.
        coefficient: Non-negative scalar diffusion coefficient.
        grid: Optional strictly increasing cell-center coordinates of shape
            ``(n,)``. Exactly one of ``dx`` and ``grid`` must be supplied.
        bc: Boundary mode: neumann, reflecting, absorbing, or periodic.
        reference: Optional array whose namespace and floating dtype control the
            result.

    Returns:
        Dense ``(n, n)`` operator in the selected Array-API namespace.

    Raises:
        ValueError: If the geometry, coefficient, or boundary mode is invalid.
    """
    lower, main, upper, periodic_corner = _diffusion_diagonals(
        n,
        dx,
        grid=grid,
        bc=bc,
    )

    namespace_source: object = reference if reference is not None else coefficient
    try:
        _namespace_of(namespace_source)
    except TypeError:
        namespace_source = np.asarray(namespace_source)
    xp = _namespace_of(namespace_source)
    source_dtype = cast("Any", namespace_source).dtype
    dtype = (
        source_dtype
        if xp.isdtype(source_dtype, "real floating")
        else xp.result_type(xp.asarray(0.0).dtype, source_dtype)
    )
    coefficient_array = cast("Array", xp.asarray(coefficient, dtype=dtype))
    if coefficient_array.shape != ():
        msg = (
            "diffusion coefficient must be scalar; "
            f"got shape {coefficient_array.shape}."
        )
        raise ValueError(msg)
    if isinstance(coefficient_array, np.ndarray):
        if not np.isfinite(coefficient_array).all():
            msg = "diffusion coefficient must be finite."
            raise ValueError(msg)
        if bool(coefficient_array < 0.0):
            msg = "diffusion coefficient must be non-negative."
            raise ValueError(msg)

    stencil = np.diag(main) + np.diag(lower, k=-1) + np.diag(upper, k=1)
    if periodic_corner:
        stencil[0, -1] += periodic_corner
        stencil[-1, 0] += periodic_corner

    stencil_array = xp.asarray(stencil, dtype=dtype)
    return cast("Array", xp.multiply(stencil_array, coefficient_array))

build_laplacian_tridiag(n, dx, coeff, dtype=np.float64, bc='neumann', *, grid=None)

Build a sparse 1D finite-volume Laplacian.

The resulting operator corresponds to coeff * Δ_h. Supply either a uniform cell width with dx or strictly increasing cell-center coordinates with grid. No time-step scaling is applied here.

Parameters:

Name	Type	Description	Default
n	int	Number of finite-volume cells.	required
dx	float | None	Positive uniform cell width, or None when grid is supplied.	required
coeff	float	Physical diffusion coefficient D (units length^2 / time).	required
dtype	DTypeLike	Floating dtype (e.g. np.float64).	float64
bc	str	Boundary condition; either "neumann" or "absorbing".	'neumann'
grid	ArrayLike | None	Optional strictly increasing cell-center coordinates of shape (n,). Exactly one of dx and grid must be supplied.	None

Returns:

Type	Description
csr_matrix	Sparse CSR matrix representing the Laplacian operator.

Raises:

Type	Description
ValueError	If the geometry or boundary condition is invalid.

Source code in src/op_engine/matrix_ops.py

def build_laplacian_tridiag(  # noqa: PLR0913
    n: int,
    dx: float | None,
    coeff: float,
    dtype: DTypeLike = np.float64,
    bc: str = "neumann",
    *,
    grid: ArrayLike | None = None,
) -> csr_matrix:
    """Build a sparse 1D finite-volume Laplacian.

    The resulting operator corresponds to ``coeff * Δ_h``. Supply either a
    uniform cell width with ``dx`` or strictly increasing cell-center
    coordinates with ``grid``. No time-step scaling is applied here.

    Args:
        n: Number of finite-volume cells.
        dx: Positive uniform cell width, or ``None`` when ``grid`` is supplied.
        coeff: Physical diffusion coefficient D (units length^2 / time).
        dtype: Floating dtype (e.g. np.float64).
        bc: Boundary condition; either ``"neumann"`` or ``"absorbing"``.
        grid: Optional strictly increasing cell-center coordinates of shape
            ``(n,)``. Exactly one of ``dx`` and ``grid`` must be supplied.

    Returns:
        Sparse CSR matrix representing the Laplacian operator.

    Raises:
        ValueError: If the geometry or boundary condition is invalid.
    """
    bc_normalized = str(bc).strip().lower()
    if bc_normalized not in {"neumann", "absorbing"}:
        msg = _UNKNOWN_BC_ERROR.format(bc=bc)
        raise ValueError(msg)

    dtype_obj = np.dtype(dtype)
    lower, main, upper, _periodic_corner = _diffusion_diagonals(
        n,
        dx,
        grid=grid,
        bc=bc_normalized,
        dtype=dtype_obj,
    )
    laplacian = diags(
        [lower.tolist(), main.tolist(), upper.tolist()],
        [-1, 0, 1],
        shape=(n, n),
        dtype=dtype_obj,
    )
    return (laplacian * coeff).tocsr()

build_predictor_corrector(base_matrix)

Build predictor-corrector matrices with dense/sparse autodispatch.

Parameters:

Name	Type	Description	Default
base_matrix	DenseOperator | csr_matrix	Base linear operator A.	required

Returns:

Type	Description
tuple[Operator, Operator, Operator]	Tuple of (predictor, L, R) operators for predictor-corrector scheme.

Source code in src/op_engine/matrix_ops.py

def build_predictor_corrector(
    base_matrix: DenseOperator | csr_matrix,
) -> tuple[Operator, Operator, Operator]:
    """
    Build predictor-corrector matrices with dense/sparse autodispatch.

    Args:
        base_matrix: Base linear operator A.

    Returns:
        Tuple of (predictor, L, R) operators for predictor-corrector scheme.
    """
    n = base_matrix.shape[0]
    if issparse(base_matrix) and n >= _DISPATCH_THRESHOLD:
        return _build_predictor_corrector_sparse(base_matrix)

    if issparse(base_matrix):
        dense_base = np.asarray(base_matrix.toarray())
    else:
        dense_base = np.asarray(base_matrix)

    return _build_predictor_corrector_dense(cast("DenseOperator", dense_base))

clear_implicit_solver_cache()

Clear the internal implicit solver cache.

Source code in src/op_engine/matrix_ops.py

def clear_implicit_solver_cache() -> None:
    """Clear the internal implicit solver cache."""
    _IMPLICIT_SOLVER_CACHE.clear()

compile_reaction_network(reactions, *, template_shapes, axis_sizes, params, n_state=None, reaction_names=None)

Compile op_system reaction artifacts into flat channels.

Each source-cell propensity becomes one channel. Consequently, a collapsed axis is represented by distinct columns that share one destination row; summing simultaneous firings is then exactly the stoichiometric matrix multiplication performed by the core stochastic solvers. An offset axis moves each channel to its shifted coordinate; when that leaves the axis, the column only removes the donor.

Parameters:

Name	Type	Description	Default
reactions	Sequence[ReactionArtifact]	Typed reaction artifacts, such as CompiledRhs.reactions.	required
template_shapes	Mapping[str, Sequence[int]]	Shape of each state template, in the flat state's order, such as CompiledRhs.template_shapes.	required
axis_sizes	Mapping[str, int]	Number of coordinates on each axis.	required
params	Mapping[str, object]	Parameter values passed to every propensity.	required
n_state	int | None	Number of cells in the flat state. Defaults to the total that template_shapes describes; when given, the two must agree.	None
reaction_names	Sequence[str] | None	Optional parent reaction names to select for a hybrid jump partition. None selects every artifact.	None

Returns:

Type	Description
CompiledReactionNetwork	Validated flat reaction network.

Raises:

Type	Description
TypeError	If an argument or artifact field has the wrong type.
ValueError	If the artifacts are inconsistent with each other or with the state layout.

Source code in src/op_engine/reactions.py

def compile_reaction_network(  # noqa: PLR0913
    reactions: Sequence[ReactionArtifact],
    *,
    template_shapes: Mapping[str, Sequence[int]],
    axis_sizes: Mapping[str, int],
    params: Mapping[str, object],
    n_state: int | None = None,
    reaction_names: Sequence[str] | None = None,
) -> CompiledReactionNetwork:
    """Compile op_system reaction artifacts into flat channels.

    Each source-cell propensity becomes one channel. Consequently, a collapsed
    axis is represented by distinct columns that share one destination row;
    summing simultaneous firings is then exactly the stoichiometric matrix
    multiplication performed by the core stochastic solvers. An offset axis
    moves each channel to its shifted coordinate; when that leaves the axis,
    the column only removes the donor.

    Args:
        reactions: Typed reaction artifacts, such as ``CompiledRhs.reactions``.
        template_shapes: Shape of each state template, in the flat state's
            order, such as ``CompiledRhs.template_shapes``.
        axis_sizes: Number of coordinates on each axis.
        params: Parameter values passed to every propensity.
        n_state: Number of cells in the flat state. Defaults to the total
            that ``template_shapes`` describes; when given, the two must
            agree.
        reaction_names: Optional parent reaction names to select for a hybrid
            jump partition. ``None`` selects every artifact.

    Returns:
        Validated flat reaction network.

    Raises:
        TypeError: If an argument or artifact field has the wrong type.
        ValueError: If the artifacts are inconsistent with each other or
            with the state layout.
    """
    if not isinstance(reactions, tuple | list):
        msg = "'reactions' must be a sequence of reaction artifacts."
        raise TypeError(msg)
    selected = tuple(cast("ReactionArtifact", item) for item in reactions)
    if reaction_names is not None:
        requested = tuple(reaction_names)
        if len(requested) != len(set(requested)):
            msg = "stochastic_reactions must not contain duplicates."
            raise ValueError(msg)
        available = {reaction.name for reaction in selected}
        missing = sorted(set(requested) - available)
        if missing:
            msg = f"Unknown stochastic reaction names: {', '.join(missing)}."
            raise ValueError(msg)
        wanted = set(requested)
        selected = tuple(reaction for reaction in selected if reaction.name in wanted)
    if not selected:
        msg = "No typed reactions are available for stochastic execution."
        raise ValueError(msg)
    reactions = selected

    if not isinstance(template_shapes, Mapping):
        msg = "'template_shapes' must map state template names to shapes."
        raise TypeError(msg)
    if n_state is None:
        n_state = sum(
            int(np.prod(tuple(shape), dtype=np.int64)) if shape else 1
            for shape in template_shapes.values()
            if isinstance(shape, tuple | list)
        )
    blocks_tuple = _state_blocks(
        cast("Mapping[object, object]", template_shapes), n_state=n_state
    )
    blocks = {block.base: block for block in blocks_tuple}

    if not isinstance(axis_sizes, Mapping):
        msg = "'axis_sizes' must be a mapping."
        raise TypeError(msg)
    raw_axis_sizes = axis_sizes
    axis_sizes = {}
    for axis, size in raw_axis_sizes.items():
        if (
            not isinstance(axis, str)
            or not isinstance(size, int)
            or isinstance(size, bool)
            or size < 1
        ):
            msg = "'axis_sizes' contains an invalid entry."
            raise TypeError(msg)
        axis_sizes[axis] = size
    base_axes = _reaction_axes(reactions, blocks, axis_sizes)

    columns: list[np.ndarray] = []
    reactant_columns: list[np.ndarray] = []
    dependency_columns: list[np.ndarray] = []
    channel_orders: list[int] = []
    channel_reactions: list[str] = []
    channel_names: list[str] = []
    event_shapes: list[tuple[int, ...]] = []
    reactants_complete = True
    incomplete: list[str] = []

    for reaction in reactions:
        complete = getattr(reaction, "reactants_complete", False)
        if not isinstance(complete, bool):
            msg = f"Reaction {reaction.name!r} reactants_complete must be boolean."
            raise TypeError(msg)
        if complete and getattr(reaction, "reactants", None) is None:
            msg = (
                f"Reaction {reaction.name!r} claims complete reactants without "
                "publishing reactant metadata."
            )
            raise ValueError(msg)
        # Complete reactants take precedence; otherwise dependencies and a
        # propensity order can describe the reaction instead.
        order = 0 if complete else _dependency_order(reaction)
        covered = complete or order > 0
        reactants_complete = reactants_complete and covered
        if not covered:
            incomplete.append(reaction.name)
        from_axes = _require_string_tuple(reaction.from_axes, field="from_axes")
        target_axes = base_axes[reaction.to_base]
        pinned, from_pinned, offsets, routed = _reaction_axis_metadata(
            reaction,
            from_axes=from_axes,
            source_axes=(
                target_axes
                if reaction.from_base is None
                else base_axes[reaction.from_base]
            ),
            target_axes=target_axes,
        )

        # A routed target coordinate is a trailing channel dimension.
        event_shape = tuple(axis_sizes[axis] for axis in from_axes + routed)
        event_shapes.append(event_shape)
        coordinates = np.ndindex(event_shape) if event_shape else iter(((),))
        for coordinate in coordinates:
            varying = dict(zip(from_axes, coordinate[: len(from_axes)], strict=True))
            routed_to = dict(zip(routed, coordinate[len(from_axes) :], strict=True))
            column = np.zeros(n_state, dtype=np.int64)
            if reaction.from_base is not None:
                source_coordinates = {**varying, **from_pinned}
                source = _flat_cell(
                    blocks[reaction.from_base],
                    base_axes[reaction.from_base],
                    source_coordinates,
                )
                column[source] -= 1

            reactants = _expanded_reactants(
                reaction,
                varying=varying,
                from_axes=from_axes,
                base_axes=base_axes,
                blocks=blocks,
                n_state=n_state,
            )

            read = np.zeros(n_state, dtype=np.int64)
            if order > 0:
                read = _expanded_entries(
                    reaction,
                    getattr(reaction, "dependencies", ()),
                    label="dependency",
                    varying=varying,
                    from_axes=from_axes,
                    base_axes=base_axes,
                    blocks=blocks,
                    n_state=n_state,
                )
                # A consumed species is read too, whether or not listed.
                read = ((read > 0) | (reactants > 0)).astype(np.int64)

            destination = _destination_cell(
                blocks[reaction.to_base],
                base_axes[reaction.to_base],
                {**varying, **pinned, **routed_to},
                offsets=offsets,
            )
            if destination is not None:
                column[destination] += 1
            columns.append(column)
            reactant_columns.append(reactants)
            dependency_columns.append(read)
            channel_orders.append(order)
            channel_reactions.append(reaction.name)
            coordinate_text = ",".join(
                [f"{axis}={index}" for axis, index in varying.items()]
                + [f"to:{axis}={index}" for axis, index in routed_to.items()]
            )
            channel_names.append(
                reaction.name
                if not coordinate_text
                else f"{reaction.name}[{coordinate_text}]"
            )

    return CompiledReactionNetwork(
        stoichiometry=np.stack(columns, axis=1),
        reactant_stoichiometry=np.stack(reactant_columns, axis=1),
        channel_reactions=tuple(channel_reactions),
        channel_names=tuple(channel_names),
        _blocks=blocks_tuple,
        _reactions=reactions,
        _event_shapes=tuple(event_shapes),
        _params=dict(params),
        reactants_complete=reactants_complete,
        incomplete_reactions=tuple(incomplete),
        dependency_incidence=(
            np.stack(dependency_columns, axis=1) if any(channel_orders) else None
        ),
        propensity_orders=(
            np.asarray(channel_orders, dtype=np.int64) if any(channel_orders) else None
        ),
    )

conserved_quantities(stoichiometry, *, tol=1e-10)

Return the structural conservation laws of a reaction network.

Rows w satisfy w @ stoichiometry = 0, so w @ y is unchanged by every reaction, whatever the rates. Totals conserved only because rates balance (births mu * N against deaths) are not structural; use :func:linear_invariants for those.

Parameters:

Name	Type	Description	Default
stoichiometry	object	(n_state, n_reactions) net change matrix, such as CompiledReactionNetwork.stoichiometry.	required
tol	float	Relative singular-value threshold.	1e-10

Returns:

Type	Description
ndarray	(k, n_state) orthonormal rows, possibly empty.

Examples:

>>> import numpy as np
>>> transfers = np.array([[-1.0], [1.0]])
>>> laws = conserved_quantities(transfers)
>>> np.allclose(laws @ transfers, 0.0)
True
>>> laws.shape
(1, 2)

Source code in src/op_engine/steady_state.py

def conserved_quantities(stoichiometry: object, *, tol: float = 1e-10) -> np.ndarray:
    """Return the structural conservation laws of a reaction network.

    Rows ``w`` satisfy ``w @ stoichiometry = 0``, so ``w @ y`` is unchanged by
    every reaction, whatever the rates. Totals conserved only because rates
    balance (births ``mu * N`` against deaths) are not structural; use
    :func:`linear_invariants` for those.

    Args:
        stoichiometry: ``(n_state, n_reactions)`` net change matrix, such as
            ``CompiledReactionNetwork.stoichiometry``.
        tol: Relative singular-value threshold.

    Returns:
        ``(k, n_state)`` orthonormal rows, possibly empty.

    Examples:
        >>> import numpy as np
        >>> transfers = np.array([[-1.0], [1.0]])
        >>> laws = conserved_quantities(transfers)
        >>> np.allclose(laws @ transfers, 0.0)
        True
        >>> laws.shape
        (1, 2)
    """
    return _left_null_space(np.asarray(stoichiometry, dtype=np.float64), tol=tol)

diagnose_trajectory(trajectory, config=None)

Check an output trajectory without modifying it or extracting host scalars.

Parameters:

Name	Type	Description	Default
trajectory	Array	Array shaped (output, state, *subgroups_or_field_axes).	required
config	OutputStateCheckConfig | None	Component selectors and absolute tolerances.	None

Returns:

Type	Description
OutputStateDiagnostics	Diagnostics for the complete supplied trajectory. valid describes
OutputStateDiagnostics	these recorded outputs, not intermediate-step positivity or accuracy.

Raises:

Type	Description
TypeError	If config has the wrong type.
ValueError	If the trajectory shape or selectors are invalid.

Examples:

>>> import numpy as np
>>> path = np.array([[[1.0, 1.0]], [[0.9, -1.0]], [[0.81, 1.0]]])
>>> result = diagnose_trajectory(path)
>>> result.n_negative_outputs.tolist()
[0, 1]
>>> result.valid.tolist()
[True, False]

Source code in src/op_engine/state_diagnostics.py

def diagnose_trajectory(
    trajectory: Array,
    config: OutputStateCheckConfig | None = None,
) -> OutputStateDiagnostics:
    """Check an output trajectory without modifying it or extracting host scalars.

    Args:
        trajectory: Array shaped ``(output, state, *subgroups_or_field_axes)``.
        config: Component selectors and absolute tolerances.

    Returns:
        Diagnostics for the complete supplied trajectory. ``valid`` describes
        these recorded outputs, not intermediate-step positivity or accuracy.

    Raises:
        TypeError: If config has the wrong type.
        ValueError: If the trajectory shape or selectors are invalid.

    Examples:
        >>> import numpy as np
        >>> path = np.array([[[1.0, 1.0]], [[0.9, -1.0]], [[0.81, 1.0]]])
        >>> result = diagnose_trajectory(path)
        >>> result.n_negative_outputs.tolist()
        [0, 1]
        >>> result.valid.tolist()
        [True, False]
    """
    cfg = OutputStateCheckConfig() if config is None else config
    if not isinstance(cfg, OutputStateCheckConfig):
        msg = "config must be an OutputStateCheckConfig"
        raise TypeError(msg)
    if len(trajectory.shape) < 2 or trajectory.shape[0] < 1:
        msg = "Trajectory must have nonempty output and state axes"
        raise ValueError(msg)
    _validate_components(cfg, trajectory.shape[1])
    xp = array_namespace(trajectory)
    if not xp.isdtype(trajectory.dtype, "real floating"):
        msg = "Output-state diagnostics require real floating states"
        raise TypeError(msg)
    values = xp.asarray(trajectory)
    selected = (
        values
        if cfg.nonnegative_indices is None
        else values[:, xp.asarray(cfg.nonnegative_indices, dtype=xp.int32)]
    )
    minimum = xp.min(xp.min(selected, axis=1), axis=0)
    negatives = xp.sum(
        xp.astype(xp.any(selected < -cfg.negative_tolerance, axis=1), xp.int32), axis=0
    )
    nonfinite = xp.sum(
        xp.astype(xp.any(~xp.isfinite(values), axis=1), xp.int32), axis=0
    )
    min_increment = decreasing = None
    if cfg.monotone_indices:
        monotone = values[:, xp.asarray(cfg.monotone_indices, dtype=xp.int32)]
        if trajectory.shape[0] == 1:
            min_increment = xp.full_like(minimum, float("inf"))
            decreasing = xp.zeros_like(negatives)
        else:
            increments = monotone[1:] - monotone[:-1]
            min_increment = xp.min(xp.min(increments, axis=1), axis=0)
            decreasing = xp.sum(
                xp.astype(
                    xp.any(increments < -cfg.monotone_tolerance, axis=1), xp.int32
                ),
                axis=0,
            )
    return OutputStateDiagnostics(
        n_outputs=int(trajectory.shape[0]),
        complete=True,
        min_state=minimum,
        n_negative_outputs=negatives,
        n_nonfinite_outputs=nonfinite,
        min_increment=min_increment,
        n_decreasing_intervals=decreasing,
    )

encode_groups(group_ids, n_groups, *, prefer_sparse=None, dtype=np.float64)

Encode group IDs into a one-hot group membership matrix.

Parameters:

Name	Type	Description	Default
group_ids	NDArray[integer]	1D array of integer group IDs for each item.	required
n_groups	int	Total number of groups.	required
prefer_sparse	bool | None	If True, always return a sparse matrix; if False, always return a dense array; if None, autodispatch based on n_groups.	None
dtype	DTypeLike	Data type for the output matrix.	float64

Returns:

Type	Description
csr_matrix | DenseOperator	A (n_groups, n_items) one-hot encoded group membership matrix.

Source code in src/op_engine/matrix_ops.py

def encode_groups(
    group_ids: NDArray[np.integer],
    n_groups: int,
    *,
    prefer_sparse: bool | None = None,
    dtype: DTypeLike = np.float64,
) -> csr_matrix | DenseOperator:
    """
    Encode group IDs into a one-hot group membership matrix.

    Args:
        group_ids: 1D array of integer group IDs for each item.
        n_groups: Total number of groups.
        prefer_sparse: If True, always return a sparse matrix; if False, always
            return a dense array; if None, autodispatch based on n_groups.
        dtype: Data type for the output matrix.

    Returns:
        A (n_groups, n_items) one-hot encoded group membership matrix.
    """
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)

    if prefer_sparse is True:
        return _encode_sparse_groups(group_ids_arr, n_groups, dtype=dtype)
    if prefer_sparse is False:
        return _encode_dense_groups(group_ids_arr, n_groups, dtype=dtype)

    if n_groups >= _DISPATCH_THRESHOLD:
        return _encode_sparse_groups(group_ids_arr, n_groups, dtype=dtype)
    return _encode_dense_groups(group_ids_arr, n_groups, dtype=dtype)

from_compiled_replicates(compiled, params, *, batch_axis)

Build a local network from compiler-certified independent reactions.

Requires the independent-reaction contract from op_system 0.8.0+. Every evolving state must carry the batch axis. Operators, gaps, explicit-time hazards and nonconstant linear tables are rejected. Time-varying parameters must explicitly carry the batch axis to represent asynchronous current values.

Returns:

Type	Description
ReplicateReactionNetwork	Local network for BatchedDirectSSASolver.

Raises:

Type	Description
ValueError	Independence, completeness or time assumptions are unmet.
TypeError	Required compiler artifacts are unavailable.

Source code in src/op_engine/replicate_network.py

def from_compiled_replicates(  # noqa: C901
    compiled: CompiledRhsLike,
    params: Mapping[str, object],
    *,
    batch_axis: str,
) -> ReplicateReactionNetwork:
    """Build a local network from compiler-certified independent reactions.

    Requires the independent-reaction contract from op_system 0.8.0+. Every
    evolving state must carry the batch axis. Operators, gaps, explicit-time
    hazards and nonconstant linear tables are rejected. Time-varying parameters
    must explicitly carry the batch axis to represent asynchronous current values.

    Returns:
        Local network for BatchedDirectSSASolver.

    Raises:
        ValueError: Independence, completeness or time assumptions are unmet.
        TypeError: Required compiler artifacts are unavailable.
    """
    if compiled.reaction_gaps or compiled.meta.get("operators"):
        msg = (
            "Independent SSA requires complete reactions and no unrepresented operators"
        )
        raise ValueError(msg)
    shapes = compiled.template_shapes
    if not shapes or not compiled.reactions:
        msg = "Independent SSA requires compiled reaction and template metadata"
        raise TypeError(msg)
    blocks = tuple(getattr(compiled, "block_axes", ()))
    block = next((b for b in blocks if b.name == batch_axis), None)
    if block is None or set(block.state_axis_pos) != set(shapes):
        msg = f"Every state must carry certified factorize axis {batch_axis!r}"
        raise ValueError(msg)
    sizes = compiled_rhs_axis_sizes(compiled)
    time_axis = str(compiled.meta.get("time_axis", "time"))
    coordinates = np.asarray(compiled.meta.get("time_coordinates", ()), dtype=float)
    local_shapes = {
        name: tuple(n for i, n in enumerate(shape) if i != block.state_axis_pos[name])
        for name, shape in shapes.items()
    }
    local_reactions = []
    tables: dict[str, tuple[str, ...]] = {}
    for reaction in compiled.reactions:
        if batch_axis not in getattr(reaction, "independent_axes", ()):
            msg = (
                f"Reaction {reaction.name!r} is not certified independent "
                f"on {batch_axis!r}"
            )
            raise ValueError(msg)
        if getattr(reaction, "uses_time", True):
            msg = (
                f"Reaction {reaction.name!r} has explicit or "
                "unqualified time dependence"
            )
            raise ValueError(msg)
        if not callable(getattr(reaction, "resolved_propensity_fn", None)):
            msg = (
                "Independent SSA requires op_system's resolved "
                "propensity contract (0.8.0+)"
            )
            raise TypeError(msg)
        tables.update(dict(getattr(reaction, "time_varying_params", ())))
        fields: dict[str, Any] = {
            name: tuple(a for a in getattr(reaction, name) if a != batch_axis)
            for name in ("from_axes", "full_axes", "to_axes", "sum_axes", "routed_axes")
        }
        destination_axes = getattr(reaction, "to_full_axes", None)
        if destination_axes is not None:
            fields["to_full_axes"] = tuple(
                a for a in destination_axes if a != batch_axis
            )
        fields.update(
            reactants=(),
            reactants_complete=False,
            dependencies=(),
            dependencies_complete=False,
            propensity_order=None,
        )
        local_reactions.append(replace(cast("Any", reaction), **fields))
    if tables and (
        coordinates.size == 0
        or (
            coordinates.size > 1
            and compiled.meta.get("time_interpolation") != "previous"
        )
    ):
        msg = (
            "Independent exact SSA requires previous interpolation "
            "or a constant one-knot table"
        )
        raise ValueError(msg)
    for name, axes in tables.items():
        if batch_axis not in axes:
            msg = f"Time parameter {name!r} must declare replicate axis {batch_axis!r}"
            raise ValueError(msg)
    local = compile_reaction_network(
        tuple(local_reactions),
        template_shapes=local_shapes,
        axis_sizes={a: n for a, n in sizes.items() if a != batch_axis},
        params={},
    )
    return ReplicateReactionNetwork(
        local.stoichiometry,
        local.channel_names,
        local.channel_reactions,
        int(sizes[batch_axis]),
        tuple(float(t) for t in coordinates[1:]) if tables else (),
        batch_axis,
        tuple(
            (name, tuple(shape), block.state_axis_pos[name])
            for name, shape in shapes.items()
        ),
        tuple(cast("_IndependentReaction", r) for r in compiled.reactions),
        dict(params),
        time_axis,
        coordinates,
        tables,
        sizes,
    )

from_compiled_rhs(compiled, params, *, reaction_names=None, allow_gaps=False)

Compile the reaction network of a compiled op_system RHS.

The flat state follows compiled.template_shapes: each template's cells in C order, templates in declaration order. That is also the layout of compiled.eval_fn's state vector when the RHS is vectorized, so network.mean_drift(t, y) should equal compiled.eval_fn(t, y, **params) for every y. Check that identity once for a new model: a mismatch means some dynamics have no reaction.

Parameters:

Name	Type	Description	Default
compiled	CompiledRhsLike	A compiled RHS, such as op_system.compile_spec(spec).	required
params	Mapping[str, object]	Parameter values passed to every propensity.	required
reaction_names	Sequence[str] | None	Optional reaction names to compile, for a hybrid jump partition. None compiles every reaction.	None
allow_gaps	bool	Compile even when compiled.reaction_gaps lists dynamics without a reaction. Selecting reaction_names also skips the check, since a partition is partial by design.	False

Returns:

Type	Description
CompiledReactionNetwork	Validated flat reaction network.

Raises:

Type	Description
TypeError	If the RHS has no vectorized state layout.
ValueError	If the RHS has reaction gaps and neither allow_gaps nor reaction_names is given.

Source code in src/op_engine/reactions.py

def from_compiled_rhs(
    compiled: CompiledRhsLike,
    params: Mapping[str, object],
    *,
    reaction_names: Sequence[str] | None = None,
    allow_gaps: bool = False,
) -> CompiledReactionNetwork:
    """Compile the reaction network of a compiled ``op_system`` RHS.

    The flat state follows ``compiled.template_shapes``: each template's
    cells in C order, templates in declaration order. That is also the
    layout of ``compiled.eval_fn``'s state vector when the RHS is
    vectorized, so ``network.mean_drift(t, y)`` should equal
    ``compiled.eval_fn(t, y, **params)`` for every ``y``. Check that
    identity once for a new model: a mismatch means some dynamics have no
    reaction.

    Args:
        compiled: A compiled RHS, such as ``op_system.compile_spec(spec)``.
        params: Parameter values passed to every propensity.
        reaction_names: Optional reaction names to compile, for a hybrid
            jump partition. ``None`` compiles every reaction.
        allow_gaps: Compile even when ``compiled.reaction_gaps`` lists
            dynamics without a reaction. Selecting ``reaction_names`` also
            skips the check, since a partition is partial by design.

    Returns:
        Validated flat reaction network.

    Raises:
        TypeError: If the RHS has no vectorized state layout.
        ValueError: If the RHS has reaction gaps and neither ``allow_gaps``
            nor ``reaction_names`` is given.
    """
    template_shapes = compiled.template_shapes
    if template_shapes is None:
        msg = (
            "The compiled RHS has no vectorized state layout "
            "(template_shapes is None), so its reactions cannot be indexed."
        )
        raise TypeError(msg)
    gaps = tuple(compiled.reaction_gaps)
    if gaps and not allow_gaps and reaction_names is None:
        described = ", ".join(
            f"{getattr(gap, 'name', None) or getattr(gap, 'origin', '?')} "
            f"({getattr(gap, 'reason', 'unknown')})"
            for gap in gaps
        )
        msg = (
            "The compiled RHS has dynamics without a reaction artifact: "
            f"{described}. A network of its reactions alone would drop them; "
            "pass allow_gaps=True to compile it anyway."
        )
        raise ValueError(msg)
    return compile_reaction_network(
        tuple(compiled.reactions),
        template_shapes=template_shapes,
        axis_sizes=compiled_rhs_axis_sizes(compiled),
        params=params,
        reaction_names=reaction_names,
    )

generate_adaptive_grid(profile, domain, n_points, *, epsilon=0.001, sampling_points=4097, smoothing_window=None, minimum_spacing=None)

Generate a curvature-weighted grid from a vectorized profile callable.

Parameters:

Name	Type	Description	Default
profile	Callable[[NDArray[float64]], ArrayLike]	Callable evaluated on a one-dimensional NumPy sample grid. It must return one finite value per sample.	required
domain	tuple[float, float]	Finite increasing (start, end) coordinate interval.	required
n_points	int	Number of output points, including both domain endpoints.	required
epsilon	float	Positive curvature-density floor.	0.001
sampling_points	int	Number of uniform samples used to estimate curvature. This must be at least n_points and at least three.	4097
smoothing_window	int | None	Optional odd moving-average width for curvature.	None
minimum_spacing	float | None	Optional positive lower bound on output spacing.	None

Returns:

Type	Description
NDArray[float64]	Strictly increasing curvature-weighted coordinates with shape
NDArray[float64]	(n_points,).

Raises:

Type	Description
TypeError	If profile is not callable.
ValueError	If the domain, sampled profile, or controls are invalid.

Source code in src/op_engine/spatial_grid.py

def generate_adaptive_grid(  # noqa: PLR0913
    profile: Callable[[NDArray[np.float64]], ArrayLike],
    domain: tuple[float, float],
    n_points: int,
    *,
    epsilon: float = 1e-3,
    sampling_points: int = 4097,
    smoothing_window: int | None = None,
    minimum_spacing: float | None = None,
) -> NDArray[np.float64]:
    """Generate a curvature-weighted grid from a vectorized profile callable.

    Args:
        profile: Callable evaluated on a one-dimensional NumPy sample grid. It
            must return one finite value per sample.
        domain: Finite increasing ``(start, end)`` coordinate interval.
        n_points: Number of output points, including both domain endpoints.
        epsilon: Positive curvature-density floor.
        sampling_points: Number of uniform samples used to estimate curvature.
            This must be at least ``n_points`` and at least three.
        smoothing_window: Optional odd moving-average width for curvature.
        minimum_spacing: Optional positive lower bound on output spacing.

    Returns:
        Strictly increasing curvature-weighted coordinates with shape
        ``(n_points,)``.

    Raises:
        TypeError: If ``profile`` is not callable.
        ValueError: If the domain, sampled profile, or controls are invalid.
    """
    if not callable(profile):
        msg = "profile must be callable."
        raise TypeError(msg)
    count = _positive_integer(n_points, name="n_points", minimum=2)
    sample_count = _positive_integer(
        sampling_points,
        name="sampling_points",
        minimum=3,
    )
    if sample_count < count:
        msg = "sampling_points must be greater than or equal to n_points."
        raise ValueError(msg)
    start, end = domain
    if not np.isfinite(start) or not np.isfinite(end) or start >= end:
        msg = "domain must contain two finite, strictly increasing bounds."
        raise ValueError(msg)

    coordinates = np.linspace(start, end, sample_count, dtype=np.float64)
    values = np.asarray(profile(coordinates), dtype=np.float64)
    return generate_adaptive_grid_from_data(
        coordinates,
        values,
        count,
        epsilon=epsilon,
        smoothing_window=smoothing_window,
        minimum_spacing=minimum_spacing,
    )

generate_adaptive_grid_from_data(coordinates, values, n_points, *, epsilon=0.001, smoothing_window=None, minimum_spacing=None)

Equidistribute points according to sampled profile curvature.

The curvature monitor is abs(f'') / (1 + f'**2)**(3/2) + epsilon. Its trapezoidal cumulative integral defines a monotone distribution whose evenly spaced quantiles are inverted with linear interpolation. The result contains exactly n_points coordinates, including both input endpoints.

This is an eager geometry-preprocessing utility: it always returns a NumPy array. The resulting static coordinates can be passed to numerical operators running in any supported Array-API namespace.

Parameters:

Name	Type	Description	Default
coordinates	ArrayLike	Strictly increasing one-dimensional sample coordinates.	required
values	ArrayLike	Finite profile values at coordinates.	required
n_points	int	Number of output points, including both endpoints.	required
epsilon	float	Positive curvature-density floor.	0.001
smoothing_window	int | None	Optional odd moving-average width applied to the sampled curvature before integration.	None
minimum_spacing	float | None	Optional positive lower bound on adjacent output spacing. It must be feasible within the sampled domain.	None

Returns:

Type	Description
NDArray[float64]	Strictly increasing curvature-weighted coordinates with shape
NDArray[float64]	(n_points,).

Raises:

Type	Description
ValueError	If samples or controls are invalid.

Source code in src/op_engine/spatial_grid.py

def generate_adaptive_grid_from_data(  # noqa: PLR0913, PLR0914
    coordinates: ArrayLike,
    values: ArrayLike,
    n_points: int,
    *,
    epsilon: float = 1e-3,
    smoothing_window: int | None = None,
    minimum_spacing: float | None = None,
) -> NDArray[np.float64]:
    """Equidistribute points according to sampled profile curvature.

    The curvature monitor is
    ``abs(f'') / (1 + f'**2)**(3/2) + epsilon``. Its trapezoidal cumulative
    integral defines a monotone distribution whose evenly spaced quantiles
    are inverted with linear interpolation. The result contains exactly
    ``n_points`` coordinates, including both input endpoints.

    This is an eager geometry-preprocessing utility: it always returns a NumPy
    array. The resulting static coordinates can be passed to numerical
    operators running in any supported Array-API namespace.

    Args:
        coordinates: Strictly increasing one-dimensional sample coordinates.
        values: Finite profile values at ``coordinates``.
        n_points: Number of output points, including both endpoints.
        epsilon: Positive curvature-density floor.
        smoothing_window: Optional odd moving-average width applied to the
            sampled curvature before integration.
        minimum_spacing: Optional positive lower bound on adjacent output
            spacing. It must be feasible within the sampled domain.

    Returns:
        Strictly increasing curvature-weighted coordinates with shape
        ``(n_points,)``.

    Raises:
        ValueError: If samples or controls are invalid.
    """
    count = _positive_integer(n_points, name="n_points", minimum=2)
    x = np.asarray(coordinates, dtype=np.float64)
    y = np.asarray(values, dtype=np.float64)
    if x.ndim != 1 or x.size < 3:
        msg = "coordinates must be one-dimensional with at least three samples."
        raise ValueError(msg)
    if y.shape != x.shape:
        msg = f"values must have shape {x.shape}; got {y.shape}."
        raise ValueError(msg)
    if not np.isfinite(x).all() or not np.isfinite(y).all():
        msg = "coordinates and values must be finite."
        raise ValueError(msg)
    if not np.all(np.diff(x) > 0.0):
        msg = "coordinates must be strictly increasing."
        raise ValueError(msg)
    if not np.isfinite(epsilon) or epsilon <= 0.0:
        msg = "epsilon must be finite and positive."
        raise ValueError(msg)

    first_derivative = np.gradient(y, x, edge_order=2)
    second_derivative = np.gradient(first_derivative, x, edge_order=2)
    curvature = np.abs(second_derivative) / np.power(
        1.0 + first_derivative * first_derivative,
        1.5,
    )
    sample_spacing = float(np.min(np.diff(x)))
    value_scale = max(1.0, float(np.max(np.abs(y))))
    curvature_noise = 64.0 * np.finfo(np.float64).eps * value_scale / sample_spacing**2
    curvature = np.where(curvature <= curvature_noise, 0.0, curvature)
    curvature = _smooth_curvature(
        np.asarray(curvature, dtype=np.float64),
        smoothing_window,
    )
    if not np.any(curvature):
        uniform = np.linspace(x[0], x[-1], count, dtype=np.float64)
        return _enforce_minimum_spacing(uniform, minimum_spacing)

    density = curvature + epsilon
    increments = 0.5 * (density[:-1] + density[1:]) * np.diff(x)
    cumulative = np.concatenate((np.asarray([0.0]), np.cumsum(increments)))
    cumulative /= cumulative[-1]
    quantiles = np.linspace(0.0, 1.0, count, dtype=np.float64)
    grid = np.asarray(np.interp(quantiles, cumulative, x), dtype=np.float64)
    grid[0] = x[0]
    grid[-1] = x[-1]
    return _enforce_minimum_spacing(grid, minimum_spacing)

grouped_count_ids(group_ids, n_groups)

Perform grouped count using group IDs.

Parameters:

Name	Type	Description	Default
group_ids	NDArray[integer]	1D array of integer group IDs.	required
n_groups	int	Total number of groups.	required

Returns:

Type	Description
NDArray[floating]	A 1D array of length n_groups where each element contains the count of
NDArray[floating]	occurrences of the corresponding group ID.

Source code in src/op_engine/matrix_ops.py

def grouped_count_ids(
    group_ids: NDArray[np.integer],
    n_groups: int,
) -> NDArray[np.floating]:
    """
    Perform grouped count using group IDs.

    Args:
        group_ids: 1D array of integer group IDs.
        n_groups: Total number of groups.

    Returns:
        A 1D array of length n_groups where each element contains the count of
        occurrences of the corresponding group ID.
    """
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)
    counts = np.bincount(group_ids_arr, minlength=n_groups)
    return counts.astype(float)

grouped_sum_ids(values, group_ids, n_groups)

Perform grouped sum over 1D values array using group IDs.

Parameters:

Name	Type	Description	Default
values	NDArray[floating]	1D array of values.	required
group_ids	NDArray[integer]	1D array of integer group IDs.	required
n_groups	int	Total number of groups.	required

Returns:

Type	Description
NDArray[floating]	A 1D array of length n_groups where each element contains the sum of values
NDArray[floating]	for the corresponding group ID.

Source code in src/op_engine/matrix_ops.py

def grouped_sum_ids(
    values: NDArray[np.floating],
    group_ids: NDArray[np.integer],
    n_groups: int,
) -> NDArray[np.floating]:
    """
    Perform grouped sum over 1D values array using group IDs.

    Args:
        values: 1D array of values.
        group_ids: 1D array of integer group IDs.
        n_groups: Total number of groups.

    Returns:
        A 1D array of length n_groups where each element contains the sum of values
        for the corresponding group ID.
    """
    values_arr = np.asarray(values)
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)
    sums = np.bincount(group_ids_arr, weights=values_arr, minlength=n_groups)
    return sums.astype(float)

grouped_sum_ids_2d(values, group_ids, n_groups)

Perform grouped sum over 2D values array using group IDs.

Parameters:

Name	Type	Description	Default
values	NDArray[floating]	2D (N, K) array where N is num of items and K is num of features.	required
group_ids	NDArray[integer]	1D array of integer group IDs of length N.	required
n_groups	int	Total number of groups.	required

Returns:

Type	Description
NDArray[floating]	A 2D array of shape (n_groups, K) where each row contains the sum of values
NDArray[floating]	for the corresponding group ID.

Raises:

Type	Description
ValueError	If values is not 2D or if group_ids length does not match the number of items in values.

Source code in src/op_engine/matrix_ops.py

def grouped_sum_ids_2d(
    values: NDArray[np.floating],
    group_ids: NDArray[np.integer],
    n_groups: int,
) -> NDArray[np.floating]:
    """
    Perform grouped sum over 2D values array using group IDs.

    Args:
        values: 2D (N, K) array where N is num of items and K is num of features.
        group_ids: 1D array of integer group IDs of length N.
        n_groups: Total number of groups.

    Returns:
        A 2D array of shape (n_groups, K) where each row contains the sum of values
        for the corresponding group ID.

    Raises:
        ValueError: If values is not 2D or if group_ids length does not match
            the number of items in values.
    """
    values_arr = np.asarray(values)
    group_ids_arr = np.asarray(group_ids, dtype=np.int64)

    if values_arr.ndim != 2:
        raise ValueError(_VALUES_2D_ERROR)

    n_items, n_features = values_arr.shape
    if group_ids_arr.shape[0] != n_items:
        raise ValueError(_GROUP_IDS_LENGTH_ERROR)

    out = np.zeros((n_groups, n_features), dtype=values_arr.dtype)
    for feature_idx in range(n_features):
        out[:, feature_idx] = np.bincount(
            group_ids_arr,
            weights=values_arr[:, feature_idx],
            minlength=n_groups,
        )
    return out

implicit_solve(left_op, right_op, x)

Perform an implicit solve with dense fallback and cached sparse dispatch.

Parameters:

Name	Type	Description	Default
left_op	object	Left operator L in the equation L @ y = R @ x.	required
right_op	object	Right operator R in the equation L @ y = R @ x.	required
x	Array	1D or 2D array representing the input vector(s).	required

Returns:

Type	Description
Array	A 1D or 2D array containing the solution vector(s) y.

Source code in src/op_engine/matrix_ops.py

def implicit_solve(
    left_op: object,
    right_op: object,
    x: Array,
) -> Array:
    """
    Perform an implicit solve with dense fallback and cached sparse dispatch.

    Args:
        left_op: Left operator L in the equation L @ y = R @ x.
        right_op: Right operator R in the equation L @ y = R @ x.
        x: 1D or 2D array representing the input vector(s).

    Returns:
        A 1D or 2D array containing the solution vector(s) y.
    """
    _validate_solve_dimensions(left_op, right_op, x)

    adapter = _find_sparse_adapter(left_op, right_op)
    if adapter is not None:
        return _sparse_implicit_solve(adapter, left_op, right_op, x)

    xp = _namespace_of(x)
    left_dense = _dense_operator_array(left_op, xp=xp, dtype=x.dtype)
    right_dense = _dense_operator_array(right_op, xp=xp, dtype=x.dtype)
    rhs = cast("Array", xp.matmul(right_dense, x))
    return cast("Array", xp.linalg.solve(left_dense, rhs))

kron_prod(a, b)

Compute the Kronecker product of two operators.

Parameters:

Name	Type	Description	Default
a	Operator	First operator.	required
b	Operator	Second operator.	required

Returns:

Type	Description
Operator	The Kronecker product operator.

Source code in src/op_engine/matrix_ops.py

def kron_prod(a: Operator, b: Operator) -> Operator:
    """
    Compute the Kronecker product of two operators.

    Args:
        a: First operator.
        b: Second operator.

    Returns:
        The Kronecker product operator.
    """
    if issparse(a) or issparse(b):
        a_csr = a if issparse(a) else csr_matrix(np.asarray(a))
        b_csr = b if issparse(b) else csr_matrix(np.asarray(b))
        return kron(
            a_csr,
            b_csr,
            format="csr",
        )
    return cast("DenseOperator", np.kron(np.asarray(a), np.asarray(b)))

kron_sum(ops)

Compute a Kronecker sum of square operators.

Parameters:

Name	Type	Description	Default
ops	list[Operator]	List of 2D square operators.	required

Returns:

Type	Description
Operator	The Kronecker sum operator.

Raises:

Type	Description
ValueError	If ops is empty or if operators are not square or have incompatible shapes.

Source code in src/op_engine/matrix_ops.py

def kron_sum(ops: list[Operator]) -> Operator:
    """
    Compute a Kronecker sum of square operators.

    Args:
        ops: List of 2D square operators.

    Returns:
        The Kronecker sum operator.

    Raises:
        ValueError: If ops is empty or if operators are not square or
            have incompatible shapes.
    """
    if not ops:
        raise ValueError(_KRON_EMPTY_ERROR)

    shapes = [
        tuple(np.asarray(op).shape) if not issparse(op) else op.shape for op in ops
    ]

    if any(s[0] != s[1] for s in shapes):
        raise ValueError(_KRON_INCOMPATIBLE_ERROR.format(shapes=shapes))

    any_sparse = any(issparse(op) for op in ops)
    sizes = [s[0] for s in shapes]
    dtype_obj = np.result_type(*[
        (op.dtype if issparse(op) else np.asarray(op).dtype) for op in ops
    ])

    def _eye(n: int) -> Operator:
        if any_sparse:
            return identity(n, format="csr", dtype=dtype_obj)
        return cast("DenseOperator", np.eye(n, dtype=dtype_obj))

    total: Operator | None = None
    n_ops = len(ops)

    for i, op_i in enumerate(ops):
        term: Operator = op_i
        for j in range(i - 1, -1, -1):
            term = kron_prod(_eye(sizes[j]), term)
        for j in range(i + 1, n_ops):
            term = kron_prod(term, _eye(sizes[j]))
        total = term if total is None else cast("Operator", total + term)

    if total is None:
        raise ValueError(_KRON_EMPTY_ERROR)
    return total

linear_invariants(rhs, states, *, t=0.0, jacobian=None, fixed=None, tol=1e-06)

Find totals w @ y that the dynamics conserve at sampled states.

A row w is returned when w @ J(y) = 0 and w @ rhs(t, y) = 0 at every sampled y, which includes rate-balanced totals that :func:conserved_quantities misses. Sample a few varied states away from special points such as the disease-free state.

The test is numerical. A mode decaying more slowly than tol times the fastest rate looks conserved, and each row's accuracy is limited by the Jacobian's error relative to the slowest real rate: a slow mode tilts the rows toward itself. The default Jacobian uses central differences for that reason; an exact one (jax.jacfwd, for example) is better still. When you know an invariant, such as a row of ones over the living compartments, pass that exact row to :func:steady_state and use this function to confirm that nothing else is conserved.

Parameters:

Name	Type	Description	Default
rhs	SteadyStateRhs	rhs(t, y) -> dy/dt.	required
states	Sequence[Array]	Sample states, each a one-dimensional array.	required
t	float	Evaluation time.	0.0
jacobian	SteadyStateJacobian | None	Dense Jacobian callback; central differences by default.	None
fixed	Sequence[int] | Array | None	Indices, or a boolean mask, of states that :func:steady_state will hold fixed. They are left out, so the rows have zeros there.	None
tol	float	Singular-value threshold relative to the fastest rate.	1e-06

Returns:

Type	Description
ndarray	(k, n_state) orthonormal rows, possibly empty.

Raises:

Type	Description
ValueError	If states is empty.

Source code in src/op_engine/steady_state.py

def linear_invariants(  # noqa: PLR0913
    rhs: SteadyStateRhs,
    states: Sequence[Array],
    *,
    t: float = 0.0,
    jacobian: SteadyStateJacobian | None = None,
    fixed: Sequence[int] | Array | None = None,
    tol: float = 1e-6,
) -> np.ndarray:
    """Find totals ``w @ y`` that the dynamics conserve at sampled states.

    A row ``w`` is returned when ``w @ J(y) = 0`` and ``w @ rhs(t, y) = 0`` at
    every sampled ``y``, which includes rate-balanced totals that
    :func:`conserved_quantities` misses. Sample a few varied states away from
    special points such as the disease-free state.

    The test is numerical. A mode decaying more slowly than ``tol`` times the
    fastest rate looks conserved, and each row's accuracy is limited by the
    Jacobian's error relative to the slowest real rate: a slow mode tilts the
    rows toward itself. The default Jacobian uses central differences for
    that reason; an exact one (``jax.jacfwd``, for example) is better still.
    When you know an invariant, such as a row of ones over the living
    compartments, pass that exact row to :func:`steady_state` and use this
    function to confirm that nothing else is conserved.

    Args:
        rhs: ``rhs(t, y) -> dy/dt``.
        states: Sample states, each a one-dimensional array.
        t: Evaluation time.
        jacobian: Dense Jacobian callback; central differences by default.
        fixed: Indices, or a boolean mask, of states that
            :func:`steady_state` will hold fixed. They are left out, so the
            rows have zeros there.
        tol: Singular-value threshold relative to the fastest rate.

    Returns:
        ``(k, n_state)`` orthonormal rows, possibly empty.

    Raises:
        ValueError: If ``states`` is empty.
    """
    jacobian = jacobian or _central_difference_jacobian(rhs)
    blocks: list[np.ndarray] = []
    free: np.ndarray | None = None
    for y in states:
        values = np.asarray(y, dtype=np.float64)
        if free is None:
            free = _free_mask(fixed, values.shape[0])
        jac = np.asarray(jacobian(t, y), dtype=np.float64)[np.ix_(free, free)]
        # Divide by the state's size so both blocks are rates, comparable
        # with one threshold.
        rate = np.asarray(rhs(t, y), dtype=np.float64)[free] / max(
            float(np.abs(values).max()), 1.0
        )
        blocks.extend((jac, rate.reshape(-1, 1)))
    if free is None:
        msg = "linear_invariants needs at least one sample state"
        raise ValueError(msg)
    rows = _left_null_space(np.concatenate(blocks, axis=1), tol=tol)
    embedded = np.zeros((rows.shape[0], free.shape[0]))
    embedded[:, free] = rows
    return embedded

local_stability(rhs, candidate, *, config, t=0.0, coordinates=None, coordinate_system='supplied state', fixed=None, invariants=None, nonnegative=False, derivatives=None, parameters=None, jacobian=None, jacobian_reference=None, time_unit=None)

Assess a candidate equilibrium through an eager NumPy linearization.

The classification applies only to an interior, residual-qualified point and, when constraints are declared, only to their tangent manifold. A boundary status deliberately suppresses an interior stability conclusion.

Parameters:

Name	Type	Description	Default
rhs	StabilityRhs	Dense NumPy right-hand side in the supplied coordinates.	required
candidate	NumpyArray | SteadyStateResult	NumPy state or NumPy-backed steady-state result.	required
config	LocalStabilityConfig	Explicit residual, spectral, boundary and difference controls.	required
t	float	Frozen evaluation time. A time-dependent RHS is only diagnosed there.	0.0
coordinates	Sequence[str] | None	Ordered state labels. An op_system artifact supplies these.	None
coordinate_system	str	Human-readable coordinate-system declaration.	'supplied state'
fixed	Sequence[int] | NumpyArray | None	Coordinates held fixed, as indices or a boolean mask.	None
invariants	ArrayLike | None	Declared invariant rows in the supplied coordinates.	None
nonnegative	bool | Sequence[bool] | NumpyArray	Coordinates with a one-sided non-negative domain.	False
derivatives	SymbolicDerivativeArtifact | None	Optional op_system symbolic derivative artifact.	None
parameters	Mapping[str, object] | None	Exact runtime bindings for derivatives.	None
jacobian	StabilityJacobian | None	Optional dense full-state Jacobian callback.	None
jacobian_reference	str | None	Required provenance label for jacobian.	None
time_unit	str | None	Unit attached to any reported e-folding timescale.	None

Returns:

Type	Description
LocalStabilityResult	Partial or complete numerical evidence with an explicit status.

Raises:

Type	Description
TypeError	The state, solver flag or configuration is not eager NumPy.
ValueError	Input declarations, shapes, dtypes or controls conflict.

Source code in src/op_engine/equilibrium_stability.py

def local_stability(  # noqa: C901,PLR0911,PLR0912,PLR0913,PLR0914,PLR0915
    rhs: StabilityRhs,
    candidate: NumpyArray | SteadyStateResult,
    *,
    config: LocalStabilityConfig,
    t: float = 0.0,
    coordinates: Sequence[str] | None = None,
    coordinate_system: str = "supplied state",
    fixed: Sequence[int] | NumpyArray | None = None,
    invariants: ArrayLike | None = None,
    nonnegative: bool | Sequence[bool] | NumpyArray = False,
    derivatives: SymbolicDerivativeArtifact | None = None,
    parameters: Mapping[str, object] | None = None,
    jacobian: StabilityJacobian | None = None,
    jacobian_reference: str | None = None,
    time_unit: str | None = None,
) -> LocalStabilityResult:
    """Assess a candidate equilibrium through an eager NumPy linearization.

    The classification applies only to an interior, residual-qualified point
    and, when constraints are declared, only to their tangent manifold.  A
    boundary status deliberately suppresses an interior stability conclusion.

    Args:
        rhs: Dense NumPy right-hand side in the supplied coordinates.
        candidate: NumPy state or NumPy-backed steady-state result.
        config: Explicit residual, spectral, boundary and difference controls.
        t: Frozen evaluation time. A time-dependent RHS is only diagnosed there.
        coordinates: Ordered state labels. An op_system artifact supplies these.
        coordinate_system: Human-readable coordinate-system declaration.
        fixed: Coordinates held fixed, as indices or a boolean mask.
        invariants: Declared invariant rows in the supplied coordinates.
        nonnegative: Coordinates with a one-sided non-negative domain.
        derivatives: Optional op_system symbolic derivative artifact.
        parameters: Exact runtime bindings for ``derivatives``.
        jacobian: Optional dense full-state Jacobian callback.
        jacobian_reference: Required provenance label for ``jacobian``.
        time_unit: Unit attached to any reported e-folding timescale.

    Returns:
        Partial or complete numerical evidence with an explicit status.

    Raises:
        TypeError: The state, solver flag or configuration is not eager NumPy.
        ValueError: Input declarations, shapes, dtypes or controls conflict.
    """
    if not isinstance(config, LocalStabilityConfig):
        msg = "config must be an explicit LocalStabilityConfig"
        raise TypeError(msg)
    solver_converged: bool | None = None
    state_value: object = candidate
    if isinstance(candidate, SteadyStateResult):
        state_value = candidate.state
    if not isinstance(state_value, np.ndarray):
        msg = "local_stability requires an eager NumPy state; transfer explicitly"
        raise TypeError(msg)
    if isinstance(candidate, SteadyStateResult):
        if not isinstance(candidate.converged, np.ndarray):
            msg = "local_stability requires an eager NumPy convergence flag"
            raise TypeError(msg)
        if candidate.converged.shape != ():
            msg = "SteadyStateResult.converged must be a scalar array"
            raise ValueError(msg)
        solver_converged = bool(candidate.converged.item())
    state = state_value
    if state.ndim != 1 or state.size < 1:
        msg = f"candidate must be a non-empty vector; got shape {state.shape}"
        raise ValueError(msg)
    if state.dtype not in {np.dtype(np.float32), np.dtype(np.float64)}:
        msg = "candidate must have NumPy float32 or float64 dtype"
        raise TypeError(msg)
    if not isinstance(t, (int, float)) or isinstance(t, bool) or not math.isfinite(t):
        msg = "t must be a finite real scalar"
        raise ValueError(msg)
    if not isinstance(coordinate_system, str) or not coordinate_system:
        msg = "coordinate_system must be a non-empty string"
        raise ValueError(msg)
    if time_unit is not None and (not isinstance(time_unit, str) or not time_unit):
        msg = "time_unit must be None or a non-empty string"
        raise ValueError(msg)
    n = state.size
    residual_tolerance = _vector_control(
        config.residual_tolerance, n, "residual_tolerance", positive=True
    )
    boundary_tolerance = _vector_control(
        config.boundary_tolerance, n, "boundary_tolerance", positive=False
    )
    fixed_mask = _selection_mask(fixed, n, "fixed")
    nonnegative_mask = _nonnegative_mask(nonnegative, n)
    invariant_rows = _invariant_rows(invariants, n)
    source, names, bound_parameters, model_digest = _source_and_coordinates(
        derivatives,
        jacobian,
        parameters,
        coordinates,
        n,
        jacobian_reference,
    )

    evidence: dict[str, Any] = {
        "state": state.copy(),
        "time": float(t),
        "coordinates": names,
        "coordinate_system": coordinate_system,
        "time_unit": time_unit,
        "jacobian_source": source,
        "jacobian_reference": jacobian_reference,
        "model_digest": model_digest,
        "solver_converged": solver_converged,
        "residual_tolerance": residual_tolerance,
        "boundary_tolerance": boundary_tolerance,
        "fixed_mask": fixed_mask,
        "nonnegative_mask": nonnegative_mask,
        "invariant_rows": invariant_rows,
    }

    def result(
        status: LocalStabilityStatus,
        message: str,
        classification: LocalSpectrum | None = None,
        **updates: object,
    ) -> LocalStabilityResult:
        return LocalStabilityResult(
            status=status,
            classification=classification,
            message=message,
            **(evidence | updates),
        )

    if not np.isfinite(state).all():
        return result(LocalStabilityStatus.NONFINITE_STATE, "candidate is non-finite")
    infeasible = nonnegative_mask & (state < -boundary_tolerance)
    if np.any(infeasible):
        return result(
            LocalStabilityStatus.INFEASIBLE,
            "candidate violates a declared non-negative coordinate",
            domain=EquilibriumDomain.INFEASIBLE,
        )
    try:
        raw_residual = _strict_numpy_vector(rhs(float(t), state), state, "rhs")
    except (
        ArithmeticError,
        AttributeError,
        LookupError,
        RuntimeError,
        TypeError,
        ValueError,
    ) as error:
        return result(
            LocalStabilityStatus.RESIDUAL_FAILURE,
            f"RHS evaluation failed: {type(error).__name__}: {error}",
        )
    evidence["raw_residual"] = raw_residual
    active = ~fixed_mask
    active_residual = raw_residual[active]
    evidence["active_residual"] = active_residual
    if not np.isfinite(raw_residual).all():
        return result(
            LocalStabilityStatus.NONFINITE_RESIDUAL,
            "candidate residual is non-finite",
        )
    if active_residual.size:
        ratios = active_residual / residual_tolerance[active]
        residual_scaled_rms = float(np.sqrt(np.mean(ratios * ratios)))
    else:
        residual_scaled_rms = 0.0
    evidence["residual_scaled_rms"] = residual_scaled_rms
    if solver_converged is False:
        return result(
            LocalStabilityStatus.SOLVE_FAILURE,
            "input SteadyStateResult did not converge",
        )
    if residual_scaled_rms > 1:
        return result(
            LocalStabilityStatus.NOT_EQUILIBRIUM,
            "active residual exceeds its declared tolerance",
        )

    constrained = nonnegative_mask & active
    boundary = constrained & (np.abs(state) <= boundary_tolerance)
    boundary_coordinates = tuple(
        name for name, hit in zip(names, boundary, strict=True) if hit
    )
    evidence["boundary_coordinates"] = boundary_coordinates
    domain = (
        EquilibriumDomain.BOUNDARY
        if np.any(boundary)
        else EquilibriumDomain.INTERIOR
        if np.any(constrained)
        else EquilibriumDomain.UNSPECIFIED
    )
    evidence["domain"] = domain
    if (
        domain is EquilibriumDomain.BOUNDARY
        and source is StabilityJacobianSource.FINITE_DIFFERENCE
    ):
        return result(
            LocalStabilityStatus.BOUNDARY,
            "central finite differences are not evaluated across a declared boundary",
        )

    try:
        finite_difference_steps = None
        if derivatives is not None:
            values = derivatives.evaluate(float(t), state, **bound_parameters)
            matrix = _strict_numpy_matrix(
                values.state, state, "symbolic derivative artifact"
            )
        elif jacobian is not None:
            matrix = _strict_numpy_matrix(jacobian(float(t), state), state, "jacobian")
        else:
            matrix, finite_difference_steps = _finite_difference_jacobian(
                rhs, float(t), state, config.finite_difference_step
            )
    except (
        ArithmeticError,
        AttributeError,
        LookupError,
        RuntimeError,
        TypeError,
        ValueError,
    ) as error:
        return result(
            LocalStabilityStatus.JACOBIAN_FAILURE,
            f"Jacobian evaluation failed: {type(error).__name__}: {error}",
        )
    evidence["finite_difference_steps"] = finite_difference_steps
    evidence["jacobian"] = matrix
    if not np.isfinite(matrix).all():
        return result(
            LocalStabilityStatus.NONFINITE_JACOBIAN,
            "Jacobian contains non-finite entries",
        )

    free = np.flatnonzero(active)
    active_jacobian = matrix[np.ix_(free, free)]
    evidence["active_jacobian"] = active_jacobian
    try:
        full_eigenvalues = np.linalg.eigvals(matrix)
        active_eigenvalues = np.linalg.eigvals(active_jacobian)
    except np.linalg.LinAlgError as error:
        return result(
            LocalStabilityStatus.EIGENSOLVE_FAILED,
            f"full or active eigensolve failed: {error}",
        )
    evidence["full_eigenvalues"] = full_eigenvalues
    evidence["active_eigenvalues"] = active_eigenvalues
    if (
        not np.isfinite(full_eigenvalues).all()
        or not np.isfinite(active_eigenvalues).all()
    ):
        return result(
            LocalStabilityStatus.NONFINITE_EIGENVALUES,
            "full or active eigenspectrum is non-finite",
        )

    try:
        tangent_basis, constraint_rank, constraint_svd_tolerance = _constraint_basis(
            fixed_mask, invariant_rows
        )
    except (np.linalg.LinAlgError, ValueError) as error:
        return result(
            LocalStabilityStatus.INCONSISTENT_INVARIANTS,
            f"constraint reduction failed: {error}",
        )
    evidence["constraint_rank"] = constraint_rank
    evidence["constraint_svd_tolerance"] = constraint_svd_tolerance
    evidence["tangent_basis"] = tangent_basis

    invariant_rank = invariant_rows.shape[0]
    evidence["invariant_rank"] = invariant_rank
    evidence["expected_conservation_neutral_modes"] = invariant_rank
    if tangent_basis.shape[1] == 0:
        evidence["observed_active_neutral_modes"] = 0
        return result(
            LocalStabilityStatus.NO_DYNAMIC_DIRECTIONS,
            "declared constraints leave no dynamic tangent directions",
        )
    active_spectral_tolerance = (
        config.spectral_absolute_tolerance
        + config.spectral_relative_tolerance * float(np.max(np.abs(active_eigenvalues)))
    )
    evidence["observed_active_neutral_modes"] = int(
        np.count_nonzero(np.abs(active_eigenvalues) <= active_spectral_tolerance)
    )
    if invariant_rank:
        invariant_free = invariant_rows[:, active]
        try:
            active_jacobian_scale = float(np.linalg.norm(active_jacobian, ord=2))
            conservation_error = float(
                np.linalg.norm(invariant_free @ active_jacobian, ord=2)
            )
        except np.linalg.LinAlgError as error:
            return result(
                LocalStabilityStatus.INCONSISTENT_INVARIANTS,
                f"invariant norm evaluation failed: {error}",
            )
        conservation_tolerance = (
            config.spectral_absolute_tolerance
            + config.spectral_relative_tolerance * active_jacobian_scale
        )
        invariant_flow_error = float(np.max(np.abs(invariant_free @ active_residual)))
        invariant_flow_tolerance = float(
            np.max(np.abs(invariant_free) @ residual_tolerance[active])
        )
        evidence.update({
            "conservation_error": conservation_error,
            "conservation_tolerance": conservation_tolerance,
            "invariant_flow_error": invariant_flow_error,
            "invariant_flow_tolerance": invariant_flow_tolerance,
        })
        if conservation_error > conservation_tolerance:
            return result(
                LocalStabilityStatus.INCONSISTENT_INVARIANTS,
                "declared invariant rows do not annihilate the local Jacobian",
            )

    tangent_jacobian = tangent_basis.T @ matrix.astype(np.float64) @ tangent_basis
    evidence["tangent_jacobian"] = tangent_jacobian
    try:
        tangent_eigenvalues = np.linalg.eigvals(tangent_jacobian)
    except np.linalg.LinAlgError as error:
        return result(
            LocalStabilityStatus.EIGENSOLVE_FAILED,
            f"tangent eigensolve failed: {error}",
        )
    evidence["tangent_eigenvalues"] = tangent_eigenvalues
    if not np.isfinite(tangent_eigenvalues).all():
        return result(
            LocalStabilityStatus.NONFINITE_EIGENVALUES,
            "tangent eigenspectrum is non-finite",
        )
    effective_tolerance = (
        config.spectral_absolute_tolerance
        + config.spectral_relative_tolerance
        * float(np.max(np.abs(tangent_eigenvalues)))
    )
    real_parts = np.real(tangent_eigenvalues)
    spectral_abscissa = float(np.max(real_parts))
    evidence["effective_spectral_tolerance"] = effective_tolerance
    evidence["spectral_abscissa"] = spectral_abscissa
    if domain is EquilibriumDomain.BOUNDARY:
        return result(
            LocalStabilityStatus.BOUNDARY,
            "active boundary requires a one-sided tangent-cone analysis",
        )
    if np.any(real_parts > effective_tolerance):
        classification = LocalSpectrum.UNSTABLE
        leading_index = int(np.argmax(real_parts))
        leading_eigenvalue = complex(tangent_eigenvalues[leading_index])
        evidence.update({
            "leading_eigenvalue": leading_eigenvalue,
            "leading_timescale": 1 / spectral_abscissa,
            "timescale_kind": StabilityTimescaleKind.GROWTH,
        })
    elif np.all(real_parts < -effective_tolerance):
        classification = LocalSpectrum.ASYMPTOTICALLY_STABLE
        leading_index = int(np.argmax(real_parts))
        leading_eigenvalue = complex(tangent_eigenvalues[leading_index])
        evidence.update({
            "leading_eigenvalue": leading_eigenvalue,
            "leading_timescale": -1 / spectral_abscissa,
            "timescale_kind": StabilityTimescaleKind.DECAY,
        })
    elif np.all(np.abs(real_parts) <= effective_tolerance):
        classification = LocalSpectrum.NEUTRAL_SPECTRUM
    else:
        classification = LocalSpectrum.NONHYPERBOLIC
    return result(
        LocalStabilityStatus.ANALYZED,
        (
            "classification is local to the declared constraint manifold"
            if invariant_rank or np.any(fixed_mask)
            else "classification is a local interior linearization"
        ),
        classification,
    )

matrix_grouped_count(group_matrix)

Perform grouped count using a group matrix.

Parameters:

Name	Type	Description	Default
group_matrix	csr_matrix | DenseOperator	2D group matrix (csr_matrix or dense ndarray).	required

Returns:

Type	Description
NDArray[floating]	A 1D array containing the counts for each group.

Source code in src/op_engine/matrix_ops.py

def matrix_grouped_count(
    group_matrix: csr_matrix | DenseOperator,
) -> NDArray[np.floating]:
    """
    Perform grouped count using a group matrix.

    Args:
        group_matrix: 2D group matrix (csr_matrix or dense ndarray).

    Returns:
        A 1D array containing the counts for each group.
    """
    n_groups = group_matrix.shape[0]
    if issparse(group_matrix) and n_groups >= _DISPATCH_THRESHOLD:
        return _matrix_grouped_count_sparse(group_matrix)
    if issparse(group_matrix):
        dense_matrix = group_matrix.toarray()
    else:
        dense_matrix = np.asarray(group_matrix)
    return _matrix_grouped_count_dense(cast("DenseOperator", dense_matrix))

matrix_grouped_sum(group_matrix, values)

Perform grouped sum using a group matrix.

Parameters:

Name	Type	Description	Default
group_matrix	csr_matrix | DenseOperator	2D group matrix (csr_matrix or dense ndarray).	required
values	NDArray[floating]	1D array of values to be summed.	required

Returns:

Type	Description
NDArray[floating]	A 1D array containing the grouped sums.

Source code in src/op_engine/matrix_ops.py

def matrix_grouped_sum(
    group_matrix: csr_matrix | DenseOperator,
    values: NDArray[np.floating],
) -> NDArray[np.floating]:
    """
    Perform grouped sum using a group matrix.

    Args:
        group_matrix: 2D group matrix (csr_matrix or dense ndarray).
        values: 1D array of values to be summed.

    Returns:
        A 1D array containing the grouped sums.
    """
    n_groups = group_matrix.shape[0]
    if issparse(group_matrix) and n_groups >= _DISPATCH_THRESHOLD:
        return _matrix_grouped_sum_sparse(group_matrix, values)
    if issparse(group_matrix):
        dense_matrix = np.asarray(group_matrix.toarray())
    else:
        dense_matrix = np.asarray(group_matrix)
    return _matrix_grouped_sum_dense(cast("DenseOperator", dense_matrix), values)

matrix_masked_sum(mask_matrix, data)

Perform masked sum using a mask matrix and data array.

Parameters:

Name	Type	Description	Default
mask_matrix	csr_matrix | DenseOperator	2D mask matrix (csr_matrix or dense ndarray).	required
data	NDArray[floating]	1D or 2D data array to be masked and summed.	required

Returns:

Type	Description
NDArray[floating]	A 1D or 2D array containing the masked sums.

Source code in src/op_engine/matrix_ops.py

def matrix_masked_sum(
    mask_matrix: csr_matrix | DenseOperator,
    data: NDArray[np.floating],
) -> NDArray[np.floating]:
    """
    Perform masked sum using a mask matrix and data array.

    Args:
        mask_matrix: 2D mask matrix (csr_matrix or dense ndarray).
        data: 1D or 2D data array to be masked and summed.

    Returns:
        A 1D or 2D array containing the masked sums.
    """
    n_masks = mask_matrix.shape[0]
    if issparse(mask_matrix) and n_masks >= _DISPATCH_THRESHOLD:
        return _matrix_masked_sum_sparse(mask_matrix, data)
    if issparse(mask_matrix):
        dense_matrix = np.asarray(mask_matrix.toarray())
    else:
        dense_matrix = np.asarray(mask_matrix)
    return _matrix_masked_sum_dense(dense_matrix, data)

observation_from_compiled(observable)

Adapt a compiled observable structurally, without a runtime compiler dependency.

The input trajectory must retain the compiler's flattened state order. The adapter reshapes each state to a vector and delegates all spatial reductions and coordinate selection to the compiler, applying no additional weights.

Parameters:

Name	Type	Description	Default
observable	CompiledObservableLike	Artifact exposing op_system's compiled-observable surface.	required

Returns:

Type	Description
ObservationSpec	Native evaluator and detached scientific metadata.

Source code in src/op_engine/observations.py

def observation_from_compiled(observable: CompiledObservableLike) -> ObservationSpec:
    """Adapt a compiled observable structurally, without a runtime compiler dependency.

    The input trajectory must retain the compiler's flattened state order. The
    adapter reshapes each state to a vector and delegates all spatial reductions
    and coordinate selection to the compiler, applying no additional weights.

    Args:
        observable: Artifact exposing op_system's compiled-observable surface.

    Returns:
        Native evaluator and detached scientific metadata.
    """
    evaluate = observable.eval_fn

    def eval_fn(t: object, y: object, **params: object) -> object:
        xp = array_namespace(y)
        return evaluate(t, xp.reshape(y, (-1,)), **params)

    return ObservationSpec(
        observable.definition.name,
        eval_fn,
        ObservationMetadata(
            tuple(observable.output_shape),
            tuple((axis, tuple(labels)) for axis, labels in observable.coordinates),
            observable.definition.units,
            tuple(observable.state_dependencies),
            tuple((name, tuple(shape)) for name, shape in observable.parameter_shapes),
            tuple(
                (
                    axis,
                    ObservationMeasure(
                        measure.kind, tuple(measure.weights), measure.units
                    ),
                )
                for axis, measure in observable.measure_requirements
            ),
            observable.time_dependent,
            observable.definition.description,
        ),
    )

observe_model(core, observations, *, params=None)

Evaluate a completed ModelCore's stored history using its output grid.

Parameters:

Name	Type	Description	Default
core	ModelCore	Completed state manager configured to store its full trajectory.	required
observations	Sequence[ObservationSpec]	Named evaluators and their output metadata.	required
params	Mapping[str, object] | None	Optional observation parameters.	None

Returns:

Type	Description
ObservationResult	Native named series at the model's stored output times.

Raises:

Type	Description
ValueError	If history is unavailable or the model has not completed.

Source code in src/op_engine/observations.py

def observe_model(
    core: ModelCore,
    observations: Sequence[ObservationSpec],
    *,
    params: Mapping[str, object] | None = None,
) -> ObservationResult:
    """Evaluate a completed ModelCore's stored history using its output grid.

    Args:
        core: Completed state manager configured to store its full trajectory.
        observations: Named evaluators and their output metadata.
        params: Optional observation parameters.

    Returns:
        Native named series at the model's stored output times.

    Raises:
        ValueError: If history is unavailable or the model has not completed.
    """
    if not core.store_history or core.state_array is None:
        msg = "observations require stored history; streaming is not available"
        raise ValueError(msg)
    if core.current_step != core.n_timesteps - 1:
        msg = "observations require a completed model trajectory"
        raise ValueError(msg)
    return observe_solution(
        core.time_grid, core.state_array, observations, params=params
    )

observe_solution(times, trajectory, observations, *, params=None)

Evaluate native named outputs at every supplied stored output time.

Time is a static, finite, strictly increasing host grid. Each callback must return an array in the trajectory namespace and dtype with its declared shape. Native values are stacked without conversion or gradient detachment. Inputs and returned arrays should be treated as immutable.

Parameters:

Name	Type	Description	Default
times	Sequence[float] | ndarray	Stored output times, independent of internal solver step times.	required
trajectory	Array	Native state history with a leading output-time dimension.	required
observations	Sequence[ObservationSpec]	Unique named evaluators with static output metadata.	required
params	Mapping[str, object] | None	Parameters passed to each evaluator, including observation-only ones.	None

Returns:

Type	Description
ObservationResult	Native series and evaluation coverage. No interpolation/error estimate.

Raises:

Type	Description
ValueError	If grids, trajectory shape, metadata or output shapes disagree.
TypeError	If states/outputs violate the real floating native-array contract.

Source code in src/op_engine/observations.py

def observe_solution(
    times: Sequence[float] | np.ndarray,
    trajectory: Array,
    observations: Sequence[ObservationSpec],
    *,
    params: Mapping[str, object] | None = None,
) -> ObservationResult:
    """Evaluate native named outputs at every supplied stored output time.

    Time is a static, finite, strictly increasing host grid. Each callback must
    return an array in the trajectory namespace and dtype with its declared
    shape. Native values are stacked without conversion or gradient detachment.
    Inputs and returned arrays should be treated as immutable.

    Args:
        times: Stored output times, independent of internal solver step times.
        trajectory: Native state history with a leading output-time dimension.
        observations: Unique named evaluators with static output metadata.
        params: Parameters passed to each evaluator, including observation-only ones.

    Returns:
        Native series and evaluation coverage. No interpolation/error estimate.

    Raises:
        ValueError: If grids, trajectory shape, metadata or output shapes disagree.
        TypeError: If states/outputs violate the real floating native-array contract.
    """
    grid = np.asarray(times)
    if (
        grid.ndim != 1
        or grid.size == 0
        or grid.dtype.kind not in "fiu"
        or not np.all(np.isfinite(grid))
        or not np.all(grid[1:] > grid[:-1])
    ):
        msg = "observation times must be a nonempty finite strictly increasing 1D grid"
        raise ValueError(msg)
    if not trajectory.shape or trajectory.shape[0] != grid.size:
        msg = "trajectory must have one leading state row per observation time"
        raise ValueError(msg)
    xp = array_namespace(trajectory)
    if not xp.isdtype(trajectory.dtype, "real floating"):
        msg = "observation trajectory must have a real floating dtype"
        raise TypeError(msg)
    _validate_specs(observations)
    static_times = tuple(float(time) for time in grid)
    states = cast("_Indexable", trajectory)
    series = {}
    for spec in observations:
        values = []
        for index, time in enumerate(static_times):
            value = spec.eval_fn(
                time, states[index], **({} if params is None else params)
            )
            if (
                array_namespace(value) is not xp
                or getattr(value, "dtype", None) != trajectory.dtype
            ):
                msg = (
                    f"observation {spec.name!r} must preserve trajectory "
                    "namespace and dtype"
                )
                raise TypeError(msg)
            if getattr(value, "shape", None) != spec.metadata.output_shape:
                msg = (
                    f"observation {spec.name!r} must have shape "
                    f"{spec.metadata.output_shape}"
                )
                raise ValueError(msg)
            values.append(value)
        series[spec.name] = ObservationSeries(
            spec, cast("Array", xp.stack(values, axis=0))
        )
    return ObservationResult(
        static_times,
        MappingProxyType(series),
        ObservationDiagnostics(len(observations) * len(static_times)),
    )

reduce_observations(samples, observations, *, params=None)

Reduce a finite native sample stream in one pass without storing a trajectory.

Samples define the exact window endpoints. Native initial/final values, changes and composite-trapezoid integrals preserve evaluator shapes and gradients. A single sample has zero duration, change and integral. Callback outputs are snapshotted so a producer may reuse its numerical state buffer.

The reducer retains a constant number of output arrays and one current state reference. Producer storage and reverse-mode computation graphs are separate: no constant-memory gradient guarantee is implied. There is no persistent checkpoint or partial-result recovery; an interrupted call must restart from a producer-owned replay/checkpoint. Callbacks must be pure.

Parameters:

Name	Type	Description	Default
samples	Iterable[tuple[float, Array]]	Finite, nonempty iterable of static times and native states. Times must increase strictly; state shape, namespace and real floating dtype stay fixed. Only the supplied times define quadrature samples.	required
observations	Sequence[ObservationSpec]	Unique named evaluators and scientific output metadata.	required
params	Mapping[str, object] | None	Observation parameters passed to every callback.	None

Returns:

Type	Description
ObservationSummary	Bounded-size native reductions with window, sample count and diagnostics.
ObservationSummary	Integrals inherit trajectory and sampled-quadrature error; no dense-output
ObservationSummary	order, event reconstruction or separate error estimate is asserted.

Raises:

Type	Description
ValueError	If the stream is empty, intervals are invalid or shapes change.
TypeError	If states or outputs violate the native namespace/dtype contract.

Source code in src/op_engine/observation_reductions.py

def reduce_observations(
    samples: Iterable[tuple[float, Array]],
    observations: Sequence[ObservationSpec],
    *,
    params: Mapping[str, object] | None = None,
) -> ObservationSummary:
    """Reduce a finite native sample stream in one pass without storing a trajectory.

    Samples define the exact window endpoints. Native initial/final values,
    changes and composite-trapezoid integrals preserve evaluator shapes and
    gradients. A single sample has zero duration, change and integral. Callback
    outputs are snapshotted so a producer may reuse its numerical state buffer.

    The reducer retains a constant number of output arrays and one current
    state reference. Producer storage and reverse-mode computation graphs are
    separate: no constant-memory gradient guarantee is implied. There is no
    persistent checkpoint or partial-result recovery; an interrupted call must
    restart from a producer-owned replay/checkpoint. Callbacks must be pure.

    Args:
        samples: Finite, nonempty iterable of static times and native states.
            Times must increase strictly; state shape, namespace and real floating
            dtype stay fixed. Only the supplied times define quadrature samples.
        observations: Unique named evaluators and scientific output metadata.
        params: Observation parameters passed to every callback.

    Returns:
        Bounded-size native reductions with window, sample count and diagnostics.
        Integrals inherit trajectory and sampled-quadrature error; no dense-output
        order, event reconstruction or separate error estimate is asserted.

    Raises:
        ValueError: If the stream is empty, intervals are invalid or shapes change.
        TypeError: If states or outputs violate the native namespace/dtype contract.
    """
    selection = tuple(observations)
    iterator = iter(samples)
    try:
        time, state = next(iterator)
    except StopIteration as error:
        msg = "observation summary requires at least one sample"
        raise ValueError(msg) from error
    first = _sample(_static_time(time), state, selection, params)
    layout = _StateLayout(state.shape, state.dtype, array_namespace(state))
    values = _initial_values(first)
    last_time = first.times[0]
    count = 1
    for sample_time, state in iterator:
        time = _static_time(sample_time)
        duration = time - last_time
        if not np.isfinite(duration) or duration <= 0:
            msg = "summary sample times must increase with finite intervals"
            raise ValueError(msg)
        if state.shape != layout.shape:
            msg = "summary samples must preserve state shape"
            raise ValueError(msg)
        if (
            state.dtype != layout.dtype
            or array_namespace(state) is not layout.namespace
        ):
            msg = "summary samples must preserve state namespace and dtype"
            raise TypeError(msg)
        current = _sample(time, state, selection, params)
        values = _advance_values(values, current, duration)
        last_time = time
        count += 1
    return ObservationSummary(
        (first.times[0], last_time),
        count,
        MappingProxyType(values),
        ObservationDiagnostics(count * len(selection)),
    )

register_sparse_adapter(adapter, *, replace=False)

Register an operator adapter in the existing sparse dispatch registry.

Replacement clears that ecosystem's cached factors. Predicates must be pure; overlapping pair matches are rejected when dispatching a solve.

Raises:

Type	Description
ValueError	If an identifier already exists without explicit replacement.

Source code in src/op_engine/matrix_ops.py

def register_sparse_adapter(adapter: SparseAdapter, *, replace: bool = False) -> None:
    """Register an operator adapter in the existing sparse dispatch registry.

    Replacement clears that ecosystem's cached factors. Predicates must be
    pure; overlapping pair matches are rejected when dispatching a solve.

    Raises:
        ValueError: If an identifier already exists without explicit replacement.
    """
    if adapter.ecosystem_id in _SPARSE_ADAPTERS and not replace:
        msg = f"adapter {adapter.ecosystem_id!r} is already registered"
        raise ValueError(msg)
    _clear_adapter_cache(adapter.ecosystem_id)
    _SPARSE_ADAPTERS[adapter.ecosystem_id] = adapter

require_converged(result)

Return a converged result or raise at an eager host boundary.

This function intentionally converts the scalar convergence array to a Python boolean. Do not call it from jax.jit or another traced region; compiled integrations must return diagnostics and validate them afterward.

Parameters:

Name	Type	Description	Default
result	NonlinearSolveResult	Result to validate.	required

Returns:

Type	Description
NonlinearSolveResult	The unchanged converged result.

Raises:

Type	Description
NonlinearConvergenceError	If the result did not converge.

Source code in src/op_engine/nonlinear_solver.py

def require_converged(result: NonlinearSolveResult) -> NonlinearSolveResult:
    """Return a converged result or raise at an eager host boundary.

    This function intentionally converts the scalar convergence array to a
    Python boolean. Do not call it from ``jax.jit`` or another traced region;
    compiled integrations must return diagnostics and validate them afterward.

    Args:
        result: Result to validate.

    Returns:
        The unchanged converged result.

    Raises:
        NonlinearConvergenceError: If the result did not converge.
    """
    if not bool(result.diagnostics.converged.item()):
        raise NonlinearConvergenceError(result)
    return result

require_linear_converged(result)

Return a successful eager solution or raise with per-column diagnostics.

Raises:

Type	Description
LinearSolveError	If any RHS column failed the checked solve.

Source code in src/op_engine/linear_operator.py

def require_linear_converged(result: LinearSolveResult) -> NDArray[np.floating]:
    """Return a successful eager solution or raise with per-column diagnostics.

    Raises:
        LinearSolveError: If any RHS column failed the checked solve.
    """
    if not np.all(result.diagnostics.converged):
        raise LinearSolveError(result)
    return result.value

require_steady_state(result)

Return a converged result or raise.

This converts the convergence flag to a Python boolean; do not call it inside jax.jit.

Parameters:

Name	Type	Description	Default
result	SteadyStateResult	Result to check.	required

Returns:

Type	Description
SteadyStateResult	result unchanged.

Raises:

Type	Description
SteadyStateConvergenceError	If the solve did not converge.

Source code in src/op_engine/steady_state.py

def require_steady_state(result: SteadyStateResult) -> SteadyStateResult:
    """Return a converged result or raise.

    This converts the convergence flag to a Python boolean; do not call it
    inside ``jax.jit``.

    Args:
        result: Result to check.

    Returns:
        ``result`` unchanged.

    Raises:
        SteadyStateConvergenceError: If the solve did not converge.
    """
    if not bool(result.converged.item()):
        raise SteadyStateConvergenceError(result)
    return result

scipy_banded_adapter(*, ecosystem_id='scipy-banded', options=None)

Build an explicitly registered, uncached eager banded direct adapter.

Returns:

Type	Description
SparseAdapter	Existing-registry adapter recognizing a BandedOperator left operand.

Source code in src/op_engine/scipy_operator.py

def scipy_banded_adapter(
    *,
    ecosystem_id: str = "scipy-banded",
    options: LinearSolveOptions | None = None,
) -> SparseAdapter:
    """Build an explicitly registered, uncached eager banded direct adapter.

    Returns:
        Existing-registry adapter recognizing a BandedOperator left operand.
    """
    return SparseAdapter(
        ecosystem_id,
        lambda op: isinstance(op, BandedOperator),
        linear_solver=_banded_solve,
        solve_options=options if options is not None else LinearSolveOptions(),
        accepts_right=_accepted_right,
    )

scipy_gmres_adapter(*, ecosystem_id='scipy-gmres', restart=20, options=None, preconditioner=None, context=None)

Build an explicit eager matrix-free adapter in the existing registry.

Returns:

Type	Description
SparseAdapter	Adapter for ActionOperator left operands, with left inverse preconditioning.

Source code in src/op_engine/scipy_operator.py

def scipy_gmres_adapter(
    *,
    ecosystem_id: str = "scipy-gmres",
    restart: int = 20,
    options: LinearSolveOptions | None = None,
    preconditioner: ActionOperator | None = None,
    context: object | None = None,
) -> SparseAdapter:
    """Build an explicit eager matrix-free adapter in the existing registry.

    Returns:
        Adapter for ActionOperator left operands, with left inverse preconditioning.
    """
    return SparseAdapter(
        ecosystem_id,
        lambda op: isinstance(op, ActionOperator),
        linear_solver=ScipyGmresSolver(restart),
        solve_options=options if options is not None else LinearSolveOptions(),
        preconditioner=preconditioner,
        context=context,
        accepts_right=_accepted_right,
    )

sink_states(rhs, states, *, t=0.0, jacobian=None, tol=1e-09)

Return a mask of states no derivative depends on.

Cumulative counters and absorbing compartments that feed nothing back have a zero Jacobian column. They never reach a zero derivative while their inflow is positive, so hold them fixed in :func:steady_state.

Parameters:

Name	Type	Description	Default
rhs	SteadyStateRhs	rhs(t, y) -> dy/dt.	required
states	Sequence[Array]	Sample states.	required
t	float	Evaluation time.	0.0
jacobian	SteadyStateJacobian | None	Dense Jacobian callback; central differences by default.	None
tol	float	Largest column entry, relative to the largest Jacobian entry, still treated as zero.	1e-09

Returns:

Type	Description
ndarray	Boolean (n_state,) mask.

Source code in src/op_engine/steady_state.py

def sink_states(
    rhs: SteadyStateRhs,
    states: Sequence[Array],
    *,
    t: float = 0.0,
    jacobian: SteadyStateJacobian | None = None,
    tol: float = 1e-9,
) -> np.ndarray:
    """Return a mask of states no derivative depends on.

    Cumulative counters and absorbing compartments that feed nothing back have
    a zero Jacobian column. They never reach a zero derivative while their
    inflow is positive, so hold them fixed in :func:`steady_state`.

    Args:
        rhs: ``rhs(t, y) -> dy/dt``.
        states: Sample states.
        t: Evaluation time.
        jacobian: Dense Jacobian callback; central differences by default.
        tol: Largest column entry, relative to the largest Jacobian entry,
            still treated as zero.

    Returns:
        Boolean ``(n_state,)`` mask.
    """
    jacobian = jacobian or _central_difference_jacobian(rhs)
    columns = np.zeros(0)
    for y in states:
        magnitude = np.abs(np.asarray(jacobian(t, y), dtype=np.float64)).max(axis=0)
        columns = magnitude if columns.size == 0 else np.maximum(columns, magnitude)
    return columns <= tol * max(float(columns.max(initial=0.0)), 1e-300)

smooth(x, alpha=0.02, out=None)

Apply simple smoothing along the last axis.

Parameters:

Name	Type	Description	Default
x	NDArray[floating]	Input array to smooth.	required
alpha	float	Smoothing factor between 0 and 1.	0.02
out	NDArray[floating] | None	Optional output array to store the result.	None

Returns:

Type	Description
NDArray[floating]	Smoothed array with the same shape as x.

Source code in src/op_engine/matrix_ops.py

def smooth(
    x: NDArray[np.floating],
    alpha: float = 0.02,
    out: NDArray[np.floating] | None = None,
) -> NDArray[np.floating]:
    """
    Apply simple smoothing along the last axis.

    Args:
        x: Input array to smooth.
        alpha: Smoothing factor between 0 and 1.
        out: Optional output array to store the result.

    Returns:
        Smoothed array with the same shape as x.
    """
    x_arr = np.asarray(x)
    smoothed = (1.0 - alpha) * x_arr + alpha * x_arr.mean(axis=-1, keepdims=True)
    if out is not None:
        np.copyto(out, smoothed)
        return out
    return cast("NDArray[np.floating]", smoothed)

solve_linear(request, solver)

Call an eager solver and verify every original, unpreconditioned residual.

Returns:

Type	Description
LinearSolveResult	Candidate solution and per-column diagnostics, including failures.

Raises:

Type	Description
TypeError	If namespace, dtype or callback metadata is invalid.
ValueError	If operator/RHS/result dimensions or iteration counts disagree.

Source code in src/op_engine/linear_operator.py

def solve_linear(
    request: LinearSolveRequest, solver: LinearSolver
) -> LinearSolveResult:
    """Call an eager solver and verify every original, unpreconditioned residual.

    Returns:
        Candidate solution and per-column diagnostics, including failures.

    Raises:
        TypeError: If namespace, dtype or callback metadata is invalid.
        ValueError: If operator/RHS/result dimensions or iteration counts disagree.
    """
    rhs = request.rhs
    if not isinstance(rhs, np.ndarray) or rhs.dtype not in (np.float32, np.float64):  # noqa: PLR6201 - dtype equality is not hash equality
        msg = "checked linear solves require eager NumPy float32/float64 RHS arrays"
        raise TypeError(msg)
    shape = request.operator.shape
    if request.operator.dtype != rhs.dtype or (
        isinstance(request.operator, ActionOperator)
        and request.operator.namespace is not array_namespace(rhs)
    ):
        msg = "operator must match the NumPy RHS dtype and namespace"
        raise TypeError(msg)
    if len(shape) != 2 or shape[0] != shape[1] or len(rhs.shape) not in {1, 2}:
        msg = "checked linear solves require a square operator and vector/matrix RHS"
        raise ValueError(msg)
    if shape[0] != rhs.shape[0]:
        msg = "operator and RHS dimensions disagree"
        raise ValueError(msg)
    preconditioner = request.preconditioner
    if preconditioner is not None and (
        preconditioner.shape != shape
        or preconditioner.dtype != rhs.dtype
        or preconditioner.namespace is not array_namespace(rhs)
    ):
        msg = "inverse preconditioner must match operator shape and RHS dtype/namespace"
        raise ValueError(msg)
    candidate = solver(request)
    _numpy_like(candidate.value, rhs, "solve result")
    if not isinstance(candidate.success, bool):
        msg = "solver success must be a Python bool"
        raise TypeError(msg)
    if (
        not isinstance(candidate.iterations, Integral)
        or isinstance(candidate.iterations, bool)
        or not 0 <= candidate.iterations <= request.options.max_iterations
    ):
        msg = "solver iterations must be an integer within the requested budget"
        raise ValueError(msg)
    product = request.operator @ candidate.value
    _numpy_like(product, rhs, "operator result")
    # Accumulate float32 norms in float64 to avoid avoidable sum-of-squares overflow.
    residual = rhs.astype(np.float64) - cast("np.ndarray", product).astype(np.float64)
    residual_norm = np.atleast_1d(np.hypot.reduce(np.abs(residual), axis=0))
    rhs_norm = np.atleast_1d(np.hypot.reduce(np.abs(rhs.astype(np.float64)), axis=0))
    tolerance = np.maximum(request.options.atol, request.options.rtol * rhs_norm)
    converged = (
        np.isfinite(residual_norm)
        & np.isfinite(rhs_norm)
        & (residual_norm <= tolerance)
        & np.atleast_1d(np.all(np.isfinite(candidate.value), axis=0))
        & candidate.success
    )
    diagnostics = LinearSolveDiagnostics(
        converged, residual_norm, tolerance, candidate.iterations, candidate.message
    )
    return LinearSolveResult(candidate.value, diagnostics)

solve_sensitivities(system, parameters, initial, time_grid, *, options=None)

Integrate S' = (df/dy) S + df/dtheta with existing explicit kernels.

The compiler is optional: any correctly ordered derivative artifact can satisfy the protocol. Callbacks/mesh are static under JIT; native scalar bindings and both initial arrays retain their graphs. Batch independent systems through native vmap. Forcing jumps, events, history dependence, stochastic paths and operator solvers need separately qualified methods.

Returns:

Type	Description
SensitivitySolution	Native retained outputs and measured sensitivities in physical coordinates.

Raises:

Type	Description
ValueError	Layout, mesh or declared storage limits disagree.

Source code in src/op_engine/sensitivities.py

def solve_sensitivities(
    system: SensitivitySystem,
    parameters: Mapping[str, object],
    initial: SensitivityInitial,
    time_grid: NDArray[np.floating],
    *,
    options: SensitivityOptions | None = None,
) -> SensitivitySolution:
    """Integrate ``S' = (df/dy) S + df/dtheta`` with existing explicit kernels.

    The compiler is optional: any correctly ordered derivative artifact can
    satisfy the protocol. Callbacks/mesh are static under JIT; native scalar
    bindings and both initial arrays retain their graphs. Batch independent
    systems through native vmap. Forcing jumps, events, history dependence,
    stochastic paths and operator solvers need separately qualified methods.

    Returns:
        Native retained outputs and measured sensitivities in physical coordinates.

    Raises:
        ValueError: Layout, mesh or declared storage limits disagree.
    """
    opts = options or SensitivityOptions()
    bindings = _validate_inputs(system, parameters, initial)
    xp = array_namespace(initial.state)
    n, p, m = (
        len(system.derivatives.state_names),
        len(system.parameter_names),
        len(system.derivatives.observable_names),
    )
    grid = np.asarray(time_grid, dtype=np.float64)
    if (
        grid.ndim != 1
        or not grid.size
        or not np.all(np.isfinite(grid))
        or np.any(np.diff(grid) <= 0)
    ):
        msg = "Sensitivity mesh must be finite and strictly increasing"
        raise ValueError(msg)
    count = grid.size if opts.store_history else 1
    dimension = n * (1 + p)
    if (
        dimension > opts.max_augmented_states
        or count * (n + n * p + m * p) * (xp.finfo(initial.state.dtype).bits // 8)
        > opts.max_output_bytes
    ):
        msg = "Sensitivity solve exceeds its augmented/output storage limit"
        raise ValueError(msg)
    core = ModelCore(
        dimension,
        1,
        grid,
        options=ModelCoreOptions(
            dtype=np.dtype(str(initial.state.dtype).removeprefix("torch.")),
            store_history=opts.store_history,
        ),
    )
    core.set_initial_state(
        xp.reshape(
            xp.concat((
                initial.state,
                xp.reshape(initial.parameter_derivatives, (-1,)),
            )),
            (-1, 1),
        )
    )
    CoreSolver(core).run(_augmented_rhs(system, bindings), config=opts.solver)
    return _solution(system, bindings, core, opts)

unregister_sparse_adapter(ecosystem_id)

Remove an adapter and release all of its cached factors.

Raises:

Type	Description
KeyError	If the adapter is not registered.

Source code in src/op_engine/matrix_ops.py

def unregister_sparse_adapter(ecosystem_id: str) -> None:
    """Remove an adapter and release all of its cached factors.

    Raises:
        KeyError: If the adapter is not registered.
    """  # noqa: DOC502 - dictionary removal raises KeyError
    del _SPARSE_ADAPTERS[ecosystem_id]
    _clear_adapter_cache(ecosystem_id)