fix vartheta bug

desc/grid.py

@@ -199,8 +199,11 @@ def _create_nodes(self, nodes):
 
         """
         nodes = np.atleast_2d(nodes).reshape((-1, 3)).astype(float)
-        nodes[nodes[:, 1] != 2 * np.pi, 1] %= 2 * np.pi
-        nodes[nodes[:, 2] != 2 * np.pi / self.NFP, 2] %= 2 * np.pi / self.NFP
+        # Do not alter nodes given by the user for custom grids.
+        # In particular, do not modulo nodes by 2pi or 2pi/NFP.
+        # This may cause the surface_integrals() function to fail recognizing
+        # surfaces outside the interval [0, 2pi] as duplicates. However, most
+        # surface integral computations are done with LinearGrid anyway.
         spacing = (  # make weights sum to 4pi^2
             np.ones_like(nodes) * np.array([1, 2 * np.pi, 2 * np.pi]) / nodes.shape[0]
         )

desc/plotting.py

@@ -1410,7 +1410,7 @@ def plot_surfaces(eq, rho=8, theta=8, zeta=None, ax=None, return_data=False, **k
     }
     if plot_theta:
         # Note: theta* (also known as vartheta) is the poloidal straight field-line
-        # anlge in PEST-like flux coordinates
+        # angle in PEST-like flux coordinates
         t_grid = _get_grid(**grid_kwargs)
         v_grid = Grid(eq.compute_theta_coords(t_grid.nodes))
     rows = np.floor(np.sqrt(nzeta)).astype(int)

tests/test_examples.py

@@ -112,7 +112,7 @@ def test_precise_QH_results(precise_QH):
     eq2 = EquilibriaFamily.load(load_from=str(precise_QH["output_path"]))[-1]
     rho_err, theta_err = area_difference_desc(eq1, eq2)
     np.testing.assert_allclose(rho_err, 0, atol=1e-2)
-    np.testing.assert_allclose(theta_err, 0, atol=3e-2)
+    np.testing.assert_allclose(theta_err, 0, atol=1e-2)
 
 
 @pytest.mark.regression
@@ -123,7 +123,7 @@ def test_HELIOTRON_vac2_results(HELIOTRON_vac, HELIOTRON_vac2):
     eq2 = EquilibriaFamily.load(load_from=str(HELIOTRON_vac2["desc_h5_path"]))[-1]
     rho_err, theta_err = area_difference_desc(eq1, eq2)
     np.testing.assert_allclose(rho_err[:, 3:], 0, atol=1e-2)
-    np.testing.assert_allclose(theta_err, 0, atol=0.003)
+    np.testing.assert_allclose(theta_err, 0, atol=1e-4)
     curr1 = eq1.get_profile("current")
     curr2 = eq2.get_profile("current")
     iota1 = eq1.get_profile("iota")
@@ -369,7 +369,7 @@ def test_ATF_results(tmpdir_factory):
     eqf = load(output_dir.join("ATF.h5"))
     rho_err, theta_err = area_difference_desc(eq0, eqf[-1])
     np.testing.assert_allclose(rho_err[:, 4:], 0, atol=4e-2)
-    np.testing.assert_allclose(theta_err, 0, atol=0.03)
+    np.testing.assert_allclose(theta_err, 0, atol=5e-4)
 
 
 @pytest.mark.regression
@@ -402,7 +402,7 @@ def test_ESTELL_results(tmpdir_factory):
     eqf = load(output_dir.join("ESTELL.h5"))
     rho_err, theta_err = area_difference_desc(eq0, eqf[-1])
     np.testing.assert_allclose(rho_err[:, 3:], 0, atol=5e-2)
-    np.testing.assert_allclose(theta_err, 0, atol=1e-3)
+    np.testing.assert_allclose(theta_err, 0, atol=1e-4)
 
 
 @pytest.mark.regression

tests/test_grid.py

@@ -775,11 +775,11 @@ def test(grid, midpoint_rule=False):
                 dr[1:-1] = (r_unique[2:] - r_unique[:-2]) / 2
                 dr[-1] = 1 - (r_unique[-2] + r_unique[-1]) / 2
             else:
-                dr = np.ones(grid.num_rho) / grid.num_rho
+                dr = 1 / grid.num_rho
             expected_integral = np.sum(dr * compress(grid, function_of_rho))
             true_integral = np.log(1.35 / 0.35)
-            midpoint_rule_error_bound = max(dr) ** 2 / 24 * (2 / 0.35**3)
-            right_riemann_error_bound = dr[0] * (1 / 0.35 - 1 / 1.35)
+            midpoint_rule_error_bound = np.max(dr) ** 2 / 24 * (2 / 0.35**3)
+            right_riemann_error_bound = dr * (1 / 0.35 - 1 / 1.35)
             np.testing.assert_allclose(
                 expected_integral,
                 true_integral,

tests/test_plotting.py

@@ -1,8 +1,11 @@
 """Regression tests for plotting functions, by comparing to saved baseline images."""
 
 import numpy as np
+import matplotlib.pyplot as plt
 import pytest
 
+from scipy.interpolate import interp1d
+
 from desc.basis import (
     DoubleFourierSeries,
     FourierSeries,
@@ -12,7 +15,7 @@
 from desc.coils import CoilSet, FourierXYZCoil
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.examples import get
-from desc.grid import ConcentricGrid, LinearGrid, QuadratureGrid
+from desc.grid import ConcentricGrid, Grid, LinearGrid, QuadratureGrid
 from desc.plotting import (
     _find_idx,
     plot_1d,
@@ -763,3 +766,47 @@ def flatten_coils(coilset):
         assert string in data.keys()
         assert len(data[string]) == len(coil_list)
     return fig
+
+
+@pytest.mark.unit
+@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d)
+def test_plot_b_mag():
+    """Test plot of |B| on longer field lines for gyrokinetic simulations."""
+    psi = 0.5
+    npol = 2
+    nzgrid = 128
+    alpha = 0
+    # compute and fit iota profile
+    eq = get("W7-X")
+    data = eq.compute("iota")
+    fi = interp1d(data["rho"], data["iota"])
+
+    # get flux tube coordinate system
+    rho = np.sqrt(psi)
+    iota = fi(rho)
+    zeta = np.linspace(
+        -np.pi * npol / np.abs(iota), np.pi * npol / np.abs(iota), 2 * nzgrid + 1
+    )
+    thetas = alpha * np.ones_like(zeta) + iota * zeta
+
+    rhoa = rho * np.ones_like(zeta)
+    c = np.vstack([rhoa, thetas, zeta]).T
+    coords = eq.compute_theta_coords(c)
+    grid = Grid(coords)
+
+    # compute |B| normalized in the usual flux tube way
+    psib = np.abs(eq.compute("psi")["psi"][-1])
+    Lref = eq.compute("a")["a"]
+    Bref = 2 * psib / Lref**2
+    bmag = eq.compute("|B|", grid=grid)["|B|"] / Bref
+    fig, ax = plt.subplots()
+    ax.plot(bmag)
+    return fig
+
+
+@pytest.mark.unit
+@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_2d)
+def test_plot_surfaces_HELIOTRON():
+    """Test plot surfaces of equilibrium for correctness of vartheta lines."""
+    fig, ax = plot_surfaces(get("HELIOTRON"))
+    return fig