Aggregators

Aggregators implement the stepping logic for simulating jump processes.

Abstract Base Classes

jumpax.AbstractAggregator

Abstract base class for all aggregators (solvers).

Aggregators implement the stepping logic for jump processes, determining when jumps occur and which reaction(s) fire(s).

init(jumps, u0: Shaped[Array, '?*u'], args: PyTree[Any], key: Key[Array, '']) -> ~_SolverState

Initialize solver-specific state.

Arguments:

• jumps: the jump problem
• u0: initial state
• args: static arguments
• key: random key

Returns:

Initial solver state for this aggregator.

step(jumps, t: Real[ArrayLike, ''], u: Shaped[Array, '?*u'], args: PyTree[Any], key: Key[Array, ''], t1: Real[ArrayLike, ''], jump_state: PyTree[Any], solver_state: ~_SolverState) -> tuple

Advance one step of the simulation.

Arguments:

• jumps: the jump problem
• t: current time
• u: current state
• args: static arguments
• key: random key
• t1: end time (step will not exceed this)
• jump_state: current jump state
• solver_state: current solver state

Returns:

A tuple (u_new, t_next, made_jump, jump_index, jump_state_new, solver_state_new, info) , if made_jump is False then u_new == u and jump_index == -1.

jumpax.AbstractHybridAggregator(jumpax.AbstractAggregator)

Abstract base class for hybrid aggregators that combine diffrax integration with jumps.

init(jumps, u0: Shaped[Array, '?*u'], args: PyTree[Any], key: Key[Array, '']) -> ~_SolverState

Initialize solver-specific state.

Arguments:

• jumps: the jump problem
• u0: initial state
• args: static arguments
• key: random key

Returns:

Initial solver state for this aggregator.

step(jumps, t: Real[ArrayLike, ''], u: Shaped[Array, '?*u'], args: PyTree[Any], key: Key[Array, ''], t1: Real[ArrayLike, ''], jump_state: PyTree[Any], solver_state: ~_SolverState) -> tuple

Advance one step of the simulation.

Arguments:

• jumps: the jump problem
• t: current time
• u: current state
• args: static arguments
• key: random key
• t1: end time (step will not exceed this)
• jump_state: current jump state
• solver_state: current solver state

Returns:

A tuple (u_new, t_next, made_jump, jump_index, jump_state_new, solver_state_new, info) , if made_jump is False then u_new == u and jump_index == -1.

Pure Jump Aggregators

---

jumpax.SSA(jumpax.AbstractAggregator)

Stochastic Simulation Algorithm (Gillespie's direct method).

Samples the next reaction time \(\tau\) from an exponential distribution with rate \(a_0 = \sum_j a_j(u)\), where \(a_j(u)\) is the propensity of reaction \(j\). The reaction channel is selected proportionally to its rate.

This is an exact method for simulating continuous-time Markov chains and chemical reaction networks.

Reference

@article{gillespie1977exact,
    title={Exact stochastic simulation of coupled chemical reactions},
    author={Gillespie, Daniel T},
    journal={The Journal of Physical Chemistry},
    volume={81},
    number={25},
    pages={2340--2361},
    year={1977},
    publisher={ACS Publications}
}

---

jumpax.SimpleTauLeaping(jumpax.AbstractAggregator)

Simple \(\tau\)-leaping aggregator with fixed step size.

At each step, draws reaction counts from Poisson distributions:

\[K_j \sim \text{Poisson}(a_j(u) \cdot \tau)\]

where \(a_j(u)\) is the propensity of reaction \(j\) and \(\tau\) is the step size. The state is then updated by applying all reactions simultaneously.

This is an approximate method that trades exactness for computational efficiency, particularly useful for systems with many reactions per unit time.

Reference

@article{gillespie2001approximate,
    title={Approximate accelerated stochastic simulation of chemically
           reacting systems},
    author={Gillespie, Daniel T},
    journal={The Journal of Chemical Physics},
    volume={115},
    number={4},
    pages={1716--1733},
    year={2001},
    publisher={AIP Publishing}
}

Hybrid Aggregators

Hybrid aggregators combine jump processes with continuous dynamics (ODEs/SDEs) using diffrax.

---

jumpax.HybridSSA(jumpax.AbstractHybridAggregator)

Hybrid SSA + diffrax ODE integration for CRJ/MAJ with piecewise-constant rates.

Assumptions:

• Jump rates do not depend on ODE-evolved components (constant between jumps).
• ode_fn(t, y, args) returns dy with same shape as y.
• Only ODE is integrated between jumps; at a jump time we apply the discrete affect.

__init__(ode_fn, *, solver: diffrax.AbstractSolver, dt0: Real[ArrayLike, ''] | None = None, stepsize_controller: diffrax.AbstractStepSizeController | None = None, ode_args: PyTree[Any] | None = None, max_steps: int | None = None)

Arguments:

• ode_fn: ODE drift function f(t, u, args) -> du/dt
• solver: diffrax solver (e.g., diffrax.Tsit5())
• dt0: initial step size for diffrax
• stepsize_controller: diffrax step size controller
• ode_args: separate args for the ODE
• max_steps: maximum diffrax steps per integration interval

---

jumpax.HazardSSA(jumpax.AbstractHybridAggregator)

Hazard (time-change) SSA that supports ODE and SDEs.

Uses the random time change representation: for each reaction channel \(j\), maintain an integrated hazard \(A_j(t) = \int_0^t \lambda_j(s, u(s)) ds\) and fire when \(A_j\) exceeds an exponential threshold \(z_j\).

This allows coupling jump processes with continuous dynamics via diffrax.

__init__(drift_fn: Callable | None, *, solver: diffrax.AbstractSolver, dt0: Real[ArrayLike, ''] | None = None, stepsize_controller: diffrax.AbstractStepSizeController | None = None, diffusion_fn: Callable | None = None, brownian_tol: Real[ArrayLike, ''] = 0.001, ode_args: PyTree[Any] | None = None, max_steps: int | None = None)

Arguments:

• drift_fn: drift function f(t, u, args) -> du/dt
• solver: diffrax solver (e.g., diffrax.Tsit5())
• dt0: initial step size for diffrax
• stepsize_controller: diffrax step size controller
• diffusion_fn: SDE diffusion function g(t, u, args), or None for ODE
• brownian_tol: tolerance for virtual Brownian tree (SDE only)
• ode_args: separate args for the ODE/SDE (defaults to jump args)
• max_steps: maximum diffrax steps per integration interval