diff --git a/desc/equilibrium/coords.py b/desc/equilibrium/coords.py index a89742b40f..bb9b5b8be9 100644 --- a/desc/equilibrium/coords.py +++ b/desc/equilibrium/coords.py @@ -91,13 +91,28 @@ def map_coordinates( # noqa: C901 inbasis = tuple(inbasis) outbasis = tuple(outbasis) + basis_derivs = tuple(f"{X}_{d}" for X in inbasis for d in ("r", "t", "z")) + for key in basis_derivs: + errorif( + key not in data_index["desc.equilibrium.equilibrium.Equilibrium"], + NotImplementedError, + f"don't have recipe to compute partial derivative {key}", + ) + + profiles = get_profiles(inbasis + basis_derivs, eq) + # TODO: make this work for permutations of in/out basis if outbasis == ("rho", "theta", "zeta"): - # TODO: get iota if not supplied using below logic - if inbasis == ("rho", "alpha", "zeta") and "iota" in kwargs: + if inbasis == ("rho", "alpha", "zeta"): + if "iota" in kwargs: + iota = kwargs.pop("iota") + else: + if profiles["iota"] is None: + profiles["iota"] = eq.get_profile(["iota", "iota_r"], params=params) + iota = profiles["iota"].compute(Grid(coords, sort=False, jitable=True)) return _map_clebsch_coordinates( coords, - kwargs.pop("iota"), + iota, params["L_lmn"], eq.L_basis, guess[:, 1] if guess is not None else None, @@ -118,29 +133,20 @@ def map_coordinates( # noqa: C901 **kwargs, ) - basis_derivs = tuple(f"{X}_{d}" for X in inbasis for d in ("r", "t", "z")) - for key in basis_derivs: - errorif( - key not in data_index["desc.equilibrium.equilibrium.Equilibrium"], - NotImplementedError, - f"don't have recipe to compute partial derivative {key}", - ) + # do surface average to get iota once + if "iota" in profiles and profiles["iota"] is None: + profiles["iota"] = eq.get_profile(["iota", "iota_r"], params=params) + params["i_l"] = profiles["iota"].params rhomin = kwargs.pop("rhomin", tol / 10) warnif(period is None, msg="Assuming no periodicity.") period = np.asarray(setdefault(period, (np.inf, np.inf, np.inf))) coords = _periodic(coords, period) - profiles = get_profiles(inbasis + basis_derivs, eq) p = "desc.equilibrium.equilibrium.Equilibrium" names = inbasis + basis_derivs + outbasis deps = list(set(get_data_deps(names, obj=p) + list(names))) - # do surface average to get iota once - if "iota" in profiles and profiles["iota"] is None: - profiles["iota"] = eq.get_profile(["iota", "iota_r"], params=params) - params["i_l"] = profiles["iota"].params - @functools.partial(jit, static_argnums=1) def compute(y, basis): grid = Grid(y, sort=False, jitable=True) diff --git a/desc/grid.py b/desc/grid.py index ee471e5d1b..4f318afcaf 100644 --- a/desc/grid.py +++ b/desc/grid.py @@ -1295,7 +1295,10 @@ def _create_nodes(self, L=1, M=1, N=1, NFP=1): self._N = check_nonnegint(N, "N", False) self._NFP = check_posint(NFP, "NFP", False) self._period = (np.inf, 2 * np.pi, 2 * np.pi / self._NFP) - L = L + 1 + # floor divide (L+2) by 2 bc only need (L+1)/2 points to + # integrate L-th order jacobi polynomial exactly, so this + # ensures we have enough pts for both odd and even L + L = (L + 2) // 2 M = 2 * M + 1 N = 2 * N + 1 diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png index 9f59728c33..f86a3b1830 100644 Binary files a/tests/baseline/test_plot_b_mag.png and b/tests/baseline/test_plot_b_mag.png differ diff --git a/tests/baseline/test_plot_grid_quad.png b/tests/baseline/test_plot_grid_quad.png index dd9185a7ce..a20bc963a0 100644 Binary files a/tests/baseline/test_plot_grid_quad.png and b/tests/baseline/test_plot_grid_quad.png differ diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 5e54166e05..f371d0291e 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index 811480bbaf..d4e49016a1 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -1196,7 +1196,7 @@ def test_BootstrapRedlConsistency_normalization(self): ) # Results are not perfectly identical because ln(Lambda) is not quite invariant. - np.testing.assert_allclose(results, expected, rtol=2e-3) + np.testing.assert_allclose(results, expected, rtol=3e-3) @pytest.mark.regression def test_bootstrap_consistency_iota(self, TmpDir): diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index 52958a6a9b..471bfa61ae 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -15,7 +15,6 @@ from desc.io import InputReader from desc.objectives import ForceBalance, ObjectiveFunction, get_equilibrium_objective from desc.profiles import PowerSeriesProfile -from desc.utils import errorif from .utils import area_difference, compute_coords @@ -50,33 +49,16 @@ def test_map_coordinates(): """Test root finding for (rho,theta,zeta) for common use cases.""" # finding coordinates along a single field line eq = get("NCSX") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(3, 3, 3, 6, 6, 6) n = 100 coords = np.array([np.ones(n), np.zeros(n), np.linspace(0, 10 * np.pi, n)]).T - grid = LinearGrid(rho=1, M=eq.M_grid, N=eq.N_grid, sym=eq.sym, NFP=eq.NFP) - iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) - iota = np.broadcast_to(iota, shape=(n,)) - - tol = 1e-5 - out_1 = eq.map_coordinates( - coords, inbasis=["rho", "alpha", "zeta"], iota=iota, tol=tol - ) - assert np.isfinite(out_1).all() - out_2 = eq.map_coordinates( + out = eq.map_coordinates( coords, inbasis=["rho", "alpha", "zeta"], period=(np.inf, 2 * np.pi, np.inf), - tol=tol, - ) - assert np.isfinite(out_2).all() - diff = out_1 - out_2 - errorif( - not np.all( - np.isclose(diff, 0, atol=2 * tol) - | np.isclose(np.abs(diff), 2 * np.pi, atol=2 * tol) - ), - AssertionError, - f"diff: {diff}", ) + assert np.isfinite(out).all() eq = get("DSHAPE") @@ -89,9 +71,9 @@ def test_map_coordinates(): grid = Grid(np.vstack([rho, theta, zeta]).T, sort=False) in_data = eq.compute(inbasis, grid=grid) - in_coords = np.stack([in_data[k] for k in inbasis], axis=-1) + in_coords = np.column_stack([in_data[k] for k in inbasis]) out_data = eq.compute(outbasis, grid=grid) - out_coords = np.stack([out_data[k] for k in outbasis], axis=-1) + out_coords = np.column_stack([out_data[k] for k in outbasis]) out = eq.map_coordinates( in_coords, diff --git a/tests/test_grid.py b/tests/test_grid.py index 160c6aac9c..051ba1b89f 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -482,14 +482,16 @@ def test_quadrature_grid(self): N = 0 NFP = 1 grid_quad = QuadratureGrid(L, M, N, NFP) - roots, weights = special.js_roots(3, 2, 2) + roots, weights = special.js_roots(2, 2, 2) quadrature_nodes = np.stack( [ - np.array([roots[0]] * 5 + [roots[1]] * 5 + [roots[2]] * 5), np.array( - [0, 2 * np.pi / 5, 4 * np.pi / 5, 6 * np.pi / 5, 8 * np.pi / 5] * 3 + [roots[0]] * grid_quad.num_theta + [roots[1]] * grid_quad.num_theta ), - np.zeros(15), + np.array( + [0, 2 * np.pi / 5, 4 * np.pi / 5, 6 * np.pi / 5, 8 * np.pi / 5] * 2 + ), + np.zeros(10), ] ).T np.testing.assert_allclose(grid_quad.spacing.prod(axis=1), grid_quad.weights) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 6de443dad0..ff5c17a415 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -208,8 +208,8 @@ def test(eq): obj.build() f = obj.compute_unscaled(*obj.xs(eq)) f_scaled = obj.compute_scaled_error(*obj.xs(eq)) - np.testing.assert_allclose(f, 1.3 / 0.7, rtol=5e-3) - np.testing.assert_allclose(f_scaled, 2 * (1.3 / 0.7), rtol=5e-3) + np.testing.assert_allclose(f, 1.3 / 0.7, rtol=8e-3) + np.testing.assert_allclose(f_scaled, 2 * (1.3 / 0.7), rtol=8e-3) test(get("HELIOTRON")) test(get("HELIOTRON").surface) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index c1bd782ac5..16b89c1543 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -5,7 +5,6 @@ import matplotlib.pyplot as plt import numpy as np import pytest -from scipy.interpolate import interp1d from desc.basis import ( DoubleFourierSeries, @@ -823,17 +822,15 @@ def flatten_coils(coilset): def test_plot_b_mag(): """Test plot of |B| on longer field lines for gyrokinetic simulations.""" psi = 0.5 + rho = np.sqrt(psi) npol = 2 nzgrid = 128 alpha = 0 - # compute and fit iota profile + # compute iota eq = get("W7-X") - data = eq.compute("iota") - fi = interp1d(data["rho"], data["iota"]) + iota = eq.compute("iota", grid=LinearGrid(rho=rho, NFP=eq.NFP))["iota"][0] # 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 )