Merge branch 'master' into yge/cpu

.github/workflows/nbtests.yml

@@ -39,7 +39,7 @@ jobs:
           lscpu
           export PYTHONPATH=$(pwd)
           pytest -v --nbmake "./docs/notebooks" \
-            --nbmake-timeout=1000 \
+            --nbmake-timeout=2000 \
             --ignore=./docs/notebooks/zernike_eval.ipynb \
             --splits 2 \
             --group ${{ matrix.group }} \

CHANGELOG.md

@@ -1,14 +1,69 @@
 Changelog
 =========
 
+
+v0.12.0
+-------
+
+[Github Commits](https://github.com/PlasmaControl/DESC/compare/v0.11.1...v0.12.0)
+
 New Features
 
-- Add method ``from_values`` to ``FourierRZCurve`` to allow fitting of data points
-to a ``FourierRZCurve`` object, and ``to_FourierRZCurve`` methods to ``Curve`` class.
-- Adds the objective `CoilsetMinDistance`, which returns the minimum distance to another
-coil for each coil in a coilset.
-- Adds the objective `PlasmaCoilsetMinDistance`, which returns the minimum distance to the
-plasma surface for each coil in a coilset.
+- Coil optimization is now possible in DESC using various filamentary coils. This includes
+a number of new objectives:
+    - ``desc.objectives.QuadraticFlux``
+    - ``desc.objectives.ToroidalFlux``
+    - ``desc.objectives.CoilLength``
+    - ``desc.objectives.CoilCurvature``
+    - ``desc.objectives.CoilTorsion``
+    - ``desc.objectives.CoilCurrentLength``
+    - ``desc.objectives.CoilSetMinDistance``
+    - ``desc.objectives.PlasmaCoilSetMinDistance``
+    - ``desc.objectives.FixCoilCurrent``
+    - ``desc.objectives.FixSumCoilCurrent``
+- Add Normal Field Error ``"B*n"`` as a plot quantity to ``desc.plotting.{plot_2d, plot_3d}``.
+- New function ``desc.plotting.poincare_plot`` for creating Poincare plots by tracing
+field lines from coils or other external fields.
+- New profile type ``desc.profiles.TwoPowerProfile``.
+- Add ``desc.geometry.FourierRZCurve.from_values`` method to fit curve with data.
+- Add ``desc.geometry.FourierRZToroidalSurface.from_shape_parameters`` method for generating a surface
+with specified elongation, triangularity, squareness, etc.
+- New class ``desc.magnetic_fields.MagneticFieldFromUser`` for user defined B(R,phi,Z).
+- All vector variables are now computed in toroidal (R,phi,Z) coordinates by default.
+Cartesian (X,Y,Z) coordinates can be requested with the compute keyword ``basis='xyz'``.
+- Add method ``desc.coils.CoilSet.is_self_intersecting``, which checks if any coils
+intersect each other in the coilset.
+
+Minor changes
+
+- Improved heuristic initial guess for ``Equilibrium.map_coordinates``.
+- Add documentation for default grid and target/bounds for objectives.
+- Add documentation for compute function keyword arguments.
+- Loading a coilset from a MAKEGRID file will now return a nested ``MixedCoilSet`` if there
+are coil groups present in the MAKEGRID file.
+- Users must now pass in spacing/weights to custom ``Grid``s (the previous defaults were
+often wrong, leading to incorrect results)
+- The ``normal`` and ``center`` parameters of a ``FourierPlanarCurve`` can now be specified
+in either cartesian or cylindrical coordinates, as determined by the ``basis`` parameter.
+- Misc small changes to reduce compile time and memory consumption (more coming soon!)
+- Linear constraint factorization has been refactored to improve efficiency and reduce
+floating point error.
+- ``desc.objectives.{GenericObjective, ObjectiveFromUser}`` can now work with other objects
+besides an ``Equilibrium`` (such as surfaces, curves, etc.)
+- Improve warning for missing attributes when loading desc objects.
+
+Bug Fixes
+
+- Several small fixes to ensure things that should be ``int``s are ``int``s
+- Fix incorrect toroidal components of surface basis vectors.
+- Fix a regression in performance in evaluating Zernike polynomials.
+- Fix errors in ``Equilibrium.map_coordinates`` for prescribed current equilibria.
+- Fix definition of ``b0`` in VMEC output.
+- Fix a bug where calling ``Equilibrium.compute(..., data=data)`` would lead to excessive
+recalculation and potentially wrong results.
+- Fixes a bug causing NaN in reverse mode AD for ``Omnigenity`` objective.
+- Fix a bug where ``"A(z)"`` would be zero if the grid doesn't contain nodes at rho=1.
+
 
 v0.11.1
 -------

README.rst

@@ -46,6 +46,7 @@ The best place to start learning about DESC is our tutorials:
 - `Basic optimization`_: specifying objectives, fixing degrees of freedom.
 - `Advanced optimization`_: advanced constraints, precise quasi-symmetry, constrained optimization.
 - `Near axis constraints`_: loading solutions from QSC/QIC and fixing near axis expansion.
+- `Coil optimization`_: "second stage" optimization of magnetic coils.
 
 For details on the various objectives, constraints, optimizable objects and more, see
 the full `api documentation`_.
@@ -73,6 +74,7 @@ The equilibrium solution is output in a HDF5 binary file, whose format is detail
 .. _Basic optimization: https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/basic_optimization.html
 .. _Advanced optimization: https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/advanced_optimization.html
 .. _Near axis constraints: https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/nae_constraint.html
+.. _Coil optimization: https://desc-docs.readthedocs.io/en/latest/notebooks/tutorials/coil_stage_two_optimization.html
 .. _api documentation: https://desc-docs.readthedocs.io/en/latest/api.html
 
 Repository Contents

desc/coils.py

@@ -19,6 +19,7 @@
 from desc.compute import get_params, rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec
 from desc.compute.geom_utils import reflection_matrix
 from desc.compute.utils import _compute as compute_fun
+from desc.compute.utils import safenorm
 from desc.geometry import (
     FourierPlanarCurve,
     FourierRZCurve,
@@ -157,6 +158,11 @@ def current(self, new):
         assert jnp.isscalar(new) or new.size == 1
         self._current = float(np.squeeze(new))
 
+    @property
+    def num_coils(self):
+        """int: Number of coils."""
+        return 1
+
     def _compute_position(self, params=None, grid=None, **kwargs):
         """Compute coil positions accounting for stellarator symmetry.
 
@@ -294,6 +300,8 @@ def to_FourierXYZ(self, N=10, grid=None, s=None, name=""):
         """
         if (grid is None) and (s is not None) and (not isinstance(s, str)):
             grid = LinearGrid(zeta=s)
+        if grid is None:
+            grid = LinearGrid(N=2 * N + 1)
         coords = self.compute("x", grid=grid, basis="xyz")["x"]
         return FourierXYZCoil.from_values(
             self.current, coords, N=N, s=s, basis="xyz", name=name
@@ -797,13 +805,22 @@ class CoilSet(OptimizableCollection, _Coil, MutableSequence):
         assuming 'virtual' coils from the other half field period. Default = False.
     name : str
         Name of this CoilSet.
+    check_intersection: bool
+        Whether or not to check the coils in the coilset for intersections.
 
     """
 
     _io_attrs_ = _Coil._io_attrs_ + ["_coils", "_NFP", "_sym"]
     _io_attrs_.remove("_current")
 
-    def __init__(self, *coils, NFP=1, sym=False, name=""):
+    def __init__(
+        self,
+        *coils,
+        NFP=1,
+        sym=False,
+        name="",
+        check_intersection=True,
+    ):
         coils = flatten_list(coils, flatten_tuple=True)
         assert all([isinstance(coil, (_Coil)) for coil in coils])
         [_check_type(coil, coils[0]) for coil in coils]
@@ -812,6 +829,9 @@ def __init__(self, *coils, NFP=1, sym=False, name=""):
         self._sym = bool(sym)
         self._name = str(name)
 
+        if check_intersection:
+            self.is_self_intersecting()
+
     @property
     def name(self):
         """str: Name of the curve."""
@@ -899,7 +919,10 @@ def compute(
 
         """
         if params is None:
-            params = [get_params(names, coil) for coil in self]
+            params = [
+                get_params(names, coil, basis=kwargs.get("basis", "rpz"))
+                for coil in self
+            ]
         if data is None:
             data = [{}] * len(self)
 
@@ -940,9 +963,9 @@ def _compute_position(self, params=None, grid=None, **kwargs):
             Coil positions, in [R,phi,Z] or [X,Y,Z] coordinates.
 
         """
-        if params is None:
-            params = [get_params("x", coil) for coil in self]
         basis = kwargs.pop("basis", "xyz")
+        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)
         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)
@@ -1009,7 +1032,9 @@ def compute_magnetic_field(
         assert basis.lower() in ["rpz", "xyz"]
         coords = jnp.atleast_2d(jnp.asarray(coords))
         if params is None:
-            params = [get_params(["x_s", "x", "s", "ds"], coil) for coil in self]
+            params = [
+                get_params(["x_s", "x", "s", "ds"], coil, basis=basis) for coil in self
+            ]
             for par, coil in zip(params, self):
                 par["current"] = coil.current
 
@@ -1140,7 +1165,7 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
 
         Parameters
         ----------
-        coils : Coil, CoilGroup, Coilset
+        coils : Coil, CoilSet
             Coil or collection of coils in one field period or half field period.
         NFP : int (optional)
             Number of field periods for enforcing field period symmetry.
@@ -1169,21 +1194,6 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
             # only need to check this for a CoilSet, not MixedCoilSet
             [_check_type(coil, coils[0]) for coil in coils]
 
-        # check toroidal extent of coils to be repeated
-        maxphi = 2 * np.pi / NFP / (sym + 1)
-        data = coils.compute("phi")
-        for i, cdata in enumerate(data):
-            errorif(
-                np.any(cdata["phi"] > maxphi),
-                ValueError,
-                f"coil {i} exceeds the toroidal extent for NFP={NFP} and sym={sym}",
-            )
-            warnif(
-                sym and np.any(cdata["phi"] < np.finfo(cdata["phi"].dtype).eps),
-                UserWarning,
-                f"coil {i} is on the symmetry plane phi=0",
-            )
-
         coilset = []
         if sym:
             # first reflect/flip original coilset
@@ -1498,6 +1508,100 @@ def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""):
         coils = [coil.to_SplineXYZ(knots, grid, method) for coil in self]
         return self.__class__(*coils, NFP=self.NFP, sym=self.sym, name=name)
 
+    def is_self_intersecting(self, grid=None, tol=None):
+        """Check if any coils in the CoilSet intersect.
+
+        By default, checks intersection by checking that for each point on a given coil
+        the closest point in the coilset is on that same coil. If the closest point is
+        on another coil, that indicates that the coils may be close to intersecting.
+
+        If instead the ``tol`` argument is provided, then the function will
+        check the minimum distance from each coil to each other coil against
+        that tol and if it finds the minimum distance is less than the ``tol``,
+        it will take it as intersecting coils, returning True and raising a warning.
+
+        NOTE: If grid resolution used is too low, this function may fail to return
+        the correct answer.
+
+        Parameters
+        ----------
+        grid : Grid, optional
+            Collocation grid containing the nodes to evaluate the coil positions at.
+            If a list, must have the same structure as the coilset. Defaults to a
+            LinearGrid(N=100)
+        tol : float, optional
+            the tolerance (in meters) to check the intersections to, if points on any
+            two coils are closer than this tolerance, then the function will return
+            True and a warning will be raised. If not passed, then the method used
+            to determine coilset intersection will be based off of checking that
+            each point on a coil is closest to a point on the same coil, which does
+            not rely on a ``tol`` parameter.
+
+        Returns
+        -------
+        is_self_intersecting : bool
+            Whether or not any coils in the CoilSet come close enough to each other to
+            possibly be intersecting.
+
+        """
+        from desc.objectives._coils import CoilSetMinDistance
+
+        grid = grid if grid else LinearGrid(N=100)
+        obj = CoilSetMinDistance(self, grid=grid)
+        obj.build(verbose=0)
+        if tol:
+            min_dists = obj.compute(self.params_dict)
+            is_nearly_intersecting = np.any(min_dists < tol)
+            warnif(
+                is_nearly_intersecting,
+                UserWarning,
+                "Found coils which are nearly intersecting according to the given tol "
+                + "(min coil-coil distance = "
+                + f"{np.min(min_dists):1.3e} m < {tol:1.3e} m)"
+                + " in the coilset, it is recommended to check coils closely.",
+            )
+            return is_nearly_intersecting
+        else:
+
+            pts = obj._constants["coilset"]._compute_position(
+                params=self.params_dict, grid=obj._constants["grid"], basis="xyz"
+            )
+            pts = np.array(pts)
+            num_nodes = pts.shape[1]
+            bad_coil_inds = []
+            # We will raise the warning if the jth point on the
+            # kth coil is closer to a point on a different coil than
+            # it is to the neighboring points on itself
+            for k in range(self.num_coils):
+                # dist[i,j,n] is the distance from the jth point on the kth coil
+                # to the nth point on the ith coil
+                dist = np.asarray(
+                    safenorm(pts[k][None, :, None] - pts[:, None, :], axis=-1)
+                )
+                for j in range(num_nodes):
+                    dists_for_this_pt = dist[:, j, :].copy()
+                    dists_for_this_pt[k][
+                        j
+                    ] = np.inf  # Set the dist from the pt to itself to inf to ignore
+                    ind_min = np.argmin(dists_for_this_pt)
+                    # check if the index returned corresponds to a point on the same
+                    # coil. if it does not, then this jth pt on the kth coil is closer
+                    #  to a point on another coil than it is to pts on its own coil,
+                    # which means it may be intersecting it.
+                    if ind_min not in np.arange((num_nodes) * k, (num_nodes) * (k + 1)):
+                        bad_coil_inds.append(k)
+            bad_coil_inds = set(bad_coil_inds)
+            is_nearly_intersecting = True if bad_coil_inds else False
+            warnif(
+                is_nearly_intersecting,
+                UserWarning,
+                "Found coils which are nearly intersecting according to the given grid"
+                + " it is recommended to check coils closely or run function "
+                + "again with a higher resolution grid."
+                + f" Offending coil indices are {bad_coil_inds}.",
+            )
+            return is_nearly_intersecting
+
     def __add__(self, other):
         if isinstance(other, (CoilSet)):
             return CoilSet(*self.coils, *other.coils)
@@ -1548,23 +1652,27 @@ class MixedCoilSet(CoilSet):
         Collection of coils.
     name : str
         Name of this CoilSet.
+    check_intersection: bool
+        Whether or not to check the coils in the coilset for intersections.
 
     """
 
     _io_attrs_ = CoilSet._io_attrs_
 
-    def __init__(self, *coils, name=""):
+    def __init__(self, *coils, name="", check_intersection=True):
         coils = flatten_list(coils, flatten_tuple=True)
         assert all([isinstance(coil, (_Coil)) for coil in coils])
         self._coils = list(coils)
         self._NFP = 1
         self._sym = False
         self._name = str(name)
+        if check_intersection:
+            self.is_self_intersecting()
 
     @property
     def num_coils(self):
         """int: Number of coils."""
-        return sum([c.num_coils if hasattr(c, "num_coils") else 1 for c in self])
+        return sum([c.num_coils for c in self])
 
     def compute(
         self,

desc/compute/__init__.py

@@ -63,10 +63,10 @@ def _build_data_index():
     for p in data_index:
         for key in data_index[p]:
             full = {
-                "data": get_data_deps(key, p, has_axis=False),
-                "transforms": get_derivs(key, p, has_axis=False),
-                "params": get_params(key, p, has_axis=False),
-                "profiles": get_profiles(key, p, has_axis=False),
+                "data": get_data_deps(key, p, has_axis=False, basis="rpz"),
+                "transforms": get_derivs(key, p, has_axis=False, basis="rpz"),
+                "params": get_params(key, p, has_axis=False, basis="rpz"),
+                "profiles": get_profiles(key, p, has_axis=False, basis="rpz"),
             }
             data_index[p][key]["full_dependencies"] = full
 
@@ -81,9 +81,9 @@ def _build_data_index():
             else:
                 full_with_axis = {
                     "data": full_with_axis_data,
-                    "transforms": get_derivs(key, p, has_axis=True),
-                    "params": get_params(key, p, has_axis=True),
-                    "profiles": get_profiles(key, p, has_axis=True),
+                    "transforms": get_derivs(key, p, has_axis=True, basis="rpz"),
+                    "params": get_params(key, p, has_axis=True, basis="rpz"),
+                    "profiles": get_profiles(key, p, has_axis=True, basis="rpz"),
                 }
                 for _key, val in full_with_axis.items():
                     if full[_key] == val: