Merge branch 'master' into clebsh_basis

.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,20 +1,69 @@
 Changelog
 =========
 
+
+v0.12.0
+-------
+
+[Github Commits](https://github.com/PlasmaControl/DESC/compare/v0.11.1...v0.12.0)
+
 New Features
 
+- 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 ``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.
-- Add method ``is_self_intersecting`` to ``CoilSet``, which checks if any coils intersect eachother in the coilset.
-- Removes error in ``from_symmetry`` method of ``CoilSet`` when a coil crosses the symmetry plane,
-and instead adds a check for intersection, to allow for valid coilsets which may cross the
-symmetry plane but not be self-intersecting after rotation/reflection.
+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/backend.py

@@ -71,8 +71,10 @@
     switch = jax.lax.switch
     while_loop = jax.lax.while_loop
     vmap = jax.vmap
-    scan = jax.lax.scan
     bincount = jnp.bincount
+    repeat = jnp.repeat
+    take = jnp.take
+    scan = jax.lax.scan
     from jax import custom_jvp
     from jax.experimental.ode import odeint
     from jax.scipy.linalg import block_diag, cho_factor, cho_solve, qr, solve_triangular
@@ -635,6 +637,13 @@ def bincount(x, weights=None, minlength=None, length=None):
         """Same as np.bincount but with a dummy parameter to match jnp.bincount API."""
         return np.bincount(x, weights, minlength)
 
+    def repeat(a, repeats, axis=None, total_repeat_length=None):
+        """A numpy implementation of jnp.repeat."""
+        out = np.repeat(a, repeats, axis)
+        if total_repeat_length is not None:
+            out = out[:total_repeat_length]
+        return out
+
     def custom_jvp(fun, *args, **kwargs):
         """Dummy function for custom_jvp without JAX."""
         fun.defjvp = lambda *args, **kwargs: None
@@ -744,3 +753,37 @@ def root(
         """
         out = scipy.optimize.root(fun, x0, args, jac=jac, tol=tol)
         return out.x, out
+
+    def take(
+        a,
+        indices,
+        axis=None,
+        out=None,
+        mode="fill",
+        unique_indices=False,
+        indices_are_sorted=False,
+        fill_value=None,
+    ):
+        """A numpy implementation of jnp.take."""
+        if mode == "fill":
+            if fill_value is None:
+                # copy jax logic
+                # https://jax.readthedocs.io/en/latest/_modules/jax/_src/lax/slicing.html#gather
+                if np.issubdtype(a.dtype, np.inexact):
+                    fill_value = np.nan
+                elif np.issubdtype(a.dtype, np.signedinteger):
+                    fill_value = np.iinfo(a.dtype).min
+                elif np.issubdtype(a.dtype, np.unsignedinteger):
+                    fill_value = np.iinfo(a.dtype).max
+                elif a.dtype == np.bool_:
+                    fill_value = True
+                else:
+                    raise ValueError(f"Unsupported dtype {a.dtype}.")
+            out = np.where(
+                (-a.size <= indices) & (indices < a.size),
+                np.take(a, indices, axis, out, mode="wrap"),
+                fill_value,
+            )
+        else:
+            out = np.take(a, indices, axis, out, mode)
+        return out

desc/coils.py

@@ -158,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.
 
@@ -227,12 +232,10 @@ def compute_magnetic_field(
             current = self.current
         else:
             current = params.pop("current", self.current)
-        if source_grid is None and hasattr(self, "NFP"):
-            # NFP=1 to ensure we have points along whole grid
-            # multiply by NFP in case the coil has NFP>1
-            # to ensure whole coil gets counted for the
-            # biot savart integration
-            source_grid = LinearGrid(N=2 * self.N * self.NFP + 5, NFP=1, endpoint=False)
+        if source_grid is None:
+            # NFP=1 to ensure points span the entire length of the coil
+            # multiply resolution by NFP to ensure Biot-Savart integration is accurate
+            source_grid = LinearGrid(N=2 * self.N * getattr(self, "NFP", 1) + 5)
 
         if not params or not transforms:
             data = self.compute(
@@ -295,6 +298,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
@@ -1665,7 +1670,7 @@ def __init__(self, *coils, name="", check_intersection=True):
     @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/_bootstrap.py

@@ -31,6 +31,7 @@
     coordinates="r",
     data=["sqrt(g)", "V_r(r)", "|B|", "<|B|^2>", "max_tz |B|"],
     axis_limit_data=["sqrt(g)_r", "V_rr(r)"],
+    resolution_requirement="tz",
     n_gauss="int: Number of quadrature points to use for estimating trapped fraction. "
     + "Default 20.",
 )