Add fft grid and raz grid to test against master

desc/equilibrium/equilibrium.py

@@ -946,6 +946,8 @@ def compute(  # noqa: C901
                 inbasis=[inbasis[char] for char in grid.coordinates],
                 outbasis=("rho", "theta", "zeta"),
                 period=grid.period,
+                tol=kwargs.pop("tol", 1e-6),
+                maxiter=kwargs.pop("maxiter", 30),
             )
             grid = Grid(
                 nodes=rtz_nodes,

tests/test_compute_everything.py

@@ -24,13 +24,14 @@
     FourierXYZCurve,
     ZernikeRZToroidalSection,
 )
-from desc.grid import LinearGrid
+from desc.grid import Grid, LinearGrid
+from desc.integrals import Bounce2D
 from desc.magnetic_fields import (
     CurrentPotentialField,
     FourierCurrentPotentialField,
     OmnigenousField,
 )
-from desc.utils import ResolutionWarning, errorif, xyz2rpz, xyz2rpz_vec
+from desc.utils import ResolutionWarning, apply, errorif, xyz2rpz, xyz2rpz_vec
 
 
 def _compare_against_master(
@@ -259,6 +260,11 @@ def need_special(name):
                 f"Parameterization: {p}. Can't compute "
                 + f"{names - this_branch_data_rpz[p].keys()}."
             )
+
+            # now we get the special grid data
+            this_branch_data_rpz[p].update(fft_grid_data(p))
+            this_branch_data_rpz[p].update(raz_grid_data(p))
+
             # compare data against master branch
             error_rpz, update_master_data_rpz = _compare_against_master(
                 p,
@@ -297,3 +303,60 @@ def need_special(name):
             pickle.dump(this_branch_data_rpz, file)
 
     assert not error_rpz
+
+
+def fft_grid_data(p):
+    """Compute fft grid quantities."""
+    if p != "desc.equilibrium.equilibrium.Equilibrium":
+        return {}
+
+    # TODO: can automate this later to add omngeneity, boozer transform, etc.
+    fft_names = ["effective ripple", "Gamma_c", "Gamma_c Velasco"]
+
+    eq = get("W7-X")
+    # ci and my laptop differ a bunch at rho = 0, so skip that
+    rho = np.linspace(1e-2, 1, 10)
+    grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False)
+
+    kwargs = dict(
+        angle=Bounce2D.angle(eq, X=32, Y=32, rho=rho, tol=1e-10),
+        Y_B=grid.num_zeta * grid.NFP,
+        num_transit=5,
+        num_well=20 * 5,
+        nufft_eps=1e-10,
+    )
+    data = eq.compute(fft_names, grid, surf_batch_size=1, **kwargs)
+
+    # check vectorization too
+    d = data.copy()
+    for name in fft_names:
+        d.pop(name)
+    d = eq.compute(fft_names, grid, surf_batch_size=2, data=d, **kwargs)
+    for name in fft_names:
+        np.testing.assert_allclose(
+            d[name], data[name], rtol=1e-9, atol=1e-9, err_msg=name
+        )
+
+    data = apply(data, grid.compress, fft_names)
+    return data
+
+
+def raz_grid_data(p):
+    """Compute field line grid quantities."""
+    if p != "desc.equilibrium.equilibrium.Equilibrium":
+        return {}
+
+    # TODO: can automate this later to add ballooning etc.
+    raz_names = ["old effective ripple", "old Gamma_c", "old Gamma_c Velasco"]
+
+    eq = get("W7-X")
+    num_transit = 2
+    Y_B = eq.N_grid * 2 * eq.NFP
+    rho = np.linspace(1e-2, 1, 10)
+    alpha = np.array([0])
+    zeta = np.linspace(0, num_transit * 2 * np.pi, num_transit * Y_B)
+    grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz")
+
+    data = eq.compute(raz_names, grid, num_well=20 * num_transit, tol=1e-10)
+    data = apply(data, grid.compress, raz_names)
+    return data

tests/test_compute_funs.py

@@ -2049,6 +2049,45 @@ def test_parallel_grad_fd(DummyStellarator):
     )
 
 
+@pytest.mark.unit
+@pytest.mark.slow
+def test_fieldline_average():
+    """Test that fieldline average converges to surface average."""
+    rho = np.array([1])
+    alpha = np.array([0])
+    eq = get("DSHAPE")
+    iota_grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym)
+    iota = iota_grid.compress(eq.compute("iota", grid=iota_grid)["iota"]).item()
+    # For axisymmetric devices, one poloidal transit must be exact.
+    zeta = np.linspace(0, 2 * np.pi / iota, 25)
+    grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz")
+    data = eq.compute(
+        ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid
+    )
+    np.testing.assert_allclose(
+        data["fieldline length"] / data["fieldline length/volume"],
+        data["V_r(r)"] / (4 * np.pi**2),
+        rtol=1e-3,
+    )
+    assert np.all(data["fieldline length"] > 0)
+    assert np.all(data["fieldline length/volume"] > 0)
+
+    # Otherwise, many toroidal transits are necessary to sample surface.
+    eq = get("W7-X")
+    zeta = np.linspace(0, 40 * np.pi, 300)
+    grid = Grid.create_meshgrid([rho, alpha, zeta], coordinates="raz")
+    data = eq.compute(
+        ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid
+    )
+    np.testing.assert_allclose(
+        data["fieldline length"] / data["fieldline length/volume"],
+        data["V_r(r)"] / (4 * np.pi**2),
+        rtol=2e-3,
+    )
+    assert np.all(data["fieldline length"] > 0)
+    assert np.all(data["fieldline length/volume"] > 0)
+
+
 @pytest.mark.unit
 def test_compute_deprecation_warning():
     """Test DeprecationWarning for deprecated compute names."""