Testing and CI

NEO_JAX is tested at several levels: low-level numerics, solver behavior on reference fixtures, CLI behavior, performance guardrails, and optional GPU smoke tests.

Testing layers

The test suite is organized around complementary questions:

  • do low-level spline, Fourier, and I/O routines produce the expected arrays?

  • do solver outputs remain stable on curated reference cases?

  • do the user-facing Python API and CLI behave as documented?

  • do optional GPU and performance checks stay within acceptable bounds?

Representative test files include:

Test file

Coverage

tests/unit/test_api.py

Public API behavior, result accessors, surface mapping, and JAX-path return types.

tests/unit/test_control.py

Control-file parsing.

tests/regression/test_constellaration_guard.py

Low-|iota| safeguards and approximate fallback behavior.

tests/regression/test_cli_legacy.py

CLI file generation, progress logging, and parity against committed xneo reference fixtures for the file-based workflow.

tests/regression/test_landreman_qa_lowres_parity.py

Dense QA fixture comparison.

tests/regression/test_orbits_parity.py

ORBITS reference behavior.

tests/regression/test_ncsx_parity.py

NCSX comparison case.

tests/regression/test_gpu_smoke.py

Optional CPU-versus-GPU agreement checks.

Local test commands

Common local validation commands are:

pytest -q tests/unit/test_api.py
pytest -q tests/regression/test_constellaration_guard.py
pytest -q tests/regression/test_cli_legacy.py
pytest -q tests/regression/test_landreman_qa_lowres_parity.py
pytest -q tests/regression/test_orbits_parity.py
pytest -q tests/regression/test_ncsx_parity.py

The documentation build is also part of the release workflow:

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

Continuous integration

The repository CI runs on GitHub Actions. The workflow installs the package, pulls the VMEC and Boozer dependencies used by the pipeline tests, runs the full pytest suite on CPU, and executes a small performance regression check.

name: ci

on:
  push:
  pull_request:

jobs:
  tests:
    name: Tests (Python ${{ matrix.python-version }})
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.10", "3.11", "3.12"]
    env:
      JAX_PLATFORM_NAME: cpu
      JAX_ENABLE_X64: "1"
      NEO_JAX_FETCH_EXTERNAL_FIXTURES: "1"
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
          cache: pip
      - name: Install package and test dependencies
        run: |
          python -m pip install --upgrade pip
          python -m pip install ".[dev]"
          python -m pip install vmec_jax booz_xform_jax
      - name: Run tests
        run: pytest -q
      - name: Smoke test installed CLI
        run: |
          mkdir -p /tmp/neo_jax_cli_smoke
          cp tests/fixtures/orbits/boozmn_ORBITS_FAST.nc /tmp/neo_jax_cli_smoke/boozmn_ORBITS_SMOKE.nc
          cat > /tmp/neo_jax_cli_smoke/neo_in.ORBITS_SMOKE <<'EOF'
          ! NEO control file (auto-generated)
          ! line2
          ! line3
          boozmn
          neo_out.ORBITS_SMOKE
          1
          96
          2
          2
          0
          0
          2
          1
          0.05
          4
          2
          4
          6
          0
          1
          0
          0
          2
          0
          0
          0
          0
          0
          0
          0
          0
          0
          neo_cur.ORBITS_SMOKE
          200
          2
          0
          EOF
          cd /tmp/neo_jax_cli_smoke
          neo-jax ORBITS_SMOKE --quiet
          test -f neo_out.ORBITS_SMOKE

  docs:
    name: Docs
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"
          cache: pip
      - name: Install package and docs dependencies
        run: |
          python -m pip install --upgrade pip
          python -m pip install ".[docs]"
          python -m pip install vmec_jax booz_xform_jax
      - name: Build docs
        run: python -m sphinx -W -b html docs docs/_build/html

What the CI does not do by default

Some checks are intentionally opt-in:

  • full slow reference cases behind NEO_JAX_RUN_SLOW=1

  • GPU smoke tests behind NEO_JAX_RUN_GPU=1

  • external NCSX fixture consumers behind NEO_JAX_FETCH_EXTERNAL_FIXTURES=1

That separation keeps standard CI fast while still preserving the heavier validation workflows for release and benchmarking. The GitHub Actions CPU test job sets NEO_JAX_FETCH_EXTERNAL_FIXTURES=1 so the NCSX regression path still runs in CI even though the large Boozer file is not committed in the repository.

Performance guardrails

The CI workflow also runs benchmarks/ci_perf_check.py with explicit limits on compile and reuse times. This is not a substitute for full benchmarking, but it catches large regressions in the compiled JAX path early.

Reference comparisons

When NEO_JAX is compared against established external reference outputs, the goal is to verify correctness of the current implementation on representative geometries. Those comparisons are summarized in Validation. The default CLI parity suite uses committed xneo reference fixtures so the checks stay fully reproducible on CI and on developer machines without a local STELLOPT build. The testing language in NEO_JAX is therefore evidence-driven: reference agreement is part of the acceptance story, while the user-facing solver remains NEO_JAX itself.