API Reference

Coordinates

Generic coordinate computation using Jacobian-based architecture.

This module provides a generic function to compute magnetic coordinates for any coordinate system by using the appropriate Jacobian function.

pycocos.coordinates.compute_coordinates.compute_magnetic_coordinates(Rgrid, zgrid, br, bz, bphi, raxis, zaxis, psigrid, ltheta=256, phiclockwise=True, jacobian_func=None, R_at_psi=None, coordinate_system='boozer')

Compute magnetic coordinates using a generic Jacobian function.

Parameters:

  • Rgrid (np.ndarray) – Radial grid where (br, bz, bphi) are defined

  • zgrid (np.ndarray) – Vertical grid where (br, bz, bphi) are defined

  • br (np.ndarray) – Radial component of the magnetic field

  • bz (np.ndarray) – Vertical component of the magnetic field

  • bphi (np.ndarray) – Toroidal component of the magnetic field

  • raxis (float) – Radial position of the magnetic axis

  • zaxis (float) – Vertical position of the magnetic axis

  • psigrid (np.ndarray) – Poloidal flux grid where coordinates are defined

  • ltheta (int, optional) – Number of points in the poloidal direction. Default is 256

  • phiclockwise (bool, optional) – Whether toroidal angle increases clockwise. Default is True

  • jacobian_func (Callable, optional) – Function to compute Jacobian: jacobian_func(context) -> J If None, uses Boozer Jacobian

  • R_at_psi (np.ndarray, optional) – Radial positions corresponding to psigrid at midplane. If None, will be computed from psigrid

  • coordinate_system (str, optional) – Name of coordinate system, used for Jacobian context construction. Default is ‘boozer’

Returns:

  • qprof (np.ndarray) – Safety factor profile

  • Fprof (np.ndarray) – F(psi) = R*B_phi profile

  • Iprof (np.ndarray) – Toroidal current profile

  • thtable (np.ndarray) – Magnetic poloidal angle table (psi x theta)

  • nutable (np.ndarray) – Magnetic toroidal angle table (psi x theta)

  • jacobian (np.ndarray) – Jacobian table (psi x theta)

  • Rtransform (np.ndarray) – Inverse transformation R(psi, theta)

  • ztransform (np.ndarray) – Inverse transformation z(psi, theta)

Return type:

Tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]

Jacobian builders for magnetic coordinate systems.

This module keeps Python orchestration and data validation separate from Numba kernels used for heavy per-surface numerical work.

pycocos.coordinates.jacobian_builders.make_jacobian_context(*, coordinate_system, R, B, Bpol, dlp, I, F, q)

Build a Numba-friendly Jacobian context for a single flux surface.

Parameters:

  • coordinate_system (str)

  • R (ndarray)

  • B (ndarray)

  • Bpol (ndarray)

  • dlp (ndarray)

  • I (float)

  • F (float)

  • q (float)

Return type:

Dict[str, Any]

pycocos.coordinates.jacobian_builders.build_boozer_jacobian_from_context(context)
Boozer Jacobian:

J = h(psi) / B^2, with h = I + qF

Parameters:

context (Mapping[str, Any])

Return type:

ndarray

pycocos.coordinates.jacobian_builders.boozer_consistency_residual(context, jacobian)
Diagnostic residual for Boozer relation:

residual = max_theta |J*B^2 - (I + qF)|

Parameters:

  • context (Mapping[str, Any])

  • jacobian (ndarray)

Return type:

float

pycocos.coordinates.jacobian_builders.build_power_family_jacobian_from_context(context, *, i_power, j_power, k_power)

Build + normalize Jacobian from family J = R^i / |grad(psi)|^j / B^k.

Parameters:

  • context (Mapping[str, Any])

  • i_power (int)

  • j_power (int)

  • k_power (int)

Return type:

ndarray

pycocos.coordinates.jacobian_builders.build_coordinate_jacobian(context)

Main builder for registered coordinate systems.

Parameters:

context (Mapping[str, Any])

Return type:

ndarray

