Chemical Master Equation

Exact Chemical Master Equation solver via matrix exponentiation.

For time-independent reaction systems with finite (or finitely truncated) state spaces, the probability distribution \(P(t)\) evolves according to:

\[\frac{dP}{dt} = Q P\]

where \(Q\) is the generator (rate) matrix. The exact solution is:

\[P(t) = e^{Qt} P(0)\]

This module provides utilities to:

1. Build the generator matrix \(Q\) from a MassActionJump problem on a truncated state space
2. Solve the CME exactly using matrix exponentiation

Functions

jumpax.build_generator(problem: jumpax.MassActionJump, max_counts: Int[Array, 'S'], *, args: PyTree[Any] = None, t0: float = 0.0) -> tuple

Build the generator matrix Q for a jumpax.MassActionJump on a truncated state space.

The state space is the rectangular lattice \(\{0, ..., \text{max\_counts}_i\}\) for each species \(i\). Transitions leaving this box are dropped.

Arguments:

• problem: A jumpax.MassActionJump defining the reaction system
• max_counts: array of maximum population per species
• args: Optional arguments passed to rate function
• t0: Time at which to evaluate rates (only matters for time-dependent rates)

Returns:

• Q: generator matrix where N = prod(max_counts + 1)
• cme_state: CMEState containing state space information

---

jumpax.solve_cme(Q: Float[Array, 'N N'], p0: Float[Array, 'N'], t: float) -> Float[Array, 'N']

Solve the CME exactly via matrix exponentiation.

Computes \(P(t) = \exp(Q t) P(0)\).

Arguments:

• Q: generator matrix from jumpax.build_generator
• p0: initial probability distribution (should sum to 1)
• t: Time at which to evaluate the distribution

Returns:

• p_t: probability distribution at time t

---

jumpax.marginal_distribution(p: Float[Array, 'N'], cme_state: jumpax.CMEState, species: int) -> Float[Array, 'M']

Compute the marginal distribution for a single species.

Arguments:

• p: probability distribution over states
• cme_state: State space information
• species: Index of species to marginalize over

Returns:

• marginal: marginal distribution where M = max_counts[species] + 1

Types

jumpax.CMEState

State space information for Chemical Master Equation (CME) solver.

Attributes:

Name	Type	Description
states		array of all states in the truncated state space
strides		array for mixed-radix encoding
max_counts		array of maximum count per species