diff --git a/desc/grid.py b/desc/grid.py index dc3eef6bea..cd99b72467 100644 --- a/desc/grid.py +++ b/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] ) diff --git a/desc/plotting.py b/desc/plotting.py index 3381009dc1..bcfda442a6 100644 --- a/desc/plotting.py +++ b/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) 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_examples.py b/tests/test_examples.py index d9f5bac45b..c993699706 100644 --- a/tests/test_examples.py +++ b/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 diff --git a/tests/test_grid.py b/tests/test_grid.py index 93cce3b9f9..05b932b566 100644 --- a/tests/test_grid.py +++ b/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, diff --git a/tests/test_plotting.py b/tests/test_plotting.py index c014fb6d34..6f5506cddb 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 matplotlib.pyplot as plt import numpy as np 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