Adding adjoint ideal-ballooning stability solver to DESC

desc/backend.py

@@ -58,12 +58,15 @@
                 desc.__version__, np.__version__, y.dtype
             )
         )
+
+
 print(
     "Using device: {}, with {:.2f} GB available memory".format(
         desc_config.get("device"), desc_config.get("avail_mem")
     )
 )
 
+
 if use_jax:  # noqa: C901 - FIXME: simplify this, define globally and then assign?
     jit = jax.jit
     fori_loop = jax.lax.fori_loop
@@ -73,6 +76,12 @@
     vmap = jax.vmap
     scan = jax.lax.scan
     bincount = jnp.bincount
+    from functools import partial
+
+    import jax
+    import jax.numpy as jnp
+    import jaxlib
+    from jax import config as jax_config
     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
@@ -158,6 +167,94 @@ def tree_unstack(tree):
         leaves, treedef = jtu.tree_flatten(tree)
         return [treedef.unflatten(leaf) for leaf in zip(*leaves)]
 
+    _eigvals_cpu = jax.jit(jnp.linalg.eigvals, device=jax.devices("cpu")[0])
+
+    @jax.custom_jvp
+    def eigvals(A):
+        """
+        Eigenvalue solver.
+
+        Returns the eigenvalues of the square matrix A.
+        eigvals only run on CPUs
+        """
+        u = jax.pure_callback(
+            _eigvals_cpu, jnp.zeros_like(A[..., -1]) + 1j, A, vectorized=True
+        )
+        return u
+
+    @eigvals.defjvp
+    def _eigvals_jvp(primals, tangents):
+
+        u = eigvals(primals[0])
+
+        @partial(jnp.vectorize, signature="(n,n),(n,n)->(n)")
+        def jvpfun(primals, tangents):
+            u, du = jax.jvp(_eigvals_cpu, (primals,), (tangents,))
+            return du.squeeze()
+
+        du = jax.pure_callback(jvpfun, u, *primals, *tangents, vectorized=True)
+        return u, du
+
+    _gen_eigval_cpu = jax.jit(jax.scipy.linalg.eigh, device=jax.devices("cpu")[0])
+
+    @jax.custom_jvp
+    def gen_eigval(A):
+        """
+        Generalize eigenvalue solver.
+
+        Returns the top n eigenvalues of the square matrix A. Calculation is
+        being performed on a CPU. If the CPU version can provide the top eigenvalue,
+        the calculation should be faster on a CPU.
+        Currently doesn't work because of the limitations of the jax functionality.
+        """
+        neigs, N, _ = jnp.shape(A)
+        u = jnp.zeros((N,))
+        i = jnp.arange(N)
+
+        u = u.at[i].set(
+            jax.pure_callback(
+                _gen_eigval_cpu,
+                jnp.zeros_like(A[i, :, :]),
+                A[i, :, :],
+                k=1,
+                sigma=0.42,
+                vectorized=True,
+            )
+        )
+        return u
+
+    @gen_eigval.defjvp
+    def _gen_eigval_jvp(primals, tangents):
+
+        u = gen_eigval(primals[0])
+
+        @partial(jnp.vectorize, signature="(n,n),(n,n)->(n)")
+        def jvpfun(primals, tangents):
+            u, du = jax.jvp(_gen_eigval_cpu, (primals,), (tangents,))
+            return du.squeeze()
+
+        du = jax.pure_callback(jvpfun, u, *primals, *tangents, vectorized=True)
+        return u, du
+
+    @jit
+    def simspson_integrator(y, dx):
+        """Simpsons integrations scheme for high-order accurate integrals."""
+        if len(y[..., :]) % 2 == 1:
+            raise ValueError("n must be even")
+
+        S = (
+            dx
+            / 3
+            * (
+                y[..., 0]
+                + y[..., -1]
+                + 4 * jnp.sum(y[..., 1:-1:2], axis=-1)
+                + 2 * jnp.sum(y[..., 2:-2:2], axis=-1)
+            )
+        )
+
+        return S
+
     def root_scalar(
         fun,
         x0,
@@ -373,6 +470,7 @@ def tangent_solve(g, y):
 # for coverage purposes
 else:  # pragma: no cover
     jit = lambda func, *args, **kwargs: func
+    import numpy as np
     import scipy.optimize
     from scipy.integrate import odeint  # noqa: F401
     from scipy.linalg import (  # noqa: F401
@@ -636,6 +734,15 @@ def custom_jvp(fun, *args, **kwargs):
         fun.defjvps = lambda *args, **kwargs: None
         return fun
 
+    def eigvals(A):
+        """
+        Eigenvalue solver.
+
+        Returns the eigenvalues of the square matrix A.
+        """
+        u = np.linalg.eigvals(A)
+        return u
+
     def root_scalar(
         fun,
         x0,

desc/compute/_basis_vectors.py

@@ -4370,6 +4370,233 @@ def _grad_alpha(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="grad(alpha)_t",
+    label="\\partial_{\\theta} \\nabla \\alpha",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description="Poloidal derivative of the gradient of the field line label",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "e^rho",
+        "e^theta",
+        "e^zeta",
+        "e^rho_t",
+        "e^theta_t",
+        "e^zeta_t",
+        "alpha_r",
+        "alpha_t",
+        "alpha_z",
+        "alpha_rt",
+        "alpha_tt",
+        "alpha_tz",
+    ],
+)
+def _grad_alpha_t(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha)_t"] = (
+        data["alpha_r"] * data["e^rho_t"].T
+        + data["alpha_rt"] * data["e^rho"].T
+        + data["alpha_t"] * data["e^theta_t"].T
+        + data["alpha_tt"] * data["e^theta"].T
+        + data["alpha_z"] * data["e^zeta_t"].T
+        + data["alpha_tz"] * data["e^zeta"].T
+    ).T
+    return data
+
+
+@register_compute_fun(
+    name="grad(alpha)_z",
+    label="\\partial_{\\zeta} (\\nabla \\alpha)",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description="Toroidal derivative of the gradient of the field line label",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "e^rho",
+        "e^theta",
+        "e^zeta",
+        "e^rho_z",
+        "e^theta_z",
+        "e^zeta_z",
+        "alpha_r",
+        "alpha_t",
+        "alpha_z",
+        "alpha_rz",
+        "alpha_tz",
+        "alpha_zz",
+    ],
+)
+def _grad_alpha_z(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha)_z"] = (
+        data["alpha_r"] * data["e^rho_z"].T
+        + data["alpha_rz"] * data["e^rho"].T
+        + data["alpha_t"] * data["e^theta_z"].T
+        + data["alpha_tz"] * data["e^theta"].T
+        + data["alpha_z"] * data["e^zeta_z"].T
+        + data["alpha_zz"] * data["e^zeta"].T
+    ).T
+    return data
+
+
+@register_compute_fun(
+    name="grad(alpha)_tt",
+    label="\\partial_{\\theta \\theta} \\nabla \\alpha",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description="Poloidal derivative of the gradient of the field line label",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "e^rho",
+        "e^theta",
+        "e^zeta",
+        "e^rho_t",
+        "e^theta_t",
+        "e^zeta_t",
+        "e^rho_tt",
+        "e^theta_tt",
+        "e^zeta_tt",
+        "alpha_r",
+        "alpha_t",
+        "alpha_z",
+        "alpha_rt",
+        "alpha_tt",
+        "alpha_tz",
+        "alpha_rtt",
+        "alpha_ttt",
+        "alpha_ttz",
+    ],
+)
+def _grad_alpha_tt(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha)_tt"] = (
+        data["alpha_r"] * data["e^rho_tt"].T
+        + data["alpha_rtt"] * data["e^rho"].T
+        + 2 * data["alpha_rt"] * data["e^rho_t"].T
+        + data["alpha_t"] * data["e^theta_tt"].T
+        + data["alpha_ttt"] * data["e^theta"].T
+        + 2 * data["alpha_tt"] * data["e^theta_t"].T
+        + data["alpha_z"] * data["e^zeta_tt"].T
+        + data["alpha_ttz"] * data["e^zeta"].T
+        + 2 * data["alpha_tz"] * data["e^zeta_t"].T
+    ).T
+    return data
+
+
+@register_compute_fun(
+    name="grad(alpha)_zz",
+    label="\\partial_{\\zeta \\zeta} (\\nabla \\alpha)",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description="Toroidal derivative of the gradient of the field line label",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "e^rho",
+        "e^theta",
+        "e^zeta",
+        "e^rho_z",
+        "e^theta_z",
+        "e^zeta_z",
+        "e^rho_zz",
+        "e^theta_zz",
+        "e^zeta_zz",
+        "alpha_r",
+        "alpha_t",
+        "alpha_z",
+        "alpha_rz",
+        "alpha_tz",
+        "alpha_zz",
+        "alpha_rzz",
+        "alpha_tzz",
+        "alpha_zzz",
+    ],
+)
+def _grad_alpha_zz(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha)_zz"] = (
+        data["alpha_r"] * data["e^rho_zz"].T
+        + data["alpha_rzz"] * data["e^rho"].T
+        + 2 * data["alpha_rz"] * data["e^rho_z"].T
+        + data["alpha_t"] * data["e^theta_zz"].T
+        + data["alpha_tzz"] * data["e^theta"].T
+        + 2 * data["alpha_tz"] * data["e^theta_z"].T
+        + data["alpha_z"] * data["e^zeta_zz"].T
+        + data["alpha_zzz"] * data["e^zeta"].T
+        + 2 * data["alpha_zz"] * data["e^zeta_z"].T
+    ).T
+    return data
+
+
+@register_compute_fun(
+    name="grad(alpha)_tz",
+    label="\\partial_{\\theta \\zeta} (\\nabla \\alpha)",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description=(
+        "Mixed toroidal and poloidal derivative of the gradient"
+        " of the field line label"
+    ),
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "e^rho",
+        "e^theta",
+        "e^zeta",
+        "e^rho_z",
+        "e^theta_z",
+        "e^zeta_z",
+        "e^rho_t",
+        "e^theta_t",
+        "e^zeta_t",
+        "e^rho_tz",
+        "e^theta_tz",
+        "e^zeta_tz",
+        "alpha_r",
+        "alpha_t",
+        "alpha_z",
+        "alpha_rz",
+        "alpha_tz",
+        "alpha_tt",
+        "alpha_rt",
+        "alpha_rtz",
+        "alpha_ttz",
+        "alpha_tzz",
+    ],
+)
+def _grad_alpha_tz(params, transforms, profiles, data, **kwargs):
+    data["grad(alpha)_tz"] = (
+        data["alpha_r"] * data["e^rho_tz"].T
+        + data["alpha_rtz"] * data["e^rho"].T
+        + data["alpha_rz"] * data["e^rho_t"].T
+        + data["alpha_rt"] * data["e^rho_z"].T
+        + data["alpha_t"] * data["e^theta_tz"].T
+        + data["alpha_ttz"] * data["e^theta"].T
+        + data["alpha_tz"] * data["e^theta_t"].T
+        + data["alpha_tt"] * data["e^theta_z"].T
+        + data["alpha_z"] * data["e^zeta_tz"].T
+        + data["alpha_tzz"] * data["e^zeta"].T
+        + data["alpha_tz"] * data["e^zeta_t"].T
+        + data["alpha_tt"] * data["e^zeta_z"].T
+    ).T
+    return data
+
+
 @register_compute_fun(
     name="grad(psi)",
     label="\\nabla\\psi",