Numba kernels for Jacobian construction hot paths.

These functions are intentionally array-centric and object-free so they can be used inside performance-critical loops.

pycocos.coordinates.jacobian_numba_kernels.compute_grad_psi_abs(R, bpol)
Compute |grad(psi)| using axisymmetric identity:

|grad(psi)| = R * Bpol

Parameters:

  • R (ndarray)

  • bpol (ndarray)

Return type:

ndarray

pycocos.coordinates.jacobian_numba_kernels.build_boozer_jacobian(B, h)

Build Boozer Jacobian: J = h / B^2.

Parameters:

  • B (ndarray)

  • h (float)

Return type:

ndarray

pycocos.coordinates.jacobian_numba_kernels.build_power_law_jacobian(R, grad_psi, B, i_power, j_power, k_power, prefactor=1.0)
Build Jacobian from power-law family:

J = prefactor * R^i / |grad(psi)|^j / B^k

Parameters:

  • R (ndarray)

  • grad_psi (ndarray)

  • B (ndarray)

  • i_power (int)

  • j_power (int)

  • k_power (int)

  • prefactor (float)

Return type:

ndarray

pycocos.coordinates.jacobian_numba_kernels.compute_theta_span(R, jacobian, grad_psi, dlp)
Compute unnormalized poloidal-angle span implied by Jacobian:

span = integral ( R / (|J| * |grad(psi)|) ) dlp

Parameters:

  • R (ndarray)

  • jacobian (ndarray)

  • grad_psi (ndarray)

  • dlp (ndarray)

Return type:

float

pycocos.coordinates.jacobian_numba_kernels.apply_scalar_scale(values, scale)

Multiply a vector by a scalar and return a new array.

Parameters:

  • values (ndarray)

  • scale (float)

Return type:

ndarray

Jacobian computation functions for magnetic coordinate systems.

Primary API (used by pycocos internals) is context-based:

jacobian_func(context_dict) -> J(theta)

Legacy Boozer API is kept for compatibility:

compute_boozer_jacobian(I, F, q, B) -> J

pycocos.coordinates.jacobians.compute_boozer_jacobian(*args, **kwargs)

Compute Boozer coordinate Jacobian.

Supported signatures

  • Context API: compute_boozer_jacobian(context) where context contains at least: coordinate_system, R, B, Bpol, dlp, I, F, q.

  • Legacy API: compute_boozer_jacobian(I, F, q, B)

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

ndarray

pycocos.coordinates.jacobians.compute_hamada_jacobian(*args, **kwargs)

Compute Hamada coordinate Jacobian from context.

Requires context-based API because Hamada construction depends on geometric quantities beyond (I, F, q, B).

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

ndarray

pycocos.coordinates.jacobians.compute_pest_jacobian(*args, **kwargs)

Compute PEST coordinate Jacobian from context.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

ndarray

pycocos.coordinates.jacobians.compute_equal_arc_jacobian(*args, **kwargs)

Compute equal-arc coordinate Jacobian from context.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

ndarray

Registry for coordinate-system Jacobian callables.

pycocos.coordinates.registry.get_jacobian_function(name)

Get a context-based Jacobian function for a coordinate system.

Parameters:

name (str)

Return type:

Callable

pycocos.coordinates.registry.register_coordinate_system(name, jacobian_func)

Register a coordinate system Jacobian callable.

Accepted signatures: - New API: jacobian_func(context) -> J - Legacy API: jacobian_func(I, F, q, B, **kwargs) -> J

Parameters:

  • name (str)

  • jacobian_func (Callable)

Return type:

None

pycocos.coordinates.registry.list_coordinate_systems()

List all available coordinate systems.

Return type:

list

Core

Library for the magnetic coordinates class handler that eases the access and control of the magnetic coordinates.

class pycocos.core.magnetic_coordinates.magnetic_coordinates(coords, deriv, Raxis, zaxis, pad=0)

Container for magnetic coordinates and transformation routines.

This class stores magnetic coordinate transformations and provides methods to transform between cylindrical (R, z, phi) and magnetic (psi, theta, nu) coordinate systems.

