Installation

Requirements

JAX-in-Cell is a pure Python package. It depends on jax, jax_tqdm (a progress bar inside compiled loops) and matplotlib. Python 3.9 or newer is required; the continuous-integration matrix tests 3.9 through 3.12 on Linux, and the documentation is built with 3.12.

The package enables 64-bit floating point in JAX when it is imported (jax.config.update("jax_enable_x64", True)). This applies to the whole Python process, including any other JAX code you run afterwards.

From PyPI

pip install jaxincell

From source

git clone https://github.com/uwplasma/JAX-in-Cell
cd JAX-in-Cell
pip install -e .

The editable install (-e) lets you modify the code under jaxincell/ and run it without reinstalling. The version string is generated by setuptools_scm from the git tags, so a source install reports the tag of the last release plus the number of commits since.

GPU and TPU

pip install jaxincell pulls in the CPU build of JAX. For a GPU, install the JAX wheel that matches your CUDA or ROCm stack first, following the JAX installation instructions, then install jaxincell on top. For example, with CUDA 12:

pip install -U "jax[cuda12]"
pip install jaxincell

Nothing in JAX-in-Cell is device specific. The simulation runs on whatever default device JAX selects; jax.devices() shows which one that is.

Checking the install

python -c "import jaxincell, jax; print(jaxincell.__file__); print(jax.devices())"
jaxincell

The second command runs the built-in default configuration (a two-stream instability with 500 pseudo-particles per species) and opens an animation window. It prints a summary of the derived plasma parameters before the run starts.

Optional dependencies

• scipy is used by some example scripts (least-squares optimisation) and by the scripts that generate the figures in this documentation (plasma dispersion function).

• ffmpeg on the PATH is needed to save animations to MP4 with jaxincell.plot().

• pytest and pytest-cov run the test suite, see Testing.