Merge branch 'yge/rpz2xyz' into dd/curvature

desc/coils.py

@@ -224,7 +224,7 @@ def compute_magnetic_field(
             current = params.pop("current", self.current)
 
         data = self.compute(
-            ["x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz"
+            ["phi", "x", "x_s", "ds"], grid=source_grid, params=params, basis="xyz"
         )
         B = biot_savart_quad(
             coords, data["x"], data["x_s"] * data["ds"][:, None], current

desc/compute/_omnigenity.py

@@ -443,15 +443,25 @@ def _omni_map(params, transforms, profiles, data, **kwargs):
     iota = kwargs.get("iota", 1)
 
     # coordinate mapping matrix from (alpha,h) to (theta_B,zeta_B)
+    # need a bunch of wheres to avoid division by zero causing NaN in backward pass
+    # this is fine since the incorrect values get ignored later, except in OT or OH
+    # where fieldlines are exactly parallel to |B| contours, but this is a degenerate
+    # case of measure 0 so this kludge shouldn't affect things too much.
+    mat_OP = jnp.array(
+        [[N, iota / jnp.where(N == 0, 1, N)], [0, 1 / jnp.where(N == 0, 1, N)]]
+    )
+    mat_OT = jnp.array([[0, -1], [M, -1 / jnp.where(iota == 0, 1.0, iota)]])
+    den = jnp.where((N - M * iota) == 0, 1.0, (N - M * iota))
+    mat_OH = jnp.array([[N, M * iota / den], [M, M / den]])
     matrix = jnp.where(
         M == 0,
-        jnp.array([N, iota / N, 0, 1 / N]),  # OP
+        mat_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
+            mat_OT,
+            mat_OH,
         ),
-    ).reshape((2, 2))
+    )
 
     # solve for (theta_B,zeta_B) corresponding to (eta,alpha)
     booz = matrix @ jnp.vstack((data["alpha"], data["h"]))

desc/compute/utils.py

@@ -9,6 +9,7 @@
 
 from desc.backend import cond, fori_loop, jnp, put
 from desc.grid import ConcentricGrid, Grid, LinearGrid
+from desc.utils import errorif
 
 from .data_index import allowed_kwargs, data_index
 
@@ -105,7 +106,12 @@ def compute(parameterization, names, params, transforms, profiles, data=None, **
             if name == "x":
                 data[name] = rpz2xyz(data[name])
             else:
-                # TODO: check if phi in data?
+                errorif(
+                    "phi" not in data.keys(),
+                    ValueError,
+                    "'phi' must be included in the compute data "
+                    + "to convert to 'xyz' basis.",
+                )
                 data[name] = rpz2xyz_vec(data[name], phi=data["phi"])
 
     return data

desc/objectives/_omnigenity.py

@@ -666,7 +666,7 @@ def __init__(
         normalize=True,
         normalize_target=True,
         loss_function=None,
-        deriv_mode="fwd",  # FIXME: get it working with rev mode (see GH issue #943)
+        deriv_mode="auto",
         eq_grid=None,
         field_grid=None,
         M_booz=None,
@@ -891,20 +891,26 @@ def compute(self, params_1=None, params_2=None, constants=None):
             # update theta_B and zeta_B with new iota from the equilibrium
             M, N = constants["helicity"]
             iota = jnp.mean(eq_data["iota"])
+            # see comment in desc.compute._omnigenity for the explanation of these
+            # wheres
+            mat_OP = jnp.array(
+                [[N, iota / jnp.where(N == 0, 1, N)], [0, 1 / jnp.where(N == 0, 1, N)]]
+            )
+            mat_OT = jnp.array([[0, -1], [M, -1 / jnp.where(iota == 0, 1.0, iota)]])
+            den = jnp.where((N - M * iota) == 0, 1.0, (N - M * iota))
+            mat_OH = jnp.array([[N, M * iota / den], [M, M / den]])
             matrix = jnp.where(
                 M == 0,
-                jnp.array([N, iota / N, 0, 1 / N]),  # OP
+                mat_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
-                    ),
+                    mat_OT,
+                    mat_OH,
                 ),
-            ).reshape((2, 2))
+            )
             booz = matrix @ jnp.vstack((field_data["alpha"], field_data["h"]))
