JAX BSSN

JAX-BSSN is an autodifferentiable implementation of the BSSN evolution system for numerical relativity. JAX-BSSN features Kreiss–Oliger dissipation, momentum constraint damping, and fixed mesh refinement.

Installation

From the repository root:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e ".[demos,test]"

Install a JAX build appropriate for the intended CPU or accelerator when the default dependency does not match the target hardware.

Supported demos

The supported physical examples are:

1. Gauge wave
2. Linear wave
3. Spherical Cartoon puncture
4. Axisymmetric black hole head-on collision
5. Full 3D single charged black hole

Each demo owns its initial data and run configuration. There is no generic simulation CLI or installed initial-data dispatcher.

Run the periodic analytic gauge wave:

python demos/gravitational_waves/gauge_wave.py

Run the linear wave through a nested fixed-refinement hierarchy and write composite openPMD output:

python demos/gravitational_waves/linear_wave.py

Run the compact spherical Cartoon puncture and write its complete reflected axis to openPMD:

python demos/spherically_symmetric_cartoon_puncture/cartoon_puncture.py

Run the axisymmetric Cartoon head-on collision of two black holes:

python demos/axisymmetric_blackhole_collision/axisymmetric_blackhole_collision.py

Run the Cartesian Einstein--Maxwell evolution of one charged puncture on one explicitly selected GPU:

python demos/single_charged_blackhole/run_single_charged_blackhole.py

The charged demo defaults to a coarse 96^3 uniform grid with M=1 and Q=0.5. Edit its configuration block before a production run. It is an initial-data and evolution setup, not a validated long-time calculation; it does not include mesh refinement, restart, or wave extraction. Run its CPU-scale analytic, one-step, and output checks with:

JAX_PLATFORMS=cpu python -m pytest tests/test_charged_puncture.py

The default production runs compile substantial JAX kernels. The demo guide describes the evolution and configuration options.

Package structure

JAX_BSSN/
├── bssn/          # state, tensor algebra, geometry, equations, raw constraints
├── cartoon/       # spherical and axisymmetric reconstruction/evolution
├── EM/            # Cartesian first-order Maxwell evolution and BSSN coupling
├── evolution/     # derivatives, boundaries, RHS assembly, projections, RK4
├── fmr/           # nested-patch geometry, transfers, stage-synchronous RK4
├── diagnostics/   # reductions, reporting, plotting, openPMD
└── utilities/     # reusable initial-data and numerical helper routines

Demos own their physical initial data and configuration; reusable Bowen–York helpers live in JAX_BSSN.utilities. Analytic reference data lives under tests/, while integration tests also exercise demo runners.

Tests

Install test dependencies and run the complete suite with:

python -m pip install -e ".[test]"
python -m pytest

Tests use double precision and write no simulation output except dedicated openPMD round-trip tests in temporary directories.