diff --git a/desc/batching.py b/desc/batching.py index 7b2a18f7b..129ef7ec3 100644 --- a/desc/batching.py +++ b/desc/batching.py @@ -1,8 +1,24 @@ """Utility functions for the ``batched_vectorize`` function.""" import functools +from functools import partial from typing import Callable, Optional +from jax._src.api import ( + _check_input_dtype_jacfwd, + _check_input_dtype_jacrev, + _check_output_dtype_jacfwd, + _check_output_dtype_jacrev, + _jacfwd_unravel, + _jacrev_unravel, + _jvp, + _std_basis, + _vjp, +) +from jax._src.api_util import _ensure_index, argnums_partial, check_callable +from jax._src.tree_util import tree_map, tree_structure, tree_transpose +from jax._src.util import wraps + from desc.backend import jax, jnp if jax.__version_info__ >= (0, 4, 16): @@ -320,3 +336,164 @@ def wrapped(*args, **kwargs): return jnp.expand_dims(result, axis=dims_to_expand) return wrapped + + +# The following section of this code is derived from JAX +# https://github.com/jax-ml/jax/blob/ff0a98a2aef958df156ca149809cf532efbbcaf4/ +# jax/_src/api.py +# +# The original copyright notice is as follows +# Copyright 2018 The JAX Authors. +# Licensed under the Apache License, Version 2.0 (the "License"); + + +def jacfwd_chunked( + fun, + argnums=0, + has_aux=False, + holomorphic=False, + *, + chunk_size=None, +): + """Jacobian of ``fun`` evaluated column-by-column using forward-mode AD. + + Parameters + ---------- + fun: callable + Function whose Jacobian is to be computed. + argnums: Optional, integer or sequence of integers. + Specifies which positional argument(s) to differentiate with respect to + (default ``0``). + has_aux: Optional, bool. + Indicates whether ``fun`` returns a pair where the first element is considered + the output of the mathematical function to be differentiated and the second + element is auxiliary data. Default False. + holomorphic: Optional, bool. + Indicates whether ``fun`` is promised to be holomorphic. Default False. + chunk_size: int + The size of the batches to pass to vmap. If None, defaults to the largest + possible chunk_size. + + Returns + ------- + jac: callable + A function with the same arguments as ``fun``, that evaluates the Jacobian of + ``fun`` using forward-mode automatic differentiation. If ``has_aux`` is True + then a pair of (jacobian, auxiliary_data) is returned. + + """ + check_callable(fun) + argnums = _ensure_index(argnums) + + docstr = ( + "Jacobian of {fun} with respect to positional argument(s) " + "{argnums}. Takes the same arguments as {fun} but returns the " + "jacobian of the output with respect to the arguments at " + "positions {argnums}." + ) + + @wraps(fun, docstr=docstr, argnums=argnums) + def jacfun(*args, **kwargs): + f = lu.wrap_init(fun, kwargs) + f_partial, dyn_args = argnums_partial( + f, argnums, args, require_static_args_hashable=False + ) + tree_map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args) + if not has_aux: + pushfwd: Callable = partial(_jvp, f_partial, dyn_args) + y, jac = vmap_chunked(pushfwd, chunk_size=chunk_size)(_std_basis(dyn_args)) + y = tree_map(lambda x: x[0], y) + jac = tree_map(lambda x: jnp.moveaxis(x, 0, -1), jac) + else: + pushfwd: Callable = partial(_jvp, f_partial, dyn_args, has_aux=True) + y, jac, aux = vmap_chunked(pushfwd, chunk_size=chunk_size)( + _std_basis(dyn_args) + ) + y = tree_map(lambda x: x[0], y) + jac = tree_map(lambda x: jnp.moveaxis(x, 0, -1), jac) + aux = tree_map(lambda x: x[0], aux) + tree_map(partial(_check_output_dtype_jacfwd, holomorphic), y) + example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args + jac_tree = tree_map(partial(_jacfwd_unravel, example_args), y, jac) + if not has_aux: + return jac_tree + else: + return jac_tree, aux + + return jacfun + + +def jacrev_chunked( + fun, + argnums=0, + has_aux=False, + holomorphic=False, + allow_int=False, + *, + chunk_size=None, +): + """Jacobian of ``fun`` evaluated row-by-row using reverse-mode AD. + + Parameters + ---------- + fun: callable + Function whose Jacobian is to be computed. + argnums: Optional, integer or sequence of integers. + Specifies which positional argument(s) to differentiate with respect to + (default ``0``). + has_aux: Optional, bool. + Indicates whether ``fun`` returns a pair where the first element is considered + the output of the mathematical function to be differentiated and the second + element is auxiliary data. Default False. + holomorphic: Optional, bool. + Indicates whether ``fun`` is promised to be holomorphic. Default False. + allow_int: Optional, bool. + Whether to allow differentiating with respect to integer valued inputs. The + gradient of an integer input will have a trivial vector-space dtype (float0). + Default False. + chunk_size: int + The size of the batches to pass to vmap. If None, defaults to the largest + possible chunk_size. + + Returns + ------- + jac: callable + A function with the same arguments as ``fun``, that evaluates the Jacobian of + ``fun`` using reverse-mode automatic differentiation. If ``has_aux`` is True + then a pair of (jacobian, auxiliary_data) is returned. + + """ + check_callable(fun) + + docstr = ( + "Jacobian of {fun} with respect to positional argument(s) " + "{argnums}. Takes the same arguments as {fun} but returns the " + "jacobian of the output with respect to the arguments at " + "positions {argnums}." + ) + + @wraps(fun, docstr=docstr, argnums=argnums) + def jacfun(*args, **kwargs): + f = lu.wrap_init(fun, kwargs) + f_partial, dyn_args = argnums_partial( + f, argnums, args, require_static_args_hashable=False + ) + tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args) + if not has_aux: + y, pullback = _vjp(f_partial, *dyn_args) + else: + y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True) + tree_map(partial(_check_output_dtype_jacrev, holomorphic), y) + jac = vmap_chunked(pullback, chunk_size=chunk_size)(_std_basis(y)) + jac = jac[0] if isinstance(argnums, int) else jac + example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args + jac_tree = tree_map(partial(_jacrev_unravel, y), example_args, jac) + jac_tree = tree_transpose( + tree_structure(example_args), tree_structure(y), jac_tree + ) + if not has_aux: + return jac_tree + else: + return jac_tree, aux + + return jacfun diff --git a/desc/coils.py b/desc/coils.py index be954c115..ea0da348a 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -1,5 +1,6 @@ """Classes for magnetic field coils.""" +import functools import numbers from abc import ABC from collections.abc import MutableSequence @@ -28,7 +29,7 @@ from desc.grid import LinearGrid from desc.magnetic_fields import _MagneticField from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter -from desc.utils import equals, errorif, flatten_list, safenorm, warnif +from desc.utils import cross, dot, equals, errorif, flatten_list, safenorm, warnif @jit @@ -245,7 +246,7 @@ def num_coils(self): """int: Number of coils.""" return 1 - def _compute_position(self, params=None, grid=None, **kwargs): + def _compute_position(self, params=None, grid=None, dx1=False, **kwargs): """Compute coil positions accounting for stellarator symmetry. Parameters @@ -255,18 +256,31 @@ def _compute_position(self, params=None, grid=None, **kwargs): grid : Grid or int, optional Grid of coordinates to evaluate at. Defaults to a Linear grid. If an integer, uses that many equally spaced points. + dx1 : bool + If True, also return dx/ds for the curve. Returns ------- x : ndarray, shape(len(self),source_grid.num_nodes,3) Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates. + x_s : ndarray, shape(len(self),source_grid.num_nodes,3) + Coil position derivatives, in [R,phi,Z] or [X,Y,Z] coordinates. + Only returned if dx1=True. """ - x = self.compute("x", grid=grid, params=params, **kwargs)["x"] - x = jnp.transpose(jnp.atleast_3d(x), [2, 0, 1]) # shape=(1,num_nodes,3) - basis = kwargs.pop("basis", "xyz") + kwargs.setdefault("basis", "xyz") + keys = ["x", "x_s"] if dx1 else ["x"] + data = self.compute(keys, grid=grid, params=params, **kwargs) + x = jnp.transpose(jnp.atleast_3d(data["x"]), [2, 0, 1]) # shape=(1,num_nodes,3) + if dx1: + x_s = jnp.transpose( + jnp.atleast_3d(data["x_s"]), [2, 0, 1] + ) # shape=(1,num_nodes,3) + basis = kwargs.get("basis", "xyz") if basis.lower() == "rpz": x = x.at[:, :, 1].set(jnp.mod(x[:, :, 1], 2 * jnp.pi)) + if dx1: + return x, x_s return x def _compute_A_or_B( @@ -1359,7 +1373,7 @@ def flip(self, *args, **kwargs): """Flip the coils across a plane.""" [coil.flip(*args, **kwargs) for coil in self.coils] - def _compute_position(self, params=None, grid=None, **kwargs): + def _compute_position(self, params=None, grid=None, dx1=False, **kwargs): """Compute coil positions accounting for stellarator symmetry. Parameters @@ -1369,25 +1383,35 @@ def _compute_position(self, params=None, grid=None, **kwargs): grid : Grid or int, optional Grid of coordinates to evaluate at. Defaults to a Linear grid. If an integer, uses that many equally spaced points. + dx1 : bool + If True, also return dx/ds for each curve. Returns ------- x : ndarray, shape(len(self),source_grid.num_nodes,3) Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates. + x_s : ndarray, shape(len(self),source_grid.num_nodes,3) + Coil position derivatives, in [R,phi,Z] or [X,Y,Z] coordinates. + Only returned if dx1=True. """ basis = kwargs.pop("basis", "xyz") + keys = ["x", "x_s"] if dx1 else ["x"] if params is None: - params = [get_params("x", coil, basis=basis) for coil in self] - data = self.compute("x", grid=grid, params=params, basis=basis, **kwargs) + params = [get_params(keys, coil, basis=basis) for coil in self] + data = self.compute(keys, grid=grid, params=params, basis=basis, **kwargs) data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict)) x = jnp.dstack([d["x"].T for d in data]).T # shape=(ncoils,num_nodes,3) - + if dx1: + x_s = jnp.dstack([d["x_s"].T for d in data]).T # shape=(ncoils,num_nodes,3) # stellarator symmetry is easiest in [X,Y,Z] coordinates - if basis.lower() == "rpz": - xyz = rpz2xyz(x) - else: - xyz = x + xyz = rpz2xyz(x) if basis.lower() == "rpz" else x + if dx1: + xyz_s = ( + rpz2xyz_vec(x_s, xyz[:, :, 0], xyz[:, :, 1]) + if basis.lower() == "rpz" + else x_s + ) # if stellarator symmetric, add reflected coils from the other half field period if self.sym: @@ -1396,27 +1420,64 @@ def _compute_position(self, params=None, grid=None, **kwargs): ) xyz_sym = xyz @ reflection_matrix(normal).T @ reflection_matrix([0, 0, 1]).T xyz = jnp.vstack((xyz, jnp.flipud(xyz_sym))) + if dx1: + xyz_s_sym = ( + xyz_s @ reflection_matrix(normal).T @ reflection_matrix([0, 0, 1]).T + ) + xyz_s = jnp.vstack((xyz_s, jnp.flipud(xyz_s_sym))) # field period rotation is easiest in [R,phi,Z] coordinates rpz = xyz2rpz(xyz) + if dx1: + rpz_s = xyz2rpz_vec(xyz_s, xyz[:, :, 0], xyz[:, :, 1]) # if field period symmetry, add rotated coils from other field periods - if self.NFP > 1: - rpz0 = rpz - for k in range(1, self.NFP): - rpz = jnp.vstack( - (rpz, rpz0 + jnp.array([0, 2 * jnp.pi * k / self.NFP, 0])) - ) + rpz0 = rpz + for k in range(1, self.NFP): + rpz = jnp.vstack((rpz, rpz0 + jnp.array([0, 2 * jnp.pi * k / self.NFP, 0]))) + if dx1: + rpz_s = jnp.tile(rpz_s, (self.NFP, 1, 1)) # ensure phi in [0, 2pi) rpz = rpz.at[:, :, 1].set(jnp.mod(rpz[:, :, 1], 2 * jnp.pi)) - if basis.lower() == "xyz": - x = rpz2xyz(rpz) - else: - x = rpz + x = rpz2xyz(rpz) if basis.lower() == "xyz" else rpz + if dx1: + x_s = ( + rpz2xyz_vec(rpz_s, phi=rpz[:, :, 1]) + if basis.lower() == "xyz" + else rpz_s + ) + return x, x_s return x + def _compute_linking_number(self, params=None, grid=None): + """Calculate linking numbers for coils in the coilset. + + Parameters + ---------- + params : dict or array-like of dict, optional + Parameters to pass to coils, either the same for all coils or one for each. + grid : Grid or int, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid. + If an integer, uses that many equally spaced points. + + Returns + ------- + link : ndarray, shape(num_coils, num_coils) + Linking number of each coil with each other coil. link=0 means they are not + linked, +/- 1 means the coils link each other in one direction or another. + + """ + if grid is None: + grid = LinearGrid(N=50) + dx = grid.spacing[:, 2] + x, x_s = self._compute_position(params, grid, dx1=True, basis="xyz") + link = _linking_number( + x[:, None], x[None, :], x_s[:, None], x_s[None, :], dx, dx + ) + return link / (4 * jnp.pi) + def _compute_A_or_B( self, coords, @@ -2307,7 +2368,7 @@ def compute( ) ] - def _compute_position(self, params=None, grid=None, **kwargs): + def _compute_position(self, params=None, grid=None, dx1=False, **kwargs): """Compute coil positions accounting for stellarator symmetry. Parameters @@ -2318,11 +2379,16 @@ def _compute_position(self, params=None, grid=None, **kwargs): Grid of coordinates to evaluate at. Defaults to a Linear grid. If an integer, uses that many equally spaced points. If array-like, should be 1 value per coil. + dx1 : bool + If True, also return dx/ds for each curve. Returns ------- x : ndarray, shape(len(self),source_grid.num_nodes,3) Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates. + x_s : ndarray, shape(len(self),source_grid.num_nodes,3) + Coil position derivatives, in [R,phi,Z] or [X,Y,Z] coordinates. + Only returned if dx1=True. """ errorif( @@ -2331,15 +2397,17 @@ def _compute_position(self, params=None, grid=None, **kwargs): "grid must be supplied to MixedCoilSet._compute_position, since the " + "default grid for each coil could have a different number of nodes.", ) + kwargs.setdefault("basis", "xyz") params = self._make_arraylike(params) grid = self._make_arraylike(grid) - x = jnp.vstack( - [ - coil._compute_position(par, grd, **kwargs) - for coil, par, grd in zip(self.coils, params, grid) - ] - ) - return x + out = [] + for coil, par, grd in zip(self.coils, params, grid): + out.append(coil._compute_position(par, grd, dx1, **kwargs)) + if dx1: + x = jnp.vstack([foo[0] for foo in out]) + x_s = jnp.vstack([foo[1] for foo in out]) + return x, x_s + return jnp.vstack(out) def _compute_A_or_B( self, @@ -2795,3 +2863,18 @@ def flatten_coils(coilset): if ignore_groups: cset = cls(*flatten_coils(cset), check_intersection=check_intersection) return cset + + +@functools.partial(jnp.vectorize, signature="(m,3),(n,3),(m,3),(n,3),(m),(n)->()") +def _linking_number(x1, x2, x1_s, x2_s, dx1, dx2): + """Linking number between curves x1 and x2 with tangents x1_s, x2_s.""" + x1_s *= dx1[:, None] + x2_s *= dx2[:, None] + dx = x1[:, None, :] - x2[None, :, :] # shape(m,n,3) + dx_norm = safenorm(dx, axis=-1) # shape(m,n) + den = dx_norm**3 + dr1xdr2 = cross(x1_s[:, None, :], x2_s[None, :, :], axis=-1) # shape(m,n,3) + num = dot(dx, dr1xdr2, axis=-1) # shape(m,n) + small = dx_norm < jnp.finfo(x1.dtype).eps + ratio = jnp.where(small, 0.0, num / jnp.where(small, 1.0, den)) + return ratio.sum() diff --git a/desc/derivatives.py b/desc/derivatives.py index 684ea37d4..d851bf7b5 100644 --- a/desc/derivatives.py +++ b/desc/derivatives.py @@ -6,10 +6,13 @@ from termcolor import colored from desc.backend import jnp, put, use_jax +from desc.utils import ensure_tuple if use_jax: import jax + from desc.batching import jacfwd_chunked, jacrev_chunked + class _Derivative(ABC): """_Derivative is an abstract base class for derivative matrix calculations. @@ -123,11 +126,11 @@ class AutoDiffDerivative(_Derivative): """ - def __init__(self, fun, argnum=0, mode="fwd", **kwargs): + def __init__(self, fun, argnum=0, mode="fwd", chunk_size=None, **kwargs): self._fun = fun self._argnum = argnum - + self._chunk_size = chunk_size self._set_mode(mode) def compute(self, *args, **kwargs): @@ -205,7 +208,7 @@ def compute_jvp(cls, fun, argnum, v, *args, **kwargs): """ _ = kwargs.pop("rel_step", None) # unused by autodiff argnum = (argnum,) if jnp.isscalar(argnum) else tuple(argnum) - v = (v,) if not isinstance(v, (tuple, list)) else v + v = ensure_tuple(v) def _fun(*x): _args = list(args) @@ -241,14 +244,14 @@ def compute_jvp2(cls, fun, argnum1, argnum2, v1, v2, *args, **kwargs): """ if np.isscalar(argnum1): - v1 = (v1,) if not isinstance(v1, (tuple, list)) else v1 + v1 = ensure_tuple(v1) argnum1 = (argnum1,) else: v1 = tuple(v1) if np.isscalar(argnum2): argnum2 = (argnum2 + 1,) - v2 = (v2,) if not isinstance(v2, (tuple, list)) else v2 + v2 = ensure_tuple(v2) else: argnum2 = tuple([i + 1 for i in argnum2]) v2 = tuple(v2) @@ -284,21 +287,21 @@ def compute_jvp3(cls, fun, argnum1, argnum2, argnum3, v1, v2, v3, *args, **kwarg """ if np.isscalar(argnum1): - v1 = (v1,) if not isinstance(v1, (tuple, list)) else v1 + v1 = ensure_tuple(v1) argnum1 = (argnum1,) else: v1 = tuple(v1) if np.isscalar(argnum2): argnum2 = (argnum2 + 1,) - v2 = (v2,) if not isinstance(v2, (tuple, list)) else v2 + v2 = ensure_tuple(v2) else: argnum2 = tuple([i + 1 for i in argnum2]) v2 = tuple(v2) if np.isscalar(argnum3): argnum3 = (argnum3 + 2,) - v3 = (v3,) if not isinstance(v3, (tuple, list)) else v3 + v3 = ensure_tuple(v3) else: argnum3 = tuple([i + 2 for i in argnum3]) v3 = tuple(v3) @@ -323,9 +326,13 @@ def _set_mode(self, mode) -> None: self._mode = mode if self._mode == "fwd": - self._compute = jax.jacfwd(self._fun, self._argnum) + self._compute = jacfwd_chunked( + self._fun, self._argnum, chunk_size=self._chunk_size + ) elif self._mode == "rev": - self._compute = jax.jacrev(self._fun, self._argnum) + self._compute = jacrev_chunked( + self._fun, self._argnum, chunk_size=self._chunk_size + ) elif self._mode == "grad": self._compute = jax.grad(self._fun, self._argnum) elif self._mode == "hess": @@ -512,7 +519,7 @@ def compute_jvp(cls, fun, argnum, v, *args, **kwargs): argnum = (argnum,) else: nargs = len(argnum) - v = (v,) if not isinstance(v, tuple) else v + v = ensure_tuple(v) f = np.array( [ @@ -549,14 +556,14 @@ def compute_jvp2(cls, fun, argnum1, argnum2, v1, v2, *args, **kwargs): """ if np.isscalar(argnum1): - v1 = (v1,) if not isinstance(v1, tuple) else v1 + v1 = ensure_tuple(v1) argnum1 = (argnum1,) else: v1 = tuple(v1) if np.isscalar(argnum2): argnum2 = (argnum2 + 1,) - v2 = (v2,) if not isinstance(v2, tuple) else v2 + v2 = ensure_tuple(v2) else: argnum2 = tuple([i + 1 for i in argnum2]) v2 = tuple(v2) @@ -592,21 +599,21 @@ def compute_jvp3(cls, fun, argnum1, argnum2, argnum3, v1, v2, v3, *args, **kwarg """ if np.isscalar(argnum1): - v1 = (v1,) if not isinstance(v1, tuple) else v1 + v1 = ensure_tuple(v1) argnum1 = (argnum1,) else: v1 = tuple(v1) if np.isscalar(argnum2): argnum2 = (argnum2 + 1,) - v2 = (v2,) if not isinstance(v2, tuple) else v2 + v2 = ensure_tuple(v2) else: argnum2 = tuple([i + 1 for i in argnum2]) v2 = tuple(v2) if np.isscalar(argnum3): argnum3 = (argnum3 + 2,) - v3 = (v3,) if not isinstance(v3, tuple) else v3 + v3 = ensure_tuple(v3) else: argnum3 = tuple([i + 2 for i in argnum3]) v3 = tuple(v3) diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py index 2f74200aa..817eee4a1 100644 --- a/desc/geometry/surface.py +++ b/desc/geometry/surface.py @@ -350,6 +350,8 @@ def from_qp_model( ): """Create a surface from a near-axis model for quasi-poloidal symmetry. + Model is based off of section III of Goodman et. al. [1]_ + Parameters ---------- major_radius : float @@ -376,6 +378,11 @@ def from_qp_model( surface : FourierRZToroidalSurface Surface with given geometric properties. + References + ---------- + .. [1] Goodman, Alan, et al. "Constructing Precisely Quasi-Isodynamic + Magnetic Fields." Journal of Plasma Physics 89.5 (2023): 905890504. + """ assert mirror_ratio <= 1 a = major_radius * np.sqrt(elongation) / aspect_ratio # major axis diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 1504dbdd7..00c959582 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -5,6 +5,7 @@ CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, PlasmaCoilSetMinDistance, diff --git a/desc/objectives/_bootstrap.py b/desc/objectives/_bootstrap.py index c76a003b2..9f6850cb2 100644 --- a/desc/objectives/_bootstrap.py +++ b/desc/objectives/_bootstrap.py @@ -9,7 +9,7 @@ from desc.utils import Timer, errorif, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class BootstrapRedlConsistency(_Objective): @@ -31,31 +31,6 @@ class BootstrapRedlConsistency(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` @@ -64,22 +39,13 @@ class BootstrapRedlConsistency(_Objective): First entry must be M=1. Second entry is the toroidal mode number N, used for evaluating the Redl bootstrap current formula. Set to 0 for axisymmetry or quasi-axisymmetry; set to +/-NFP for quasi-helical symmetry. - name : str, optional - Name of the objective function. - jac_chunk_size : int, optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "r" _units = "(T A m^-2)" _print_value_fmt = "Bootstrap current self-consistency error: " diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 40c53bd85..d8b68604e 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -17,7 +17,7 @@ from desc.utils import Timer, errorif, safenorm, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class _CoilObjective(_Objective): @@ -29,50 +29,16 @@ class _CoilObjective(_Objective): Coil for which the data keys will be optimized. data_keys : list of str data keys that will be optimized when this class is inherited. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, list, optional Collocation grid containing the nodes to evaluate at. If a list, must have the same structure as coil. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + coil=True, + ) + def __init__( self, coil, @@ -235,52 +201,18 @@ class CoilLength(_CoilObjective): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``target=2*np.pi``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=2*np.pi``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=2*np.pi``.", + bounds_default="``target=2*np.pi``.", + coil=True, + ) + _scalar = False # Not always a scalar, if a coilset is passed in _units = "(m)" _print_value_fmt = "Coil length: " @@ -373,52 +305,18 @@ class CoilCurvature(_CoilObjective): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``bounds=(0,1)``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(0,1)``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(0,1).``", + bounds_default="``bounds=(0,1).``", + coil=True, + ) + _scalar = False _units = "(m^-1)" _print_value_fmt = "Coil curvature: " @@ -506,52 +404,18 @@ class CoilTorsion(_CoilObjective): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``target=0``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + coil=True, + ) + _scalar = False _units = "(m^-1)" _print_value_fmt = "Coil torsion: " @@ -639,52 +503,18 @@ class CoilCurrentLength(CoilLength): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``target=0``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + coil=True, + ) + _scalar = False _units = "(A*m)" _print_value_fmt = "Coil current length: " @@ -780,52 +610,19 @@ class CoilSetMinDistance(_Objective): ---------- coil : CoilSet Coil(s) that are to be optimized. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individial coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, list, optional Collocation grid used to discretize each coil. Defaults to the default grid for the given coil-type, see ``coils.py`` and ``curve.py`` for more details. If a list, must have the same structure as coils. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + coil=True, + ) + _scalar = False _units = "(m)" _print_value_fmt = "Minimum coil-coil distance: " @@ -951,32 +748,6 @@ class PlasmaCoilSetMinDistance(_Objective): to satisfy the Objective. coil : CoilSet Coil(s) that are to be optimized. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individial coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. plasma_grid : Grid, optional Collocation grid containing the nodes to evaluate plasma geometry at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. @@ -997,22 +768,15 @@ class PlasmaCoilSetMinDistance(_Objective): during optimization, and self.things = [eq] only. If False, the coil coordinates are computed at every iteration. False by default, so that self.things = [coil, eq]. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + coil=True, + ) + _scalar = False _units = "(m)" _print_value_fmt = "Minimum plasma-coil distance: " @@ -1209,22 +973,6 @@ class QuadraticFlux(_Objective): field : MagneticField External field produced by coils or other source, which will be optimized to minimize the normal field error on the provided equilibrium's surface. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - len(target) must be equal to Objective.dim_f. - Default target is zero. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - len(bounds[0]) and len(bounds[1]) must be equal to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - len(weight) must be equal to Objective.dim_f - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. source_grid : Grid, optional Collocation grid containing the nodes for plasma source terms. Default grid is detailed in the docs for ``compute_B_plasma`` @@ -1240,22 +988,14 @@ class QuadraticFlux(_Objective): vacuum : bool If true, B_plasma (the contribution to the normal field on the boundary from the plasma currents) is set to zero. - name : str - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + ) + _scalar = False _linear = False _print_value_fmt = "Boundary normal field error: " @@ -1430,30 +1170,6 @@ class ToroidalFlux(_Objective): field : MagneticField MagneticField object, the parameters of this will be optimized to minimize the objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Defaults to eq.Psi. Must be broadcastable to Objective.dim_f. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: has no effect for this objective - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. field_grid : Grid, optional Grid containing the nodes to evaluate field source at on the winding surface. (used if e.g. field is a CoilSet or @@ -1463,23 +1179,15 @@ class ToroidalFlux(_Objective): Collocation grid containing the nodes to evaluate the normal magnetic field at plasma geometry at. Defaults to a LinearGrid(L=eq.L_grid, M=eq.M_grid, zeta=jnp.array(0.0), NFP=eq.NFP). - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. - """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=eq.Psi``.", + bounds_default="``target=eq.Psi``.", + loss_detail=" Note: has no effect for this objective.", + ) + _coordinates = "rtz" _units = "(Wb)" _print_value_fmt = "Toroidal Flux: " @@ -1654,3 +1362,118 @@ def compute(self, field_params=None, constants=None): ) return Psi + + +class CoilSetLinkingNumber(_Objective): + """Prevents coils from becoming interlinked. + + The linking number of 2 curves is (approximately) 0 if they are not linked, and + (approximately) +/-1 if they are (with the sign indicating the helicity of the + linking). + + This objective returns a single value for each coil in the coilset, with that number + being the sum of the absolute value of the linking numbers of that coil with every + other coil in the coilset, approximating the number of other coils that are linked + + Parameters + ---------- + coil : CoilSet + Coil(s) that are to be optimized. + grid : Grid, list, optional + Collocation grid used to discretize each coil. Defaults to + ``LinearGrid(N=50)`` + + """ + + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + coil=True, + ) + + _scalar = False + _units = "(dimensionless)" + _print_value_fmt = "Coil linking number: " + + def __init__( + self, + coil, + grid=None, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + loss_function=None, + deriv_mode="auto", + jac_chunk_size=None, + name="coil-coil linking number", + ): + from desc.coils import CoilSet + + if target is None and bounds is None: + target = 0 + self._grid = grid + errorif( + not isinstance(coil, CoilSet), + ValueError, + "coil must be of type CoilSet, not an individual Coil", + ) + super().__init__( + things=coil, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + loss_function=loss_function, + deriv_mode=deriv_mode, + jac_chunk_size=jac_chunk_size, + name=name, + ) + + def build(self, use_jit=True, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + coilset = self.things[0] + grid = self._grid or LinearGrid(N=50) + + self._dim_f = coilset.num_coils + self._constants = {"coilset": coilset, "grid": grid, "quad_weights": 1.0} + + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params, constants=None): + """Compute linking numbers between coils. + + Parameters + ---------- + params : dict + Dictionary of coilset degrees of freedom, eg CoilSet.params_dict + constants : dict + Dictionary of constant data, eg transforms, profiles etc. + Defaults to self._constants. + + Returns + ------- + f : array of floats + For each coil, the sum of the absolute value of the linking numbers between + that coil and every other coil in the coilset, which approximates the + number of coils linked with that coil. + + """ + if constants is None: + constants = self.constants + link = constants["coilset"]._compute_linking_number( + params=params, grid=constants["grid"] + ) + + return jnp.abs(link).sum(axis=0) diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index 0124a84b1..7be04509e 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -7,7 +7,7 @@ from desc.utils import Timer from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class ForceBalance(_Objective): @@ -33,50 +33,16 @@ class ForceBalance(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(N)" @@ -222,46 +188,16 @@ class ForceBalanceAnisotropic(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : float, ndarray, optional - Target value(s) of the objective. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. grid : Grid, ndarray, optional - Collocation grid containing the nodes to evaluate at. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. grid : Grid, ndarray, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - name : str - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _units = "(N)" _coordinates = "rtz" _equilibrium = True @@ -389,50 +325,16 @@ class RadialForceBalance(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(N)" @@ -560,50 +462,16 @@ class HelicalForceBalance(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(N)" @@ -727,52 +595,18 @@ class Energy(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` gamma : float, optional Adiabatic (compressional) index. Default = 0. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = True _coordinates = "" _equilibrium = True @@ -909,50 +743,16 @@ class CurrentDensity(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(A*m)" diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index cd75899f4..9429b328e 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -11,8 +11,8 @@ from desc.grid import LinearGrid from desc.integrals import DFTInterpolator, FFTInterpolator, virtual_casing_biot_savart from desc.nestor import Nestor -from desc.objectives.objective_funs import _Objective -from desc.utils import PRINT_WIDTH, Timer, errorif, warnif +from desc.objectives.objective_funs import _Objective, collect_docs +from desc.utils import PRINT_WIDTH, Timer, errorif, parse_argname_change, warnif from .normalization import compute_scaling_factors @@ -38,32 +38,7 @@ class VacuumBoundaryError(_Objective): Equilibrium that will be optimized to satisfy the Objective. field : MagneticField External field produced by coils or other sources outside the plasma. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. - grid : Grid, optional + eval_grid : Grid, optional Collocation grid containing the nodes to evaluate error at. Should be at rho=1. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` field_grid : Grid, optional @@ -71,22 +46,13 @@ class VacuumBoundaryError(_Objective): field_fixed : bool Whether to assume the field is fixed. For free boundary solve, should be fixed. For single stage optimization, should be False (default). - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = False _linear = False _print_value_fmt = "Boundary Error: " @@ -104,15 +70,17 @@ def __init__( normalize_target=True, loss_function=None, deriv_mode="auto", - grid=None, + eval_grid=None, field_grid=None, field_fixed=False, name="Vacuum boundary error", jac_chunk_size=None, + **kwargs, ): + eval_grid = parse_argname_change(eval_grid, kwargs, "grid", "eval_grid") if target is None and bounds is None: target = 0 - self._grid = grid + self._eval_grid = eval_grid self._eq = eq self._field = field self._field_grid = field_grid @@ -146,12 +114,12 @@ def build(self, use_jit=True, verbose=1): """ eq = self.things[0] - if self._grid is None: + if self._eval_grid is None: grid = LinearGrid( rho=np.array([1.0]), M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False ) else: - grid = self._grid + grid = self._eval_grid pres = np.max(np.abs(eq.compute("p")["p"])) curr = np.max(np.abs(eq.compute("current")["current"])) @@ -379,31 +347,6 @@ class BoundaryError(_Objective): Equilibrium that will be optimized to satisfy the Objective. field : MagneticField External field produced by coils. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. s, q : integer Hyperparameters for singular integration scheme, s is roughly equal to the size of the local singular grid with respect to the global grid, q is the order of @@ -421,21 +364,13 @@ class BoundaryError(_Objective): loop : bool If True, evaluate integral using loops, as opposed to vmap. Slower, but uses less memory. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. + """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + __doc__ += """ Examples -------- Assigning a surface current to the equilibrium: @@ -860,53 +795,19 @@ class BoundaryErrorNESTOR(_Objective): Equilibrium that will be optimized to satisfy the Objective. field : MagneticField External field produced by coils. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. mf, nf : integer maximum poloidal and toroidal mode numbers to use for NESTOR scalar potential. ntheta, nzeta : int number of grid points in poloidal, toroidal directions to use in NESTOR. field_grid : Grid, optional Grid used to discretize field. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = False _linear = False _print_value_fmt = "Boundary magnetic pressure error: " diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index b2e568fd2..8e16e61cf 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -14,7 +14,7 @@ from desc.utils import errorif, parse_argname_change from .linear_objectives import _FixedObjective -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class GenericObjective(_Objective): @@ -26,51 +26,16 @@ class GenericObjective(_Objective): Name of the quantity to compute. thing : Optimizable Object that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f. - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect on this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _print_value_fmt = "Generic objective value: " def __init__( @@ -200,28 +165,13 @@ class LinearObjectiveFromUser(_FixedObjective): Custom objective function. thing : Optimizable Object whose degrees of freedom are being constrained. - target : dict of {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of dict {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : dict of {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Should be a scalar or have the same tree structure as thing.params. - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Has no effect for this objective. - name : str, optional - Name of the objective function. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = False _linear = True _fixed = True @@ -326,49 +276,16 @@ class ObjectiveFromUser(_Objective): Custom objective function. thing : Optimizable Object that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. + """ + + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + __doc__ += """ Examples -------- .. code-block:: python diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 6c8c1376d..65fa1cf82 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -9,7 +9,7 @@ from desc.utils import Timer, errorif, parse_argname_change, safenorm, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs from .utils import check_if_points_are_inside_perimeter, softmin @@ -21,52 +21,21 @@ class AspectRatio(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=2``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=2``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=2``.", + bounds_default="``target=2``.", + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(dimensionless)" _print_value_fmt = "Aspect ratio: " @@ -196,52 +165,20 @@ class Elongation(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1``.", + bounds_default="``target=1``.", + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(dimensionless)" _print_value_fmt = "Elongation: " @@ -370,52 +307,19 @@ class Volume(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1``.", + bounds_default="``target=1``.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(m^3)" _print_value_fmt = "Plasma volume: " @@ -570,31 +474,6 @@ class PlasmaVesselDistance(_Objective): Equilibrium that will be optimized to satisfy the Objective. surface : Surface Bounding surface to penalize distance to. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``bounds=(1,np.inf)``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(1,np.inf)``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. surface_grid : Grid, optional Collocation grid containing the nodes to evaluate surface geometry at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. @@ -623,22 +502,14 @@ class PlasmaVesselDistance(_Objective): the array by 2/min_val to ensure that softmin_alpha*array>1. Making softmin_alpha larger than this minimum value will make the softmin a more accurate approximation of the true min. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + ) + _coordinates = "rtz" _units = "(m)" _print_value_fmt = "Plasma-vessel distance: " @@ -946,51 +817,18 @@ class MeanCurvature(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``bounds=(-np.inf, 0)``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(-np.inf, 0)``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(-np.inf,0)``.", + bounds_default="``bounds=(-np.inf,0)``.", + ) + _coordinates = "rtz" _units = "(m^-1)" _print_value_fmt = "Mean curvature: " @@ -1119,51 +957,18 @@ class PrincipalCurvature(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1``.", + bounds_default="``target=1``.", + ) + _coordinates = "rtz" _units = "(m^-1)" _print_value_fmt = "Principal curvature: " @@ -1288,50 +1093,17 @@ class BScaleLength(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``bounds=(1,np.inf)``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(1,np.inf)``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + ) + _coordinates = "rtz" _units = "(m)" _print_value_fmt = "Magnetic field scale length: " @@ -1453,49 +1225,16 @@ class GoodCoordinates(_Objective): Equilibrium that will be optimized to satisfy the Objective. sigma : float Relative weight between the Jacobian and radial terms. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + ) + _scalar = False _units = "(dimensionless)" _print_value_fmt = "Coordinate goodness : " diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index 415716c2e..8ac2faa9e 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -10,7 +10,7 @@ from desc.vmec_utils import ptolemy_linear_transform from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class QuasisymmetryBoozer(_Objective): @@ -20,31 +20,6 @@ class QuasisymmetryBoozer(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Must be a LinearGrid with sym=False. @@ -55,22 +30,13 @@ class QuasisymmetryBoozer(_Objective): Poloidal resolution of Boozer transformation. Default = 2 * eq.M. N_booz : int, optional Toroidal resolution of Boozer transformation. Default = 2 * eq.N. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _units = "(T)" _print_value_fmt = "Quasi-symmetry Boozer error: " @@ -163,7 +129,7 @@ def build(self, use_jit=True, verbose=1): M_booz=M_booz, N_booz=N_booz, ) - matrix, modes, idx = ptolemy_linear_transform( + matrix, _, idx = ptolemy_linear_transform( transforms["B"].basis.modes, helicity=self.helicity, NFP=transforms["B"].basis.NFP, @@ -180,7 +146,7 @@ def build(self, use_jit=True, verbose=1): if verbose > 1: timer.disp("Precomputing transforms") - self._dim_f = idx.size + self._dim_f = idx.size * grid.num_rho if self._normalize: scales = compute_scaling_factors(eq) @@ -250,52 +216,18 @@ class QuasisymmetryTwoTerm(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. helicity : tuple, optional Type of quasi-symmetry (M, N). - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(T^3)" _print_value_fmt = "Quasi-symmetry two-term error: " @@ -450,50 +382,16 @@ class QuasisymmetryTripleProduct(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(T^4/m^2)" _print_value_fmt = "Quasi-symmetry error: " @@ -614,31 +512,6 @@ class Omnigenity(_Objective): Equilibrium to be optimized to satisfy the Objective. field : OmnigenousField Omnigenous magnetic field to be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. eq_grid : Grid, optional Collocation grid containing the nodes to evaluate at for equilibrium data. Defaults to a linearly space grid on the rho=1 surface. @@ -669,22 +542,13 @@ class Omnigenity(_Objective): computation time during optimization and self.things = [eq] only. If False, the field is allowed to change during the optimization and its associated data are re-computed at every iteration (Default). - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(T)" _print_value_fmt = "Omnigenity error: " @@ -988,51 +852,16 @@ class Isodynamicity(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(dimensionless)" _print_value_fmt = "Isodynamicity error: " diff --git a/desc/objectives/_power_balance.py b/desc/objectives/_power_balance.py index 313aa8734..bc291cb8a 100644 --- a/desc/objectives/_power_balance.py +++ b/desc/objectives/_power_balance.py @@ -6,7 +6,7 @@ from desc.utils import Timer, errorif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class FusionPower(_Objective): @@ -24,52 +24,20 @@ class FusionPower(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1e9``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1e9``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. fuel : str, optional Fusion fuel, assuming a 50/50 mix. One of {'DT'}. Default = 'DT'. grid : Grid, optional Collocation grid used to compute the intermediate quantities. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP)``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1e9``.", + bounds_default="``target=1e9``.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(W)" _print_value_fmt = "Fusion power: " @@ -220,31 +188,6 @@ class HeatingPowerISS04(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. H_ISS04 : float, optional ISS04 confinement enhancement factor. Default = 1. gamma : float, optional @@ -252,22 +195,15 @@ class HeatingPowerISS04(_Objective): grid : Grid, optional Collocation grid used to compute the intermediate quantities. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP)``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(W)" _print_value_fmt = "Heating power: " diff --git a/desc/objectives/_profiles.py b/desc/objectives/_profiles.py index 1103421d4..0d0a522e4 100644 --- a/desc/objectives/_profiles.py +++ b/desc/objectives/_profiles.py @@ -8,65 +8,43 @@ from desc.utils import Timer, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs from .utils import _parse_callable_target_bounds - -class Pressure(_Objective): - """Target pressure profile. - - Parameters - ---------- - eq : Equilibrium - Equilibrium that will be optimized to satisfy the Objective. +profile_overwrite = { + "target": """ target : {float, ndarray, callable}, optional Target value(s) of the objective. Only used if bounds is None. Must be broadcastable to Objective.dim_f. If a callable, should take a single argument `rho` and return the desired value of the profile at those locations. Defaults to ``target=0``. + """, + "bounds": """ bounds : tuple of {float, ndarray, callable}, optional Lower and upper bounds on the objective. Overrides target. Both bounds must be broadcastable to to Objective.dim_f If a callable, each should take a single argument `rho` and return the desired bound (lower or upper) of the profile at those locations. Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. + """, +} + + +class Pressure(_Objective): + """Target pressure profile. + + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid)``. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs(overwrite=profile_overwrite) + _coordinates = "r" _units = "(Pa)" _print_value_fmt = "Pressure: " @@ -188,56 +166,20 @@ class RotationalTransform(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``target=0``. - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + overwrite=profile_overwrite, + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _coordinates = "r" _units = "(dimensionless)" _print_value_fmt = "Rotational transform: " @@ -372,56 +314,20 @@ class Shear(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``target=0``. - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + overwrite=profile_overwrite, + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _coordinates = "r" _units = "(dimensionless)" _print_value_fmt = "Shear: " @@ -552,56 +458,16 @@ class ToroidalCurrent(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``target=0``. - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs(overwrite=profile_overwrite) + _coordinates = "r" _units = "(A)" _print_value_fmt = "Toroidal current: " diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index 600359b40..02da27459 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -9,9 +9,27 @@ from desc.utils import Timer, errorif, setdefault, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs from .utils import _parse_callable_target_bounds +overwrite_stability = { + "target": """ + target : {float, ndarray, callable}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. If a callable, should take a + single argument `rho` and return the desired value of the profile at those + locations. Defaults to ``bounds=(0, np.inf)`` + """, + "bounds": """ + bounds : tuple of {float, ndarray, callable}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to to Objective.dim_f + If a callable, each should take a single argument `rho` and return the + desired bound (lower or upper) of the profile at those locations. + Defaults to ``bounds=(0, np.inf)`` + """, +} + class MercierStability(_Objective): """The Mercier criterion is a fast proxy for MHD stability. @@ -28,56 +46,16 @@ class MercierStability(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``bounds=(0, np.inf)`` - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``bounds=(0, np.inf)`` - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs(overwrite=overwrite_stability) + _coordinates = "r" _units = "(Wb^-2)" _print_value_fmt = "Mercier Stability: " @@ -222,56 +200,20 @@ class MagneticWell(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``bounds=(0, np.inf)`` - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``bounds=(0, np.inf)`` - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - name : str, optional - Name of the objective function. - jac_chunk_size : int , optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + overwrite=overwrite_stability, + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _coordinates = "r" _units = "(dimensionless)" _print_value_fmt = "Magnetic Well: " diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py index 727c6f183..7d4772a8b 100644 --- a/desc/objectives/getters.py +++ b/desc/objectives/getters.py @@ -1,6 +1,6 @@ """Utilities for getting standard groups of objectives and constraints.""" -from desc.utils import flatten_list, is_any_instance, unique_list +from desc.utils import flatten_list, is_any_instance, isposint, unique_list from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance from .linear_objectives import ( @@ -86,7 +86,10 @@ def get_equilibrium_objective(eq, mode="force", normalize=True, jac_chunk_size=" objectives = (RadialForceBalance(**kwargs), HelicalForceBalance(**kwargs)) else: raise ValueError("got an unknown equilibrium objective type '{}'".format(mode)) - return ObjectiveFunction(objectives, jac_chunk_size=jac_chunk_size) + deriv_mode = "batched" if isposint(jac_chunk_size) else "auto" + return ObjectiveFunction( + objectives, jac_chunk_size=jac_chunk_size, deriv_mode=deriv_mode + ) def get_fixed_axis_constraints(eq, profiles=True, normalize=True): diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 6ef14660c..6f302f2df 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -3065,7 +3065,7 @@ def build(self, use_jit=True, verbose=1): np.logical_and((basis.modes[:, 1] % 2 == 0), basis.modes[:, 1] > 0), ) )[0] - mm = basis.modes[idx_m, 2] + mm = basis.modes[idx_m, 1] self._A[i, idx_0] = 1 self._A[i, idx_m] = (mm % 2 - 1) * (mm % 4 - 1) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 9a540d0b5..28d7530d1 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -13,6 +13,7 @@ jit, jnp, tree_flatten, + tree_map, tree_unflatten, use_jax, ) @@ -23,6 +24,7 @@ from desc.utils import ( PRINT_WIDTH, Timer, + ensure_tuple, errorif, flatten_list, is_broadcastable, @@ -31,6 +33,153 @@ unique_list, ) +doc_target = """ + target : {float, ndarray}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. +""" +doc_bounds = """ + bounds : tuple of {float, ndarray}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to Objective.dim_f +""" +doc_weight = """ + weight : {float, ndarray}, optional + Weighting to apply to the Objective, relative to other Objectives. + Must be broadcastable to Objective.dim_f +""" +doc_normalize = """ + normalize : bool, optional + Whether to compute the error in physical units or non-dimensionalize. +""" +doc_normalize_target = """ + normalize_target : bool, optional + Whether target and bounds should be normalized before comparing to computed + values. If `normalize` is `True` and the target is in physical units, + this should also be set to True. +""" +doc_loss_function = """ + loss_function : {None, 'mean', 'min', 'max'}, optional + Loss function to apply to the objective values once computed. This loss function + is called on the raw compute value, before any shifting, scaling, or + normalization. +""" +doc_deriv_mode = """ + deriv_mode : {"auto", "fwd", "rev"} + Specify how to compute Jacobian matrix, either forward mode or reverse mode AD. + "auto" selects forward or reverse mode based on the size of the input and output + of the objective. Has no effect on self.grad or self.hess which always use + reverse mode and forward over reverse mode respectively. +""" +doc_name = """ + name : str, optional + Name of the objective. +""" +doc_jac_chunk_size = """ + jac_chunk_size : int or "auto", optional + Will calculate the Jacobian + ``jac_chunk_size`` columns at a time, instead of all at once. + The memory usage of the Jacobian calculation is roughly + ``memory usage = m0 + m1*jac_chunk_size``: the smaller the chunk size, + the less memory the Jacobian calculation will require (with some baseline + memory usage). The time it takes to compute the Jacobian is roughly + ``t= t0 + t1/jac_chunk_size` so the larger the ``jac_chunk_size``, the faster + the calculation takes, at the cost of requiring more memory. + If None, it will use the largest size i.e ``obj.dim_x``. + Defaults to ``chunk_size=None``. +""" +docs = { + "target": doc_target, + "bounds": doc_bounds, + "weight": doc_weight, + "normalize": doc_normalize, + "normalize_target": doc_normalize_target, + "loss_function": doc_loss_function, + "deriv_mode": doc_deriv_mode, + "name": doc_name, + "jac_chunk_size": doc_jac_chunk_size, +} + + +def collect_docs( + overwrite=None, + target_default="", + bounds_default="", + normalize_detail=None, + normalize_target_detail=None, + loss_detail=None, + coil=False, +): + """Collect default parameters for the docstring of Objective. + + Parameters + ---------- + overwrite : dict, optional + Dict of strings to overwrite from the _Objective's docstring. If None, + all default parameters are included as they are. Use this argument if + you want to specify a special docstring for a specific parameter in + your objective definition. + target_default : str, optional + Default value for the target parameter. + bounds_default : str, optional + Default value for the bounds parameter. + normalize_detail : str, optional + Additional information about the normalize parameter. + normalize_target_detail : str, optional + Additional information about the normalize_target parameter. + loss_detail : str, optional + Additional information about the loss function. + coil : bool, optional + Whether the objective is a coil objective. If True, adds extra docs to + target and loss_function. + + Returns + ------- + doc_params : str + String of default parameters for the docstring. + + """ + doc_params = "" + for key in docs.keys(): + if overwrite is not None and key in overwrite.keys(): + doc_params += overwrite[key].rstrip() + else: + if key == "target": + target = "" + if coil: + target += ( + "If array, it has to be flattened according to the " + + "number of inputs." + ) + if target_default != "": + target = target + " Defaults to " + target_default + doc_params += docs[key].rstrip() + target + elif key == "bounds" and bounds_default != "": + doc_params = ( + doc_params + docs[key].rstrip() + " Defaults to " + bounds_default + ) + elif key == "loss_function": + loss = "" + if coil: + loss = " Operates over all coils, not each individual coil." + if loss_detail is not None: + loss += loss_detail + doc_params += docs[key].rstrip() + loss + elif key == "normalize": + norm = "" + if normalize_detail is not None: + norm += normalize_detail + doc_params += docs[key].rstrip() + norm + elif key == "normalize_target": + norm_target = "" + if normalize_target_detail is not None: + norm_target = normalize_target_detail + doc_params += docs[key].rstrip() + norm_target + else: + doc_params += docs[key].rstrip() + + return doc_params + class ObjectiveFunction(IOAble): """Objective function comprised of one or more Objectives. @@ -111,14 +260,6 @@ def __init__( self._compiled = False self._name = name - def _set_derivatives(self): - """Choose derivative mode based on mode of sub-objectives.""" - if self._deriv_mode == "auto": - if all((obj._deriv_mode == "fwd") for obj in self.objectives): - self._deriv_mode = "batched" - else: - self._deriv_mode = "blocked" - def _unjit(self): """Remove jit compiled methods.""" methods = [ @@ -176,35 +317,34 @@ def build(self, use_jit=None, verbose=1): else: self._scalar = False - self._set_derivatives() + self._set_things() + + # setting derivative mode and chunking. + errorif( + isposint(self._jac_chunk_size) and self._deriv_mode in ["auto", "blocked"], + ValueError, + "'jac_chunk_size' was passed into ObjectiveFunction, but the " + "ObjectiveFunction is not using 'batched' deriv_mode", + ) sub_obj_jac_chunk_sizes_are_ints = [ isposint(obj._jac_chunk_size) for obj in self.objectives ] errorif( - any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode != "blocked", + any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode == "batched", ValueError, "'jac_chunk_size' was passed into one or more sub-objectives, but the" - " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " + " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " "'jac_chunk_size' will be ignored in favor of the ObjectiveFunction's " f"'jac_chunk_size' of {self._jac_chunk_size}." " Specify 'blocked' deriv_mode if each sub-objective is desired to have a " "different 'jac_chunk_size' for its Jacobian computation.", ) - errorif( - self._jac_chunk_size not in ["auto", None] - and self._deriv_mode == "blocked", - ValueError, - "'jac_chunk_size' was passed into ObjectiveFunction, but the " - "ObjectiveFunction is using 'blocked' deriv_mode, so sub-objective " - "'jac_chunk_size' are used to compute each sub-objective's Jacobian, " - "`ignoring the ObjectiveFunction's 'jac_chunk_size'.", - ) - - if not self.use_jit: - self._unjit() - self._set_things() - self._built = True + if self._deriv_mode == "auto": + if all((obj._deriv_mode == "fwd") for obj in self.objectives): + self._deriv_mode = "batched" + else: + self._deriv_mode = "blocked" if self._jac_chunk_size == "auto": # Heuristic estimates of fwd mode Jacobian memory usage, @@ -216,6 +356,15 @@ def build(self, use_jit=None, verbose=1): * self.dim_x ) self._jac_chunk_size = max([1, max_chunk_size]) + if self._deriv_mode == "blocked": + for obj in self.objectives: + if obj._jac_chunk_size is None: + obj._jac_chunk_size = self._jac_chunk_size + + if not self.use_jit: + self._unjit() + + self._built = True timer.stop("Objective build") if verbose > 1: @@ -500,15 +649,36 @@ def hess(self, x, constants=None): Derivative(self.compute_scalar, mode="hess")(x, constants).squeeze() ) - def _jac_blocked(self, op, x, constants=None): - # could also do something similar for grad and hess, but probably not - # worth it. grad is already super cheap to eval all at once, and blocked - # hess would only be block diag which may miss important interactions. + @jit + def jac_scaled(self, x, constants=None): + """Compute Jacobian matrix of self.compute_scaled wrt x.""" + v = jnp.eye(x.shape[0]) + return self.jvp_scaled(v, x, constants).T + + @jit + def jac_scaled_error(self, x, constants=None): + """Compute Jacobian matrix of self.compute_scaled_error wrt x.""" + v = jnp.eye(x.shape[0]) + return self.jvp_scaled_error(v, x, constants).T + + @jit + def jac_unscaled(self, x, constants=None): + """Compute Jacobian matrix of self.compute_unscaled wrt x.""" + v = jnp.eye(x.shape[0]) + return self.jvp_unscaled(v, x, constants).T + + def _jvp_blocked(self, v, x, constants=None, op="scaled"): + v = ensure_tuple(v) + if len(v) > 1: + # using blocked for higher order derivatives is a pain, and only really + # is needed for perturbations. Just pass that to jvp_batched for now + return self._jvp_batched(v, x, constants, op) if constants is None: constants = self.constants xs_splits = np.cumsum([t.dim_x for t in self.things]) xs = jnp.split(x, xs_splits) + vs = jnp.split(v[0], xs_splits, axis=-1) J = [] assert len(self.objectives) == len(self.constants) # basic idea is we compute the jacobian of each objective wrt each thing @@ -518,63 +688,18 @@ def _jac_blocked(self, op, x, constants=None): # get the xs that go to that objective thing_idx = self._things_per_objective_idx[k] xi = [xs[i] for i in thing_idx] - Ji_ = getattr(obj, op)(*xi, constants=const) # jac wrt to just those things - Ji = [] # jac wrt all things - for i, thing in enumerate(self.things): - if i in thing_idx: # dfi/dxj != 0 - Ji += [Ji_[thing_idx.index(i)]] - else: # dfi/dxj == 0 - Ji += [jnp.zeros((obj.dim_f, thing.dim_x))] - Ji = jnp.hstack(Ji) # something like [df1/dx1, df1/dx2, 0] - J += [Ji] - # something like [df1/dx1, df1/dx2, 0] - # [df2/dx1, 0, df2/dx3] # noqa:E800 - J = jnp.vstack(J) + vi = [vs[i] for i in thing_idx] + Ji_ = getattr(obj, "jvp_" + op)(vi, xi, constants=const) + J += [Ji_] + # this is the transpose of the jvp when v is a matrix, for consistency with + # jvp_batched + J = jnp.hstack(J) return J - @jit - def jac_scaled(self, x, constants=None): - """Compute Jacobian matrix of self.compute_scaled wrt x.""" - if constants is None: - constants = self.constants + def _jvp_batched(self, v, x, constants=None, op="scaled"): + v = ensure_tuple(v) - if self._deriv_mode == "batched": - J = Derivative(self.compute_scaled, mode="fwd")(x, constants) - if self._deriv_mode == "blocked": - J = self._jac_blocked("jac_scaled", x, constants) - - return jnp.atleast_2d(J.squeeze()) - - @jit - def jac_scaled_error(self, x, constants=None): - """Compute Jacobian matrix of self.compute_scaled_error wrt x.""" - if constants is None: - constants = self.constants - - if self._deriv_mode == "batched": - J = Derivative(self.compute_scaled_error, mode="fwd")(x, constants) - if self._deriv_mode == "blocked": - J = self._jac_blocked("jac_scaled_error", x, constants) - - return jnp.atleast_2d(J.squeeze()) - - @jit - def jac_unscaled(self, x, constants=None): - """Compute Jacobian matrix of self.compute_unscaled wrt x.""" - if constants is None: - constants = self.constants - - if self._deriv_mode == "batched": - J = Derivative(self.compute_unscaled, mode="fwd")(x, constants) - if self._deriv_mode == "blocked": - J = self._jac_blocked("jac_unscaled", x, constants) - - return jnp.atleast_2d(J.squeeze()) - - def _jvp(self, v, x, constants=None, op="compute_scaled"): - v = v if isinstance(v, (tuple, list)) else (v,) - - fun = lambda x: getattr(self, op)(x, constants) + fun = lambda x: getattr(self, "compute_" + op)(x, constants) if len(v) == 1: jvpfun = lambda dx: Derivative.compute_jvp(fun, 0, dx, x) return batched_vectorize( @@ -612,7 +737,11 @@ def jvp_scaled(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._jvp(v, x, constants, "compute_scaled") + if self._deriv_mode == "batched": + J = self._jvp_batched(v, x, constants, "scaled") + if self._deriv_mode == "blocked": + J = self._jvp_blocked(v, x, constants, "scaled") + return J @jit def jvp_scaled_error(self, v, x, constants=None): @@ -629,7 +758,11 @@ def jvp_scaled_error(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._jvp(v, x, constants, "compute_scaled_error") + if self._deriv_mode == "batched": + J = self._jvp_batched(v, x, constants, "scaled_error") + if self._deriv_mode == "blocked": + J = self._jvp_blocked(v, x, constants, "scaled_error") + return J @jit def jvp_unscaled(self, v, x, constants=None): @@ -646,10 +779,14 @@ def jvp_unscaled(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._jvp(v, x, constants, "compute_unscaled") + if self._deriv_mode == "batched": + J = self._jvp_batched(v, x, constants, "unscaled") + if self._deriv_mode == "blocked": + J = self._jvp_blocked(v, x, constants, "unscaled") + return J - def _vjp(self, v, x, constants=None, op="compute_scaled"): - fun = lambda x: getattr(self, op)(x, constants) + def _vjp(self, v, x, constants=None, op="scaled"): + fun = lambda x: getattr(self, "compute_" + op)(x, constants) return Derivative.compute_vjp(fun, 0, v, x) @jit @@ -666,7 +803,7 @@ def vjp_scaled(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._vjp(v, x, constants, "compute_scaled") + return self._vjp(v, x, constants, "scaled") @jit def vjp_scaled_error(self, v, x, constants=None): @@ -682,7 +819,7 @@ def vjp_scaled_error(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._vjp(v, x, constants, "compute_scaled_error") + return self._vjp(v, x, constants, "scaled_error") @jit def vjp_unscaled(self, v, x, constants=None): @@ -698,7 +835,7 @@ def vjp_unscaled(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._vjp(v, x, constants, "compute_unscaled") + return self._vjp(v, x, constants, "unscaled") def compile(self, mode="auto", verbose=1): """Call the necessary functions to ensure the function is compiled. @@ -809,7 +946,7 @@ def dim_x(self): @property def dim_f(self): """int: Number of objective equations.""" - if not self.built: + if not hasattr(self, "_dim_f"): raise RuntimeError("ObjectiveFunction must be built first.") return self._dim_f @@ -978,11 +1115,13 @@ def __init__( def _set_derivatives(self): """Choose derivative mode based on size of inputs/outputs.""" if self._deriv_mode == "auto": - # choose based on shape of jacobian. fwd mode is more memory efficient - # so we prefer that unless the jacobian is really wide + # choose based on shape of jacobian. dim_x is usually an overestimate of + # the true number of DOFs because linear constraints remove some. Also + # fwd mode is more memory efficient so we prefer that unless the jacobian + # is really wide self._deriv_mode = ( "fwd" - if self.dim_f >= 0.5 * sum(t.dim_x for t in self.things) + if self.dim_f >= 0.2 * sum(t.dim_x for t in self.things) else "rev" ) @@ -1167,37 +1306,57 @@ def hess(self, *args, **kwargs): def jac_scaled(self, *args, **kwargs): """Compute Jacobian matrix of self.compute_scaled wrt x.""" argnums = tuple(range(len(self.things))) - return Derivative(self.compute_scaled, argnums, mode=self._deriv_mode)( - *args, **kwargs - ) + return Derivative( + self.compute_scaled, + argnums, + mode=self._deriv_mode, + chunk_size=self._jac_chunk_size, + )(*args, **kwargs) @jit def jac_scaled_error(self, *args, **kwargs): """Compute Jacobian matrix of self.compute_scaled_error wrt x.""" argnums = tuple(range(len(self.things))) - return Derivative(self.compute_scaled_error, argnums, mode=self._deriv_mode)( - *args, **kwargs - ) + return Derivative( + self.compute_scaled_error, + argnums, + mode=self._deriv_mode, + chunk_size=self._jac_chunk_size, + )(*args, **kwargs) @jit def jac_unscaled(self, *args, **kwargs): """Compute Jacobian matrix of self.compute_unscaled wrt x.""" argnums = tuple(range(len(self.things))) - return Derivative(self.compute_unscaled, argnums, mode=self._deriv_mode)( - *args, **kwargs - ) - - def _jvp(self, v, x, constants=None, op="compute_scaled"): - v = v if isinstance(v, (tuple, list)) else (v,) - x = x if isinstance(x, (tuple, list)) else (x,) + return Derivative( + self.compute_unscaled, + argnums, + mode=self._deriv_mode, + chunk_size=self._jac_chunk_size, + )(*args, **kwargs) + + def _jvp(self, v, x, constants=None, op="scaled"): + v = ensure_tuple(v) + x = ensure_tuple(x) assert len(x) == len(v) - fun = lambda *x: getattr(self, op)(*x, constants=constants) - jvpfun = lambda *dx: Derivative.compute_jvp(fun, tuple(range(len(x))), dx, *x) - sig = ",".join(f"(n{i})" for i in range(len(x))) + "->(k)" - return batched_vectorize( - jvpfun, signature=sig, chunk_size=self._jac_chunk_size - )(*v) + if self._deriv_mode == "fwd": + fun = lambda *x: getattr(self, "compute_" + op)(*x, constants=constants) + jvpfun = lambda *dx: Derivative.compute_jvp( + fun, tuple(range(len(x))), dx, *x + ) + sig = ",".join(f"(n{i})" for i in range(len(x))) + "->(k)" + return batched_vectorize( + jvpfun, signature=sig, chunk_size=self._jac_chunk_size + )(*v) + else: # rev mode. We compute full jacobian and manually do mv. In this case + # the jacobian should be wide so this isn't very expensive. + jac = getattr(self, "jac_" + op)(*x, constants=constants) + # jac is a tuple, 1 array for each thing. Transposes here and below make it + # equivalent to fwd mode above, which batches over the first axis + Jv = tree_map(lambda a, b: jnp.dot(a, b.T), jac, v) + # sum over different things. + return jnp.sum(jnp.asarray(Jv), axis=0).T @jit def jvp_scaled(self, v, x, constants=None): @@ -1213,7 +1372,7 @@ def jvp_scaled(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._jvp(v, x, constants, "compute_scaled") + return self._jvp(v, x, constants, "scaled") @jit def jvp_scaled_error(self, v, x, constants=None): @@ -1229,7 +1388,7 @@ def jvp_scaled_error(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._jvp(v, x, constants, "compute_scaled_error") + return self._jvp(v, x, constants, "scaled_error") @jit def jvp_unscaled(self, v, x, constants=None): @@ -1245,7 +1404,7 @@ def jvp_unscaled(self, v, x, constants=None): Constant parameters passed to sub-objectives. """ - return self._jvp(v, x, constants, "compute_unscaled") + return self._jvp(v, x, constants, "unscaled") def print_value(self, args, args0=None, **kwargs): """Print the value of the objective.""" diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py index 604609fc8..bbc5daf4e 100644 --- a/desc/optimize/_constraint_wrappers.py +++ b/desc/optimize/_constraint_wrappers.py @@ -289,12 +289,6 @@ def hess(self, x_reduced, constants=None): def _jac(self, x_reduced, constants=None, op="scaled"): x = self.recover(x_reduced) - if self._objective._deriv_mode == "blocked": - fun = getattr(self._objective, "jac_" + op) - return fun(x, constants)[:, self._unfixed_idx] @ ( - self._Z * self._D[self._unfixed_idx, None] - ) - v = self._unfixed_idx_mat df = getattr(self._objective, "jvp_" + op)(v.T, x, constants) return df.T diff --git a/desc/plotting.py b/desc/plotting.py index 430c9ad53..12900ff41 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -864,7 +864,15 @@ def plot_3d( * ``cmap``: string denoting colormap to use. * ``levels``: array of data values where ticks on colorbar should be placed. * ``alpha``: float in [0,1.0], the transparency of the plotted surface - * ``showscale``: Bool, whether or not to show the colorbar. True by default + * ``showscale``: Bool, whether or not to show the colorbar. True by default. + * ``showgrid``: Bool, whether or not to show the coordinate grid lines. + True by default. + * ``showticklabels``: Bool, whether or not to show the coordinate tick labels. + True by default. + * ``showaxislabels``: Bool, whether or not to show the coordinate axis labels. + True by default. + * ``zeroline``: Bool, whether or not to show the zero coordinate axis lines. + True by default. * ``field``: MagneticField, a magnetic field with which to calculate Bn on the surface, must be provided if Bn is entered as the variable to plot. * ``field_grid``: MagneticField, a Grid to pass to the field as a source grid @@ -956,6 +964,10 @@ def plot_3d( data = data.reshape((grid.num_theta, grid.num_rho, grid.num_zeta), order="F") label = r"$\mathbf{B} \cdot \hat{n} ~(\mathrm{T})$" + showgrid = kwargs.pop("showgrid", True) + zeroline = kwargs.pop("zeroline", True) + showticklabels = kwargs.pop("showticklabels", True) + showaxislabels = kwargs.pop("showaxislabels", True) errorif( len(kwargs) != 0, ValueError, @@ -1024,30 +1036,48 @@ def plot_3d( if fig is None: fig = go.Figure() fig.add_trace(meshdata) + xaxis_title = ( + LatexNodes2Text().latex_to_text(_AXIS_LABELS_XYZ[0]) if showaxislabels else "" + ) + yaxis_title = ( + LatexNodes2Text().latex_to_text(_AXIS_LABELS_XYZ[1]) if showaxislabels else "" + ) + zaxis_title = ( + LatexNodes2Text().latex_to_text(_AXIS_LABELS_XYZ[2]) if showaxislabels else "" + ) fig.update_layout( scene=dict( - xaxis_title=LatexNodes2Text().latex_to_text(_AXIS_LABELS_XYZ[0]), - yaxis_title=LatexNodes2Text().latex_to_text(_AXIS_LABELS_XYZ[1]), - zaxis_title=LatexNodes2Text().latex_to_text(_AXIS_LABELS_XYZ[2]), + xaxis_title=xaxis_title, + yaxis_title=yaxis_title, + zaxis_title=zaxis_title, aspectmode="data", xaxis=dict( backgroundcolor="white", gridcolor="darkgrey", showbackground=False, zerolinecolor="darkgrey", + showgrid=showgrid, + zeroline=zeroline, + showticklabels=showticklabels, ), yaxis=dict( backgroundcolor="white", gridcolor="darkgrey", showbackground=False, zerolinecolor="darkgrey", + showgrid=showgrid, + zeroline=zeroline, + showticklabels=showticklabels, ), zaxis=dict( backgroundcolor="white", gridcolor="darkgrey", showbackground=False, zerolinecolor="darkgrey", + showgrid=showgrid, + zeroline=zeroline, + showticklabels=showticklabels, ), ), width=figsize[0] * dpi, @@ -2016,7 +2046,7 @@ def plot_boundary(eq, phi=None, plot_axis=True, ax=None, return_data=False, **kw ax.set_ylabel(_AXIS_LABELS_RPZ[2], fontsize=ylabel_fontsize) ax.tick_params(labelbottom=True, labelleft=True) - fig.legend(**legend_kw) + ax.legend(**legend_kw) _set_tight_layout(fig) plot_data = {} @@ -2120,9 +2150,7 @@ def plot_boundaries( phi = (1 if eqs[-1].N == 0 else 4) if phi is None else phi if isinstance(phi, numbers.Integral): - phi = np.linspace( - 0, 2 * np.pi / eqs[-1].NFP, phi + 1 - ) # +1 to include pi and 2pi + phi = np.linspace(0, 2 * np.pi / eqs[-1].NFP, phi, endpoint=False) phi = np.atleast_1d(phi) neq = len(eqs) @@ -2173,7 +2201,7 @@ def plot_boundaries( plot_data["R"].append(R) plot_data["Z"].append(Z) - for j in range(nz - 1): + for j in range(nz): (line,) = ax.plot( R[:, -1, j], Z[:, -1, j], color=colors[i], linestyle=ls[i], lw=lw[i] ) @@ -2190,7 +2218,7 @@ def plot_boundaries( ax.tick_params(labelbottom=True, labelleft=True) if any(labels) and kwargs.pop("legend", True): - fig.legend(**kwargs.pop("legend_kw", {})) + ax.legend(**kwargs.pop("legend_kw", {})) _set_tight_layout(fig) assert ( @@ -2416,6 +2444,14 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): * ``lw``: float, linewidth of plotted coils * ``ls``: str, linestyle of plotted coils * ``color``: str, color of plotted coils + * ``showgrid``: Bool, whether or not to show the coordinate grid lines. + True by default. + * ``showticklabels``: Bool, whether or not to show the coordinate tick labels. + True by default. + * ``showaxislabels``: Bool, whether or not to show the coordinate axis labels. + True by default. + * ``zeroline``: Bool, whether or not to show the zero coordinate axis lines. + True by default. Returns ------- @@ -2430,6 +2466,10 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): figsize = kwargs.pop("figsize", (10, 10)) color = kwargs.pop("color", "black") unique = kwargs.pop("unique", False) + showgrid = kwargs.pop("showgrid", True) + zeroline = kwargs.pop("zeroline", True) + showticklabels = kwargs.pop("showticklabels", True) + showaxislabels = kwargs.pop("showaxislabels", True) errorif( len(kwargs) != 0, ValueError, @@ -2497,28 +2537,40 @@ def flatten_coils(coilset): ) fig.add_trace(trace) + xaxis_title = "X (m)" if showaxislabels else "" + yaxis_title = "Y (m)" if showaxislabels else "" + zaxis_title = "Z (m)" if showaxislabels else "" fig.update_layout( scene=dict( - xaxis_title="X (m)", - yaxis_title="Y (m)", - zaxis_title="Z (m)", + xaxis_title=xaxis_title, + yaxis_title=yaxis_title, + zaxis_title=zaxis_title, xaxis=dict( backgroundcolor="white", gridcolor="darkgrey", showbackground=False, zerolinecolor="darkgrey", + showgrid=showgrid, + zeroline=zeroline, + showticklabels=showticklabels, ), yaxis=dict( backgroundcolor="white", gridcolor="darkgrey", showbackground=False, zerolinecolor="darkgrey", + showgrid=showgrid, + zeroline=zeroline, + showticklabels=showticklabels, ), zaxis=dict( backgroundcolor="white", gridcolor="darkgrey", showbackground=False, zerolinecolor="darkgrey", + showgrid=showgrid, + zeroline=zeroline, + showticklabels=showticklabels, ), aspectmode="data", ), @@ -2706,7 +2758,7 @@ def plot_boozer_modes( # noqa: C901 ax.set_ylabel(ylabel, fontsize=ylabel_fontsize) if kwargs.pop("legend", True): - fig.legend(**kwargs.pop("legend_kw", {"loc": "lower right"})) + ax.legend(**kwargs.pop("legend_kw", {"loc": "lower right"})) assert ( len(kwargs) == 0 @@ -3119,7 +3171,7 @@ def plot_qs_error( # noqa: 16 fxn too complex ax.set_ylabel(ylabel, fontsize=ylabel_fontsize) if kwargs.pop("legend", True): - fig.legend(**kwargs.pop("legend_kw", {"loc": "center right"})) + ax.legend(**kwargs.pop("legend_kw", {"loc": "center right"})) assert ( len(kwargs) == 0 diff --git a/desc/utils.py b/desc/utils.py index 203e29d3d..50c1db1b5 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -929,3 +929,12 @@ def tupleset(t, i, value): ) return res + + +def ensure_tuple(x): + """Returns x as a tuple of arrays.""" + if isinstance(x, tuple): + return x + if isinstance(x, list): + return tuple(x) + return (x,) diff --git a/docs/adding_objectives.rst b/docs/adding_objectives.rst index 31192c991..2fa3028ba 100644 --- a/docs/adding_objectives.rst +++ b/docs/adding_objectives.rst @@ -47,42 +47,16 @@ A full example objective with comments describing the key points is given below: ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. grid : Grid, optional Collocation grid containing the nodes to evaluate at. - name : str, optional - Name of the objective function. - jac_chunk_size : int or "auto", optional - Will calculate the Jacobian for this objective ``jac_chunk_size`` - columns at a time, instead of all at once. The memory usage of the - Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``: - the smaller the chunk size, the less memory the Jacobian calculation - will require (with some baseline memory usage). The time to compute the - Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the - ``jac_chunk_size``, the faster the calculation takes, at the cost of - requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least - memory intensive, but slowest method of calculating the Jacobian. - If None, it will use the largest possible size. """ + # Most of the documentation is shared among all objectives, so we just inherit + # the docstring from the base class and add a few details specific to this objective. + # See the documentation of `collect_docs` for more details. + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) _coordinates = "rtz" # What coordinates is this objective a function of, with r=rho, t=theta, z=zeta? # i.e. if only a profile, it is "r" , while if all 3 coordinates it is "rtz" @@ -229,12 +203,12 @@ you will have to manually convert these vectors using the geometry utility funct ``rpz2xyz`` and/or ``rpz2xyz_vec``. See the ``PlasmaVesselDistance`` objective for an example of this. -Adapting Existing Objectives with Different Loss Funtions ---------------------------------------------------------- +Adapting Existing Objectives with Different Loss Functions +---------------------------------------------------------- If your desired objective is already implemented in DESC, but not in the correct form, -a few different loss functions are available through the the ``loss_function`` kwarg -when instantiating an Objective objective to modify the objective cost in order to adapt +a few different loss functions are available through the ``loss_function`` kwarg +when instantiating an Objective, to modify the objective cost in order to adapt the objective to your desired purpose. For example, the DESC ``RotationalTransform`` objective with ``target=iota_target`` by default forms the residual by taking the target and subtracting it from the profile at the points in the grid, resulting in a residual diff --git a/docs/api.rst b/docs/api.rst index 02c7cc8c7..85c201396 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -163,6 +163,7 @@ Objective Functions desc.objectives.CoilCurrentLength desc.objectives.CoilCurvature desc.objectives.CoilLength + desc.objectives.CoilSetLinkingNumber desc.objectives.CoilSetMinDistance desc.objectives.CoilTorsion desc.objectives.CurrentDensity diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst index 273892aab..2a7c443c2 100644 --- a/docs/api_objectives.rst +++ b/docs/api_objectives.rst @@ -100,6 +100,7 @@ Coil Optimization desc.objectives.CoilLength desc.objectives.CoilCurvature desc.objectives.CoilTorsion + desc.objectives.CoilSetLinkingNumber desc.objectives.CoilSetMinDistance desc.objectives.PlasmaCoilSetMinDistance desc.objectives.CoilCurrentLength diff --git a/requirements.txt b/requirements.txt index b37f4f6c0..470158117 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,7 @@ colorama diffrax >= 0.4.1 h5py >= 3.0.0, < 4.0 interpax >= 0.3.3 -jax[cpu] >= 0.3.2, < 0.5.0 +jax[cpu] >= 0.3.2, < 0.4.34 matplotlib >= 3.5.0, < 4.0.0 mpmath >= 1.0.0, < 2.0 netcdf4 >= 1.5.4, < 2.0 diff --git a/tests/baseline/test_bounce1d_checks.png b/tests/baseline/test_bounce1d_checks.png index e43ba54cd..f3927bec6 100644 Binary files a/tests/baseline/test_bounce1d_checks.png and b/tests/baseline/test_bounce1d_checks.png differ diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png index f86a3b183..c81f89f03 100644 Binary files a/tests/baseline/test_plot_b_mag.png and b/tests/baseline/test_plot_b_mag.png differ diff --git a/tests/baseline/test_plot_boozer_modes.png b/tests/baseline/test_plot_boozer_modes.png index b72e30839..a2d1c6e37 100644 Binary files a/tests/baseline/test_plot_boozer_modes.png and b/tests/baseline/test_plot_boozer_modes.png differ diff --git a/tests/baseline/test_plot_boozer_modes_breaking_only.png b/tests/baseline/test_plot_boozer_modes_breaking_only.png index 424ba1f66..06f04654f 100644 Binary files a/tests/baseline/test_plot_boozer_modes_breaking_only.png and b/tests/baseline/test_plot_boozer_modes_breaking_only.png differ diff --git a/tests/baseline/test_plot_boozer_modes_max.png b/tests/baseline/test_plot_boozer_modes_max.png index ac59a48ed..b3d52bbe2 100644 Binary files a/tests/baseline/test_plot_boozer_modes_max.png and b/tests/baseline/test_plot_boozer_modes_max.png differ diff --git a/tests/baseline/test_plot_boozer_modes_no_norm.png b/tests/baseline/test_plot_boozer_modes_no_norm.png index b2fe1e29b..a5ccc0e32 100644 Binary files a/tests/baseline/test_plot_boozer_modes_no_norm.png and b/tests/baseline/test_plot_boozer_modes_no_norm.png differ diff --git a/tests/baseline/test_plot_boundaries.png b/tests/baseline/test_plot_boundaries.png index 44d4792c1..9e730b85a 100644 Binary files a/tests/baseline/test_plot_boundaries.png and b/tests/baseline/test_plot_boundaries.png differ diff --git a/tests/baseline/test_plot_boundary.png b/tests/baseline/test_plot_boundary.png index 8f2dd2769..ab4e97d73 100644 Binary files a/tests/baseline/test_plot_boundary.png and b/tests/baseline/test_plot_boundary.png differ diff --git a/tests/baseline/test_plot_boundary_surface.png b/tests/baseline/test_plot_boundary_surface.png index 5dfe95191..1e4f60d91 100644 Binary files a/tests/baseline/test_plot_boundary_surface.png and b/tests/baseline/test_plot_boundary_surface.png differ diff --git a/tests/baseline/test_plot_qs_error.png b/tests/baseline/test_plot_qs_error.png index c90e59fb8..10485feb2 100644 Binary files a/tests/baseline/test_plot_qs_error.png and b/tests/baseline/test_plot_qs_error.png differ diff --git a/tests/test_coils.py b/tests/test_coils.py index 71127660d..368a2fd12 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -1282,3 +1282,26 @@ def test_repr(): coils.name = "MyCoils" assert "MyCoils" in str(coils) + + +@pytest.mark.unit +def test_linking_number(): + """Test calculation of linking number.""" + coil = FourierPlanarCoil(center=[10, 1, 0]) + grid = LinearGrid(N=25) + # regular modular coilset from symmetry, so that there are 10 coils, half going + # one way and half going the other way + coilset = CoilSet(coil, NFP=5, sym=True) + coil2 = FourierRZCoil() + # add a coil along the axis that links all the other coils + coilset2 = MixedCoilSet(coilset, coil2) + link = coilset2._compute_linking_number(grid=grid) + + # modular coils don't link each other + np.testing.assert_allclose(link[:-1, :-1], 0, atol=1e-14) + # axis coil doesn't link itself + np.testing.assert_allclose(link[-1, -1], 0, atol=1e-14) + # we expect the axis coil to link all the modular coils, with alternating sign + # due to alternating orientation of the coils due to symmetry. + expected = [1, -1] * 5 + np.testing.assert_allclose(link[-1, :-1], expected, rtol=1e-3) diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index 29c2f909d..e5f50eee5 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -77,13 +77,17 @@ def test_fun(x, y, a): y = np.array([60, 1, 100, 0.02]) a = -2 - jac_AD = AutoDiffDerivative(test_fun, argnum=0) - J_AD = jac_AD.compute(x, y, a) + jacf_AD = AutoDiffDerivative(test_fun, argnum=0, mode="fwd") + Jf_AD = jacf_AD.compute(x, y, a) + jacr_AD = AutoDiffDerivative(test_fun, argnum=0, mode="rev") + Jr_AD = jacr_AD.compute(x, y, a) - jac_FD = AutoDiffDerivative(test_fun, argnum=0) + jac_FD = FiniteDiffDerivative(test_fun, argnum=0) J_FD = jac_FD.compute(x, y, a) - np.testing.assert_allclose(J_FD, J_AD, atol=1e-8) + np.testing.assert_allclose(Jf_AD, Jr_AD, atol=1e-8) + np.testing.assert_allclose(J_FD, Jf_AD, rtol=1e-2) + np.testing.assert_allclose(J_FD, Jr_AD, rtol=1e-2) @pytest.mark.unit def test_fd_hessian(self): diff --git a/tests/test_examples.py b/tests/test_examples.py index 53fb6479b..93c700b36 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1043,6 +1043,7 @@ def test_freeb_vacuum(): objective = ObjectiveFunction( VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True), jac_chunk_size=1000, + deriv_mode="batched", ) eq, _ = eq.optimize( objective, diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index ada80267d..982e98d5d 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -36,6 +36,7 @@ FixNearAxisLambda, FixNearAxisR, FixNearAxisZ, + FixOmniBmax, FixOmniMap, FixOmniWell, FixParameters, @@ -915,6 +916,26 @@ def test_fix_omni_indices(): assert constraint.dim_f == indices.size +@pytest.mark.unit +def test_fix_omni_Bmax(): + """Test that omnigenity parameters are constrained for B_max to be a straight line. + + Test for GH issue #1266. + """ + + def _test(M_x, N_x, NFP, sum): + field = OmnigenousField(L_x=2, M_x=M_x, N_x=N_x, NFP=NFP, helicity=(1, NFP)) + constraint = FixOmniBmax(field=field) + constraint.build() + assert constraint.dim_f == (2 * field.N_x + 1) * (field.L_x + 1) + # 0 - 2 + 4 - 6 + 8 ... + np.testing.assert_allclose(constraint._A @ field.x_basis.modes[:, 1], sum) + + _test(M_x=6, N_x=3, NFP=1, sum=-4) + _test(M_x=9, N_x=4, NFP=2, sum=4) + _test(M_x=12, N_x=5, NFP=3, sum=6) + + @pytest.mark.unit def test_fix_parameters_input_order(DummyStellarator): """Test that FixParameters preserves the input indices and target ordering.""" diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 68de417ea..fccfaae38 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -44,6 +44,7 @@ CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, Elongation, @@ -78,7 +79,7 @@ ) from desc.objectives._free_boundary import BoundaryErrorNESTOR from desc.objectives.normalization import compute_scaling_factors -from desc.objectives.objective_funs import _Objective +from desc.objectives.objective_funs import _Objective, collect_docs from desc.objectives.utils import softmax, softmin from desc.profiles import FourierZernikeProfile, PowerSeriesProfile from desc.utils import PRINT_WIDTH @@ -1194,6 +1195,25 @@ def test(eq, field, correct_value, rtol=1e-14, grid=None): # test on field with no vector potential test(eq, PoloidalMagneticField(1, 1, 1), 0.0) + @pytest.mark.unit + def test_coil_linking_number(self): + """Test for linking number objective.""" + coil = FourierPlanarCoil(center=[10, 1, 0]) + # regular modular coilset from symmetry, so that there are 10 coils, half going + # one way and half going the other way + coilset = CoilSet.from_symmetry(coil, NFP=5, sym=True) + coil2 = FourierRZCoil() + # add a coil along the axis that links all the other coils + coilset2 = MixedCoilSet(coilset, coil2) + + obj = CoilSetLinkingNumber(coilset2) + obj.build() + out = obj.compute_scaled_error(coilset2.params_dict) + # the modular coils all link 1 other coil (the axis) + # while the axis links all 10 modular coils + expected = np.array([1] * 10 + [10]) + np.testing.assert_allclose(out, expected, rtol=1e-3) + @pytest.mark.unit def test_signed_plasma_vessel_distance(self): """Test calculation of signed distance from plasma to vessel.""" @@ -1323,6 +1343,7 @@ def test_derivative_modes(): AspectRatio(eq), ], deriv_mode="batched", + jac_chunk_size="auto", use_jit=False, ) obj2 = ObjectiveFunction( @@ -1332,17 +1353,20 @@ def test_derivative_modes(): AspectRatio(eq, jac_chunk_size=None), ], deriv_mode="blocked", + jac_chunk_size="auto", use_jit=False, ) obj1.build() obj2.build() # check that default size works for blocked - assert obj2.objectives[1]._jac_chunk_size is None - assert obj2.objectives[2]._jac_chunk_size is None + assert obj2.objectives[0]._jac_chunk_size == 2 + assert obj2.objectives[1]._jac_chunk_size > 0 + assert obj2.objectives[2]._jac_chunk_size > 0 # hard to say what size auto will give, just check it is >0 assert obj1._jac_chunk_size > 0 obj3.build() x = obj1.x(eq, surf) + v = jnp.ones_like(x) g1 = obj1.grad(x) g2 = obj2.grad(x) g3 = obj3.grad(x) @@ -1363,6 +1387,48 @@ def test_derivative_modes(): H3 = obj3.hess(x) np.testing.assert_allclose(H1, H2, atol=1e-10) np.testing.assert_allclose(H1, H3, atol=1e-10) + j1 = obj1.jvp_scaled(v, x) + j2 = obj2.jvp_scaled(v, x) + j3 = obj3.jvp_scaled(v, x) + np.testing.assert_allclose(j1, j2, atol=1e-10) + np.testing.assert_allclose(j1, j3, atol=1e-10) + + +@pytest.mark.unit +def test_fwd_rev(): + """Test that forward and reverse mode jvps etc give same results.""" + eq = Equilibrium() + obj1 = MeanCurvature(eq, deriv_mode="fwd") + obj2 = MeanCurvature(eq, deriv_mode="rev") + obj1.build() + obj2.build() + + x = eq.pack_params(eq.params_dict) + J1 = obj1.jac_scaled(x) + J2 = obj2.jac_scaled(x) + np.testing.assert_allclose(J1, J2, atol=1e-14) + + jvp1 = obj1.jvp_scaled(x, jnp.ones_like(x)) + jvp2 = obj2.jvp_scaled(x, jnp.ones_like(x)) + np.testing.assert_allclose(jvp1, jvp2, atol=1e-14) + + surf = FourierRZToroidalSurface() + obj1 = PlasmaVesselDistance(eq, surf, deriv_mode="fwd") + obj2 = PlasmaVesselDistance(eq, surf, deriv_mode="rev") + obj1.build() + obj2.build() + + x1 = eq.pack_params(eq.params_dict) + x2 = surf.pack_params(surf.params_dict) + + J1a, J1b = obj1.jac_scaled(x1, x2) + J2a, J2b = obj2.jac_scaled(x1, x2) + np.testing.assert_allclose(J1a, J2a, atol=1e-14) + np.testing.assert_allclose(J1b, J2b, atol=1e-14) + + jvp1 = obj1.jvp_scaled((x1, x2), (jnp.ones_like(x1), jnp.ones_like(x2))) + jvp2 = obj2.jvp_scaled((x1, x2), (jnp.ones_like(x1), jnp.ones_like(x2))) + np.testing.assert_allclose(jvp1, jvp2, atol=1e-14) @pytest.mark.unit @@ -2199,6 +2265,7 @@ class TestComputeScalarResolution: CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, FusionPower, @@ -2561,7 +2628,14 @@ def test_compute_scalar_resolution_others(self, objective): @pytest.mark.regression @pytest.mark.parametrize( "objective", - [CoilLength, CoilTorsion, CoilCurvature, CoilCurrentLength, CoilSetMinDistance], + [ + CoilLength, + CoilTorsion, + CoilCurvature, + CoilCurrentLength, + CoilSetMinDistance, + CoilSetLinkingNumber, + ], ) def test_compute_scalar_resolution_coils(self, objective): """Coil objectives.""" @@ -2597,6 +2671,7 @@ class TestObjectiveNaNGrad: CoilLength, CoilCurrentLength, CoilCurvature, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, ForceBalanceAnisotropic, @@ -2788,7 +2863,14 @@ def test_objective_no_nangrad(self, objective): @pytest.mark.unit @pytest.mark.parametrize( "objective", - [CoilLength, CoilTorsion, CoilCurvature, CoilCurrentLength, CoilSetMinDistance], + [ + CoilLength, + CoilTorsion, + CoilCurvature, + CoilCurrentLength, + CoilSetMinDistance, + CoilSetLinkingNumber, + ], ) def test_objective_no_nangrad_coils(self, objective): """Coil objectives.""" @@ -2887,3 +2969,19 @@ def test_objective_print_widths(): + "change the name or increase the PRINT_WIDTH in the " + "desc/utils.py file. The former is preferred." ) + + +def test_objective_docstring(): + """Test that the objective docstring and collect_docs are consistent.""" + objective_docs = _Objective.__doc__.rstrip() + doc_header = ( + "Objective (or constraint) used in the optimization of an Equilibrium.\n\n" + + " Parameters\n" + + " ----------\n" + + " things : Optimizable or tuple/list of Optimizable\n" + + " Objects that will be optimized to satisfy the Objective.\n" + ) + collected_docs = collect_docs().strip() + collected_docs = doc_header + " " + collected_docs + + assert objective_docs == collected_docs diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 6048a7d60..8488c91a8 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -47,9 +47,9 @@ from desc.utils import isalmostequal from desc.vmec import VMECIO -tol_1d = 7.8 -tol_2d = 15 -tol_3d = 15 +tol_1d = 4.5 +tol_2d = 10 +tol_3d = 10 @pytest.mark.unit @@ -284,7 +284,15 @@ def test_3d_rt(self): def test_plot_3d_surface(self): """Test 3d plotting of surface object.""" surf = FourierRZToroidalSurface() - fig = plot_3d(surf, "curvature_H_rho") + fig = plot_3d( + surf, + "curvature_H_rho", + showgrid=False, + showscale=False, + zeroline=False, + showticklabels=False, + showaxislabels=False, + ) return fig @pytest.mark.unit @@ -507,7 +515,7 @@ def test_plot_boundary_surface(self): return fig @pytest.mark.unit - @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d) + @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) def test_plot_boundaries(self): """Test plotting boundaries.""" eq1 = get("SOLOVEV") @@ -516,6 +524,11 @@ def test_plot_boundaries(self): eq4 = get("ESTELL") with pytest.raises(ValueError, match="differing field periods"): fig, ax = plot_boundaries([eq3, eq4], theta=0) + _, _, data1 = plot_boundaries( + (eq1, eq2, eq3), + phi=4, + return_data=True, + ) fig, ax, data = plot_boundaries( (eq1, eq2, eq3), phi=np.linspace(0, 2 * np.pi / eq3.NFP, 4, endpoint=False), @@ -525,6 +538,12 @@ def test_plot_boundaries(self): assert "Z" in data.keys() assert len(data["R"]) == 3 assert len(data["Z"]) == 3 + assert ( + data["R"][-1].shape == data1["R"][-1].shape + ), "Passing phi as an integer or array results in different behavior" + assert ( + data["Z"][-1].shape == data1["Z"][-1].shape + ), "Passing phi as an integer or array results in different behavior" return fig @@ -840,6 +859,40 @@ def flatten_coils(coilset): return fig +@pytest.mark.unit +def test_plot_coils_no_grid(): + """Test 3d plotting of coils with currents without any gridlines.""" + N = 48 + NFP = 4 + I = 1 + coil = FourierXYZCoil() + coil.rotate(angle=np.pi / N) + coils = CoilSet.linspaced_angular(coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2) + with pytest.raises(ValueError, match="Expected `coils`"): + plot_coils("not a coil") + fig, data = plot_coils( + coils, + unique=True, + return_data=True, + showgrid=False, + zeroline=False, + showticklabels=False, + showaxislabels=False, + ) + + def flatten_coils(coilset): + if hasattr(coilset, "__len__"): + return [a for i in coilset for a in flatten_coils(i)] + else: + return [coilset] + + coil_list = flatten_coils(coils) + for string in ["X", "Y", "Z"]: + assert string in data.keys() + assert len(data[string]) == len(coil_list) + return fig + + @pytest.mark.unit @pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) def test_plot_b_mag():