non-singular BIE solver, much faster and effecient interpolation for singular integral, and new magnetic field API

.gitignore

@@ -23,6 +23,7 @@ tests/benchmarks/.benchmarks/
 *.swp
 *.swo
 .pymon
+.DS_store
 
 # UV
 uv.lock

CHANGELOG.md

@@ -64,6 +64,12 @@ v0.17.0
 -------
 
 New Features
+------------
+- Automatically-differentiable, non-singular Laplace BIE solver.
+- Improved performance and accuracy of FFT interpolation in singular integrals
+  ([1](https://github.com/f0uriest/interpax/pull/116), [2](https://github.com/f0uriest/interpax/pull/117)).
+  This is useful for free surface optimization.
+- [Plumbing for new magnetic field API](https://github.com/PlasmaControl/DESC/issues/1807).
 
 - Adds particle tracing capabilities in ``desc.particles`` module.
     - Particle tracing is done via ``desc.particles.trace_particles`` function.

README.rst

@@ -102,7 +102,7 @@ Contribute
 - `Contributing guidelines <https://github.com/PlasmaControl/DESC/blob/master/CONTRIBUTING.rst>`_
 - `Issue Tracker <https://github.com/PlasmaControl/DESC/issues>`_
 - `Source Code <https://github.com/PlasmaControl/DESC/>`_
-- `Documentation <https://desc-docs.readthedocs.io/>`_
+- `Documentation <https://desc-docs.readthedocs.io/en/stable/>`_
 
 .. |License| image:: https://img.shields.io/github/license/PlasmaControl/desc?color=blue&logo=open-source-initiative&logoColor=white
     :target: https://github.com/PlasmaControl/DESC/blob/master/LICENSE

desc/backend.py

@@ -65,6 +65,50 @@
     )
 
 
+def _is_converged(residual, tol):
+    return jnp.sum(residual * residual) <= tol**2
+
+
+def _lstsq(A, y):
+    """Cholesky factorized least-squares.
+
+    jnp.linalg.lstsq doesn't have JVP defined and is slower than needed,
+    so we use regularized cholesky.
+
+    For square systems, solves Ax=y directly.
+    """
+    A = jnp.atleast_2d(A)
+    y = jnp.atleast_1d(y)
+    eps = jnp.sqrt(jnp.finfo(A.dtype).eps)
+    if A.shape[-2] == A.shape[-1]:
+        return jnp.linalg.solve(A, y) if y.size > 1 else jnp.squeeze(y / A)
+    elif A.shape[-2] > A.shape[-1]:
+        P = A.T @ A + eps * jnp.eye(A.shape[-1])
+        return cho_solve(cho_factor(P), A.T @ y)
+    else:
+        P = A @ A.T + eps * jnp.eye(A.shape[-2])
+        return A.T @ cho_solve(cho_factor(P), y)
+
+
+def _tangent_solve(g, y):
+    # System is always square.
+    return _lstsq(jax.jacfwd(g)(y), y)
+
+
+def _tangent_solve_scalar(g, y):
+    return y / g(1.0)
+
+
+def _map(f, xs, *, batch_size=None, in_axes=0, out_axes=0):
+    """Generalizes jax.lax.map; uses numpy."""
+    if not isinstance(xs, np.ndarray):
+        raise NotImplementedError(
+            "Require numpy array input, or install jax to support pytrees."
+        )
+    xs = np.moveaxis(xs, source=in_axes, destination=0)
+    return np.stack([f(x) for x in xs], axis=out_axes)
+
+
 def _diag_to_full(d, e):
     j = np.arange(d.shape[-1])
     return (
@@ -377,7 +421,7 @@
 
             def condfun(state):
                 xk1, fk1, k1 = state
-                return (k1 < maxiter) & (jnp.dot(fk1, fk1) > tol**2)
+                return (k1 < maxiter) & (~_is_converged(fk1, tol))
 
             def bodyfun(state):
                 xk1, fk1, k1 = state
@@ -393,41 +437,17 @@
             else:
                 return state[0]
 
-        def tangent_solve(g, y):
-            return y / g(1.0)
-
         if full_output:
             x, (res, niter) = jax.lax.custom_root(
-                res, x0, solve, tangent_solve, has_aux=True
+                res, x0, solve, _tangent_solve_scalar, has_aux=True
             )
             return x, (abs(res), niter)
         else:
-            x = jax.lax.custom_root(res, x0, solve, tangent_solve, has_aux=False)
+            x = jax.lax.custom_root(
+                res, x0, solve, _tangent_solve_scalar, has_aux=False
+            )
             return x
 
-    def _lstsq(A, y):
-        """Cholesky factorized least-squares.
-
-        jnp.linalg.lstsq doesn't have JVP defined and is slower than needed,
-        so we use regularized cholesky.
-
-        For square systems, solves Ax=y directly.
-        """
-        A = jnp.atleast_2d(A)
-        y = jnp.atleast_1d(y)
-        eps = jnp.sqrt(jnp.finfo(A.dtype).eps)
-        if A.shape[-2] == A.shape[-1]:
-            return jnp.linalg.solve(A, y) if y.size > 1 else jnp.squeeze(y / A)
-        elif A.shape[-2] > A.shape[-1]:
-            P = A.T @ A + eps * jnp.eye(A.shape[-1])
-            return cho_solve(cho_factor(P), A.T @ y)
-        else:
-            P = A @ A.T + eps * jnp.eye(A.shape[-2])
-            return A.T @ cho_solve(cho_factor(P), y)
-
-    def _tangent_solve(g, y):
-        return _lstsq(jax.jacfwd(g)(y), y)
-
     def root(
         fun,
         x0,
@@ -512,7 +532,7 @@
 
             def condfun(state):
                 xk1, fk1, k1 = state
-                return (k1 < maxiter) & (jnp.dot(fk1, fk1) > tol**2)
+                return (k1 < maxiter) & (~_is_converged(fk1, tol))
 
             def bodyfun(state):
                 xk1, fk1, k1 = state
@@ -570,15 +590,6 @@
         excluded={"eigvals_only", "select", "select_range", "tol"},
     )
 
-    def _map(f, xs, *, batch_size=None, in_axes=0, out_axes=0):
-        """Generalizes jax.lax.map; uses numpy."""
-        if not isinstance(xs, np.ndarray):
-            raise NotImplementedError(
-                "Require numpy array input, or install jax to support pytrees."
-            )
-        xs = np.moveaxis(xs, source=in_axes, destination=0)
-        return np.stack([f(x) for x in xs], axis=out_axes)
-
     def vmap(fun, in_axes=0, out_axes=0):
         """A numpy implementation of jax.lax.map whose API is a subset of jax.vmap.
 

desc/basis.py

@@ -6,6 +6,7 @@
 
 import mpmath
 import numpy as np
+from jax.lax import stop_gradient
 
 from desc.backend import custom_jvp, fori_loop, jit, jnp, sign
 from desc.grid import Grid, _Grid
@@ -421,7 +422,7 @@
         else:
             lidx = loutidx = np.arange(len(modes))
         if (derivatives[1] != 0) or (derivatives[2] != 0):
-            return jnp.zeros((grid.num_nodes, modes.shape[0]))
+            return jnp.zeros((grid.num_nodes, self.num_modes))
         if not len(modes):
             return np.array([]).reshape((grid.num_nodes, 0))
 
@@ -537,7 +538,7 @@
         else:
             nidx = noutidx = np.arange(len(modes))
         if (derivatives[0] != 0) or (derivatives[1] != 0):
-            return jnp.zeros((grid.num_nodes, modes.shape[0]))
+            return jnp.zeros((grid.num_nodes, self.num_modes))
         if not len(modes):
             return np.array([]).reshape((grid.num_nodes, 0))
 
@@ -594,21 +595,39 @@
         * ``'cos'`` for cos(m*t-n*z) symmetry
         * ``'sin'`` for sin(m*t-n*z) symmetry
         * ``False`` for no symmetry (Default)
+    stop_gradient : bool
+        Whether to stop gradients through basis evaluations. Default is ``False``.
+        May set to ``True`` if the grid to evaluate on is independent of the
+        optimizable parameters.
 
     """
 
     _fft_poloidal = True
     _fft_toroidal = True
+    _static_attrs = _Basis._static_attrs + ["_stop_gradient"]
 
-    def __init__(self, M, N, NFP=1, sym=False):
+    def __init__(self, M, N, NFP=1, sym=False, stop_gradient=False):
         self._L = 0
         self._M = check_nonnegint(M, "M", False)
         self._N = check_nonnegint(N, "N", False)
         self._NFP = check_posint(NFP, "NFP", False)
         self._sym = bool(sym) if not sym else str(sym)
         self._spectral_indexing = "linear"
+        self._stop_gradient = bool(stop_gradient)
         self._modes = self._get_modes(M=self.M, N=self.N)
         super().__init__()
+        self._gauge_idx = np.asarray(self.get_idx(error=False)).squeeze()
+
+    def _set_up(self):
+        """Do things after loading or changing resolution."""
+        super()._set_up()
+        self._gauge_idx = np.asarray(self.get_idx(error=False)).squeeze()
+        self._stop_gradient = getattr(self, "_stop_gradient", False)
+
+    @property
+    def gauge_idx(self):
+        """ndarray: Index of constant-potential gauge mode, if present."""
+        return self._gauge_idx
 
     def _get_modes(self, M, N):
         """Get mode numbers for double Fourier series.
@@ -635,7 +654,13 @@
         z = np.zeros_like(m)
         return np.array([z, m, n]).T
 
-    def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
+    def evaluate(
+        self,
+        grid,
+        derivatives=np.array([0, 0, 0]),
+        modes=None,
+        **kwargs,
+    ):
         """Evaluate basis functions at specified nodes.
 
         Parameters
@@ -644,7 +669,7 @@
             Node coordinates, in (rho,theta,zeta).
         derivatives : ndarray of int, shape(num_derivatives,3)
             Order of derivatives to compute in (rho,theta,zeta).
-        modes : ndarray of in, shape(num_modes,3), optional
+        modes : ndarray of int, shape(num_modes,3), optional
             Basis modes to evaluate (if None, full basis is used).
 
         Returns
@@ -670,7 +695,7 @@
             midx = moutidx = np.arange(len(modes))
             nidx = noutidx = np.arange(len(modes))
         if derivatives[0] != 0:
-            return jnp.zeros((grid.num_nodes, modes.shape[0]))
+            return jnp.zeros((grid.num_nodes, self.num_modes))
         if not len(modes):
             return np.array([]).reshape((grid.num_nodes, 0))
 
@@ -688,17 +713,22 @@
         _, t, z = grid.nodes.T
         _, m, n = modes.T
 
-        t = t[tidx]
-        z = z[zidx]
+        t = kwargs["t"] if "t" in kwargs else t[tidx]
+        z = kwargs["z"] if "z" in kwargs else z[zidx]
         m = m[midx]
         n = n[nidx]
 
-        poloidal = fourier(t[:, np.newaxis], m, 1, derivatives[1])
-        toroidal = fourier(z[:, np.newaxis], n, self.NFP, derivatives[2])
-        poloidal = poloidal[toutidx][:, moutidx]
-        toroidal = toroidal[zoutidx][:, noutidx]
-
-        return poloidal * toroidal
+        poloidal = fourier(t[:, np.newaxis], m, 1, derivatives[1])[:, moutidx]
+        toroidal = fourier(z[:, np.newaxis], n, self.NFP, derivatives[2])[:, noutidx]
+        if grid.is_meshgrid and grid.num_rho == 1:
+            out = (poloidal[:, np.newaxis] * toroidal[np.newaxis]).reshape(
+                -1, self.num_modes, order="F"
+            )
+        else:
+            out = poloidal[toutidx] * toroidal[zoutidx]
+        if self._stop_gradient:
+            out = stop_gradient(out)
+        return out
 
     def change_resolution(self, M, N, NFP=None, sym=None):
         """Change resolution of the basis to the given resolutions.
@@ -881,7 +911,7 @@
             lmidx = lmoutidx = np.arange(len(modes))
             midx = moutidx = np.arange(len(modes))
         if derivatives[2] != 0:
-            return jnp.zeros((grid.num_nodes, modes.shape[0]))
+            return jnp.zeros((grid.num_nodes, self.num_modes))
         if not len(modes):
             return np.array([]).reshape((grid.num_nodes, 0))
 
@@ -1437,7 +1467,7 @@
         else:
             lidx = loutidx = np.arange(len(modes))
         if (derivatives[1] != 0) or (derivatives[2] != 0):
-            return jnp.zeros((grid.num_nodes, modes.shape[0]))
+            return jnp.zeros((grid.num_nodes, self.num_modes))
         if not len(modes):
             return np.array([]).reshape((grid.num_nodes, 0))
 

desc/coils.py

@@ -1603,7 +1603,7 @@ def compute(
 
         Returns
         -------
-        data : list of dict of ndarray
+        data : list[dict[str, jnp.ndarray]]
             Computed quantity and intermediate variables, for each coil in the set.
             List entries map to coils in coilset, each dict contains data for an
             individual coil.
@@ -1843,6 +1843,7 @@ def body(AB, x):
                 return AB, None
 
             AB += scan(body, jnp.zeros(coords_nfp.shape), tree_stack(params))[0]
+
             return AB
 
         AB = fori_loop(0, self.NFP, nfp_loop, jnp.zeros_like(coords_rpz))
@@ -2753,7 +2754,7 @@ def compute(
 
         Returns
         -------
-        data : list of dict of ndarray
+        data : list[dict[str, jnp.ndarray]]
             Computed quantity and intermediate variables, for each coil in the set.
             List entries map to coils in coilset, each dict contains data for an
             individual coil.

desc/compute/__init__.py

@@ -35,6 +35,7 @@
     _fast_ion,
     _field,
     _geometry,
+    _laplace,
     _metric,
     _neoclassical,
     _old,