flip_theta compatibility function

desc/compat.py

@@ -111,6 +111,46 @@
     return eq
 
 
+def flip_theta(eq):
+    """Change the gauge freedom of the poloidal angle of an Equilibrium.
+
+    Equivalent to redefining theta_new = theta_old + π
+
+    Parameters
+    ----------
+    eq : Equilibrium or iterable of Equilibrium
+        Equilibria to redefine the poloidal angle of.
+
+    Returns
+    -------
+    eq : Equilibrium or iterable of Equilibrium
+        Same as input, but with the poloidal angle redefined.
+
+    """
+    # maybe it's iterable:
+    if hasattr(eq, "__len__"):
+        for e in eq:
+            flip_theta(e)
+        return eq
+
+    rone = np.ones_like(eq.R_lmn)
+    rone[eq.R_basis.modes[:, 1] % 2 == 1] *= -1
+    eq.R_lmn *= rone
+
+    zone = np.ones_like(eq.Z_lmn)
+    zone[eq.Z_basis.modes[:, 1] % 2 == 1] *= -1
+    eq.Z_lmn *= zone
+
+    lone = np.ones_like(eq.L_lmn)
+    lone[eq.L_basis.modes[:, 1] % 2 == 1] *= -1
+    eq.L_lmn *= lone
+
+    eq.axis = eq.get_axis()
+    eq.surface = eq.get_surface_at(rho=1)
+
+    return eq
+
+
 def rescale(
     eq, L=("R0", None), B=("B0", None), scale_pressure=True, copy=False, verbose=0
 ):

docs/api.rst

@@ -41,6 +41,7 @@ Compatibility
 
     desc.compat.ensure_positive_jacobian
     desc.compat.flip_helicity
+    desc.compat.flip_theta
     desc.compat.rescale
 
 Continuation

docs/api_equilibrium.rst

@@ -66,4 +66,5 @@ equilibria to a given size and/or field strength.
 
     desc.compat.ensure_positive_jacobian
     desc.compat.flip_helicity
+    desc.compat.flip_theta
     desc.compat.rescale

tests/test_compat.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from desc.compat import flip_helicity, rescale
+from desc.compat import flip_helicity, flip_theta, rescale
 from desc.examples import get
 from desc.grid import Grid, LinearGrid, QuadratureGrid
 
@@ -47,7 +47,6 @@ def test_flip_helicity_axisym():
 
 
 @pytest.mark.unit
-@pytest.mark.solve
 def test_flip_helicity_iota():
     """Test flip_helicity on an Equilibrium with an iota profile."""
     eq = get("HELIOTRON")
@@ -90,7 +89,6 @@ def test_flip_helicity_iota():
 
 
 @pytest.mark.unit
-@pytest.mark.solve
 def test_flip_helicity_current():
     """Test flip_helicity on an Equilibrium with a current profile."""
     eq = get("HSX")
@@ -135,6 +133,87 @@ def test_flip_helicity_current():
     np.testing.assert_allclose(data_old["f_C"], data_flip["f_C"], atol=1e-8)
 
 
+@pytest.mark.unit
+def test_flip_theta_axisym():
+    """Test flip_theta on an axisymmetric Equilibrium."""
+    eq = get("DSHAPE")
+
+    grid = LinearGrid(
+        L=eq.L_grid,
+        theta=2 * eq.M_grid,
+        N=eq.N_grid,
+        NFP=eq.NFP,
+        sym=eq.sym,
+        axis=False,
+    )
+    data_keys = ["current", "|F|", "D_Mercier"]
+
+    data_old = eq.compute(data_keys, grid=grid)
+    eq = flip_theta(eq)
+    data_new = eq.compute(data_keys, grid=grid)
+
+    # check that Jacobian and force balance did not change
+    np.testing.assert_allclose(
+        data_old["sqrt(g)"].reshape((grid.num_rho, grid.num_theta)),
+        np.fliplr(data_new["sqrt(g)"].reshape((grid.num_rho, grid.num_theta))),
+    )
+    np.testing.assert_allclose(
+        data_old["|F|"].reshape((grid.num_rho, grid.num_theta)),
+        np.fliplr(data_new["|F|"].reshape((grid.num_rho, grid.num_theta))),
+        rtol=2e-5,
+    )
+
+    # check that profiles did not change
+    np.testing.assert_allclose(
+        grid.compress(data_old["iota"]), grid.compress(data_new["iota"])
+    )
+    np.testing.assert_allclose(
+        grid.compress(data_old["current"]), grid.compress(data_new["current"])
+    )
+    np.testing.assert_allclose(
+        grid.compress(data_old["D_Mercier"]), grid.compress(data_new["D_Mercier"])
+    )
+
+
+@pytest.mark.unit
+def test_flip_theta_nonaxisym():
+    """Test flip_theta on a non-axisymmetric Equilibrium."""
+    eq = get("HSX")
+
+    grid = QuadratureGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
+    nodes = grid.nodes.copy()
+    nodes[:, 1] = np.mod(nodes[:, 1] + np.pi, 2 * np.pi)
+    grid_flip = Grid(nodes)  # grid with flipped theta values
+    data_keys = ["current", "|F|", "D_Mercier", "f_C"]
+
+    data_old = eq.compute(data_keys, grid=grid, helicity=(1, eq.NFP))
+    eq = flip_theta(eq)
+    data_new = eq.compute(data_keys, grid=grid_flip, helicity=(1, eq.NFP))
+
+    # check that basis vectors did not change
+    np.testing.assert_allclose(data_old["e_rho"], data_new["e_rho"], atol=1e-15)
+    np.testing.assert_allclose(data_old["e_theta"], data_new["e_theta"], atol=1e-15)
+    np.testing.assert_allclose(data_old["e^zeta"], data_new["e^zeta"], atol=1e-15)
+
+    # check that Jacobian is still positive
+    np.testing.assert_array_less(0, grid.compress(data_new["sqrt(g)"]))
+
+    # check that stability did not change
+    np.testing.assert_allclose(
+        grid.compress(data_old["D_Mercier"]),
+        grid.compress(data_new["D_Mercier"]),
+        rtol=2e-2,
+    )
+
+    # check that the total force balance error on each surface did not change
+    # (equivalent collocation points now corresond to theta + pi)
+    np.testing.assert_allclose(data_old["|F|"], data_new["|F|"], rtol=1e-3)
+
+    # check that the QH helicity did not change
+    # (equivalent collocation points now corresond to theta + pi)
+    np.testing.assert_allclose(data_old["f_C"], data_new["f_C"], atol=1e-8)
+
+
 @pytest.mark.unit
 def test_rescale():
     """Test rescale function."""