Thinning Ssa

thinning_ssa

Exact bounded thinning for time-dependent reaction propensities.

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

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

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)