QUADCOIL API

Subpackages

Submodules

quadcoil.math_utils module

quadcoil.quadcoil module

quadcoil.quadcoil_params module

class quadcoil.quadcoil_params.QuadcoilParams(plasma_surface, winding_surface, net_poloidal_current_amperes: float, net_toroidal_current_amperes: float, Bnormal_plasma=None, mpol=4, ntor=4, quadpoints_phi=None, quadpoints_theta=None, stellsym=None)

Bases: _Params

A class storing all informations required to solve a quadcoil problem. These includes plasma info, winding surface info, but does not include problem-specific info such as objectives, constraints or solutions. This class is primarily intended as a concise way to pass information into objective functions. It allows functions in quadcoil.objective to have the same signature, despite requiring different info to calculate.

Parameters:

  • plasma_surface (SurfaceRZFourierJAX) – The plasma surface.

  • winding_surface (SurfaceRZFourierJAX) – The winding surface. Must have all field periods.

  • net_poloidal_current_amperes (float) – The net poloidal current.

  • net_toroidal_current_amperes (float) – The net toroidal current.

  • Bnormal_plasma (ndarray, shape (nphi, ntheta), optional, default=None) – The magnetic field distribution on the plasma surface.

  • mpol (int, optional, default=4) – The number of poloidal Fourier harmonics in the current potential \(\Phi_{sv}\).

  • ntor (int, optional, default=4) – The number of toroidal Fourier harmonics in \(\Phi_{sv}\).

  • quadpoints_phi (ndarray, shape (nphi,), optional, default=None) – The toroidal quadrature points to evaluate quantities at. Takes one field period from the winding surface by default.

  • quadpoints_theta (ndarray, shape (ntheta,), optional, default=None) – The poloidal quadrature points to evaluate quantities at. Takes the winding surface’s quadrature points by default.

plasma_surface

(Traced) The plasma surface.

Type:

SurfaceRZFourierJAX

winding_surface

(Traced) The winding surface. Must have all field periods.

Type:

SurfaceRZFourierJAX

eval_surface

(Traced) The evaluation surface. Has the same dofs as the winding surface, but uses the quadrature points given by self.quadpoints_phi and self.quadpoints_phi.

Type:

SurfaceRZFourierJAX

net_poloidal_current_amperes

(Traced) The net poloidal current.

Type:

float

net_toroidal_current_amperes

(Traced) The net toroidal current.

Type:

float

Bnormal_plasma

(Traced) The magnetic field distribution on the plasma surface. will be filled with zeros by default.

Type:

ndarray, shape (nphi, ntheta)

quadpoints_phi

(Traced) The toroidal quadrature points to evaluate quantities at.

Type:

ndarray, shape (nphi,)

quadpoints_theta

(Traced) The poloidal quadrature points to evaluate quantities at.

Type:

ndarray, shape (ntheta,)

nfp

(Static) The number of field periods.

Type:

int

stellsym

(Static) Stellarator symmetry.

Type:

bool

mpol

(Static) The number of poloidal Fourier harmonics in \(\Phi_{sv}\).

Type:

int

ntor

(Static) The number of toroidal Fourier harmonics in \(\Phi_{sv}\).

Type:

int

ndofs

(Static) The number of degrees of freedom in \(\Phi_{sv}\).

Type:

int

ndofs_half

(Static) ndof if stellsym==True, ndof//2 otherwise.

Type:

int

Kdash_helper()
change_phi_resolution()
diff_helper()
make_mn()
make_mn_helper()
tree_flatten()
tree_unflatten()
class quadcoil.quadcoil_params.QuadcoilParamsFiniteElement(plasma_surface, winding_surface, net_poloidal_current_amperes: float, net_toroidal_current_amperes: float, Bnormal_plasma=None, quadpoints_phi=None, quadpoints_theta=None, stellsym=None)

Bases: _Params

tree_flatten()
tree_unflatten()
quadcoil.quadcoil_params.cp_ndofs(stellsym, mpol, ntor)[source]

quadcoil.solver module

quadcoil.surface module

class quadcoil.surface.SurfaceJAX(nfp: int, stellsym: bool, mpol: int, ntor: int, quadpoints_phi: jax.numpy.ndarray, quadpoints_theta: jax.numpy.ndarray, dofs: jax.numpy.ndarray)[source]

Bases: object

Abstract base class for JAX-native toroidal surfaces.

Subclasses must implement gammadash() and register themselves as JAX pytrees. All geometric quantities (normals, curvatures, etc.) are derived from gammadash and defined here so that every concrete surface type shares the same interface without code duplication.

