Simulating large chemical mechanisms is vital in combustion, atmospheric chemistry and heterogeneous catalysis. JAX-reactor is a package written in python to simulate large kinetic models leveraging just-in-time (JIT) compilation, automatic differentiaion and vectorization capabilities of awesome JAX package. JAX uses XLA to JIT compile python code to CPU, GPU and TPU. JAX can automatically differentiate python and numpy functions allowing us to efficiently calculate Jacobians of large chemical systems.
JAX-reactor uses Cantera's recently developed YAML input format to read large detailed kinetic models. Currently, JAX-reactor provides a basic JAX implementation of backward differentiation formula (BDF) solver to integrate stiff chemical systems. JAX-reactor is heavily inspired by a similar package written in PyTorch called reactorch. JAX-reactor is a research project and is in early stages of development.
First set up a conda environment
conda create --name jax python=3.7 scons cython boost numpy=>1.5 ruamel_yaml scipy=>1.4 matplotlib jupyter
Install the dev build of Cantera
conda activate jax
conda install -c cantera/label/dev cantera
CPU only version of JAX can be easily installed using pip
pip install --upgrade pip
pip install --upgrade jax jaxlib # CPU-only version
For GPU version of JAX please follow the official installation instructions at JAX GPU installation
If you want to use JAX on CentOS-7 you need to build JAX from source following instructions at jax-ml/jax#2083
Once the enviroment is setup install JAX-reactor:
git clone https://github.com/comocheng/jax-reactor.git
export PYTHONPATH=<full-path-to-cloned-folder>:$PYTHONPATH
1. Weiqi Ji, Sili Deng. ReacTorch: A Differentiable Reacting Flow Simulation Package in PyTorch, https://github.com/DENG-MIT/reactorch, 2020.