non-singular BIE solver, improved FFT interpolation for singular integral, and new magnetic field API

CHANGELOG.md

@@ -49,6 +49,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 print_backend_info():
     )
 
 
+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 (
@@ -360,7 +404,7 @@ def backtrack(xk1, fk1, d):
 
             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
@@ -376,41 +420,17 @@ def bodyfun(state):
             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,
@@ -495,7 +515,7 @@ def backtrack(xk1, fk1, d):
 
             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
@@ -552,15 +572,6 @@ def bodyfun(state):
         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

@@ -421,7 +421,7 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
         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 +537,7 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
         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))
 
@@ -635,7 +635,13 @@ def _get_modes(self, M, N):
         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 +650,7 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
             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 +676,7 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
             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 +694,18 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
         _, 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:
+            return (poloidal[:, np.newaxis] * toroidal[np.newaxis]).reshape(
+                -1, self.num_modes, order="F"
+            )
+        return poloidal[toutidx] * toroidal[zoutidx]
 
     def change_resolution(self, M, N, NFP=None, sym=None):
         """Change resolution of the basis to the given resolutions.
@@ -881,7 +888,7 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
             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 +1444,7 @@ def evaluate(self, grid, derivatives=np.array([0, 0, 0]), modes=None):
         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,