Merge branch 'master' into ripple

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()

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

desc/objectives/__init__.py

@@ -5,6 +5,7 @@
     CoilCurrentLength,
     CoilCurvature,
     CoilLength,
+    CoilSetLinkingNumber,
     CoilSetMinDistance,
     CoilTorsion,
     PlasmaCoilSetMinDistance,

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 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, 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: "