add tests

desc/grid.py

@@ -199,7 +199,11 @@ def _create_nodes(self, nodes):
 
         """
         nodes = np.atleast_2d(nodes).reshape((-1, 3)).astype(float)
-        # Note: do not modulo nodes by 2pi or 2pi/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]
         )

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