Merge branch 'master' into rc/freeb_field

desc/basis.py

@@ -17,6 +17,7 @@
     "ZernikePolynomial",
     "ChebyshevDoubleFourierBasis",
     "FourierZernikeBasis",
+    "ChebyshevPolynomial",
 ]
 
 
@@ -149,8 +150,8 @@ def evaluate(
         modes : ndarray of in, shape(num_modes,3), optional
             basis modes to evaluate (if None, full basis is used)
         unique : bool, optional
-            whether to workload by only calculating for unique values of nodes, modes
-            can be faster, but doesn't work with jit or autodiff
+            whether to reduce workload by only calculating for unique values of nodes,
+            modes can be faster, but doesn't work with jit or autodiff
 
         Returns
         -------
@@ -294,8 +295,8 @@ def evaluate(
         modes : ndarray of in, shape(num_modes,3), optional
             Basis modes to evaluate (if None, full basis is used)
         unique : bool, optional
-            whether to workload by only calculating for unique values of nodes, modes
-            can be faster, but doesn't work with jit or autodiff
+            whether to reduce workload by only calculating for unique values of nodes,
+            modes can be faster, but doesn't work with jit or autodiff
 
         Returns
         -------
@@ -900,8 +901,8 @@ def evaluate(
         modes : ndarray of in, shape(num_modes,3), optional
             Basis modes to evaluate (if None, full basis is used).
         unique : bool, optional
-            whether to workload by only calculating for unique values of nodes, modes
-            can be faster, but doesn't work with jit or autodiff
+            whether to reduce workload by only calculating for unique values of nodes,
+            modes can be faster, but doesn't work with jit or autodiff
 
         Returns
         -------
@@ -1192,6 +1193,98 @@ def change_resolution(self, L, M, N, NFP=None, sym=None):
             self._set_up()
 
 
+class ChebyshevPolynomial(_Basis):
+    """Shifted Chebyshev polynomial of the first kind.
+
+    Parameters
+    ----------
+    L : int
+        Maximum radial resolution.
+
+    """
+
+    def __init__(self, L):
+        self._L = L
+        self._M = 0
+        self._N = 0
+        self._NFP = 1
+        self._sym = False
+        self._spectral_indexing = "linear"
+
+        self._modes = self._get_modes(L=self.L)
+
+        super().__init__()
+
+    def _get_modes(self, L=0):
+        """Get mode numbers for shifted Chebyshev polynomials of the first kind.
+
+        Parameters
+        ----------
+        L : int
+            Maximum radial resolution.
+
+        Returns
+        -------
+        modes : ndarray of int, shape(num_modes,3)
+            Array of mode numbers [l,m,n].
+            Each row is one basis function with modes (l,m,n).
+
+        """
+        l = np.arange(L + 1).reshape((-1, 1))
+        z = np.zeros((L + 1, 2))
+        return np.hstack([l, z])
+
+    def evaluate(
+        self, nodes, derivatives=np.array([0, 0, 0]), modes=None, unique=False
+    ):
+        """Evaluate basis functions at specified nodes.
+
+        Parameters
+        ----------
+        nodes : ndarray of float, size(num_nodes,3)
+            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
+            Basis modes to evaluate (if None, full basis is used)
+        unique : bool, optional
+            whether to reduce workload by only calculating for unique values of nodes,
+            modes can be faster, but doesn't work with jit or autodiff
+
+        Returns
+        -------
+        y : ndarray, shape(num_nodes,num_modes)
+            basis functions evaluated at nodes
+
+        """
+        if modes is None:
+            modes = self.modes
+        if (derivatives[1] != 0) or (derivatives[2] != 0):
+            return jnp.zeros((nodes.shape[0], modes.shape[0]))
+        if not len(modes):
+            return np.array([]).reshape((len(nodes), 0))
+
+        r, t, z = nodes.T
+        l, m, n = modes.T
+
+        radial = chebyshev(r[:, np.newaxis], l, dr=derivatives[0])
+        return radial
+
+    def change_resolution(self, L):
+        """Change resolution of the basis to the given resolution.
+
+        Parameters
+        ----------
+        L : int
+            Maximum radial resolution.
+
+        """
+        if L != self.L:
+            self._L = L
+            self._modes = self._get_modes(self.L)
+            self._set_up()
+
+
 def polyder_vec(p, m, exact=False):
     """Vectorized version of polyder.
 

desc/compute/__init__.py

@@ -35,8 +35,8 @@
     _field,
     _geometry,
     _metric,
+    _omnigenity,
     _profiles,
-    _qs,
     _stability,
     _surface,
 )

desc/compute/_core.py

@@ -2552,6 +2552,7 @@ def _phi_zz(params, transforms, profiles, data, **kwargs):
     parameterization=[
         "desc.equilibrium.equilibrium.Equilibrium",
         "desc.geometry.core.Surface",
+        "desc.magnetic_fields._core.OmnigenousField",
     ],
 )
 def _rho(params, transforms, profiles, data, **kwargs):

desc/compute/_field.py

@@ -3203,38 +3203,58 @@ def _B_dot_gradB_z(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="max_tz |B|",
-    label="\\max_{\\theta \\zeta} |B|",
+    name="min_tz |B|",
+    label="\\min_{\\theta \\zeta} |\\mathbf{B}|",
     units="T",
     units_long="Tesla",
-    description="Maximum field strength on each flux surface",
+    description="Minimum field strength on each flux surface",
     dim=1,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="r",
     data=["|B|"],
 )
-def _max_tz_modB(params, transforms, profiles, data, **kwargs):
-    data["max_tz |B|"] = surface_max(transforms["grid"], data["|B|"])
+def _min_tz_modB(params, transforms, profiles, data, **kwargs):
+    data["min_tz |B|"] = surface_min(transforms["grid"], data["|B|"])
     return data
 
 
 @register_compute_fun(
-    name="min_tz |B|",
-    label="\\min_{\\theta \\zeta} |B|",
+    name="max_tz |B|",
+    label="\\max_{\\theta \\zeta} |\\mathbf{B}|",
     units="T",
     units_long="Tesla",
-    description="Minimum field strength on each flux surface",
+    description="Maximum field strength on each flux surface",
     dim=1,
     params=[],
     transforms={"grid": []},
     profiles=[],
     coordinates="r",
     data=["|B|"],
 )
-def _min_tz_modB(params, transforms, profiles, data, **kwargs):
-    data["min_tz |B|"] = surface_min(transforms["grid"], data["|B|"])
+def _max_tz_modB(params, transforms, profiles, data, **kwargs):
+    data["max_tz |B|"] = surface_max(transforms["grid"], data["|B|"])
+    return data
+
+
+@register_compute_fun(
+    name="mirror ratio",
+    label="(B_{max} - B_{min}) / (B_{min} + B_{max})",
+    units="~",
+    units_long="None",
+    description="Mirror ratio on each flux surface",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="r",
+    data=["min_tz |B|", "max_tz |B|"],
+)
+def _mirror_ratio(params, transforms, profiles, data, **kwargs):
+    data["mirror ratio"] = (data["max_tz |B|"] - data["min_tz |B|"]) / (
+        data["min_tz |B|"] + data["max_tz |B|"]
+    )
     return data
 
 

desc/compute/_qs.py → desc/compute/_omnigenity.py

@@ -1,4 +1,4 @@
-"""Compute functions for quasisymmetry objectives.
+"""Compute functions for omnigenity objectives.
 
 Notes
 -----
@@ -10,8 +10,9 @@
 """
 
 import numpy as np
+from interpax import interp1d
 
-from desc.backend import jnp, put, sign
+from desc.backend import jnp, put, sign, vmap
 
 from .data_index import register_compute_fun
 from .utils import cross, dot
@@ -346,6 +347,131 @@ def _f_T(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="eta",
+    label="\\eta",
+    units="rad",
+    units_long="radians",
+    description="Intermediate omnigenity coordinate along field lines",
+    dim=1,
+    params=[],
+    transforms={"h": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=[],
+    parameterization="desc.magnetic_fields._core.OmnigenousField",
+)
+def _eta(params, transforms, profiles, data, **kwargs):
+    data["eta"] = transforms["h"].grid.nodes[:, 1]
+    return data
+
+
+@register_compute_fun(
+    name="alpha",
+    label="\\alpha",
+    units="rad",
+    units_long="radians",
+    description="Field line label, defined on [0, 2pi)",
+    dim=1,
+    params=[],
+    transforms={"h": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=[],
+    parameterization="desc.magnetic_fields._core.OmnigenousField",
+)
+def _alpha(params, transforms, profiles, data, **kwargs):
+    data["alpha"] = transforms["h"].grid.nodes[:, 2]
+    return data
+
+
+@register_compute_fun(
+    name="h",
+    label="h = \\theta + (N / M) \\zeta",
+    units="rad",
+    units_long="radians",
+    description="Omnigenity symmetry angle",
+    dim=1,
+    params=["x_lmn"],
+    transforms={"h": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=["eta"],
+    parameterization="desc.magnetic_fields._core.OmnigenousField",
+)
+def _omni_angle(params, transforms, profiles, data, **kwargs):
+    data["h"] = transforms["h"].transform(params["x_lmn"]) + 2 * data["eta"] + jnp.pi
+    return data
+
+
+@register_compute_fun(
+    name="theta_B",
+    label="(\\theta_{B},\\zeta_{B})",
+    units="rad",
+    units_long="radians",
+    description="Boozer angular coordinates",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["alpha", "h"],
+    aliases=["zeta_B"],
+    parameterization="desc.magnetic_fields._core.OmnigenousField",
+    helicity="helicity",
+    iota="iota",
+)
+def _omni_map(params, transforms, profiles, data, **kwargs):
+    M, N = kwargs.get("helicity", (1, 0))
+    iota = kwargs.get("iota", 1)
+
+    # coordinate mapping matrix from (alpha,h) to (theta_B,zeta_B)
+    matrix = jnp.where(
+        M == 0,
+        jnp.array([N, iota / N, 0, 1 / N]),  # OP
+        jnp.where(
+            N == 0,
+            jnp.array([0, -1, M, -1 / iota]),  # OT
+            jnp.array([N, M * iota / (N - M * iota), M, M / (N - M * iota)]),  # OH
+        ),
+    ).reshape((2, 2))
+
+    # solve for (theta_B,zeta_B) corresponding to (eta,alpha)
+    booz = matrix @ jnp.vstack((data["alpha"], data["h"]))
+    data["theta_B"] = booz[0, :]
+    data["zeta_B"] = booz[1, :]
+    return data
+
+
+@register_compute_fun(
+    name="|B|",
+    label="|\\mathbf{B}|",
+    units="T",
+    units_long="Tesla",
+    description="Magnitude of omnigenous magnetic field",
+    dim=1,
+    params=["B_lm"],
+    transforms={"B": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=["eta"],
+    parameterization="desc.magnetic_fields._core.OmnigenousField",
+)
+def _B_omni(params, transforms, profiles, data, **kwargs):
+    # reshaped to size (L_B, M_B)
+    B_lm = params["B_lm"].reshape((transforms["B"].basis.L + 1, -1))
+    # assuming single flux surface, so only take first row (single node)
+    B_input = vmap(lambda x: transforms["B"].transform(x))(B_lm.T)[:, 0]
+    B_input = jnp.sort(B_input)  # sort to ensure monotonicity
+    eta_input = jnp.linspace(0, jnp.pi / 2, num=B_input.size)
+
+    # |B|_omnigeneous is an even function so B(-eta) = B(+eta) = B(|eta|)
+    data["|B|"] = interp1d(
+        jnp.abs(data["eta"]), eta_input, B_input, method="monotonic-0"
+    )
+    return data
+
+
 @register_compute_fun(
     name="isodynamicity",
     label="1/|B|^2 (\\mathbf{b} \\times \\nabla B) \\cdot \\nabla \\psi",

desc/compute/data_index.py

@@ -222,6 +222,7 @@ def _decorator(func):
         "desc.geometry.curve.SplineXYZCurve",
         "desc.geometry.core.Curve",
     ],
+    "desc.magnetic_fields._core.OmnigenousField": [],
 }
 
 data_index = {p: {} for p in _class_inheritance.keys()}