CLI

NEO_JAX provides a legacy-compatible command-line interface for STELLOPT NEO. The intent is that an existing xneo workflow can be pointed at the JAX implementation and keep the same control-file and output-file conventions for both the effective-ripple and parallel-current solves.

Installed entrypoints

After pip install neo-jax (or pip install -e . from a clone), the following commands are available:

neo-jax
xneo
xneo_jax
python -m neo_jax

The legacy form is:

xneo [extension]

Examples:

xneo ORBITS
xneo_jax LandremanPaul2021_QA_lowres
python -m neo_jax ORBITS_FAST

Control-file resolution

When an extension is supplied, NEO_JAX follows the same search order as the STELLOPT executable:

  1. neo_param.<extension>

  2. neo_param.in

  3. neo_in.<extension>

If no extension is supplied, the CLI looks for neo.in.

This behavior is implemented in neo_jax.io.resolve_control_path and is used by the CLI by default. You can still override it explicitly with --control.

Boozer input resolution

The input Boozer file depends on INP_SWI:

  • inp_swi = 0: NEO_JAX follows the legacy boozmn_<extension>.nc convention used by xneo and STELLOPT’s read_booz_in path.

  • inp_swi != 0: NEO_JAX resolves the path from IN_FILE in the control file.

You can override the automatic resolution with --boozmn.

Legacy output files

When the corresponding control switches are enabled, the CLI writes the same legacy files that STELLOPT writes:

  • main output: neo_out.*

  • parallel-current summary: neo_cur.* when calc_cur = 1

  • log: neolog.*

  • diagnostics: diagnostic.dat, diagnostic_add.dat, diagnostic_bigint.dat

  • convergence history: conver.dat

  • current-history dump: current.dat when write_cur_inte = 1

  • geometry / Fourier dumps: dimension.dat, theta_arr.dat, phi_arr.dat, mn_arr.dat, rmnc_arr.dat, zmns_arr.dat, lmns_arr.dat, bmnc_arr.dat, b_s_arr.dat, r_s_arr.dat, z_s_arr.dat, l_s_arr.dat, isqrg_arr.dat, sqrg11_arr.dat, kg_arr.dat, pard_arr.dat, r_tb_arr.dat, z_tb_arr.dat, p_tb_arr.dat, b_tb_arr.dat, r_pb_arr.dat, z_pb_arr.dat, p_pb_arr.dat, b_pb_arr.dat, gtbtb_arr.dat, gpbpb_arr.dat, gtbpb_arr.dat

The Fortran-style formatting used in the text files is implemented in neo_jax.legacy so that neo_out.*, neolog.*, diagnostic*.dat, neo_cur.*, and the covered conver.dat parity cases match the STELLOPT text output. current.dat follows the same token layout and special-value formatting (including gfortran’s omitted exponent letter for some 3-digit exponents) and is validated numerically token-by-token against committed xneo reference fixtures.

Compatibility scope

Supported legacy scope:

  • calc_cur = 0 effective-ripple runs

  • calc_cur = 1 parallel-current runs

  • control-file precedence: neo_param.<extension> -> neo_param.in -> neo_in.<extension>

  • legacy Boozer input naming via boozmn_<extension>.nc for inp_swi = 0

Examples

Using the shipped ORBITS fixture:

cd tests/fixtures/orbits
python -m neo_jax ORBITS_FAST

Using the dense Landreman/Paul QA fixture:

cd tests/fixtures/landreman_qa_lowres
xneo LandremanPaul2021_QA_lowres

Using a current-enabled ORBITS case:

cd /path/to/case
xneo ORBITS_CURINT
xneo_jax ORBITS_CURINT

Both commands will read the same control file and write the same neo_out.* / neo_cur.* outputs.

Progress logging

NEO_JAX prints high-level status lines by default so long runs do not look stalled. In addition to the control file, Boozer file, solve mode, and surface count, the CLI also reports the active JAX runtime:

NEO_JAX: surfaces=10 theta_n=64 phi_n=64 npart=40 backend=JAX
NEO_JAX: jax_runtime=gpu (2 devices: NVIDIA RTX A4000, NVIDIA RTX A4000)

Each surface now also prints a preflight summary before the solve starts:

NEO_JAX: surface 1/10 index=2 s=0.005000 sqrt(s)=0.070711 iota=8.944971e-05
NEO_JAX: resolution theta_n=64 phi_n=64 npart=40 multra=2 nstep_per=20 nstep_min=200 nstep_max=500
NEO_JAX: geometry nfp=1 nmodes=2048 B00=1.171188e+01 Bmin=2.595699e+00 Bmax=1.171188e+01
NEO_JAX: preflight approx_rational_field_periods=558974 approx_substeps=11179480 approx_eta_paths=894358400 limit=100000

Use --quiet to suppress these messages for batch jobs or benchmarking.

Optional arguments

NEO_JAX keeps the legacy positional interface but also exposes a few explicit overrides:

  • --control: use a specific control file path

  • --boozmn: use a specific boozmn path

  • --output: override the main output file name

  • --jax: prefer the JAX backend when compatible

  • --no-jax: force the Python backend

  • --verbose: print extra progress information

  • --quiet: suppress the default NEO_JAX progress messages

Debugging aids

For convergence-history debugging, set NEO_JAX_WRITE_IPMAX_DEBUG=1 before running the CLI. NEO_JAX will write diagnostic_ipmax_jax.dat in the working directory with one line per trapped-amplitude update that feeds conver.dat.

For near-zero-|iota| surfaces, NEO_JAX also enforces a preflight work limit. If the estimated rational-surface correction would require too many field periods, the default policy is to abort with a detailed explanation rather than appearing to hang. The default limit is controlled by NEO_JAX_MAX_RATIONAL_FIELD_PERIODS and defaults to 100000. Set that environment variable to 0 to disable the safeguard explicitly.

You can also request a controlled fallback instead of an error by setting:

export NEO_JAX_RATIONAL_SURFACE_POLICY=approximate

In approximate mode, NEO_JAX still performs the base integration but skips the expensive rational-surface correction if the preflight estimate exceeds the configured field-period limit. The returned diagnostics record that the approximation was used and include the estimated rational workload. The CLI also prints that an approximate result is being used and points to NEO_JAX_MAX_RATIONAL_FIELD_PERIODS=0 for the full exact legacy run.

To reduce the chance of entering this regime, avoid surfaces with near-zero |iota| when possible, or loosen acc_req for exploratory runs.

This dump is intended for parity debugging of the dense WRITE_INTEGRATE=1 cases and is exercised by the CLI regression suite.

Note that control-file WRITE_PROGRESS is honored by default, so legacy runs still print progress when the control file requests it. In addition, NEO_JAX prints its own high-level status messages by default so long JAX or parity runs do not look hung. Use --quiet to suppress those messages.

Testing

The CLI parity tests live in tests/regression/test_cli_legacy.py. They run the JAX CLI against committed xneo reference fixtures on multiple geometries and control-file layouts:

  • LandremanPaul2021_QA_lowres dense legacy fixture

  • synthetic ORBITS single-surface case with diagnostics and array dumps

  • synthetic ORBITS calc_cur = 1 case with current.dat

  • NCSX mini case

  • control-file precedence checks for neo_param.* vs neo_in.*

  • optional IPMAX debug-dump coverage for the convergence logger

Optional GPU smoke coverage lives in tests/regression/test_gpu_smoke.py. When NEO_JAX_RUN_GPU=1 is set on a machine with a visible JAX GPU backend, that test file runs both:

  • a legacy CLI CPU-vs-GPU comparison on a one-surface ORBITS case

  • a public Python API CPU-vs-GPU comparison through run_neo(...)

Additional dense ORBITS and NCSX solver parity tests are available behind NEO_JAX_RUN_SLOW=1 in the regression suite.