QUADCOIL API

Subpackages

Submodules

quadcoil.math_utils module

quadcoil.math_utils.abs_lse(x, epsilon, **kwargs)[source]
quadcoil.math_utils.is_ndarray(arr, n=1)[source]
quadcoil.math_utils.linf_lse(x, epsilon, **kwargs)[source]
quadcoil.math_utils.max_lse(x, epsilon, **kwargs)[source]
quadcoil.math_utils.norm_helper(vec)

Calculates \(|v|\) and \(1/|v|\) for a vector field on a 2d surface.

Parameters:

vec (ndarray, shape (Nx, Ny, ..., 3)) – The vector field

Returns:

  • normN_prime_2d (ndarray, shape (Nx, Ny, …)) – The vector field’s length, \(|v|\)

  • inv_normN_prime_2d (ndarray, shape (Nx, Ny, …)) – 1/the vector field’s length, \(1/|v|\)

quadcoil.math_utils.project_arr_coord(operator, unit1, unit2, unit3)

Project an array of vector fields on a 2d surface in a given basis, unit1, unit2, unit3.

Parameters:

  • operator (ndarray, shape (n_phi, n_theta, 3, ...)) – An array of (n_phi, n_theta, 3) vector fields. operator.shape[:3] must be (n_phi, n_theta, 3). Otehrwise the shape is flexible.

  • unit1 (ndarray, shape (n_phi, n_theta, 3)) – Basis vector 1 where the vector field is sampled.

  • unit2 (ndarray, shape (n_phi, n_theta, 3)) – Basis vector 2 where the vector field is sampled.

  • unit3 (ndarray, shape (n_phi, n_theta, 3)) – Basis vector 3 where the vector field is sampled.

Returns:

Outputs

Return type:

ndarray, shape (n_phi, n_theta, 3, …)

quadcoil.math_utils.project_arr_cylindrical(gamma, operator)

Project a stack of vector fields onto a cylindrical coordinate for a given set of coordinate points.

Parameters:

  • gamma (ndarray, shape (n_phi, n_theta, 3)) – The location of the coordinate points where the field is sampled in x, y, z.

  • operator (ndarray, shape (n_phi, n_theta, 3, ...)) – A stack of (n_phi, n_theta, 3) vector fields. operator.shape[:3] must be (n_phi, n_theta, 3). Otherwise the shape is flexible.

Returns:

Outputs

Return type:

ndarray, shape (n_phi, n_theta, 3, …)

quadcoil.math_utils.safe_linear_solve(A, b)
quadcoil.math_utils.safe_linear_solve_bwd(res, g)[source]
quadcoil.math_utils.safe_linear_solve_fwd(A, b)[source]
quadcoil.math_utils.sin_or_cos(x, mode)[source]

Scans a pair of arrays, x and mode. Where mode==1, return jnp.sin(x). Otherwise return jnp.cos(x). Used in inverse Fourier Transforms.

Parameters:

  • x (ndarray) – The data.

  • mode (ndarray) – The choice of trigonometry functions.

Return type:

ndarray

quadcoil.math_utils.tree_len(pytree)[source]

quadcoil.quadcoil module

quadcoil.quadcoil.quadcoil(nfp: int, stellsym: bool, plasma_mpol: int, plasma_ntor: int, plasma_dofs, net_poloidal_current_amperes: float, net_toroidal_current_amperes: float = 0.0, mpol: int = 6, ntor: int = 4, quadpoints_phi=None, quadpoints_theta=None, phi_init=None, phi_unit=None, plasma_stellsym=True, plasma_quadpoints_phi=None, plasma_quadpoints_theta=None, Bnormal_plasma=None, plasma_coil_distance: float = None, winding_surface_generator=jax.jit, winding_dofs=None, winding_mpol: int = 6, winding_ntor: int = 5, winding_quadpoints_phi=None, winding_quadpoints_theta=None, winding_stellsym=True, objective_name='f_B', objective_weight=1.0, objective_unit=None, constraint_name=(), constraint_type=(), constraint_unit=(), constraint_value=jax.numpy.array, metric_name=('f_B', 'f_K'), value_only=False, smoothing='slack', smoothing_params={'lse_epsilon': 0.001}, convex: bool = False, verbose: int = 0, merge_constraints: bool = False, c_init: float = 1.0, c_growth_rate: float = 2.0, xstop_outer: float = 1e-06, ctol_outer: float = 1e-06, fstop_inner: float = 1e-06, xstop_inner: float = 1e-06, gtol_inner: float = 1e-06, fstop_inner_last: float = 0.0, xstop_inner_last: float = 1e-10, gtol_inner_last: float = 1e-10, svtol: float = 1e-06, maxiter_tot: int = 10000, maxiter_inner: int = 1000, max_linesearch_steps: int = 20, gplus_mask=<function <lambda>>, implicit_linear_solver=None)

