crn-jax

Chemical reaction networks in JAX: a lightweight, GPU-optimized Gillespie / SSA simulation library.

☝️ Wall-time to simulate 1,000,000 independent stochastic trajectories — each a full Gillespie run of the reaction network from t=0 to t=20, sampled at 200 time points (CPU vs RTX 5090 GPU).

☝️ 24 stochastic trajectories of the repressilator gene circuit, sampled in parallel from a one-line call.

Key features

• 🎯 Exact Gillespie simulations: an exact implementation of the Gillespie (SSA) algorithm.
• 🚀 GPU-optimized: run 1M+ independent trajectories on a single GPU with no Python overhead, powered by JAX.
• 🎛️ Closed-loop simulations: propensity functions take an optional input argument that can vary per-interval and per-replicate, so each of N parallel trajectories can follow its own time-dependent input schedule (useful for closed-loop experiment simulations, RL-style rollouts, …).
• 🎨 Pre-built suite of common reaction networks: a library of canonical examples already implemented with realistic default (tunable) parameters (useful for benchmarking, prototyping, and getting started).
• 🧩 Customizable: implementing a new chemical reaction network is easy.

Install

crn-jax depends on jax / jaxlib only.

pip install crn-jax

# with NVIDIA GPU support:
pip install crn-jax "jax[cuda12]"

# with plotting helpers:
pip install "crn-jax[examples]"

# for local development (uses Poetry):
git clone https://github.com/robinhenry/crn-jax && cd crn-jax
poetry install                    # main deps + dev tools
poetry install --with gpu         # add jax[cuda12] on an NVIDIA host

Quickstart

To simulate 32 stochastic trajectories of the repressilator in parallel and plot them:

import jax, jax.numpy as jnp
from crn_jax import models
from crn_jax.plotting import plot_species_trajectories

# Initial counts: (n_replicates, n_species). You always supply x0 yourself.
# We deliberately don't sample initial conditions for users because the sensible
# initial distribution is often problem-specific.
key, k_x0 = jax.random.split(jax.random.PRNGKey(0))
x0 = jax.random.uniform(
    k_x0, (32, len(models.repressilator.SPECIES)), minval=0.0, maxval=50.0,
)

ds = models.sample_trajectories(models.repressilator, key, x0, n_steps=3000, dt=0.1)

plot_species_trajectories(ds, title="Repressilator")  # see the plot at the top of this file

The returned Dataset has the following fields:

• ds.species: a tuple of species names, e.g. ("A", "B", "C").
• ds.xs: a 3D array with shape (n_replicates, n_steps, n_species) that contains the full trajectories.
• ds.X_t, ds.dX: two 2D arrays with shape (n_replicates * n_steps, n_species) that contain the timestamps and deltas of each reactions (useful to use as a training dataset for downstream ML).

Swapping models can be done by replacing models.repressilator with models.toggle_switch, models.incoherent_ffl, etc.

For a complete walkthrough of how to implement your own models, see Examples below.

Pre-built models

crn_jax.models provides a library of canonical reaction networks inspired by the literature, each exposing the same interface.

Currently, these are mostly gene reaction networks (GRNs), but this may evolve over time.

Override the defaults by passing your own Params values:

ds = models.sample_trajectories(
    models.positive_autoregulation, key, x0,
    params=models.positive_autoregulation.Params(n=3),  # default is n=2
    n_steps=3000, dt=0.1,
)

Examples

The examples/ folder walks through the main features end-to-end:

• 1_library_examples.ipynb: sample and plot from every pre-built model. Start here.
• 2_per-replicate_control.ipynb: closed-loop simulations where each replicate gets its own time-varying controller and setpoint.
• 3_implement_your_own.ipynb: define a custom reaction network from scratch, and optionally use the low-level simulate_interval / simulate_until primitives for finer-grained control (RL-style rollouts, non-uniform time grids).

See Also

• GillesPy2: C++ optimized Gillespie simulations on CPU.
• jax-smfsb: JAX implementations of algorithms from the Stochastic Modelling for Systems Biology book.
• myriad-jax: RL-style decision making fully in JAX, powered by crn-jax at its core.