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,
)

# dt is the measurement sampling frequency
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 between subsequent measurement steps (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.

model	species	reactions	shape
birth_death	X	2	minimal one-species baseline
single_gene	R, P	4	constitutive transcription-translation
negative_autoregulation	X	2	Hill-repressed self-feedback
positive_autoregulation	X	2	Hill self-activation (Params.bistable() for the bistable regime)
linear_cascade	A, B	4	A → B activation cascade
toggle_switch	A, B	4	mutual repression
incoherent_ffl	A, B, C	6	adaptive / pulse-generating FFL
repressilator	A, B, C	6	synthetic oscillator
cca_optogenetic	R, P	4	light-driven gene expression (CcaS/CcaR) — input-driven

You can verride 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.