Parameters:

  • coords (xr.Dataset) – Dataset containing magnetic coordinates (psi, theta, nu)

  • deriv (xr.Dataset) – Dataset containing derivatives of coordinate transformations

  • Raxis (float) – Radial position of the magnetic axis

  • zaxis (float) – Vertical position of the magnetic axis

  • pad (int, optional) – Number of padding points for periodic boundary conditions. Default is 0

coords

Magnetic coordinate data

Type:

xr.Dataset

deriv

Derivative data

Type:

xr.Dataset

lame_mag

Lamé factors for magnetic coordinates

Type:

xr.Dataset

metric_covariant

Covariant metric coefficients in magnetic coordinates (g_ij)

Type:

xr.Dataset

metric_contravariant

Contravariant metric coefficients in magnetic coordinates (g^ij)

Type:

xr.Dataset

Raxis

Radial position of magnetic axis

Type:

float

zaxis

Vertical position of magnetic axis

Type:

float

nthtpad

Number of theta padding points

Type:

int

Examples

>>> coords = mag_coords(R=2.0, z=0.0)  # Transform to magnetic coords
>>> cyl_coords = mag_coords.transform_inverse(psi=0.5, thetamag=0.0)
metric(i, j, tensor='covariant', return_in='Rzphi', return_psi_norm=False, return_rhopol=False)

Return a metric coefficient selected by magnetic-coordinate indices.

Parameters:

  • i (str) – Metric indices. Allowed names: psi, theta, zeta. Alias nu is accepted and mapped to zeta.

  • j (str) – Metric indices. Allowed names: psi, theta, zeta. Alias nu is accepted and mapped to zeta.

  • tensor (str, optional) – Metric tensor kind: covariant (g_ij) or contravariant (g^ij). Default is covariant.

  • return_in (str, optional) – Return space selector: Rzphi or magnetic_coordinates. Default is Rzphi.

  • return_psi_norm (bool, optional) – Passed to cyl2mag_scalar when return_in='magnetic_coordinates'. Default is False.

  • return_rhopol (bool, optional) – Passed to cyl2mag_scalar when return_in='magnetic_coordinates'. Default is False.

Return type:

DataArray

jacobian(return_in='Rzphi', inverse=False, return_psi_norm=False, return_rhopol=False)

Return the coordinate-transformation Jacobian.

Parameters:

  • return_in (str, optional) – Return space selector: Rzphi or magnetic_coordinates. Default is Rzphi.

  • inverse (bool, optional) – If True, return inverse Jacobian 1/J. Default is False.

  • return_psi_norm (bool, optional) – Passed to cyl2mag_scalar when return_in='magnetic_coordinates'. Default is False.

  • return_rhopol (bool, optional) – Passed to cyl2mag_scalar when return_in='magnetic_coordinates'. Default is False.

Return type:

DataArray

transform_inverse(psi, thetamag, grid=False, psi_is_norm=False)

Transform from magnetic to cylindrical coordinates.

Parameters:

  • psi (float or np.ndarray) – Poloidal flux values

  • thetamag (float or np.ndarray) – Magnetic poloidal angle values

  • grid (bool, optional) – If True, create a grid from psi and thetamag. Default is False

  • psi_is_norm (bool, optional) – If True, psi is normalized between 0 and 1. Default is False

Returns:

Dataset containing R_inv, z_inv coordinates

Return type:

xr.Dataset

Examples

>>> cyl = mag_coords.transform_inverse(psi=0.5, thetamag=0.0)
>>> R = cyl.R_inv.values
>>> z = cyl.z_inv.values
mag2cyl_scalar(field, psi=None, theta=None, nu=None, R=None, z=None, phi=None)

Transform a scalar field from magnetic coordinates to cylindrical.

