Stochastic Solver

stochastic_solver

Stochastic reaction-network integration methods.

This module keeps stochastic event semantics separate from deterministic ODE right-hand sides. Tau-leaping and exact direct SSA consume the same validated reaction network: reaction-channel propensities, a stoichiometric matrix, and one configured reaction axis. Random sampling is injected so the numerical methods remain independent of NumPy, JAX, or another array ecosystem's PRNG API.

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)

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),
    )

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``."""

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."""

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)