Skip to content

ReactionNetwork

stochastix.ReactionNetwork #

ReactionNetwork(
    reactions: Iterable[Reaction], volume: float = 1.0
)

A network of chemical reactions for simulation.

This class compiles a list of Reaction objects into a system ready for stochastic or deterministic simulation. It constructs the stoichiometric, reactant, and product matrices, and provides methods to compute propensities and simulate the system as an ordinary differential equation (ODE) or a stochastic differential equation (SDE).

The network can be manipulated functionally, for example, by adding reactions or creating subnetworks through slicing.

State input formats.

vector_field and log_prob accept any properly formatted pytree as state input, all other methods expect a flat array of species counts. You can use pytree_to_state and state_to_pytree to convert between the two formats.

Attributes:

  • reactions –

    A tuple of Reaction objects in the network.

  • species –

    An ordered tuple of unique species names.

  • volume –

    The system volume, used for concentration-dependent rates.

  • stoichiometry_matrix –

    The matrix S where S[i, j] is the net change in species i from reaction j. Shape: (n_species, n_reactions).

  • reactant_matrix –

    The matrix of reactant stoichiometric coefficients. Shape: (n_species, n_reactions).

  • product_matrix –

    The matrix of product stoichiometric coefficients. Shape: (n_species, n_reactions).

  • n_reactions –

    The number of reactions in the network.

  • n_species –

    The number of species in the network.

Methods:

  • propensity_fn –

    Computes the propensity vector for the network.

  • vector_field –

    The vector field for deterministic ODE integration.

  • ode_rhs –

    Alias for vector_field.

  • drift_fn –

    The drift function for the Chemical Langevin Equation (CLE).

  • cle_drift –

    Alias for drift_fn.

  • noise_coupling –

    The noise coupling matrix for the Chemical Langevin Equation (CLE).

  • cle_diffusion –

    Alias for noise_coupling.

  • diffusion_fn –

    Alias for noise_coupling.

  • diffusion_coeff_matrix –

    The diffusion coefficient matrix for the Fokker-Planck equation.

  • diffusion_coeff –

    Alias for diffusion_coeff_matrix.

  • diffrax_ode_term –

    Returns an ODETerm for ODE solving with diffrax.

  • diffrax_sde_term –

    A MultiTerm for solving the system as an SDE with diffrax.

  • log_prob –

    Calculates the log-probability of a simulation trajectory.

  • copy –

    Creates a deep copy of the reaction network.

  • to_latex –

    Generates a LaTeX representation of the reaction network.

Parameters:

  • reactions (Iterable[Reaction]) –

    An iterable of Reaction objects.

  • volume (float, default: 1.0 ) –

    The system volume. Affects concentration-dependent rates.

propensity_fn #

propensity_fn(
    x: ndarray,
    t: float = None,
    reaction_mask: ndarray | None = None,
) -> jnp.ndarray

Computes the propensity vector for the network.

This function calculates the propensity a(x, t) for each reaction. It is optimized for the common case where all propensities are calculated (reaction_mask=None). For the masked case, it uses jax.lax.cond to remain JIT-compatible with dynamic masks.

Parameters:

  • x (ndarray) –

    The current state of the system (species counts).

  • t (float, default: None ) –

    The current time.

  • reaction_mask (ndarray | None, default: None ) –

    An optional boolean array of shape (n_reactions,) indicating which propensities to compute. If None, all are computed.

Returns:

  • ndarray –

    Propensities for each reaction. Propensities for masked-out reactions will be zero.

vector_field #