Parameters:

  • field (np.ndarray or xr.DataArray) – Field defined in magnetic coordinates (psi, theta, nu) If DataArray, must have dims (‘psi’, ‘theta’, ‘nu’)

  • psi (np.ndarray, optional) – Psi coordinate grid. Required if field is not a DataArray

  • theta (np.ndarray, optional) – Theta coordinate grid. Required if field is not a DataArray

  • nu (np.ndarray, optional) – Nu coordinate grid. Required if field is not a DataArray

  • R (np.ndarray, optional) – Radial grid for output. If None, uses internal grid

  • z (np.ndarray, optional) – Vertical grid for output. If None, uses internal grid

  • phi (np.ndarray, optional) – Toroidal angle grid for output. If None, uses nu grid size

Returns:

Field transformed to cylindrical coordinates (phi, R, z)

Return type:

xr.DataArray

Raises:

ValueError – If field shape doesn’t match coordinate grids

cyl2mag_scalar(field, R=None, z=None, phi=None, return_psi_norm=False, return_rhopol=False)

Transform a scalar field from cylindrical coordinates to magnetic.

Parameters:

  • field (np.ndarray or xr.DataArray) – Field defined in cylindrical coordinates (R, z, phi) If DataArray, must have dims (‘R’, ‘z’, ‘phi’)

  • R (np.ndarray, optional) – Radial grid for input. Required if field is not a DataArray

  • z (np.ndarray, optional) – Vertical grid for input. Required if field is not a DataArray

  • phi (np.ndarray, optional) – Toroidal angle grid for input. Required if field is not a DataArray

  • return_psi_norm (bool, optional) – If True, return the first coordinate as normalized flux psi_N = (psi - psi_min) / (psi_max - psi_min). Default is False

  • return_rhopol (bool, optional) – If True, return the first coordinate as rhopol = sqrt(psi_N). Default is False

Returns:

Field transformed to magnetic coordinates (psi, theta, nu)

Return type:

xr.DataArray

Raises:

ValueError – If field shape doesn’t match coordinate grids

to_hdf5_as_xarray(fn, group_name=None)

Save to file HDF5 using the xarray package.

We use this function to save the magnetic coordinates using the xarray package. This is useful to save the magnetic coordinates with the corresponding coordinates.

Parameters:

  • fn (str)

  • group_name (str)

to_netcdf4(fn, group_name=None)

Save the magnetic coordinates object to a netcdf4 file.

Parameters:

  • fn (str) – filename where to save the magnetic coordinates.

  • group_name (str)

plot(coord_name=None, ax=None, nr=256, nz=512, **kwargs)

Plot magnetic coordinates.

If no coord_name is provided, lists available plottable coordinates. Otherwise, plots the specified coordinate.

Parameters:

  • coord_name (str, optional) – Name of the coordinate to plot. If None, lists available coordinates.

  • ax (matplotlib.axes.Axes, optional) – Axes to plot on. If None, creates new figure.

  • nr (int, optional) – Number of radial points for plot grid. Default is 256

  • nz (int, optional) – Number of vertical points for plot grid. Default is 512

  • **kwargs – Arguments passed to matplotlib contour

Returns:

  • ax (matplotlib.axes.Axes) – Axes object

  • im (matplotlib ContourSet) – Contour plot object

  • cbar (matplotlib Colorbar, optional) – Colorbar (if created)

Examples

>>> mag.plot()  # List available coordinates
>>> ax, im, cbar = mag.plot('psi')  # Plot psi coordinate
>>> ax, im, cbar = mag.plot('theta')  # Plot theta coordinate
rescale(nr, nz, ntht=None, npsi=None, rmin=None, rmax=None, zmin=None, zmax=None)

Rescale the magnetic coordinates to the input grid.

Parameters:

  • nr (int) – number of points in the radial direction.

  • nz (int) – number of points in the vertical direction.

  • ntht (int) – number of points in the poloidal angle.

  • npsi (int) – number of points in the toroidal angle.

  • rmin (float) – minimum value of the radial coordinate.

  • rmax (float) – maximum value of the radial coordinate.

  • zmin (float) – minimum value of the vertical coordinate.

  • zmax (float) – maximum value of the vertical coordinate.

property coordnames

Return the names of the coordinates.