Solves a QUADCOIL problem.

Parameters:

  • nfp (int) – (Static) The number of field periods.

  • stellsym (bool) – (Static) Whether the coils have stellarator symmetry.

  • plasma_mpol (int) – (Static) The number of poloidal Fourier harmonics in the plasma boundary.

  • plasma_ntor (int) – (Static) The number of toroidal Fourier harmonics in the plasma boundary.

  • plasma_dofs (ndarray) – (Static) The plasma surface degrees of freedom. Uses the simsopt.geo.SurfaceRZFourier.get_dofs() convention.

  • net_poloidal_current_amperes (float) – (Traced) The net poloidal current \(G\).

  • net_toroidal_current_amperes (float, optional, default=0) – (Traced) The net toroidal current \(I\).

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

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

  • quadpoints_phi (ndarray, shape (nphi,), optional, default=None) – (Traced) The poloidal quadrature points on the winding surface to evaluate the objectives at. Uses one period from the winding surface by default.

  • quadpoints_theta (ndarray, shape (ntheta,), optional, default=None) – (Traced) The toroidal quadrature points on the winding surface to evaluate the objectives at. Uses one period from the winding surface by default.

  • phi_init (ndarray, optional, default=None) – (Traced) The initial guess. All zeros by default.

  • phi_unit (float, optional, default=None) – (Traced) Current potential’s normalization constant. By default will be generated from total net current.

  • plasma_stellsym (bool, default=True) – (Static) Whether the plasma has stellarator symmetry.

  • plasma_quadpoints_phi (ndarray, shape (nphi_plasma,), optional, default=None) – (Traced) Will be set based on the shape of Bnormal_plasma if it’s provided, or default to jnp.linspace(0, 1/nfp, 32, endpoint=False) otherwise.

  • plasma_quadpoints_theta (ndarray, shape (ntheta_plasma,), optional, default=None) – (Traced) Will be set based on the shape of Bnormal_plasma if it’s provided, or default to jnp.linspace(0, 1, 34, endpoint=False) otherwise.

  • Bnormal_plasma (ndarray, shape (nphi, ntheta), optional, default=None) – (Traced) The magnetic field distribution on the plasma surface. Will be filled with zeros by default.

  • plasma_coil_distance (float, optional, default=None) – (Traced) The coil-plasma distance. Is set to None by default, but a value must be provided if winding_dofs is not provided.

  • winding_surface_generator (callable, optional, default=gen_winding_surface_atan) – (Static) The winding surface generator.

  • winding_dofs (ndarray, shape (ndof_winding,)) – (Traced) The winding surface degrees of freedom. Uses the simsopt.geo.SurfaceRZFourier.get_dofs() convention. Will be generated using winding_surface_generator if plasma_coil_distance is provided. Must be provided otherwise.

  • winding_mpol (int, optional, default=6) – (Static) The number of poloidal Fourier harmonics in the winding surface.

  • winding_ntor (int, optional, default=5) – (Static) The number of toroidal Fourier harmonics in the winding surface.

  • winding_quadpoints_phi (ndarray, shape (nphi_winding,), optional, default=None) – (Traced) Will be set to jnp.linspace(0, 1, 32*nfp, endpoint=False) by default.

  • winding_quadpoints_theta (ndarray, shape (ntheta_winding,), optional, default=None) – (Traced) Will be set to jnp.linspace(0, 1, 34, endpoint=False) by default.

  • winding_stellsym (bool, default=True) – (Static) Whether the winding surface has stellarator symmetry.

  • objective_name (tuple, optional, default='f_B_normalized_by_Bnormal_IG') – (Static) The names of the objective functions. Must be a member of quadcoil.objective that outputs a scalar.

  • objective_weight (ndarray, optional, default=None) – (Traced) The weights of the objective functions. Derivatives will be calculated w.r.t. this quantity.

  • objective_unit (ndarray, optional, default=None) – (Traced) The normalization constants of the objective terms, so that f/objective_unit is \(O(1)\). May contain None

  • constraint_name (tuple, optional, default=()) – (Static) The names of the constraint functions. Must be a member of quadcoil.objective that outputs a scalar.

  • constraint_type (tuple, optional, default=()) – (Static) The types of the constraints. Must consist of '>=', '<=', '==' only.

  • constraint_unit (ndarray, optional, default=()) – (Traced) The normalization constants of the constraints, so that f/constraint_unit is \(O(1)\) May contain None.

  • constraint_value (ndarray, optional, default=()) – (Traced) The constraint thresholds. Derivatives will be calculated w.r.t. this quantity.

  • metric_name (tuple, optional, default=('f_B', 'f_K')) – (Static) The names of the functions to diagnose the coil configurations with. Will be differentiated w.r.t. other input quantities.

  • convex (bool, optional, default=False) – (Static) Whether to assume the problem is convex. When True, QUADCOIL will apply some limited simplifications.

  • c_init (float, optional, default=1.) – (Traced) The initial \(c\) factor. Please see Constrained Optimization and Lagrange Multiplier Methods Chapter 3.

  • c_growth_rate (float, optional, default=1.2) – (Traced) The growth rate of the \(c\) factor.

  • xstop_outer (float, optional, default=1e-7) – (Traced) x convergence rate of the outer augmented Lagrangian loop. Terminates when dx falls below this.

  • ctol_outer (float, optional, default=1e-7) – (Traced) Tolerance of the constraint KKT conditions in the outer Lagrangian loop.

  • fstop_inner (float, optional, default=1e-7) – (Traced) f convergence rate of the inner LBFGS Lagrangian loop. Terminates when df falls below this.

  • fstop_inner_last (float, optional, default=1e-7) – (Traced) f convergence rate of the inner LBFGS Lagrangian loop. Terminates when df falls below this.

  • xstop_inner (float, optional, default=0) – (Traced) x convergence rate of the outer augmented Lagrangian loop. Terminates when dx falls below this.

  • xstop_inner_last (float, optional, default=0) – (Traced) x convergence rate of the outer augmented Lagrangian loop. Terminates when dx falls below this.

  • gtol_inner (float, optional, default=0.1) – (Traced) Gradient tolerance of the inner LBFGS iteration, normalized by the starting gradient.

  • gtol_inner_last (float, optional, default=0.1) – (Traced) Gradient tolerance of the inner LBFGS iteration, normalized by the starting gradient.

  • svtol (float, optional, default=0.1) – (Traced) Singular-value cut-off threshold during the pre-conditioning. Will treat singular values smaller than svtol * jnp.max(s) as 0

  • maxiter_tot (int, optional, default=50.) – (Static) The maximum of the outer iteration.

  • maxiter_inner (int, optional, default=500) – (Static) The maximum of the inner iteration.

  • maxiter_inner_last (int, optional, default=500) – (Static) The maximum of the inner iteration.

  • max_linesearch_steps (int) – (Static) The maximum steps count in the LBFGS line search.

  • gplus_mask (Callable, optional, default=quadcoil.gplus_hard) – (Static) The form of g+. Soft thresholding may improve derivative effectiveness.

  • implicit_linear_solver (lineax.AbstractLinearSolver, optional, default=lineax.AutoLinearSolver(well_posed=True)) – (Static) The lineax linear solver choice for implicit differentiation.

  • value_only (bool, optional, default=False) – (Static) When True, skip gradient calculations.

  • verbose (int, optional, default=False) – (Static) Print general info when verbose==1. Print inside the outer iteration loop, too, when verbose==2.

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()
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.solver.delta_normalized(x1, x2)[source]
quadcoil.solver.gplus_elu(x, mu, c, g_ineq, scale=1)[source]
quadcoil.solver.gplus_hard(x, mu, c, g_ineq)
quadcoil.solver.gplus_softplus(x, mu, c, g_ineq, scale=1)[source]
quadcoil.solver.run_opt_lbfgs(init_params, fun, maxiter, fstop, xstop, gtol, max_linesearch_steps, verbose)[source]

A wrapper for performing unconstrained optimization using optax.lbfgs.

Parameters:

  • init_params (ndarray, shape (N,)) – The initial condition.

  • fun (Callable) – The objective function.

  • maxiter (int) – The maximum iteration number.

  • fstop (float) – The objective function convergence rate tolerance. Terminates when any one of the tolerances is satisfied.

  • xstop (float) – The unknown convergence rate tolerance. Terminates when any one of the tolerances is satisfied.

  • gtol (float) – The gradient tolerance. Terminates when any one of the tolerances is satisfied.

  • max_linesearch_steps (int) – The maximum steps count in the LBFGS line search.

  • verbose (int) – Output levels of detail.

Returns:

  • x (ndarray, shape (N,)) – The optimum.

  • f (float) – The objective at the optimum.

  • grad (ndarray, shape (N,)) – The gradient at the optimum.

  • count (int) – The iteration number.

  • final_dx (float) – The rate of change of x at the optimum.

  • final_du (float) – The rate of change of updates at the optimum.

  • final_df (float) – The rate of change of f at the optimum.

quadcoil.solver.run_opt_optax(init_params, fun, maxiter, fstop, xstop, gtol, opt, verbose, n_history=10)[source]

A wrapper for performing unconstrained optimization using optax.base.GradientTransformationExtraArgs.

Parameters:

  • init_params (ndarray, shape (N,)) – The initial condition.

  • fun (Callable) – The objective function.

  • maxiter (int) – The maximum iteration number.

  • fstop (float) – The objective function convergence rate tolerance. Terminates when any one of the tolerances is satisfied.

  • xstop (float) – The unknown convergence rate tolerance. Terminates when any one of the tolerances is satisfied.

  • gtol (float) – The gradient tolerance. Terminates when any one of the tolerances is satisfied.

  • opt (optax.base.GradientTransformationExtraArgs) – The optimizer of choice.

Returns:

  • x (ndarray, shape (N,)) – The optimum.

  • f (float) – The objective at the optimum.

  • grad (ndarray, shape (N,)) – The gradient at the optimum.

  • count (int) – The iteration number.

  • final_dx (float) – The rate of change of x at the optimum.

  • final_du (float) – The rate of change of updates at the optimum.

  • final_df (float) – The rate of change of f at the optimum.

quadcoil.solver.solve_constrained(x_init, f_obj, c_init=1.0, c_growth_rate=1.1, lam_init=jax.numpy.zeros, h_eq=<function <lambda>>, mu_init=jax.numpy.zeros, g_ineq=<function <lambda>>, xstop_outer=1e-07, ctol_outer=1e-07, fstop_inner=1e-07, xstop_inner=1e-07, gtol_inner=1e-07, fstop_inner_last=1e-07, xstop_inner_last=1e-07, gtol_inner_last=1e-07, maxiter_tot=10000, maxiter_inner=500, max_linesearch_steps=20, verbose=0, c_k_safe=1000000000000000.0, gplus_mask=<function <lambda>>)[source]

Solves the constrained optimization problem:

\[\begin{split}\min_x f(x) \\ \text{subject to } \\ h(x) = 0, \\ g(x) \leq 0 \\\end{split}\]

Using the augmented Lagrangian method in Constrained Optimization and Lagrange Multiplier Methods Chapter 3. Please refer to the chapter for notation.

Parameters:

  • init_params (ndarray, shape (N,))

  • fun (Callable)

  • maxiter (int) – The maximum iteration number.

  • fstop (float) – The objective function convergence rate tolerance. Terminates when any one of the tolerances is satisfied.

  • xstop (float) – The unknown convergence rate tolerance. Terminates when any one of the tolerances is satisfied.

  • gtol (float) – The gradient tolerance. Terminates when any one of the tolerances is satisfied.

  • x_init (ndarray, shape (Nx,)) – The initial condition.

  • x_unit_init (ndarray, shape (Nx,)) – The initial x scale. This scaling factor ensures that x~1. Will be updated after every outer iteration.

  • f_obj (Callable) – The objective function.

  • c_init (float, optional, default=1.) – The initial \(c\) factor. Please see Constrained Optimization and Lagrange Multiplier Methods Chapter 3.

  • c_growth_rate (float, optional, default=1.1,) – The growth rate of the \(c\) factor.

  • lam_init (ndarray, shape (Nh), optional, default=jnp.zeros(1),) – The initial \(\lambda\) multiplier for equality constraints. No constraints by default.

  • h_eq (Callable, optional, default=lambda x:jnp.zeros(1),) – The equality constraint function. Must map x to an ndarray with shape (Nh). No constraints by default.

  • mu_init (ndarray, shape (Ng), optional, default=jnp.zeros(1),) – The initial \(\mu\) multiplier for inequality constraints. No constraints by default.

  • g_ineq (Callable, optional, default=lambda x:jnp.zeros(1),) – The equality constraint function. Must map x to an ndarray with shape (Ng). No constraints by default.

  • xstop_outer (float, optional, default=1e-7) – (Traced) x convergence rate of the outer augmented Lagrangian loop. Terminates when dx falls below this.

  • gtol_outer (float, optional, default=1e-7) – (Traced) Tolerance of the \(\nabla L\) KKT condition in the outer augmented Lagrangian loop.

  • ctol_outer (float, optional, default=1e-7) – (Traced) Tolerance of the constraint KKT conditions in the outer Lagrangian loop.

  • fstop_inner (float, optional, default=1e-7) – (Traced) f convergence rate of the inner LBFGS Lagrangian loop. Terminates when df falls below this.

  • xstop_inner (float, optional, default=0) – (Traced) x convergence rate of the outer augmented Lagrangian loop. Terminates when dx falls below this.

  • gtol_inner (float, optional, default=1e-7) – (Traced) Gradient tolerance of the inner LBFGS iteration. Terminates when is satisfied.

  • maxiter_tot (int, optional, default=10000) – (Static) The maximum of the outer iteration.

  • maxiter_inner (int, optional, default=500) – (Static) The maximum of the inner iteration.

  • max_linesearch_steps (int) – (Static) The maximum steps count in the LBFGS line search.

  • verbose (int, optional, default=0) – (Static) The verbosity. When >1, outputs outer iteration convergence info.

  • c_k_safe (float, optional, default=1e15,)

  • gplus_mask=gplus_hard

:param :

Returns:

status – The end state of the iteration. Contains the following entries:

init_dict = {
    'tot_niter' : int, # The outer iteration number
    'outer_dx' : float, # The L2 norm of the change in x between the last 2 outer iterations
    'inner_fin_f' : float, # The value of f at the optimum
    'inner_fin_g' : ndarray, # The value of g at the optimum
    'inner_fin_h' : ndarray, # The value of h at the optimum
    'inner_fin_x' : ndarray, # The optimum
    'inner_fin_l_aug' : float, # The value of the augmented Lagrangian objective l_k at the optimum
    'grad_l_k' : ndarray, # The gradient of the augmented Lagrangian objective l_k at the optimum
    'inner_fin_c' : float, # The final value of c
    'inner_fin_lam' : ndarray, # The final value of lambda
    'inner_fin_mu' : ndarray, # The final value of mu
    'inner_fin_niter' : int, # The number of L-BFGS iterations in the last step
    'inner_fin_dx_scaled' : float, # The L2 norm of the change in x between the last 2 inner L-BFGS iteration
    'inner_fin_du' : float, # The L2 norm of the change in update between the last 2 inner L-BFGS iteration
    'inner_fin_dl' : float, # The L2 norm of the change in f between the last 2 inner L-BFGS iteration
}

