A High-Performance JAX Framework for Biochemical Modeling, Neural ODEs, and Bayesian Inference

</div>

🚀 What is Catalax?

Catalax combines the power of JAX with advanced numerical methods for biochemical modeling. It seamlessly integrates traditional mechanistic modeling with cutting-edge neural differential equations, enabling:

⚡ Lightning-fast simulations with JIT compilation and GPU acceleration
🧠 Neural ODEs for data-driven discovery of biochemical dynamics
🎯 Bayesian parameter inference using Hamiltonian Monte Carlo
📊 Comprehensive model analysis with fit metrics and model comparison
🔬 Biochemical data integration with EnzymeML support

Check out the documentation for more details.

🌟 Key Features

🔬 Mechanistic Modeling

Intuitive model building with states, reactions, and ODEs
Automatic parameter extraction and constraint handling
Support for complex biochemical networks and pathways
Integration with experimental data and EnzymeML format

🧠 Neural Differential Equations

NeuralODE: Pure neural network dynamics learning
RateFlowODE: Reaction network discovery with learnable stoichiometry
UniversalODE: Hybrid models combining mechanistic and neural components
Advanced penalty systems for biological constraints

🎯 Bayesian Inference

MCMC sampling with NUTS (No-U-Turn Sampler)
Surrogate-accelerated inference using neural networks
Uncertainty quantification with HDI intervals
Prior specification and model comparison tools

📊 Model Analysis

Statistical model comparison (AIC, BIC, chi-square)
Phase plot visualization for rate analysis
Comprehensive plotting and visualization tools (phase plots, rate grids, etc.)

🛠️ Installation

# Install from PyPI
pip install catalax

# Or from source
git clone https://github.com/JR-1991/Catalax.git
cd Catalax
pip install .

🚀 Quick Start

1. Mechanistic Modeling

Create biochemical models using first-principles knowledge of reaction mechanisms. This approach is ideal when you understand the underlying biochemistry and want to build physically interpretable models. Catalax automatically extracts parameters from equations and enables efficient simulation with JAX compilation.

Key features:

Automatic parameter extraction from symbolic equations
Multiple initial conditions for batch simulation
Fast JAX-compiled integration with adaptive step sizes
Built-in visualization with publication-ready plots

<details open> <summary><strong>📖 Click to see mechanistic modeling code</strong></summary>

import catalax as ctx

# Create a Michaelis-Menten enzyme kinetics model
model = ctx.Model(name="Enzyme Kinetics")

# Define states
model.add_states(S="Substrate", E="Enzyme", P="Product")

# Add reaction kinetics via schemes
model.add_reaction(
    "S -> P",
    symbol="r1",
    equation="k_cat * E * S / (K_m + S)",
)

# Or as ODEs
model.add_odes(
    S="-k_cat * E * S / (K_m + S)",
    P="k_cat * E * S / (K_m + S)",
    E="0",
)

# Set parameters
model.parameters.k_cat.value = 0.1
model.parameters.K_m.value = 0.05

# Create dataset and add initial conditions
dataset = ctx.Dataset.from_model(model)
dataset.add_initial(S=100, E=10, P=0)
dataset.add_initial(S=200, E=10, P=0)

# Configure simulation
config = ctx.SimulationConfig(t1=100, nsteps=1000)

# Run simulation
results = model.simulate(dataset=dataset, config=config)

# Visualize results
results.plot()

</details>

2. Neural ODEs for Discovery

Discover unknown reaction networks directly from experimental time-series data using neural differential equations. The RateFlowODE learns both reaction rates and stoichiometric coefficients, automatically uncovering the underlying biochemical network structure without prior mechanistic knowledge.

Key features:

Learnable stoichiometry matrices that discover reaction networks
Biological constraints through penalty functions (mass balance, sparsity, integer coefficients)
Multi-step training strategies with adaptive learning rates
Rate visualization to interpret discovered reactions

<details> <summary><strong>🧠 Click to see neural ODE discovery code</strong></summary>

from catalax.neural import RateFlowODE, train_neural_ode, Strategy

# Create experimental dataset (with measurement data)
experimental_data = ctx.Dataset.from_model(model)
experimental_data.add_initial(S=100, E=10, ES=0, P=0)
# Add actual measurement data to the dataset...

# Create a neural ODE that learns reaction stoichiometry
neural_model = RateFlowODE.from_dataset(
    dataset=experimental_data,
    reaction_size=3,  # Number of reactions to discover
    width_size=64,
    depth=3
)

