Welcome to jax-sn’s documentation!

This package uses a multigroup, discontinuous Galerkin (DG) discrete ordinates (SN) method to solve a linear Boltzmann transport equation in JAX (ax), enabling differentiable neutron transport simulations.

It is written in pure functional JAX primitives and uses linear operators to represent the discretised transport operator.

Getting started

First install a JAX distribution, see here. GPU support is strongly preferred.

Then install from PyPI (Python 3.12 or newer):

pip install jax-sn

Optional extras:

  • geometry – mesh location and Cartesian domain setup (jax-sbgeom)

  • warp – NVIDIA Warp transport-sweep and mass kernels (warp-lang)

  • viz – domain plotting (plotly)

  • test – the test suite (pytest, sympy, networkx, trimesh)

  • docs – building this documentation

For example:

pip install "jax-sn[geometry,warp]"

For development, clone the repository and install it in editable mode with pip install -e ".[test,geometry,warp]".

See the example problems for some instructive examples.

Publications

This model has been described and benchmarked in the following publications. If this work has been useful to you, please cite [3]:

Contents: