API And Source Reference

High-Level Class

class booz_xform_jax.core.Booz_xform(nfp=1, asym=False, verbose=1, mpol=0, ntor=0, mnmax=0, xm=None, xn=None, xm_nyq=None, xn_nyq=None, mpol_nyq=None, ntor_nyq=None, mnmax_nyq=None, s_in=None, iota=None, rmnc=None, rmns=None, zmnc=None, zmns=None, lmnc=None, lmns=None, bmnc=None, bmns=None, bsubumnc=None, bsubumns=None, bsubvmnc=None, bsubvmns=None, Boozer_I_all=None, Boozer_G_all=None, phip=None, chi=None, pres=None, phi=None, aspect=0.0, toroidal_flux=0.0, compute_surfs=None, s_b=None, ns_in=None, ns_b=None, Boozer_I=None, Boozer_G=None, mboz=None, nboz=None, mnboz=None, xm_b=None, xn_b=None, bmnc_b=None, bmns_b=None, rmnc_b=None, rmns_b=None, zmnc_b=None, zmns_b=None, numnc_b=None, numns_b=None, gmnc_b=None, gmns_b=None, _prepared=False)[source]

Bases: object

Class implementing the Boozer coordinate transformation using JAX.

Instances of Booz_xform encapsulate all data required to convert the spectral representation of a VMEC equilibrium (in VMEC angles) to a spectral representation in Boozer coordinates.

Typical usage

>>> bx = Booz_xform()
>>> bx.read_wout("wout_mycase.nc", flux=True)   # or init_from_vmec(...)
>>> bx.register_surfaces([0.2, 0.5, 0.8])       # select surfaces in s-space
>>> bx.run()
>>> bx.write_boozmn("boozmn_mycase.nc")

After run() completes, the Boozer spectra are stored in attributes like bmnc_b, rmnc_b, zmns_b, etc., and the Boozer I/G profiles and radial grid on Boozer_I, Boozer_G, and s_b.

nfp

Field periodicity (number of field periods) of the equilibrium.

Type:

int

asym

Whether the VMEC equilibrium is non-stellarator-symmetric. If False, only the symmetric Fourier coefficients are used (cosine/sine combinations that respect stellarator symmetry). If True, additional “ns” arrays are populated and used.

Type:

bool

verbose

Controls diagnostic printing during run(). Historically this was an integer (0, 1, 2, …). In this implementation any truthy value enables basic per-surface diagnostics; setting verbose > 1 prints additional information.

Type:

int or bool

mpol, ntor

Maximum poloidal and toroidal mode numbers in the non-Nyquist VMEC spectrum, read from the wout file.

Type:

int

mnmax

Total number of non-Nyquist VMEC Fourier modes. For symmetric equilibria this is typically mpol * (2*ntor + 1).

Type:

int

xm, xn

Mode list for the non-Nyquist VMEC spectrum: poloidal and toroidal mode numbers (with xn stored as \(n n_{fp}\) to match VMEC conventions).

Type:

ndarray of int, shape (mnmax,)

xm_nyq, xn_nyq

Mode list for the Nyquist spectrum used to reconstruct w and B. Sizes and ranges mirror those in the original BOOZ_XFORM.

Type:

ndarray of int

mpol_nyq, ntor_nyq, mnmax_nyq

Nyquist resolutions and total number of Nyquist modes in the VMEC input, read from the wout file.

Type:

int

s_in

Radial coordinate values on the VMEC half grid (excluding the magnetic axis). This is stored as a NumPy array (host side) so that we can use standard Python indexing and numpy.argmin when mapping floating-point s values to nearest indices.

Type:

ndarray, shape (ns_in,)

iota

Rotational transform on the VMEC half grid.

Type:

jax.numpy.ndarray, shape (ns_in,)

rmnc, rmns, zmnc, zmns, lmnc, lmns

Non-Nyquist VMEC Fourier coefficients on the half grid, with dimensions (mnmax, ns_in). Asymmetric quantities are set to None when asym is False.

Type:

jax.numpy.ndarray

bmnc, bmns, bsubumnc, bsubumns, bsubvmnc, bsubvmns

Nyquist VMEC Fourier coefficients on the half grid, with dimensions (mnmax_nyq, ns_in). Asymmetric quantities are set to None when asym is False. These are used to reconstruct B and the covariant components B_θ, B_ζ.

Type:

jax.numpy.ndarray

Boozer_I_all, Boozer_G_all

Boozer I(s) and G(s) profiles on the full half grid. These correspond to the m=0, n=0 components of bsubumnc and bsubvmnc and are stored as NumPy arrays.

Type:

ndarray, shape (ns_in,)

phip, chi, pres, phi

Optional radial profiles read from the VMEC file when the flux flag is passed to read_wout(). They are not used directly in the Boozer transform but are convenient to have available for post-processing.

Type:

jax.numpy.ndarray, shape (ns_in,), optional

aspect

Aspect ratio of the equilibrium (copied from VMEC).

Type:

float

toroidal_flux

Total toroidal flux of the equilibrium (copied from VMEC).

Type:

float

compute_surfs

Indices of the half-grid surfaces on which to compute the Boozer transform. Indices run from 0 to ns_in-1. None (default) means “all surfaces”.

Type:

list[int] or None

s_b

Radial coordinate values on the subset of surfaces selected by compute_surfs. Populated by run() and read_boozmn().

Type:

ndarray, shape (ns_b,)

ns_in

Number of half-grid surfaces (excluding the axis) in the VMEC input.

Type:

int

ns_b

Number of surfaces selected for the Boozer transform (i.e. len(compute_surfs)).

Type:

int

Boozer_I, Boozer_G

Boozer I and G profiles restricted to the selected surfaces.

Type:

ndarray, shape (ns_b,)

mboz, nboz

Maximum poloidal and toroidal mode numbers in the Boozer spectrum. If not explicitly set by the user, these default to mpol and ntor respectively (mirroring the original BOOZ_XFORM behaviour).

Type:

int

mnboz

Total number of Boozer harmonics retained. The enumeration follows the original code:

  • m runs from 0, 1, …, mboz-1

  • for m = 0, n runs 0, 1, …, nboz

  • for m > 0, n runs -nboz, …, -1, 0, 1, …, nboz

The toroidal index is stored as xn_b = n * nfp.

Type:

int

xm_b, xn_b

Boozer mode list as described above.

Type:

ndarray of int, shape (mnboz,)

bmnc_b, bmns_b, rmnc_b, rmns_b, zmnc_b, zmns_b,
numnc_b, numns_b, gmnc_b, gmns_b

Boozer Fourier coefficients on the selected surfaces. Each has shape (mnboz, ns_b). Asymmetric arrays are None when asym is False. The “c” suffix denotes cosine-like coefficients and the “s” suffix sine-like coefficients, following the usual VMEC/BOOZ_XFORM conventions.

Type:

ndarray

_prepared

Internal flag indicating whether the angular grids and related bookkeeping (θ, ζ, grid sizes) have been initialised.

Type:

bool

init_from_vmec(*args, s_in=None)[source]

Load Fourier data from VMEC into this instance.

This method simply delegates to booz_xform_jax.vmec.init_from_vmec(). See that function for the full list of arguments and options.

Parameters:

  • *args – Passed directly to init_from_vmec().

  • s_in (ndarray | None) – Optional replacement radial grid of normalised toroidal flux. If provided, its first element should correspond to the axis; this element will be discarded so that s_in[0] on the instance is the first half-grid surface away from the axis.

Return type:

None

read_wout(filename, flux=False)[source]

Read a VMEC wout file and populate the internal arrays.

This is a thin wrapper around booz_xform_jax.vmec.read_wout(). In addition to the core Fourier coefficients needed for the Boozer transform, optional flux profile arrays can be loaded when flux=True.

Parameters:

  • filename (str) – Path to the VMEC wout file.

  • flux (bool) – If True, also read radial profile arrays (φ’, χ, p, …).

Return type:

None

read_wout_data(wout, flux=False)[source]

Populate the instance from an in-memory VMEC wout object.

This is a thin wrapper around booz_xform_jax.vmec.read_wout_data().

Parameters:

  • wout – A VMEC wout-like object (e.g. vmec_jax.WoutData).

  • flux (bool) – If True, also read radial profile arrays (φ’, χ, p, …) when available.

Return type:

None

write_boozmn(filename)[source]

Write the computed Boozer spectra to a NetCDF file.

This delegates to booz_xform_jax.io_utils.write_boozmn(). The file format (NetCDF3 vs NetCDF4) depends on the availability of the netCDF4 package and mirrors the behaviour of the original BOOZ_XFORM code.

Parameters:

filename (str)

Return type:

None

read_boozmn(filename)[source]

Read Boozer spectra from an existing boozmn file.

This delegates to booz_xform_jax.io_utils.read_boozmn() and populates the current instance with the data from that file, including mode definitions, radial profiles, and Boozer spectra.

Parameters:

filename (str)

Return type:

None

run(jit=False)[source]

Perform the Boozer coordinate transformation on selected surfaces.

Parameters:

jit (bool, optional) –

Placeholder flag (currently unused). The transform is implemented entirely in terms of JAX array operations (jax.numpy and einsum). To avoid large compile times on CPU, we do not wrap the entire run() in a single jax.jit() by default. Small helpers such as _init_trig() are jitted.

Advanced users who want full JIT compilation can wrap run() externally, but should be aware that this may lead to long compilation times for large Boozer resolutions.

Return type:

None

Notes

The implementation follows the algorithm outlined in the module docstring and in the BOOZ_XFORM documentation. The main difference from a direct translation of the Fortran/C++ code is that all loops over Fourier modes are vectorised. Only the loop over radial surfaces remains as a Python loop.

run_jax(*, jit=True)[source]

Run a JAX-native Boozer transform (no Python surface loop).

This method returns a mapping compatible with boozmn field names. The mapping includes gmnc_b and its BOOZ_XFORM-compatible gmn_b alias for Boozer Jacobian harmonics, plus asymmetric spectra when asym is true. It is intended for end-to-end JIT/differentiable workflows and does not populate the instance attributes (unlike run).

Parameters:

jit (bool)

Return type:

dict

register_surfaces(s)[source]

Register one or more surfaces on which to compute the transform.

This method mirrors the original C++ register routine. It accepts either integer half-grid indices or floating-point radial coordinate values in normalised toroidal flux space.

Parameters:

s (int, float, or iterable of these) –

Surfaces to register:

  • If an integer, it is interpreted as an index on the VMEC half grid (0 ≤ index < ns_in).

  • If a float, it should lie in [0, 1] and is interpreted as a normalised toroidal flux value. We then choose the nearest index based on self.s_in.

Return type:

None

Notes

  • Any new surfaces are appended to the existing compute_surfs list (duplicates are removed).

  • Surfaces outside the valid index range produce a ValueError.

  • The method does not perform the transform; you must call run() afterwards.

Parameters:

  • nfp (int)

  • asym (bool)

  • verbose (int | bool)

  • mpol (int)

  • ntor (int)

  • mnmax (int)

  • xm (ndarray | None)

  • xn (ndarray | None)

  • xm_nyq (ndarray | None)

  • xn_nyq (ndarray | None)

  • mpol_nyq (int | None)

  • ntor_nyq (int | None)

  • mnmax_nyq (int | None)

  • s_in (ndarray | None)

  • iota (Array | None)

  • rmnc (Array | None)

  • rmns (Array | None)

  • zmnc (Array | None)

  • zmns (Array | None)

  • lmnc (Array | None)

  • lmns (Array | None)

  • bmnc (Array | None)

  • bmns (Array | None)

  • bsubumnc (Array | None)

  • bsubumns (Array | None)

  • bsubvmnc (Array | None)

  • bsubvmns (Array | None)

  • Boozer_I_all (ndarray | None)

  • Boozer_G_all (ndarray | None)

  • phip (Array | None)

  • chi (Array | None)

  • pres (Array | None)

  • phi (Array | None)

  • aspect (float)

  • toroidal_flux (float)

  • compute_surfs (List[int] | None)

  • s_b (ndarray | None)

  • ns_in (int | None)

  • ns_b (int | None)

  • Boozer_I (ndarray | None)

  • Boozer_G (ndarray | None)

  • mboz (int | None)

  • nboz (int | None)

  • mnboz (int | None)

  • xm_b (ndarray | None)

  • xn_b (ndarray | None)

  • bmnc_b (ndarray | None)

  • bmns_b (ndarray | None)

  • rmnc_b (ndarray | None)

  • rmns_b (ndarray | None)

  • zmnc_b (ndarray | None)

  • zmns_b (ndarray | None)

  • numnc_b (ndarray | None)

  • numns_b (ndarray | None)

  • gmnc_b (ndarray | None)

  • gmns_b (ndarray | None)

  • _prepared (bool)

Functional JAX API

Pure JAX API for end-to-end Boozer transforms.

This module provides a JIT-friendly, functional interface that avoids Python loops over surfaces and keeps all arrays in JAX. It is intended for end-to-end differentiation with vmec_jax -> booz_xform_jax -> neo_jax.

class booz_xform_jax.jax_api.BoozXformConstants(nfp, mboz, nboz, asym, ntheta, nzeta, nu2_b, mmax_non, nmax_non, mmax_nyq, nmax_nyq)[source]

Bases: object

Static constants for the JAX Boozer transform.

Parameters:
class booz_xform_jax.jax_api.BoozXformGrids(theta_grid, zeta_grid, xm_b, xn_b)[source]

Bases: object

Grid arrays for the JAX Boozer transform.

Parameters:

  • theta_grid (Array)

  • zeta_grid (Array)

  • xm_b (Array)

  • xn_b (Array)

booz_xform_jax.jax_api.prepare_booz_xform_constants(*, nfp, mboz, nboz, asym, xm, xn, xm_nyq, xn_nyq)[source]

Compute static constants for the JAX Boozer transform.

This helper runs on the host and can be used before JIT compilation.

Parameters:
Return type:

tuple[BoozXformConstants, BoozXformGrids]

booz_xform_jax.jax_api.prepare_booz_xform_constants_from_inputs(*, inputs, mboz, nboz, asym)[source]

Convenience wrapper using a VMEC -> Boozer input bundle.

Parameters:
Return type:

tuple[BoozXformConstants, BoozXformGrids]

booz_xform_jax.jax_api.booz_xform_jax_impl(rmnc, zmns, lmns, bmnc, bsubumnc, bsubvmnc, iota, *, xm, xn, xm_nyq, xn_nyq, constants, grids, rmns=None, zmnc=None, lmnc=None, bmns=None, bsubumns=None, bsubvmns=None, surface_indices=None)[source]

JAX-native Boozer transform over all (or selected) surfaces.

All inputs must be JAX arrays with surface dimension first, i.e. shape (ns, mn_non) for non-Nyquist arrays and (ns, mn_nyq) for Nyquist arrays.

Parameters:

  • rmnc (Array)

  • zmns (Array)

  • lmns (Array)

  • bmnc (Array)

  • bsubumnc (Array)

  • bsubvmnc (Array)

  • iota (Array)

  • xm (Array)

  • xn (Array)

  • xm_nyq (Array)

  • xn_nyq (Array)

  • constants (BoozXformConstants)

  • grids (BoozXformGrids)

  • rmns (Array | None)

  • zmnc (Array | None)

  • lmnc (Array | None)

  • bmns (Array | None)

  • bsubumns (Array | None)

  • bsubvmns (Array | None)

  • surface_indices (Array | None)

Return type:

dict

booz_xform_jax.jax_api.booz_xform_from_inputs(*, inputs, constants, grids, surface_indices=None, jit=True)[source]

Run the JAX Boozer transform using a VMEC -> Boozer input bundle.

Parameters:
Return type:

dict

booz_xform_jax.jax_api.booz_xform_jax(*, rmnc, zmns, lmns, bmnc, bsubumnc, bsubvmnc, iota, xm, xn, xm_nyq, xn_nyq, nfp, mboz, nboz, asym=False, rmns=None, zmnc=None, lmnc=None, bmns=None, bsubumns=None, bsubvmns=None, surface_indices=None)[source]

Host-side convenience wrapper for booz_xform_jax_impl().

This wrapper computes static constants on the host (NumPy) and returns a JAX output dictionary. For full JIT, call booz_xform_jax_impl() directly with precomputed constants.