-            field_data["theta_B"] = booz[0, :]
-            field_data["zeta_B"] = booz[1, :]
+            theta_B = booz[0, :]
+            zeta_B = booz[1, :]
         else:
             field_data = compute_fun(
                 "desc.magnetic_fields._core.OmnigenousField",
@@ -915,13 +921,15 @@ def compute(self, params_1=None, params_2=None, constants=None):
                 helicity=constants["helicity"],
                 iota=jnp.mean(eq_data["iota"]),
             )
+            theta_B = field_data["theta_B"]
+            zeta_B = field_data["zeta_B"]
 
         # additional computations that cannot be part of the regular compute API
         nodes = jnp.vstack(
             (
-                jnp.zeros_like(field_data["theta_B"]),
-                field_data["theta_B"],
-                field_data["zeta_B"],
+                jnp.zeros_like(theta_B),
+                theta_B,
+                zeta_B,
             )
         ).T
         B_eta_alpha = jnp.matmul(

tests/test_curves.py

@@ -55,19 +55,15 @@ def test_torsion(self):
     def test_frenet(self):
         """Test frenet-serret frame of circular curve."""
         c = FourierRZCurve()
-        data = c.compute(
-            ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", grid=0
-        )
+        data = c.compute(["frenet_tangent", "frenet_normal", "frenet_binormal"], grid=0)
         T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"]
         np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12)
         np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12)
         np.testing.assert_allclose(B, np.array([[0, 0, 1]]), atol=1e-12)
         c.rotate(angle=np.pi)
         c.flip([0, 1, 0])
         c.translate([1, 1, 1])
-        data = c.compute(
-            ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0
-        )
+        data = c.compute(["frenet_tangent", "frenet_normal", "frenet_binormal"], grid=0)
         T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"]
         np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12)
         np.testing.assert_allclose(N, np.array([[1, 0, 0]]), atol=1e-12)
@@ -470,9 +466,9 @@ def test_rotation(self):
         datax = cx.compute("x", grid=20, basis="xyz")
         datay = cy.compute("x", grid=20, basis="xyz")
         dataz = cz.compute("x", grid=20, basis="xyz")
-        np.testing.assert_allclose(datax["x"][:, 0], 0, atol=1e-16)  # only in Y-Z plane
-        np.testing.assert_allclose(datay["x"][:, 1], 0, atol=1e-16)  # only in X-Z plane
-        np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=1e-16)  # only in X-Y plane
+        np.testing.assert_allclose(datax["x"][:, 0], 0, atol=2e-16)  # only in Y-Z plane
+        np.testing.assert_allclose(datay["x"][:, 1], 0, atol=2e-16)  # only in X-Z plane
+        np.testing.assert_allclose(dataz["x"][:, 2], 0, atol=2e-16)  # only in X-Y plane
 
     @pytest.mark.unit
     def test_length(self):

tests/test_objective_funs.py

@@ -2418,6 +2418,35 @@ def test_objective_no_nangrad_coils(self, objective):
         g = obj.grad(obj.x())
         assert not np.any(np.isnan(g)), str(objective)
 
+    @pytest.mark.unit
+    @pytest.mark.parametrize("helicity", [(1, 0), (1, 1), (0, 1)])
+    def test_objective_no_nangrad_omnigenity(self, helicity):
+        """Omnigenity."""
+        surf = FourierRZToroidalSurface.from_qp_model(
+            major_radius=1,
+            aspect_ratio=20,
+            elongation=6,
+            mirror_ratio=0.2,
+            torsion=0.1,
+            NFP=1,
+            sym=True,
+        )
+        eq = Equilibrium(Psi=6e-3, M=4, N=4, surface=surf)
+        field = OmnigenousField(
+            L_B=0,
+            M_B=2,
+            L_x=1,
+            M_x=1,
+            N_x=1,
+            NFP=eq.NFP,
+            helicity=helicity,
+            B_lm=np.array([0.8, 1.2]),
+        )
+        obj = ObjectiveFunction(Omnigenity(eq=eq, field=field))
+        obj.build()
+        g = obj.grad(obj.x())
+        assert not np.any(np.isnan(g)), str(helicity)
+
 
 @pytest.mark.unit
 def test_asymmetric_normalization():