Add fft grid and raz grid to test against master

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,59 @@ 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")
+    rho = np.linspace(0, 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),
+        Y_B=grid.num_zeta * grid.NFP,
+        num_transit=10,
+        num_well=20 * 10,
+        nufft_eps=1e-8,
+    )
+    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-8, atol=1e-8, 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 = 10
+    Y_B = eq.N_grid * 2 * eq.NFP
+    rho = np.linspace(0, 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)
+    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."""