Bounce averaging

desc/backend.py

@@ -73,6 +73,9 @@
     vmap = jax.vmap
     scan = jax.lax.scan
     bincount = jnp.bincount
+    repeat = jnp.repeat
+    flatnonzero = jnp.flatnonzero
+    take = jnp.take
     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
@@ -104,6 +107,38 @@
             return arr
         return jnp.asarray(arr).at[inds].set(vals)
 
+    def put_along_axis(arr, indices, values, axis):
+        """Put values into the destination array by matching 1d index and data slices.
+
+        This iterates over matching 1d slices oriented along the specified axis in
+        the index and data arrays, and uses the former to place values into the
+        latter.
+
+        Parameters
+        ----------
+        arr : ndarray (Ni..., M, Nk...)
+            Destination array.
+        indices : ndarray (Ni..., J, Nk...)
+            Indices to change along each 1d slice of `arr`. This must match the
+            dimension of arr, but dimensions in Ni and Nj may be 1 to broadcast
+            against `arr`.
+        values : array_like (Ni..., J, Nk...)
+            values to insert at those indices. Its shape and dimension are
+            broadcast to match that of `indices`.
+        axis : int
+            The axis to take 1d slices along. If axis is None, the destination
+            array is treated as if a flattened 1d view had been created of it.
+
+        """
+        if axis != -1:
+            raise NotImplementedError(
+                f"put_along_axis for axis={axis} not implemented yet."
+            )
+        if isinstance(arr, np.ndarray):
+            arr[..., indices] = values
+            return arr
+        return jnp.asarray(arr).at[..., indices].set(values)
+
     def sign(x):
         """Sign function, but returns 1 for x==0.
 
@@ -357,6 +392,26 @@
         )
         return x, (jnp.linalg.norm(res), niter)
 
+    def complex_sqrt(x):
+        """Compute the square root of x.
+
+        For negative input elements, a complex value is returned
+        (unlike numpy.sqrt which returns NaN).
+
+        Parameters
+        ----------
+        x : array_like
+            The input value(s).
+
+        Returns
+        -------
+        out : ndarray
+            The square root of x.
+
+        """
+        out = jnp.sqrt(x.astype("complex128"))
+        return out
+
 
 # we can't really test the numpy backend stuff in automated testing, so we ignore it
 # for coverage purposes
@@ -373,6 +428,9 @@
     )
     from scipy.special import gammaln, logsumexp  # noqa: F401
 
+    complex_sqrt = np.emath.sqrt
+    put_along_axis = np.put_along_axis
+
     def tree_stack(*args, **kwargs):
         """Stack pytree for numpy backend."""
         raise NotImplementedError
@@ -530,7 +588,7 @@
             val = body_fun(val)
         return val
 
-    def vmap(fun, out_axes=0):
+    def vmap(fun, in_axes=0, out_axes=0):
         """A numpy implementation of jax.lax.map whose API is a subset of jax.vmap.
 
         Like Python's builtin map,
@@ -541,6 +599,8 @@
         ----------
         fun: callable
             Function (A -> B)
+        in_axes: int
+            Axis to map over.
         out_axes: int
             An integer indicating where the mapped axis should appear in the output.
 
@@ -550,8 +610,16 @@
             Vectorized version of fun.
 
         """
+        if in_axes != 0:
+            raise NotImplementedError(
+                f"Backend for numpy vmap for in_axes={in_axes} not implemented yet."
+            )
 
         def fun_vmap(fun_inputs):
+            if isinstance(fun_inputs, tuple):
+                raise NotImplementedError(
+                    "Backend implementation of vmap fails for multiple arguments."
+                )
             return np.stack([fun(fun_input) for fun_input in fun_inputs], axis=out_axes)
 
         return fun_vmap
@@ -595,9 +663,21 @@
             ys.append(y)
         return carry, np.stack(ys)
 
-    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 bincount(x, weights=None, minlength=0, length=None):
+        """Numpy implementation of jnp.bincount."""
+        x = np.clip(x, 0, None)
+        if length is None:
+            length = max(minlength, x.max() + 1)
+        else:
+            minlength = max(minlength, length)
+        return np.bincount(x, weights, minlength)[:length]
+
+    def repeat(a, repeats, axis=None, total_repeat_length=None):
+        """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."""
@@ -708,3 +788,34 @@
         """
         out = scipy.optimize.root(fun, x0, args, jac=jac, tol=tol)
         return out.x, out
+
+    def flatnonzero(a, size=None, fill_value=0):
+        """Numpy implementation of jnp.flatnonzero."""
+        nz = np.flatnonzero(a)
+        if size is not None:
+            nz = np.append(nz, np.repeat(fill_value, max(size - nz.size, 0)))
+        return nz
+
+    def take(
+        a,
+        indices,
+        axis=None,
+        out=None,
+        mode="fill",
+        unique_indices=False,
+        indices_are_sorted=False,
+        fill_value=None,
+    ):
+        """Numpy implementation of jnp.take."""
+        if mode == "fill":
+            if fill_value is None:
+                # TODO: Interpret default fill value based on dtype of a.
+                fill_value = np.nan
+            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/compute/_basis_vectors.py

@@ -504,6 +504,26 @@
     return data
 
 
