Thinning Ssa

thinning_ssa

Exact bounded thinning for time-dependent reaction propensities.

NumpyThinningSampler(seed=None)

Seeded stateful thinning sampler for NumPy reaction networks.

Create a NumPy random generator with an optional seed.

Source code in src/op_engine/thinning_ssa.py

def __init__(self, seed: int | None = None) -> None:
    """Create a NumPy random generator with an optional seed."""
    self._rng = np.random.default_rng(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, *, rate_bound, config=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	Namespace-preserving candidate sampler.	required
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

Raises:

Type	Description
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,
    *,
    rate_bound: float | RateBoundFunction,
    config: ThinningSSAConfig | 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.

    Raises:
        RuntimeError: If the candidate limit is exceeded, time cannot
            advance, or a reaction produces invalid populations.
    """
    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)