diff --git a/desc/grid.py b/desc/grid.py index 543d5685b9..cd99b72467 100644 --- a/desc/grid.py +++ b/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] ) diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png new file mode 100644 index 0000000000..fc2feb35b9 Binary files /dev/null and b/tests/baseline/test_plot_b_mag.png differ diff --git a/tests/baseline/test_plot_surfaces_HELIOTRON.png b/tests/baseline/test_plot_surfaces_HELIOTRON.png new file mode 100644 index 0000000000..d7b7145884 Binary files /dev/null and b/tests/baseline/test_plot_surfaces_HELIOTRON.png differ diff --git a/tests/test_plotting.py b/tests/test_plotting.py index c014fb6d34..bc11bba350 100644 --- a/tests/test_plotting.py +++ b/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