+@register_compute_fun(
+    name="e^theta_PEST",
+    label="\\mathbf{e}^{\\theta_{PEST}}",
+    units="m^{-1}",
+    units_long="inverse meters",
+    description="Contravariant straight field line (PEST) poloidal basis vector",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e_rho", "e_phi", "sqrt(g)_PEST"],
+)
+def _e_sup_theta_PEST(params, transforms, profiles, data, **kwargs):
+    data["e^theta_PEST"] = (
+        cross(data["e_phi"], data["e_rho"]).T / data["sqrt(g)_PEST"]
+    ).T
+    return data
+
+
 @register_compute_fun(
     name="e^theta*sqrt(g)",
     label="\\mathbf{e}^{\\theta} \\sqrt{g}",
@@ -2377,7 +2397,7 @@
     data=["e_theta", "theta_PEST_t"],
 )
 def _e_sub_theta_pest(params, transforms, profiles, data, **kwargs):
-    # dX/dv at const r,z = dX/dt * dt/dv / dX/dt / dv/dt
+    # dX/dv at const r,z = dX/dt * dt/dv = dX/dt / dv/dt
     data["e_theta_PEST"] = (data["e_theta"].T / data["theta_PEST_t"]).T
     return data
 

desc/compute/_core.py

@@ -28,6 +28,29 @@
     return data
 
 
+@register_compute_fun(
+    name="1",
+    label="1",
+    units="~",
+    units_long="None",
+    description="Ones",
+    dim=1,
+    params=[],
+    transforms={"grid": []},
+    profiles=[],
+    coordinates="rtz",
+    data=[],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.core.Surface",
+        "desc.geometry.core.Curve",
+    ],
+)
+def _1(params, transforms, profiles, data, **kwargs):
+    data["1"] = jnp.ones(transforms["grid"].num_nodes)
+    return data
+
+
 @register_compute_fun(
     name="R",
     label="R",

desc/compute/_field.py

@@ -86,6 +86,24 @@
     return data
 
 
+@register_compute_fun(
+    name="B^theta_PEST",
+    label="B^{\\theta}",
+    units="T \\cdot m^{-1}",
+    units_long="Tesla / meter",
+    description="Contravariant straight field line (PEST) component of magnetic field",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["B", "e^theta_PEST"],
+)
+def _B_sup_theta_PEST(params, transforms, profiles, data, **kwargs):
+    data["B^theta_PEST"] = dot(data["B"], data["e^theta_PEST"])
+    return data
+
+
 @register_compute_fun(
     name="B^zeta",
     label="B^{\\zeta}",
@@ -2220,12 +2238,13 @@
     ],
 )
 def _B_mag_r(params, transforms, profiles, data, **kwargs):
-    data["|B|_r"] = (
+    data["|B|_r"] = safediv(
         data["B^theta_r"] * data["B_theta"]
         + data["B^theta"] * data["B_theta_r"]
         + data["B^zeta_r"] * data["B_zeta"]
-        + data["B^zeta"] * data["B_zeta_r"]
-    ) / (2 * data["|B|"])
+        + data["B^zeta"] * data["B_zeta_r"],
+        2 * data["|B|"],
+    )
     return data
 
 
@@ -2253,12 +2272,13 @@
     ],
 )
 def _B_mag_t(params, transforms, profiles, data, **kwargs):
-    data["|B|_t"] = (
+    data["|B|_t"] = safediv(
         data["B^theta_t"] * data["B_theta"]
         + data["B^theta"] * data["B_theta_t"]
         + data["B^zeta_t"] * data["B_zeta"]
-        + data["B^zeta"] * data["B_zeta_t"]
-    ) / (2 * data["|B|"])
+        + data["B^zeta"] * data["B_zeta_t"],
+        2 * data["|B|"],
+    )
     return data
 
 
@@ -2286,12 +2306,55 @@
     ],
 )
 def _B_mag_z(params, transforms, profiles, data, **kwargs):
-    data["|B|_z"] = (
+    data["|B|_z"] = safediv(
         data["B^theta_z"] * data["B_theta"]
         + data["B^theta"] * data["B_theta_z"]
         + data["B^zeta_z"] * data["B_zeta"]
-        + data["B^zeta"] * data["B_zeta_z"]
-    ) / (2 * data["|B|"])
+        + data["B^zeta"] * data["B_zeta_z"],
+        2 * data["|B|"],
+    )
+    return data
+
+
+@register_compute_fun(
+    name="|B|_alpha",
+    label="\\partial_{\\alpha} |\\mathbf{B}|",
+    units="T",
+    units_long="Tesla",
+    description="Magnitude of magnetic field, derivative wrt field line angle",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|B|_t", "alpha_t"],
+)
+def _B_mag_alpha(params, transforms, profiles, data, **kwargs):
+    # constant ρ and ζ
+    data["|B|_alpha"] = safediv(data["|B|_t"], data["alpha_t"])
+    return data
+
+
+@register_compute_fun(
+    # TODO: pick a name
+    name="|B|_z constant rho alpha",
+    label="\\(partial_{\\zeta} |\\mathbf{B}|)_{\\rho, \\alpha}",
+    units="T",
+    units_long="Tesla",
+    description="Magnitude of magnetic field, derivative along field line",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|B|_z", "|B|_alpha", "alpha_z"],
+)
+def _B_mag_z_constant_rho_alpha(params, transforms, profiles, data, **kwargs):
+    # ∂|B|/∂ζ (constant ρ and α) = ∂|B|/∂ζ (constant ρ and θ)
+    #                            - ∂|B|/∂α (constant ρ and ζ) * ∂α/∂ζ (constant ρ and θ)
+    data["|B|_z constant rho alpha"] = (
+        data["|B|_z"] - data["|B|_alpha"] * data["alpha_z"]
+    )
     return data