Parameters:

  • rmnc (Array)

  • zmns (Array)

  • lmns (Array)

  • bmnc (Array)

  • bsubumnc (Array)

  • bsubvmnc (Array)

  • iota (Array)

  • xm (Sequence[int])

  • xn (Sequence[int])

  • xm_nyq (Sequence[int])

  • xn_nyq (Sequence[int])

  • nfp (int)

  • mboz (int)

  • nboz (int)

  • asym (bool)

  • rmns (Array | None)

  • zmnc (Array | None)

  • lmnc (Array | None)

  • bmns (Array | None)

  • bsubumns (Array | None)

  • bsubvmns (Array | None)

  • surface_indices (Sequence[int] | None)

Return type:

dict

CLI Driver

Legacy-compatible command line interface for booz_xform_jax.

This module mirrors the STELLOPT xbooz_xform driver:

xbooz_xform <infile> [T|F]

where <infile> contains:

  1. mboz nboz on the first line,

  2. a VMEC extension or wout filename on the second line,

  3. an optional whitespace-separated list of full-grid surface indices.

booz_xform_jax.cli.run_from_legacy_input(input_file, *, screen_output=True)[source]

Run booz_xform_jax from an STELLOPT-style in_booz file.

Parameters:
Return type:

Path

booz_xform_jax.cli.main(argv=None)[source]

Command line entrypoint.

Parameters:

argv (list[str] | None)

Return type:

int

VMEC Input Utilities

VMEC input routines for the JAX implementation of booz_xform.

This module contains functions for loading data from VMEC output files and for initialising a Booz_xform instance with that data. The goal is to mimic the behaviour of the booz_xform C++ code while providing a Pythonic interface.

The functions defined here operate on instances of Booz_xform. They are not intended to be called standalone; instead, call the corresponding methods on a Booz_xform object (init_from_vmec and read_wout), which delegate to these functions.

booz_xform_jax.vmec.init_from_vmec(self, *args, s_in=None)[source]

Initialise a Booz_xform instance with VMEC data.

This function accepts two calling conventions for compatibility with the original C++ init_from_vmec function:

  1. init_from_vmec(ns, iotas, rmnc, rmns, zmnc, zmns, lmnc, lmns, bmnc, bmns, bsubumnc, bsubumns, bsubvmnc, bsubvmns[, phip, chi, pres, phi]) where ns is an integer giving the number of radial surfaces on the full VMEC grid (including the axis) and the remaining arguments are 1D or 2D arrays with shapes matching the VMEC documentation. The zeroth radial entry (the axis) is discarded. Optional arrays phip, chi, pres and phi may appear at the end of the argument list; if present, all four must be provided.

  2. init_from_vmec(iotas, rmnc, rmns, zmnc, zmns, lmnc, lmns, bmnc, bmns, bsubumnc, bsubumns, bsubvmnc, bsubvmns[, phip, chi, pres, phi]) omitting the ns argument. In this case the length of the iotas array determines the number of surfaces on the full grid. Optional flux arrays may appear at the end as in the first calling convention.

Parameters:

  • self (Booz_xform) – The instance to initialise.

  • *args (sequence) – Positional arguments following one of the two calling conventions described above.

  • s_in (ndarray, optional) – Optional array of length ns giving the values of normalized toroidal flux on the full radial grid. If provided, the first entry should correspond to the axis and will be discarded. When not provided, a uniform grid between 0 and 1 is used.

Return type:

None

Notes

The zeroth radial entry (corresponding to the magnetic axis) is ignored in all arrays. The remaining ns - 1 surfaces constitute the half‑grid on which the Boozer transform is performed. All input arrays are defensively copied and converted to jax.numpy.DeviceArray objects of dtype float64 for computation; however, the radial coordinate array s_in is stored as a NumPy array so that it can be indexed using Python lists in run() and read_boozmn().

booz_xform_jax.vmec.read_wout(self, filename, flux=False)[source]

Read a VMEC wout file and populate the internal arrays.

This routine loads the equilibrium data from a VMEC NetCDF file (the file whose name begins with wout_). The Fourier mode definitions, the non‑Nyquist and Nyquist Fourier coefficients, and optional flux profiles are read. Once the data are assembled, this function calls init_from_vmec() on the instance to prepare the arrays for the Boozer transformation.