Return type:

dict

quadcoil.surface module

class quadcoil.surface.SurfaceJAX(quadpoints_phi: jax.numpy.ndarray, quadpoints_theta: 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,

first_fund_form()

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

gamma()
gammadash(a: int, b: int) → jax.numpy.ndarray[source]

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)

gammadash1()
gammadash1dash1()
gammadash1dash2()
gammadash2()
gammadash2dash2()
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)

unitnormal()
unitnormaldash()

d(unitnormal)/dphi and d(unitnormal)/dtheta.

Return type:

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

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.

copy_and_set_quadpoints()
from_desc()
from_simsopt()
gammadash()
to_desc()
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

copy_and_set_quadpoints()
from_simsopt()
gammadash()
num_dofs()
to_RZFourier()
to_simsopt()
tree_flatten()
tree_unflatten()
quadcoil.surface.dof_to_gamma(dofs, phi_grid, theta_grid, nfp, stellsym, dash1_order=0, dash2_order=0, mpol: int = 10, ntor: int = 10)[source]
quadcoil.surface.dof_to_gamma_op(phi_grid, theta_grid, nfp, stellsym, dash1_order=0, dash2_order=0, mpol: int = 10, ntor: int = 10)[source]

Operator of shape (nphi, ntheta, 3, ndof) mapping dofs -> gamma.

quadcoil.surface.dof_to_rz_op(phi_grid, theta_grid, nfp: int, stellsym: bool, dash1_order=0, dash2_order=0, mpol: int = 10, ntor: int = 10)

Operator mapping DOF vector -> (R, Z) on the quadrature grid.

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.winding_surface.fit_surfacerzfourier(phi_grid, theta_grid, r_fit, z_fit, nfp: int, stellsym: bool, mpol: int = 5, ntor: int = 5, lam_tikhonov=0.0, custom_weight=None)
quadcoil.winding_surface.gen_rot_matrix(theta)
quadcoil.winding_surface.gen_winding_surface_arc(plasma_gamma, d_expand, nfp, stellsym, unitnormal=None, mpol=5, ntor=5, pol_interp=2, tor_interp=2, lam_tikhonov=1e-05, rule='self-intersection')
quadcoil.winding_surface.gen_winding_surface_atan(plasma_gamma, d_expand, nfp, stellsym, unitnormal=None, mpol=5, ntor=5, pol_interp=2, tor_interp=2, lam_tikhonov=1e-05)

Create uniform offset

quadcoil.winding_surface.gen_winding_surface_offset(plasma_gamma, d_expand, nfp, stellsym, unitnormal=None, mpol=10, ntor=10)[source]

quadcoil.wrapper module

quadcoil.wrapper.get_quantity(func_name: str)[source]

Takes a string as input and returns the function with the same name in quadcoil.quantity. throws an error if a function with the same name cannot be found. Used to parse str in quadcoil.quadcoil.

Parameters:

func_name (str) – Name of the function to find.

Returns:

A callable with the same name in quadcoil.quantity.

Return type:

callable

quadcoil.wrapper.merge_callables(callables, merge_constraints=False, smoothing=None, smoothing_params=None)[source]

Merge a tuple of callable``s into one that takes 2 arguments (all functions in the ``quadcoil.objective do), by flattening and concatenating their outputs into an 1D array. Used to construct constraints.

Parameters:

callables (tuple of callables) – The callables to merge.

Returns:

A callable that returns a 1D array

Return type:

callable

Module contents