Skip to content

Solvers#

Classic Solvers#

stochastix.DirectMethod #

DirectMethod()

Exact stochastic simulation using the Direct Method.

The Direct Method is one of the two original exact algorithms proposed by Gillespie (1977). It samples the next reaction time from an exponential distribution and selects the reaction to fire based on their relative propensities.

This implementation is JAX-compatible and designed for efficient compilation and automatic differentiation.

Attributes:

  • is_exact_solver –

    Boolean flag indicating that this is an exact solver.

  • is_pathwise_differentiable –

    Boolean flag indicating that this solver is not pathwise differentiable.

Methods:

  • propensities –

    Calculate reaction propensities for the current state.

  • step –

    Execute one step of the Direct Method algorithm.

References

Gillespie, D. T. (1977). Exact stochastic simulation of coupled chemical reactions. The Journal of Physical Chemistry, 81(25), 2340-2361.

stochastix.FirstReactionMethod #

FirstReactionMethod()

Exact stochastic simulation using the First Reaction Method.

The First Reaction Method is the second exact algorithm proposed by Gillespie (1977). It generates a candidate firing time for each reaction independently and selects the reaction with the smallest firing time.

While mathematically equivalent to the Direct Method, this approach can be computationally less efficient for large numbers of reactions due to the need to generate multiple random numbers per step.

Attributes:

  • is_exact_solver –

    Boolean flag indicating that this is an exact solver.

  • is_pathwise_differentiable –

    Boolean flag indicating that this solver is not pathwise differentiable.

Methods:

  • propensities –

    Calculate reaction propensities for the current state.

  • step –

    Execute one step of the First Reaction Method algorithm.

References

Gillespie, D. T. (1977). Exact stochastic simulation of coupled chemical reactions. The Journal of Physical Chemistry, 81(25), 2340-2361.

Pathwise Differentiable Solvers#

stochastix.DifferentiableDirect #

DifferentiableDirect(
    logits_scale: float = 1.0, exact_fwd: bool = True
)

Pathwise differentiable Direct Method using the Straight-Through Gumbel-Softmax Gradient Estimator.

This solver implements the Direct Method in a form that is pathwise differentiable with respect to reaction parameters. Differentiation is enabled by the Gumbel-max trick, which reparameterizes the discrete choice of which reaction fires next.

The forward pass can be either exact or approximate (relaxed): - If exact_fwd=True (default), it uses a straight-through estimator, making the forward pass of the simulation exact. - If exact_fwd=False, it uses a continuous relaxation of the reaction choice, making the forward pass approximate.

In both cases, the backward pass uses a continuous relaxation (softmax) to enable gradient propagation, allowing for end-to-end training of stochastic models with gradient-based optimization.

Attributes:

  • logits_scale –

    The temperature for the softmax relaxation.

  • is_exact_solver –

    Boolean flag indicating exactness based on exact_fwd.

  • is_pathwise_differentiable –

    Boolean flag indicating differentiability.

Methods:

  • propensities –

    Calculate reaction propensities for the current state.

  • step –

    Execute one step of the differentiable Direct Method.

Parameters:

  • logits_scale (float, default: 1.0 ) –

    The temperature parameter for the Gumbel-softmax relaxation. Smaller values give a better approximation but may have higher variance gradients. Defaults to 1.0.

  • exact_fwd (bool, default: True ) –

    If True, the forward pass is exact, using a straight- through gradient estimator. If False, the forward pass is relaxed and approximate.

stochastix.DifferentiableFirstReaction #

DifferentiableFirstReaction(
    logits_scale: float = 1.0, exact_fwd: bool = True
)

Pathwise differentiable First Reaction Method using the Straight-Through Gumbel-Softmax Gradient Estimator.

This solver implements the First Reaction Method in a form that is pathwise differentiable. It is the differentiable counterpart to FirstReactionMethod.

The forward pass can be either exact or approximate (relaxed), controlled by exact_fwd: - If exact_fwd=True (default), it uses a straight-through estimator for both the reaction choice and the time step dt, ensuring an exact forward pass. - If exact_fwd=False, it uses a continuous relaxation for both the state update and dt, making the forward pass approximate.

The backward pass uses a continuous relaxation (soft-argmin) of the reaction choice and a relaxed dt to enable gradient propagation.

Attributes:

  • logits_scale –

    The temperature for the softmax relaxation.

  • is_exact_solver –

    Boolean flag indicating exactness based on exact_fwd.

  • is_pathwise_differentiable –

    Boolean flag indicating differentiability.

Methods:

  • propensities –

    Calculate reaction propensities for the current state.

  • step –

    Execute one step of the differentiable First Reaction Method.

Parameters:

  • logits_scale (float, default: 1.0 ) –

    The temperature for the soft-argmin relaxation. Smaller values give a better approximation but may have higher variance gradients. Defaults to 1.0.

  • exact_fwd (bool, default: True ) –

    If True, the forward pass is exact, using a straight- through estimator. If False, the forward pass is relaxed and approximate.

stochastix.DGA #

DGA(a: float = 0.005, sigma: float = 0.1581)

Approximate simulation using the Differentiable Gillespie Algorithm.

This solver is highly unstable and sensitive to hyperparameter choices. In most cases, reparameterized versions of exact solvers should be preferred.

