Merge branch 'master' into rg/parallel-testing

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)

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
 

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):

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,

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)

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)

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
     )