Parameters:

  • self (Booz_xform) – The instance to populate.

  • filename (str) – Path to a VMEC wout NetCDF file.

  • flux (bool, optional) – If True, the flux profile arrays phipf (or phips), chi, pres and phi are read and passed to init_from_vmec(). When False, these arrays are ignored.

Return type:

None

booz_xform_jax.vmec.read_wout_data(self, wout, flux=False)[source]

Populate the instance from a VMEC wout-like object (e.g. vmec_jax.WoutData).

This mirrors read_wout() but accepts an in-memory object with the standard VMEC fields as attributes instead of reading a NetCDF file.

Parameters:

  • self (Booz_xform) – The instance to populate.

  • wout (object) – A wout-like object with VMEC attributes (e.g. vmec_jax.WoutData).

  • flux (bool, optional) – If True, attempt to populate flux profile arrays. When the required fields are unavailable, the profiles are silently skipped to match read_wout behavior.

Return type:

None

boozmn File I/O

Input/output utilities for the JAX implementation of booz_xform.

This module provides functions to read from and write to the NetCDF formats used by the original booz_xform code. In particular it handles the boozmn file format, which stores Boozer‐coordinate Fourier spectra and associated radial profiles. These functions operate on Booz_xform instances and are used by the methods booz_xform_jax.core.Booz_xform.write_boozmn() and booz_xform_jax.core.Booz_xform.read_boozmn().

booz_xform_jax.io_utils.write_boozmn(self, filename)[source]

Write the computed Boozer Fourier spectra to a boozmn NetCDF file.

The boozmn format is used by the original booz_xform package to store Boozer coordinates and related quantities. This function writes the essential information required to reconstruct the Boozer harmonics: the mode definitions, radial profiles, and spectral coefficients. If the instance was initialised from a VMEC equilibrium and booz_xform_jax.core.Booz_xform.run() has been called, this routine will create a NetCDF file that can be read by the original booz_xform or by read_boozmn() below.

Parameters:

  • self (Booz_xform) – The instance containing the results of the Boozer transform.

  • filename (str) – Path of the output NetCDF file.

Return type:

None

booz_xform_jax.io_utils.read_boozmn(self, filename)[source]

Read Boozer Fourier data from a boozmn NetCDF file.

This routine populates a Booz_xform instance with data from a file produced by the original booz_xform program or by write_boozmn(). It reads the mode definitions, radial profiles and spectral arrays, reorienting the latter into the internal (mnboz, ns_b) layout. Existing data on the instance will be overwritten.

Parameters:

  • self (Booz_xform) – The instance to populate.

  • filename (str) – Path to a boozmn NetCDF file.

Return type:

None

Plotting Utilities

Plotting utilities for the JAX version of the Boozer transform.

This module provides convenience functions for visualising the output of the booz_xform_jax.booz_xform.Booz_xform class. The API is modelled on the plotting functions in the original booz_xform package. All functions accept either a Booz_xform instance or the path to a boozmn file. Matplotlib is used for creating figures. Plotly support is not included in this version.

The following functions are provided:

surfplot

Show the magnetic field strength B on a flux surface versus the Boozer poloidal and toroidal angles.

symplot

Plot the radial variation of all Fourier modes of B grouped by symmetry type.

modeplot

Plot the radial variation of the largest Fourier modes of B.

wireplot

Make a 3D figure showing curves of constant Boozer angles on a flux surface.

These functions are optional; if matplotlib is not available they will raise an exception at runtime.

booz_xform_jax.plots.handle_b_input(b)[source]

Convert input to a Booz_xform instance.

Parameters:

b (str or Booz_xform) – If a string, it is interpreted as a path to a boozmn file which will be read into a new Booz_xform instance. If an instance of Booz_xform is provided it is returned unchanged.

Returns:

The corresponding Booz_xform instance.

Return type:

Booz_xform

booz_xform_jax.plots.surfplot(b, js=0, fill=True, ntheta=50, nphi=90, ncontours=25, *, ax=None, show=True, savefig=None, **kwargs)[source]

Plot B on a flux surface versus the Boozer angles.

This function generates a 2D contour plot of the magnetic field strength B as a function of the Boozer toroidal angle ϕ and poloidal angle θ on the specified surface.