The method has two key approximations: 1. Differentiable reaction selection: The discrete choice of which reaction fires is replaced by a "soft" selection using sigmoid functions, controlled by the steepness parameter a. 2. Differentiable state update: The state is updated using a Gaussian-based weighting scheme centered around a soft reaction index, controlled by the sharpness parameter sigma.

Attributes:

  • is_exact_solver –

    Boolean flag indicating this is an approximate method.

  • is_pathwise_differentiable –

    Boolean flag indicating differentiability.

  • a –

    Steepness parameter for the sigmoid function.

  • sigma –

    Width parameter (standard deviation) for the Gaussian function.

Methods:

  • propensities –

    Compute reaction propensities at the current state.

  • step –

    Perform a single step of the Differentiable Gillespie Algorithm.

References

Rijal, K., & Mehta, P. (2025). A differentiable Gillespie algorithm for simulating chemical kinetics, parameter estimation, and designing synthetic biological circuits. eLife, 14, RP103877. https://doi.org/10.7554/eLife.103877.2

Parameters:

  • a (float, default: 0.005 ) –

    Steepness parameter for the sigmoid function. Must be positive. Smaller values make the sigmoid steeper, approaching a true Heaviside function. Defaults to 0.005.

  • sigma (float, default: 0.1581 ) –

    Standard deviation of the Gaussian for smoothing the state update. Must be positive. Smaller values make the Gaussian sharper, approaching a true discrete update. Defaults to ~0.1581.

Raises:

  • ValueError –

    If a or sigma are not positive.

Approximate Solvers#

stochastix.TauLeaping #

TauLeaping(epsilon: float = 0.03, tau_min: float = 1e-10)

Approximate stochastic simulation using the tau-leaping method.

This solver implements the tau-leaping method described in Cao et al. (2006), which approximates the exact stochastic simulation algorithm (SSA) by allowing multiple reactions to fire in a single time step tau. This can significantly accelerate simulations when reaction propensities are large.

The method selects a time step tau such that the propensities do not change significantly during the leap, maintaining accuracy while improving computational efficiency. The leap size is determined by balancing mean and variance constraints to control the approximation error.

Attributes:

  • is_exact_solver –

    Boolean flag indicating this is an approximate method.

  • epsilon –

    Error control parameter. Smaller values give higher accuracy but smaller time steps.

  • tau_min –

    Minimum allowed time step to prevent simulation stalling.

Methods:

  • propensities –

    Compute reaction propensities at the current state.

  • step –

    Perform a single step of the tau-leaping algorithm.

References

Cao, Y., Gillespie, D. T., & Petzold, L. R. (2006). Efficient step size selection for the tau-leaping simulation method. Journal of Chemical Physics, 124(4), 044109.

Parameters:

  • epsilon (float, default: 0.03 ) –

    The error control parameter, between 0 and 1. Smaller values provide higher accuracy but require smaller time steps. Defaults to 0.03 (3% relative error tolerance).

  • tau_min (float, default: 1e-10 ) –

    The minimum allowed time step. Must be positive. Prevents the simulation from taking excessively small steps. Defaults to 1e-10.

Raises:

  • ValueError –

    If epsilon is not between 0 and 1, or if tau_min is not positive.

Abstract Base Class#

stochastix.solvers.AbstractStochasticSolver #

AbstractStochasticSolver(
    is_exact_solver: bool, is_pathwise_differentiable: bool
)

Abstract base class for stochastic solvers.

Attributes:

  • is_exact_solver –

    Boolean flag indicating whether the solver is exact.

  • is_pathwise_differentiable –

    Boolean flag indicating whether the solver is pathwise differentiable.

Methods:

  • init –

    Initialize the solver's state.

  • propensities –

    Compute the reaction propensities.

  • step –

    Perform a single step of the solver.

init #

init(
    network: ReactionNetwork,
    t: floating,
    x: ndarray,
    a: ndarray,
    *,
    key: ndarray,
)

Initialize the solver's state.

This method is called once at the beginning of a simulation.

Parameters:

  • network (ReactionNetwork) –

    The reaction network.

  • t (floating) –

    The initial time.

  • x (ndarray) –

    The initial state.

  • a (ndarray) –

    The initial propensities.

  • key (ndarray) –

    A JAX random key.

Returns:

  • –

    The initial state of the solver.

propensities #

propensities(
    network: ReactionNetwork, x: ndarray, t: floating
) -> jnp.ndarray

Compute the reaction propensities.

Parameters:

  • network (ReactionNetwork) –

    The reaction network.

  • x (ndarray) –

    The current state of the system.

  • t (floating) –

    The current time.

Returns:

  • ndarray –

    Array of propensities for each reaction.

step #

step(
    network: ReactionNetwork,
    t: floating,
    x: ndarray,
    a: ndarray,
    state,
    *,
    key: ndarray,
) -> tuple[SimulationStep, ...]

Perform a single step of the solver.

Parameters:

  • network (ReactionNetwork) –

    The reaction network.

  • t (floating) –

    The current time.

  • x (ndarray) –

    The current state.

  • a (ndarray) –

    The current propensities.

  • state –

    The current state of the solver.

  • key (ndarray) –

    A JAX random key.

Returns:

  • tuple[SimulationStep, ...] –

    Tuple containing the SimulationStep result and the new solver state.