vector_field(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The vector field for deterministic ODE integration.

This method computes the right-hand side of the ODE system, dx/dt = f(t, x), based on the deterministic rate laws of the reactions. The signature f(t, x, args) is compatible with diffrax.diffeqsolve.

Parameters:

  • t (floating) –

    The current time.

  • x (ndarray) –

    The current state vector (species counts).

  • args (Any, default: None ) –

    Additional arguments (not used).

Returns:

  • ndarray –

    The time derivative of the state, dx/dt, in molecules/time (counts/time).

ode_rhs #

ode_rhs(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The right-hand side of the ODE system. Alias for vector_field.

drift_fn #

drift_fn(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The drift function for the Chemical Langevin Equation (CLE).

The CLE is given by: dX_t = f(t, X_t) dt + g(t, X_t) dW_t. This function computes the drift term f(t, X_t) = S * a(X_t, t), where S is the stoichiometry matrix and a is the propensity vector.

Note

This is different from vector_field. The CLE drift is based on stochastic propensities, while the ODE vector field is based on deterministic rates. Better to use the cle_drift alias for clarity.

Parameters:

  • t (floating) –

    The current time.

  • x (ndarray) –

    The current state vector (species counts).

  • args (Any, default: None ) –

    Additional arguments (not used).

Returns:

  • ndarray –

    The drift vector for the CLE.

cle_drift #

cle_drift(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The drift function f(t, X_t) for the Chemical Langevin Equation (CLE). Alias for drift_fn.

noise_coupling #

noise_coupling(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The noise coupling matrix for the Chemical Langevin Equation (CLE).

The CLE is given by: dX_t = f(t, X_t) dt + g(t, X_t) dW_t. This function computes the noise coupling term g(t, X_t), which is defined as S * diag(sqrt(a(X_t, t))), where S is the stoichiometry matrix and a is the propensity vector.

Note

This is not the same as the diffusion coefficient matrix used in the Fokker-Planck equation. See diffusion_coeff_matrix.

Parameters:

  • t (floating) –

    The current time.

  • x (ndarray) –

    The current state vector (species counts).

  • args (Any, default: None ) –

    Additional arguments (not used).

Returns:

  • ndarray –

    The noise coupling matrix of shape (n_species, n_reactions).

cle_diffusion #

cle_diffusion(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The diffusion function g(t, X_t) for the Chemical Langevin Equation (CLE). Alias for noise_coupling.

diffusion_fn #

diffusion_fn(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The diffusion function g(t, X_t) for the Chemical Langevin Equation (CLE). Alias for noise_coupling.

diffrax_ode_term #

diffrax_ode_term() -> ODETerm

Returns an ODETerm for solving the system as an ODE with diffrax.

diffusion_coeff_matrix #

diffusion_coeff_matrix(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The diffusion coefficient matrix for the Fokker-Planck equation.

This matrix, D(t, X_t), is defined as g @ g.T, where g is the noise coupling matrix from the CLE. It is equivalent to S @ diag(a(X_t, t)) @ S.T.

Note

This is different from noise_coupling, which is the g matrix itself.

Parameters:

  • t (floating) –

    The current time.

  • x (ndarray) –

    The current state vector (species counts).

  • args (Any, default: None ) –

    Additional arguments (not used).

Returns:

  • ndarray –

    The diffusion coefficient matrix of shape (n_species, n_species).

diffusion_coeff #

diffusion_coeff(
    t: floating, x: ndarray, args: Any = None
) -> jnp.ndarray

The diffusion coefficient for the Fokker-Planck equation. Alias for diffusion_coeff_matrix.

diffrax_sde_term #

diffrax_sde_term(
    t1: float,
    t0: float = 0.0,
    tol: float = 0.001,
    *,
    key: ndarray = None,
) -> MultiTerm

Returns a MultiTerm for solving the system as an SDE with diffrax.

This configures the SDE with a drift term (from the ODE vector field) and a diffusion term, controlled by a VirtualBrownianTree.

Parameters:

  • t1 (float) –

    The end time of the simulation.

  • t0 (float, default: 0.0 ) –

    The start time of the simulation. Defaults to 0.0.

  • tol (float, default: 0.001 ) –

    The tolerance for the Brownian motion simulation.

  • key (ndarray, default: None ) –

    The jax.random.PRNGKey for the Brownian motion.

Returns:

  • MultiTerm –

    A diffrax.MultiTerm for SDE integration.

log_prob #

log_prob(sim_results: SimulationResults) -> jnp.ndarray

Calculates the log-probability of a simulation trajectory.

This method computes the log-likelihood of a given trajectory produced by an exact stochastic simulation algorithm (e.g., Gillespie Direct Method). It recomputes the propensities at each step of the trajectory to ensure that the computation graph is connected for gradient-based optimization.

Parameters:

  • sim_results (SimulationResults) –

    A SimulationResults object containing the trajectory.

Returns:

  • ndarray –

    Log-probability terms, one for each step in the trajectory.

copy #

copy() -> ReactionNetwork

Creates a new ReactionNetwork instance from the current one.

This method provides a way to create a mutable copy of the network. It reconstructs the network from its reactions and volume, which is more robust and efficient than a generic deep copy.

Returns:

to_latex #

to_latex(print_kinetics: bool = False) -> str

Generates a LaTeX representation of the reaction network.

This method produces a LaTeX string suitable for rendering in reports, notebooks, or publications. It automatically pairs forward and reverse reactions into reversible reactions.

Parameters:

  • print_kinetics (bool, default: False ) –

    Whether to include the kinetics types in the output.

Returns:

  • str –

    A string containing the LaTeX representation of the reaction network, wrapped in $$...$$ and an align* environment.