Parameters:

  • b (str or Booz_xform) – The Booz_xform instance to plot, or a filename of a boozmn file.

  • js (int, optional) – Index among the output surfaces to plot (default is 0).

  • fill (bool, optional) – Whether to fill contours (default True). If False, line contours are drawn instead.

  • ntheta (int, optional) – Number of grid points in the poloidal angle.

  • nphi (int, optional) – Number of grid points in the toroidal angle.

  • ncontours (int, optional) – Number of contour levels.

  • **kwargs – Additional keyword arguments passed to matplotlib’s contour functions.

  • ax (Axes | None)

  • show (bool)

  • savefig (str | None)

Return type:

None

booz_xform_jax.plots.symplot(b, max_m=20, max_n=20, ymin=None, sqrts=False, log=True, B0=True, helical_detail=False, legend_args=None, *, ax=None, show=True, savefig=None, **kwargs)[source]

Plot radial variation of all Fourier modes of B grouped by symmetry.

The modes are coloured according to whether m=0, n=0, or both. Modes with m and n exceeding the specified maxima are omitted.

Parameters:

  • b (str or Booz_xform) – The Booz_xform instance to plot, or a filename of a boozmn file.

  • max_m (int) – Maximum poloidal mode number to include.

  • max_n (int) – Maximum toroidal mode number (divided by nfp) to include.

  • ymin (float, optional) – Lower limit for the y axis in logarithmic plots. If not specified, a value is chosen automatically.

  • sqrts (bool) – Whether to plot sqrt(s) on the x axis instead of s.

  • log (bool) – Use a logarithmic y axis.

  • B0 (bool) – Include the (m,n)=(0,0) mode in the figure.

  • helical_detail (bool) – Whether to show details of helical modes separately.

  • legend_args (dict, optional) – Additional arguments to pass to plt.legend.

  • **kwargs – Additional keyword arguments passed to plt.plot.

  • ax (Axes | None)

  • show (bool)

  • savefig (str | None)

Return type:

None

booz_xform_jax.plots.modeplot(b, nmodes=10, ymin=None, sqrts=False, log=True, B0=True, legend_args=None, *, ax=None, show=True, savefig=None, **kwargs)[source]

Plot the radial variation of the largest Fourier modes of B.

Modes are sorted by amplitude at the outermost computed surface and the largest nmodes are displayed.

Parameters:

  • b (str or Booz_xform) – The Booz_xform instance or path to a boozmn file.

  • nmodes (int, optional) – Number of modes to display.

  • ymin (float, optional) – Lower limit for the y axis in logarithmic plots. If not specified, a value is chosen automatically.

  • sqrts (bool) – Whether to plot sqrt(s) on the x axis instead of s.

  • log (bool) – Use a logarithmic y axis.

  • B0 (bool) – Whether to include the (m,n)=(0,0) mode.

  • legend_args (dict, optional) – Additional arguments passed to plt.legend.

  • **kwargs – Additional keyword arguments passed to plt.plot.

  • ax (Axes | None)

  • show (bool)

  • savefig (str | None)

Return type:

None

booz_xform_jax.plots.wireplot(b, js=None, ntheta=30, nphi=80, refine=3, surf=True, orig=True, *, ax=None, show=True, savefig=None)[source]

Make a 3D wireframe of curves of constant Boozer angles.

This function recreates the 3D wireframe plot from the original booz_xform code. Curves of constant Boozer poloidal angle θ and constant toroidal angle ϕ are drawn on a specified flux surface. Optionally, the original VMEC coordinate curves are overlayed. Requires matplotlib with 3D support.

Parameters:

  • b (str or Booz_xform) – The Booz_xform instance or path to a boozmn file.

  • js (int, optional) – The index among the output surfaces to plot. If None, the outermost surface is selected.

  • ntheta (int, optional) – Number of curves in the poloidal angle.

  • nphi (int, optional) – Number of curves in the toroidal angle.

  • refine (int, optional) – Number of interpolation points between contours. A higher value yields smoother curves.

  • surf (bool, optional) – Whether to display a semi‑transparent surface representing the flux surface.

  • orig (bool, optional) – Whether to overlay curves of constant VMEC angles.

  • ax (Axes | None)

  • show (bool)

  • savefig (str | None)

Return type:

None