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/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 5752f730f..9429b328e 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -11,7 +11,7 @@ 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.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,31 +38,6 @@ 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. 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)`` @@ -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: " @@ -381,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 @@ -423,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: @@ -862,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 1dff948b2..5b6430d67 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -10,7 +10,7 @@ from ..compute.geom_utils import errorif_sym 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 @@ -22,52 +22,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: " @@ -198,52 +167,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: " @@ -375,52 +312,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: " @@ -575,31 +479,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)``. @@ -628,22 +507,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: " @@ -951,51 +822,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: " @@ -1124,51 +962,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: " @@ -1293,50 +1098,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: " @@ -1458,49 +1230,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 b720608b4..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: " @@ -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/objective_funs.py b/desc/objectives/objective_funs.py index 9cdfb61c2..28d7530d1 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -33,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. @@ -113,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 = [ @@ -178,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, @@ -218,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: @@ -799,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 diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py index ffaa507d3..161c3f057 100644 --- a/desc/objectives/utils.py +++ b/desc/objectives/utils.py @@ -92,6 +92,36 @@ def factorize_linear_constraints(objective, constraint, x_scale="auto"): # noqa A = A[:, cols] assert A.shape[1] == xp.size + # check for degenerate rows and delete if necessary + # augment A with b so that it only deletes actual degenerate constraints + # which are duplicate rows of A that also have duplicate entries of b, + # if the entries of b aren't the same then the constraints are actually + # incompatible and so we will leave those to be caught later. + A_augmented = np.hstack([A, np.reshape(b, (A.shape[0], 1))]) + row_idx_to_delete = np.array([], dtype=int) + for row_idx in range(A_augmented.shape[0]): + # find all rows equal to this row + rows_equal_to_this_row = np.where( + np.all(A_augmented[row_idx, :] == A_augmented, axis=1) + )[0] + # find the rows equal to this row that are not this row + rows_equal_to_this_row_but_not_this_row = rows_equal_to_this_row[ + rows_equal_to_this_row != row_idx + ] + # if there are rows equal to this row that aren't this row, AND this particular + # row has not already been detected as a duplicate of an earlier one and slated + # for deletion, add the duplicate row indices to the array of + # rows to be deleted + if ( + rows_equal_to_this_row_but_not_this_row.size + and row_idx not in row_idx_to_delete + ): + row_idx_to_delete = np.append(row_idx_to_delete, rows_equal_to_this_row[1:]) + # delete the affected rows, and also the corresponding rows of b + A_augmented = np.delete(A_augmented, row_idx_to_delete, axis=0) + A = A_augmented[:, :-1] + b = np.atleast_1d(A_augmented[:, -1].squeeze()) + # will store the global index of the unfixed rows, idx indices_row = np.arange(A.shape[0]) indices_idx = np.arange(A.shape[1]) @@ -161,7 +191,6 @@ def factorize_linear_constraints(objective, constraint, x_scale="auto"): # noqa Z = np.eye(A.shape[1]) xp = put(xp, unfixed_idx, A_inv @ b) xp = put(xp, fixed_idx, ((1 / D) * xp)[fixed_idx]) - # cast to jnp arrays xp = jnp.asarray(xp) A = jnp.asarray(A) 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/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_examples.py b/tests/test_examples.py index 53fb6479b..95d46a994 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -178,6 +178,8 @@ def test_1d_optimization(): constraints = ( ForceBalance(eq=eq), FixBoundaryR(eq=eq), + FixBoundaryR(eq=eq, modes=[0, 0, 0]), # add a degenerate constraint to confirm + # proximal-lsq-exact not affected by GH #1297 FixBoundaryZ(eq=eq, modes=eq.surface.Z_basis.modes[0:-1, :]), FixPressure(eq=eq), FixIota(eq=eq), @@ -644,6 +646,9 @@ def test_multiobject_optimization_al(): FixParameters(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}), FixParameters(eq, {"Psi": True, "i_l": True}), FixBoundaryR(eq, modes=[[0, 0, 0]]), + FixBoundaryR( + eq=eq, modes=[0, 0, 0] + ), # add a degenerate constraint to test fix of GH #1297 for lsq-auglag PlasmaVesselDistance(surface=surf, eq=eq, target=1), ) @@ -1043,6 +1048,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 982e98d5d..d0c17f8e3 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -205,6 +205,9 @@ def test_fixed_mode_solve(): FixIota(eq=eq), FixPsi(eq=eq), FixBoundaryR(eq=eq), + FixBoundaryR( + eq=eq, modes=[0, 0, 0] + ), # add a degenerate constraint to test fix of GH #1297 for lsq-exact FixBoundaryZ(eq=eq), fixR, fixZ, diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index cfbf09090..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,13 +1353,15 @@ 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() @@ -2242,6 +2265,7 @@ class TestComputeScalarResolution: CoilCurrentLength, CoilCurvature, CoilLength, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, FusionPower, @@ -2604,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.""" @@ -2640,6 +2671,7 @@ class TestObjectiveNaNGrad: CoilLength, CoilCurrentLength, CoilCurvature, + CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, ForceBalanceAnisotropic, @@ -2831,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.""" @@ -2930,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