Metadata-Version: 2.4
Name: crn-jax
Version: 0.3.1
Summary: Chemical reaction networks in JAX: GPU-parallel stochastic simulations.
License-Expression: MIT
License-File: LICENSE
Keywords: jax,gillespie,ssa,stochastic-simulation,chemical-reaction-networks,systems-biology,gpu
Author: Robin Henry
Author-email: robin.henry012@gmail.com
Requires-Python: >=3.11
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Provides-Extra: examples
Requires-Dist: jax (>=0.10,<0.11)
Requires-Dist: matplotlib (>=3.10) ; extra == "examples"
Requires-Dist: numpy (>=2.0) ; extra == "examples"
Project-URL: Issues, https://github.com/robinhenry/crn-jax/issues
Project-URL: Repository, https://github.com/robinhenry/crn-jax
Description-Content-Type: text/markdown

# crn-jax

[![CI](https://github.com/robinhenry/crn-jax/actions/workflows/ci.yml/badge.svg)](https://github.com/robinhenry/crn-jax/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/crn-jax.svg)](https://pypi.org/project/crn-jax/)
[![Python](https://img.shields.io/badge/python-3.11%2B-blue)](https://pypi.org/project/crn-jax/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

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

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="benchmarks/figures/throughput_speedup_birth_death_dark.svg">
    <source media="(prefers-color-scheme: light)" srcset="benchmarks/figures/throughput_speedup_birth_death.svg">
    <img alt="Birth-death benchmark: crn-jax on GPU vs GillesPy2 (C++) on CPU." src="benchmarks/figures/throughput_speedup_birth_death.svg" width="48%">
  </picture>
  &nbsp;
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="benchmarks/figures/throughput_speedup_linear_cascade_dark.svg">
    <source media="(prefers-color-scheme: light)" srcset="benchmarks/figures/throughput_speedup_linear_cascade.svg">
    <img alt="Linear-cascade benchmark: crn-jax on GPU vs GillesPy2 (C++) on CPU." src="benchmarks/figures/throughput_speedup_linear_cascade.svg" width="48%">
  </picture>
</p>

<p align="center">
  <i>☝️ 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).</i>
</p>

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="benchmarks/figures/trajectories_repressilator_dark.png">
    <source media="(prefers-color-scheme: light)" srcset="benchmarks/figures/trajectories_repressilator.png">
    <img alt="24 stochastic trajectories of the repressilator (3-species oscillator), one panel per species." src="benchmarks/figures/trajectories_repressilator.png" width="96%">
  </picture>
</p>

<p align="center">
  <i>☝️ 24 stochastic trajectories of the repressilator gene circuit, sampled in parallel from a one-line call.</i>
</p>

## 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.

```bash
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:

```python
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](#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** |

Override the defaults by passing your own `Params` values:

```python
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/`](examples/) folder walks through the main features end-to-end:

- [`1_library_examples.ipynb`](examples/1_library_examples.ipynb): sample and plot from every pre-built model. **Start here.**
- [`2_per-replicate_control.ipynb`](examples/2_per-replicate_control.ipynb): closed-loop simulations where each replicate gets its own time-varying controller and setpoint.
- [`3_implement_your_own.ipynb`](examples/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](https://github.com/StochSS/GillesPy2): C++ optimized Gillespie simulations on CPU.
* [jax-smfsb](https://github.com/darrenjw/jax-smfsb): JAX implementations of algorithms from the *Stochastic Modelling for Systems Biology* book.
* [myriad-jax](https://github.com/robinhenry/myriad-jax): RL-style decision making fully in JAX, powered by `crn-jax` at its core.

