FAX

FAX is a lean JAX code for shear-Alfvén physics in stellarators: the continuum, AE3D-form global modes and a perturbative fast-ion drive, with first-order derivatives with respect to profiles and, through VMEX, the VMEC boundary.

Left: the n = 2 CTOK continuum with the six global modes nearest 384 kHz (red). Right: the passing-only drive of the Könies et al. (2018) TAE.

Scope

The drive is perturbative and valid up to about 200 keV in the Könies case ($\gamma/\omega \le 4%$). It is about 0.6 of the CAS3D-K zero-orbit-width curve, plausibly the missing trapped share; that attribution is not established. Derivatives are first order and are NaN unless the eigensolve converged and its selection is not a tie.

Install

Python 3.10 or newer, JAX 0.4.38 or newer, solvax>=0.32.1. FAX is on PyPI as alfvex; the code is imported as fax. FAX will soon be renamed ALFVEX everywhere (repository, import and command line). Do not pip install fax: that is an unrelated package.

pip install alfvex            # core
pip install "alfvex[vmex]"    # VMEC solves and boundary gradients (Python >= 3.11)
pip install "alfvex[plot]"    # matplotlib, for the figure examples

From a clone, pip install . (with the same extras) installs the development version.

Quickstart

This runs from an installed wheel with no data files: a continuum, a global TAE and one gradient on the analytic Könies et al. (2018) torus.

import jax, jax.numpy as jnp, fax

s = (jnp.arange(200) + 0.5) / 200
geo = fax.circular_torus(s=s, major_radius=10.0, minor_radius=1.0, B0=3.0, iota=1 / (1.71 + 0.16 * s), epsilon_B=0.1)
harmonics, rho = fax.mode_family(6, 16, 5), fax.mass_density(2e19, 1.0)
f_hz = fax.continuum(geo, harmonics, rho)  # (surfaces, harmonics)
modes = fax.eigenmodes(geo, harmonics, rho, target_hz=6.46e4, count=3, ntheta=32, nzeta=4)
print(modes.frequencies_hz[-1], modes.status.converged)  # the TAE, 64.63 kHz
f = lambda r: fax.frequency_hz(geo, harmonics, r, 2, 3, target_hz=6.46e4, count=3, ntheta=32, nzeta=4)
print(jax.grad(f)(rho))  # equals -f / (2 rho)

The command line runs the continuum and global-mode tasks from a TOML case; the drive, objectives and boundary gradients are Python only.

Examples (checkout only)

These read files in the repository or need an extra; each has editable settings at the top and no command-line parsing.

Evidence

The validation page has every gate, its test and what is not done.

Ecosystem

SOLVAX provides the generalized eigensolver and its implicit derivatives; VMEX the equilibrium and its implicit derivative; booz_xform_jax the Boozer transform. ESSOS orbit transport of FAX modes is planned, not implemented.

Documentation and testing

fax.readthedocs.io: getting started and TOML reference, tutorials, theory, validation, API with its derivative contract.

pip install -e ".[dev,vmex]"
python -m pytest -q -m 'not scheduled'
ruff check . && ruff format --check . && sphinx-build -W -b html docs docs/_build/html

Citing and license

Cite the repository through CITATION.cff and the papers whose models FAX implements (references). MIT license.