Installation

NEO_JAX is distributed as a standard Python package.

Base install

Install from PyPI with:

pip install neo-jax

For development from a clone, a standard editable install is:

cd NEO_JAX
pip install -e .

Development and documentation extras

Optional development and documentation dependencies are installed with:

pip install -e ".[dev,docs]"

JAX should be installed with the correct accelerator support for your system (CPU, CUDA, or ROCm). Consult the JAX installation guide for platform-specific instructions.

The boozmn reader relies on the netCDF4 Python package, which is listed as a core dependency.

Optional pipeline dependencies

For end-to-end VMEC→Boozer→NEO workflows, install:

• vmec_jax

• booz_xform_jax

The CI workflow installs both packages for the pipeline tests.

Building the documentation

To build the documentation locally:

python -m sphinx -b html docs docs/_build/html

For a fast structural check without generating the full HTML tree:

python -m sphinx -b dummy docs docs/_build/dummy

Continuous integration

The GitHub Actions workflow installs the package, runs the test suite on CPU, and executes a small performance regression check. See Testing and CI for the full workflow.