quadpoints_phi, quadpoints_theta

Quadrature grid in [0, 1).

Type:

jnp.ndarray, shape (nphi,) / (ntheta,)

phi_mesh, theta_mesh

Meshgrid counterparts (phi varies along axis-0).

Type:

jnp.ndarray, shape (nphi, ntheta)

dphi, dtheta

Grid spacings.

Type:

float

area()
copy_and_set_quadpoints(quadpoints_phi, quadpoints_theta)[source]
da()

Area element: |N| * dphi * dtheta.

dga_inv_n_dashb()

Derivatives of (1/|N|) * (dγ/dphi) and (1/|N|) * (dγ/dtheta).

Returns:

dg2_inv_n_dash1, dg2_inv_n_dash2)

Each of shape (nphi, ntheta, 3).

Return type:

(dg1_inv_n_dash1, dg1_inv_n_dash2,

classmethod dof_to_gamma(dofs, phi_grid, theta_grid, nfp, stellsym, dash1_order=0, dash2_order=0, mpol: int = 10, ntor: int = 10)[source]

Map DOF vector to gamma (or derivatives) on the quadrature grid.

first_fund_form()

First fundamental form [E, F, G], shape (nphi, ntheta, 3).

classmethod fit(phi_target, theta_target, gamma_target, nfp: int, stellsym: bool, quadpoints_phi, quadpoints_theta, mpol: int = 7, ntor: int = 7, lam_tikhonov=0.0, custom_weight=None)
classmethod from_simsopt(surface_simsopt)[source]
gamma()
gamma_at_point(phi, theta)
gammadash(a: int, b: int) → jax.numpy.ndarray

Surface position or mixed partial derivative.

Parameters:

  • a (int) – Order of the phi derivative (0, 1, or 2).

  • b (int) – Order of the theta derivative (0, 1, or 2).

Returns:

The quantity d^(a+b) gamma / d phi^a d theta^b evaluated on the quadrature grid. Derivatives are with respect to the normalised angles in [0, 1).

Return type:

jnp.ndarray, shape (nphi, ntheta, 3)

Notes

Forwards to gammadash_at_point() with the separable meshgrid form phi=quadpoints_phi[:, None], theta=quadpoints_theta[None, :]. This is strictly a faster backend than the historical _dof_to_gamma_op @ dofs path (5x-17x for RZ/XYZ Fourier surfaces, 30x-40x for the separable XYZ tensor Fourier surface) and produces output that matches the operator approach to round-off. The operator (cls._dof_to_gamma_op) remains available and is still used internally by _fit_dofs_from_gamma.

gammadash1()
gammadash1_at_point(phi, theta)
gammadash1dash1()
gammadash1dash1_at_point(phi, theta)
gammadash1dash2()
gammadash1dash2_at_point(phi, theta)
gammadash2()
gammadash2_at_point(phi, theta)
gammadash2dash2()
gammadash2dash2_at_point(phi, theta)
gammadash_at_point(phi, theta, a: int, b: int) → jax.numpy.ndarray

Broadcastable evaluation of d^(a+b) gamma / dphi^a dtheta^b.

Unlike gammadash(), this never materialises the (nphi, ntheta, 3, ndof) operator: it computes the basis functions at the requested (phi, theta) points and contracts directly against self.dofs.

Parameters:

  • phi (jnp.ndarray) – Broadcast-compatible arrays of normalised angles in [0, 1).

  • theta (jnp.ndarray) – Broadcast-compatible arrays of normalised angles in [0, 1).

  • a (int) – Order of the phi derivative (0, 1, 2, or 3).

  • b (int) – Order of the theta derivative (0, 1, 2, or 3).

Returns:

The mixed partial of gamma at each requested point.

Return type:

jnp.ndarray, shape broadcast(phi, theta).shape + (3,)

Notes

Reproduces self.gammadash(a, b) bit-for-bit when called with the fully expanded meshgrid phi = quadpoints_phi[:, None] + 0 * quadpoints_theta[None, :], theta = quadpoints_theta[None, :] + 0 * quadpoints_phi[:, None].

Subclasses must override this method.

gen_winding_surface(d_expand, mpol=7, ntor=7, phi_interp: int = 2, theta_interp: int = 5, theta_rule_subsample: int = None, quadpoints_phi=None, quadpoints_theta=None, lam_tikhonov=1e-05, winding_surface_mode='self-intersection', theta_mode='arclen', weight_mode=None)

Generate a smoothed winding surface from a uniform normal offset.

Parameters:

  • d_expand (float) – Normal offset distance.

  • mpol (int) – Fourier resolution of the fitted surface.

  • ntor (int) – Fourier resolution of the fitted surface.

  • phi_interp (int)

  • theta_interp (int) – Toroidal/poloidal oversampling factor over self.quadpoints_theta during the Fourier fit. higher values leads to more accurate offset surfaces.

  • theta_rule_subsample (int) – Poloidal subsampling stride for removing self-intersection relative to len(self.quadpoints_theta) * theta_interp. This is necessary because the complexity of the self-intersection check is O(len(self.quadpoints_theta)**2). The default value is theta_interp.

  • quadpoints_phi (array or None) – Quadrature points for the output surface. Defaults to jnp.linspace(0, 1, self.nfp*len(self.quadpoints_phi), endpoint=False) and jnp.linspace(0, 1, len(self.quadpoints_theta), endpoint=False).

  • quadpoints_theta (array or None) – Quadrature points for the output surface. Defaults to jnp.linspace(0, 1, self.nfp*len(self.quadpoints_phi), endpoint=False) and jnp.linspace(0, 1, len(self.quadpoints_theta), endpoint=False).

  • lam_tikhonov (float) – Tikhonov regularization weight for least-square surface fit.

  • winding_surface_mode ({'uniform', 'self-intersection'}) – Type of winding surface to generate.

  • theta_mode ({'arclen', 'arctan'}) – Function to use for assigning poloidal parameterization. 'arctan' uses arctan on poloidal section, using the center-of-weight of the plasma poloidal cross sections as origin. 'arclen' uses the arc lengths of the winding surface’s poloidal cross sections. 'arctan' generates smoother surfaces, but can misbehave when the winding surface is concave. 'arctan' is robust for concave winding surfaces but less smooth.

  • weight_mode ({'poloidal_arclen', 'area'} or None) – How to weight fit points. Can be based on arclength of poloidal cross sections or the surface jacobian. 'poloidal_arclen' behaves better for compact configurations. 'area' behaves better for long-aspect ratio, helical configurations.

Returns:

fit_surface – Fitted winding surface with the given quadrature points.

Return type:

same type as self

gen_winding_surface_dofs(d_expand, mpol=7, ntor=7, phi_interp: int = 2, theta_interp: int = 5, theta_rule_subsample: int = 5, lam_tikhonov=1e-05, winding_surface_mode='self-intersection', theta_mode='arclen', weight_mode=None)

Fit winding surface DOFs from a uniform normal offset of this surface.

Generates an exact offset surface, bisects poloidal slices onto best-fit planes, computes arc-length or arc-tan parameterisation, removes self-intersecting regions, and solves a weighted least-squares fit. Please see SurfaceJAX.gen_winding_surface() for parameters.

Returns:

dofs – DOF vector for the fitted surface (same layout as self.dofs).

Return type:

jnp.ndarray

get_dofs()[source]
grad_helper()

Contravariant vectors grad-phi and grad-theta.

Return type:

(grad1, grad2) each of shape (nphi, ntheta, 3)

integrate(scalar_field)

Integrate a scalar field over the surface.

normal()
plot(**kwargs)[source]
second_fund_form()

Second fundamental form [e, f, g], shape (nphi, ntheta, 3).

surface_curvatures()

Mean (H), Gaussian (K), and principal (κ₁, κ₂) curvatures.

Returns:

Stacked [H, K, kappa1, kappa2].

Return type:

jnp.ndarray, shape (nphi, ntheta, 4)

uniform_offset(d_expand, quadpoints_phi=None, quadpoints_theta=None)[source]
unitnormal()
unitnormal_at_point(phi, theta) → jax.numpy.ndarray[source]

Convenience: unitnormaldash_at_point(phi, theta, 0, 0).

unitnormaldash(a: int, b: int) → jax.numpy.ndarray

Compute d^(a+b)(unitnormal) / dphi^a dtheta^b using autodiff.

Uses nested automatic differentiation for arbitrary-order derivatives.

Parameters:

  • a (int) – Order of derivative with respect to phi

  • b (int) – Order of derivative with respect to theta

Returns:

The derivative d^(a+b)(unitnormal) / dphi^a dtheta^b

Return type:

jnp.ndarray, shape (nphi, ntheta, 3)

Examples

>>> surf.unitnormaldash(0, 0)  # Returns unitnormal
>>> surf.unitnormaldash(1, 0)  # Returns d(unitnormal)/dphi
>>> surf.unitnormaldash(0, 1)  # Returns d(unitnormal)/dtheta
>>> surf.unitnormaldash(2, 0)  # Returns d²(unitnormal)/dphi²
unitnormaldash_at_point(phi, theta, a: int, b: int) → jax.numpy.ndarray

Broadcastable mixed derivative of the unit normal at arbitrary points.

Built on top of gammadash_at_point() and JAX forward-mode autodiff (vmap of repeated jacfwd over scalar phi/theta). Output shape: broadcast(phi, theta).shape + (3,).

Reproduces self.unitnormaldash(a, b) when called with the fully-expanded meshgrid phi=quadpoints_phi[:, None] + 0*theta_1d, theta=quadpoints_theta[None, :] + 0*phi_1d.

Parameters:

  • phi (jnp.ndarray) – Broadcast-compatible normalised angles in [0, 1).

  • theta (jnp.ndarray) – Broadcast-compatible normalised angles in [0, 1).

  • a (int) – Order of the phi / theta derivative.

  • b (int) – Order of the phi / theta derivative.

Notes

For (a, b) == (0, 0) we take an analytic fast path that evaluates cross(gammadash1_at_point, gammadash2_at_point) / norm directly on the broadcasted arrays (no autodiff, no vmap). For all other orders we vmap a per-point scalar function and apply jacfwd a + b times.

unitnormaldash_legacy()

d(unitnormal)/dphi and d(unitnormal)/dtheta (legacy hard-coded implementation).

This is the original hard-coded gradient implementation, preserved for performance comparison. For production code, use unitnormaldash(a, b) instead.

Return type:

(unitnormaldash1, unitnormaldash2), each (nphi, ntheta, 3)

class quadcoil.surface.SurfaceOffsetJAX(base_surface, d_expand, quadpoints_phi=None, quadpoints_theta=None)

Bases: SurfaceJAX

Subclass that applies a fixed normal offset to any SurfaceJAX subclass.

This class inherits from SurfaceJAX and offsets all geometric quantities by a fixed distance along the surface normal. It maintains the full SurfaceJAX interface and can be used anywhere a SurfaceJAX is expected.

Parameters:

  • base_surface (SurfaceJAX) – The underlying surface to offset.

  • d_expand (float) – Distance to offset along the unit normal (positive = outward).

Examples

>>> plasma_surf = SurfaceRZFourierJAX(...)
>>> winding_surf = SurfaceOffsetJAX(plasma_surf, d_expand=0.2)
>>> gamma_offset = winding_surf.gamma()
>>> isinstance(winding_surf, SurfaceJAX)  # Returns True
True

Notes

DOF-related methods (e.g., get_dofs(), from_simsopt(), _fit_dofs_from_gamma()) raise NotImplementedError since offset surfaces don’t have independent DOFs.

copy_and_set_quadpoints()
dof_to_gamma()
from_simsopt()
gammadash()
gammadash_at_point()
gen_winding_surface()
gen_winding_surface_dofs()
get_dofs()
plot()
to_simsopt()
tree_flatten()
tree_unflatten()
uniform_offset()
class quadcoil.surface.SurfaceRZFourierJAX(nfp: int, stellsym: bool, mpol: int, ntor: int, quadpoints_phi: jax.numpy.ndarray, quadpoints_theta: jax.numpy.ndarray, dofs: jax.numpy.ndarray)

Bases: SurfaceJAX

JAX-native surface in cylindrical Fourier (RZ) coordinates.

Representation:

r(phi, theta) = sum_{m,n} [rc_{mn} cos(m*theta - nfp*n*phi)
                          + rs_{mn} sin(m*theta - nfp*n*phi)]
z(phi, theta) = sum_{m,n} [zc_{mn} cos(m*theta - nfp*n*phi)
                          + zs_{mn} sin(m*theta - nfp*n*phi)]

The DOF vector is [rc, zs] for stellarator-symmetric surfaces and [rc, rs, zc, zs] otherwise, matching simsopt’s convention exactly.

dof_to_rz_op()
from_desc()
from_simsopt()
gammadash_at_point()
to_desc()
to_simsopt()
tree_flatten()
tree_unflatten()
class quadcoil.surface.SurfaceXYZFourierJAX(nfp: int, stellsym: bool, mpol: int, ntor: int, quadpoints_phi: jax.numpy.ndarray, quadpoints_theta: jax.numpy.ndarray, dofs: jax.numpy.ndarray)

Bases: SurfaceJAX

JAX-native surface in Cartesian Fourier (XYZ) coordinates.

Matches simsopt.geo.SurfaceXYZFourier exactly.

Representation:

x_hat(phi, theta) = sum_{m,n} [xc_{mn} cos(m*theta - nfp*n*phi)
                              + xs_{mn} sin(m*theta - nfp*n*phi)]
y_hat(phi, theta) = sum_{m,n} [yc_{mn} cos(m*theta - nfp*n*phi)
                              + ys_{mn} sin(m*theta - nfp*n*phi)]
z(phi, theta)     = sum_{m,n} [zc_{mn} cos(m*theta - nfp*n*phi)
                              + zs_{mn} sin(m*theta - nfp*n*phi)]
x = x_hat * cos(2*pi*phi) - y_hat * sin(2*pi*phi)
y = x_hat * sin(2*pi*phi) + y_hat * cos(2*pi*phi)

Under stellarator symmetry the xs, yc, and zc terms are zero.

The DOF vector is [xc, ys, zs] for stellarator-symmetric surfaces and [xc, xs, yc, ys, zc, zs] otherwise, matching simsopt’s convention exactly. The (m, n) mode indexing follows make_rzfourier_mc_ms_nc_ns().

dof_to_xhatz_op()
from_simsopt()
gammadash_at_point()
to_simsopt()
tree_flatten()
tree_unflatten()
class quadcoil.surface.SurfaceXYZTensorFourierJAX(nfp: int, stellsym: bool, mpol: int, ntor: int, quadpoints_phi: jax.numpy.ndarray, quadpoints_theta: jax.numpy.ndarray, dofs: jax.numpy.ndarray)

Bases: SurfaceJAX

JAX-native surface in Cartesian tensor-product Fourier coordinates.

Matches simsopt.geo.SurfaceXYZTensorFourier exactly.

Representation:

x_hat(theta, phi) = sum_{i,j} x_{ij} w_i(theta) v_j(phi)
y_hat(theta, phi) = sum_{i,j} y_{ij} w_i(theta) v_j(phi)
x(phi, theta) = x_hat * cos(phi_rad) - y_hat * sin(phi_rad)
y(phi, theta) = x_hat * sin(phi_rad) + y_hat * cos(phi_rad)
z(theta, phi) = sum_{i,j} z_{ij} w_i(theta) v_j(phi)

where phi_rad = 2*pi*phi_normalised, and the toroidal basis is:

v_j : j=0..ntor  -> cos(nfp*j*phi_rad)
      j=ntor+1..2*ntor -> sin(nfp*(j-ntor)*phi_rad)

and the poloidal basis is:

w_i : i=0..mpol  -> cos(i*theta_rad)
      i=mpol+1..2*mpol -> sin((i-mpol)*theta_rad)

The DOF vector is [x_active, y_active, z_active] where the active coefficients follow simsopt’s get_dofs() ordering (row-major over (m, n), skipping stellarator-symmetric zeros).

Stellarator symmetry rules

  • x: keep (n <= ntor and m <= mpol) OR (n > ntor and m > mpol)

  • y, z: keep (n <= ntor and m > mpol) OR (n > ntor and m <= mpol)

param nfp:

type nfp:

int

param stellsym:

type stellsym:

bool

param mpol:

type mpol:

int

param ntor:

type ntor:

int

param quadpoints_phi:

type quadpoints_phi:

array-like 1-D, values in [0, 1)

param quadpoints_theta:

type quadpoints_theta:

array-like 1-D, values in [0, 1)

param dofs:

Active Fourier coefficients in simsopt ordering.

type dofs:

1-D array

from_simsopt()
gammadash_at_point()
num_dofs()
to_RZFourier()
to_simsopt()
tree_flatten()
tree_unflatten()
quadcoil.surface.make_rzfourier_mc_ms_nc_ns(mpol: int, ntor: int)
quadcoil.surface.xyztensor_gammadash(dofs, quadpoints_phi, quadpoints_theta, nfp: int, stellsym: bool, a: int, b: int, mpol: int, ntor: int)

Compute d^(a+b) gamma / dphi^a dtheta^b for XYZ tensor Fourier surface.

Uses the Leibniz product rule to differentiate x = x_hat * cos(phi_rad) - y_hat * sin(phi_rad) and similarly for y, then combines with the theta derivative contained in the W basis.

Parameters:

  • dofs (1-D jax array)

  • quadpoints_phi (1-D jax arrays in [0, 1))

  • quadpoints_theta (1-D jax arrays in [0, 1))

  • nfp (static)

  • stellsym (static)

  • a (static)

  • b (static)

  • mpol (static)

  • ntor (static)

Return type:

jnp.ndarray, shape (nphi, ntheta, 3)

quadcoil.winding_surface module

quadcoil.wrapper module

Module contents