Adaptive Tau

adaptive_tau

Bounded adaptive tau-leaping for stochastic reaction networks.

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, ssa_sampler, *, config=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	Backend-specific noncritical Poisson sampler.	required
ssa_sampler	SSASampler	Backend-specific exact-event sampler.	required
config	AdaptiveTauLeapingConfig | None	Optional adaptive, retry, and forcing controls.	None

Raises:

Type	Description
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,
    ssa_sampler: SSASampler,
    *,
    config: AdaptiveTauLeapingConfig | 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.

    Raises:
        RuntimeError: If a step guard, retry guard, or non-negativity
            invariant fails.
    """
    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)