# Set up training strategy
strategy = Strategy()
strategy.add_step(lr=1e-3, steps=1000, batch_size=32)
strategy.add_step(lr=1e-4, steps=2000, batch_size=64)

# Train the model
trained_model = neural_model.train(
    model=neural_model,
    dataset=experimental_data,
    strategy=strategy
)

# Visualize learned reactions
trained_model.plot_learned_rates(experimental_data, original_model)

</details>

3. Bayesian Parameter Inference

Quantify parameter uncertainty and obtain credible intervals using Hamiltonian Monte Carlo (HMC) sampling. This approach goes beyond point estimates to provide full posterior distributions, enabling robust uncertainty quantification and model comparison through Bayesian statistics.

Key features:

NUTS sampler (No-U-Turn Sampler) for efficient exploration of parameter space
Multiple chains for convergence diagnostics and robust sampling
Flexible priors supporting various probability distributions
HDI intervals (Highest Density Intervals) for credible parameter ranges
Integrated visualization with corner plots, trace plots, and credibility intervals

<details> <summary><strong>🎯 Click to see Bayesian inference code</strong></summary>

from catalax.mcmc import HMC
from catalax.mcmc.priors import Normal, LogNormal
import numpyro.distributions as dist

# Set parameter priors
model.parameters.k1.prior = LogNormal(mu=0.0, sigma=1.0)
model.parameters.k2.prior = LogNormal(mu=0.0, sigma=1.0)
model.parameters.k3.prior = LogNormal(mu=0.0, sigma=1.0)

# Create HMC sampler with configuration
hmc = HMC(
    num_warmup=1000,
    num_samples=2000,
    likelihood=dist.SoftLaplace,  # or dist.Normal
    num_chains=4,
    chain_method="parallel"
)

# Create experimental dataset
experimental_data = ctx.Dataset.from_model(model)
experimental_data.add_initial(S=100, E=10, ES=0, P=0)
# Add more experimental measurements as needed...

# Run inference
results = hmc.run(
    model=model,
    dataset=experimental_data,
    yerrs=0.1  # Measurement uncertainty
)

# Analyze results with integrated visualization
samples = results.get_samples()
results.print_summary()

# Create publication-ready plots
fig1 = results.plot_corner()
fig2 = results.plot_posterior()
fig3 = results.plot_trace()
fig4 = results.plot_credibility_interval(
    initial_condition={"S": 100, "E": 10, "ES": 0, "P": 0},
    time=jnp.linspace(0, 100, 200)
)

# Get summary statistics
summary_stats = results.summary(hdi_prob=0.95)

</details>

📝 Documentation

Catalax documentation is available at https://catalax.mintlify.app/.

📚 Examples

Explore comprehensive examples in the examples/ directory:

Optimization - Parameter estimation and model fitting
HMC - Bayesian parameter inference
NeuralODE - Learning dynamics from data
UniversalODE - Hybrid mechanistic-neural models
SurrogateHMC - Accelerated inference with surrogates
Data Import - Working with experimental data

🛠️ Developer Guide

Prerequisites

Python 3.11+
uv for dependency management
Git for version control

Development Setup

Clone the repository

git clone https://github.com/JR-1991/Catalax.git
cd Catalax

Install development dependencies

# Install all dependencies including dev dependencies
uv sync --dev

Set up pre-commit hooks

# Install pre-commit hooks
uv run pre-commit install

Testing

Catalax uses pytest for testing with multiple test categories:

# Run all tests (excluding expensive/slow tests)
uv run pytest -vv -m "not expensive"

# Run all tests including expensive ones
uv run pytest -vv

# Run specific test categories
uv run pytest -vv tests/unit/          # Unit tests
uv run pytest -vv tests/integration/   # Integration tests

# Run tests with coverage
uv run pytest -vv --cov=catalax --cov-report=html

The test suite includes:

Unit tests: Fast tests for individual components
Integration tests: End-to-end testing of workflows
Expensive tests: Computationally intensive tests (marked with @pytest.mark.expensive)

Code Quality

The project uses several tools to maintain code quality:

Ruff: Fast Python linter and formatter
Pyright: Type checking
Pre-commit: Git hooks for automa

Catalax

Install / Use

README

🚀 What is Catalax?

🌟 Key Features

🔬 Mechanistic Modeling

🧠 Neural Differential Equations

🎯 Bayesian Inference

📊 Model Analysis

🛠️ Installation

🚀 Quick Start

1. Mechanistic Modeling

2. Neural ODEs for Discovery

3. Bayesian Parameter Inference

📝 Documentation

📚 Examples

🛠️ Developer Guide

Prerequisites

Development Setup

Testing

Code Quality