Overview

JAX-in-Cell is an open-source project in Python that uses JAX to speedup simulations, leading to a simple to use, fast and concise code. It can be imported in a Python script using the jaxincell package, or run directly in the command line as jaxincell. To install it, use

pip install jaxincell

Alternatively, install dependencies jax, jax_tqdm and matplotlib, and run:

git clone https://github.com/uwplasma/JAX-in-Cell
python example_script.py

The project can be downloaded in its GitHub repository: https://github.com/uwplasma/JAX-in-Cell