diff --git a/CHANGELOG.md b/CHANGELOG.md index aab80a4173..85ab3a47ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,14 @@ Changelog New Features - Add ``use_signed_distance`` flag to ``PlasmaVesselDistance`` which will use a signed distance as the target, which is positive when the plasma is inside of the vessel surface and negative if the plasma is outside of the vessel surface, to allow optimizer to distinguish if the equilbrium surface exits the vessel surface and guard against it by targeting a positive signed distance. +- Allow specification of Nyquist spectrum maximum modenumbers when using ``VMECIO.save`` to save a DESC .h5 file as a VMEC-format wout file + +Bug Fixes + +- Fixes bugs that occur when saving asymmetric equilibria as wout files +- Fixes bug that occurs when using ``VMECIO.plot_vmec_comparison`` to compare to an asymmetric wout file + + v0.12.1 ------- diff --git a/desc/vmec.py b/desc/vmec.py index fc6fc5498f..17e7bf3b30 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -25,7 +25,7 @@ from desc.objectives.utils import factorize_linear_constraints from desc.profiles import PowerSeriesProfile, SplineProfile from desc.transform import Transform -from desc.utils import Timer +from desc.utils import Timer, warnif from desc.vmec_utils import ( fourier_to_zernike, ptolemy_identity_fwd, @@ -158,7 +158,7 @@ def load( zax_cs = file.variables["zaxis_cs"][:].filled() try: rax_cs = file.variables["raxis_cs"][:].filled() - rax_cc = file.variables["zaxis_cc"][:].filled() + zax_cc = file.variables["zaxis_cc"][:].filled() except KeyError: rax_cs = np.zeros_like(rax_cc) zax_cc = np.zeros_like(zax_cs) @@ -208,7 +208,9 @@ def load( return eq @classmethod - def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify + def save( # noqa: C901 - FIXME - simplify + cls, eq, path, surfs=128, verbose=1, M_nyq=None, N_nyq=None + ): """Save an Equilibrium as a netCDF file in the VMEC format. Parameters @@ -224,6 +226,10 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify * 0: no output * 1: status of quantities computed * 2: as above plus timing information + M_nyq, N_nyq: int + The max poloidal and toroidal modenumber to use in the + Nyquist spectrum that the derived quantities are Fourier + fit with. Defaults to M+4 and N+2. Returns ------- @@ -242,8 +248,14 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify NFP = eq.NFP M = eq.M N = eq.N - M_nyq = M + 4 - N_nyq = N + 2 if N > 0 else 0 + M_nyq = M + 4 if M_nyq is None else M_nyq + warnif( + N_nyq is not None and int(N) == 0, + UserWarning, + "Passed in N_nyq but equilibrium is axisymmetric, setting N_nyq to zero", + ) + N_nyq = N + 2 if N_nyq is None else N_nyq + N_nyq = 0 if int(N) == 0 else N_nyq # VMEC radial coordinate: s = rho^2 = Psi / Psi(LCFS) s_full = np.linspace(0, 1, surfs) @@ -807,6 +819,14 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify lmnc.long_name = "cos(m*t-n*p) component of lambda, on half mesh" lmnc.units = "rad" l1 = np.ones_like(eq.L_lmn) + # should negate lambda coefs bc theta_DESC + lambda = theta_PEST, + # since we are reversing the theta direction (and the theta_PEST direction), + # so -theta_PEST = -theta_DESC - lambda, so the negative of lambda is what + # should be saved, so that would be negating all of eq.L_lmn + # BUT since we are also reversing the poloidal angle direction, which + # would negate only the coeffs of L_lmn corresponding to m<0 + # (sin theta modes in DESC), the effective result is to only + # negate the cos(theta) (m>0) lambda modes l1[eq.L_basis.modes[:, 1] >= 0] *= -1 m, n, x_mn = zernike_to_fourier(l1 * eq.L_lmn, basis=eq.L_basis, rho=r_half) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) @@ -823,7 +843,7 @@ def save(cls, eq, path, surfs=128, verbose=1): # noqa: C901 - FIXME - simplify sin_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym="sin") cos_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym="cos") - full_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym=None) + full_basis = DoubleFourierSeries(M=M_nyq, N=N_nyq, NFP=NFP, sym=False) if eq.sym: sin_transform = Transform( grid=grid_lcfs, basis=sin_basis, build=False, build_pinv=True @@ -932,7 +952,7 @@ def fullfit(x): if eq.sym: x_mn[i, :] = cosfit(data[i, :]) else: - x_mn[i, :] = full_transform.fit(data[i, :]) + x_mn[i, :] = fullfit(data[i, :]) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) bmnc[0, :] = 0 bmnc[1:, :] = c @@ -975,7 +995,7 @@ def fullfit(x): if eq.sym: x_mn[i, :] = cosfit(data[i, :]) else: - x_mn[i, :] = full_transform.fit(data[i, :]) + x_mn[i, :] = fullfit(data[i, :]) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) bsupumnc[0, :] = 0 bsupumnc[1:, :] = -c # negative sign for negative Jacobian @@ -1018,7 +1038,7 @@ def fullfit(x): if eq.sym: x_mn[i, :] = cosfit(data[i, :]) else: - x_mn[i, :] = full_transform.fit(data[i, :]) + x_mn[i, :] = fullfit(data[i, :]) xm, xn, s, c = ptolemy_identity_rev(m, n, x_mn) bsupvmnc[0, :] = 0 bsupvmnc[1:, :] = c @@ -1641,13 +1661,15 @@ def vmec_interpolate(Cmn, Smn, xm, xn, theta, phi, s=None, si=None, sym=True): return C + S @classmethod - def compute_theta_coords(cls, lmns, xm, xn, s, theta_star, zeta, si=None): + def compute_theta_coords( + cls, lmns, xm, xn, s, theta_star, zeta, si=None, lmnc=None + ): """Find theta such that theta + lambda(theta) == theta_star. Parameters ---------- lmns : array-like - fourier coefficients for lambda + sin(mt-nz) Fourier coefficients for lambda xm : array-like poloidal mode numbers xn : array-like @@ -1662,6 +1684,8 @@ def compute_theta_coords(cls, lmns, xm, xn, s, theta_star, zeta, si=None): si : ndarray values of radial coordinates where lmns are defined. Defaults to linearly spaced on half grid between (0,1) + lmnc : array-like, optional + cos(mt-nz) Fourier coefficients for lambda Returns ------- @@ -1672,19 +1696,30 @@ def compute_theta_coords(cls, lmns, xm, xn, s, theta_star, zeta, si=None): if si is None: si = np.linspace(0, 1, lmns.shape[0]) si[1:] = si[0:-1] + 0.5 / (lmns.shape[0] - 1) - lmbda_mn = interpolate.CubicSpline(si, lmns) + lmbda_mns = interpolate.CubicSpline(si, lmns) + if lmnc is None: + lmbda_mnc = lambda s: 0 + else: + lmbda_mnc = interpolate.CubicSpline(si, lmnc) # Note: theta* (also known as vartheta) is the poloidal straight field line # angle in PEST-like flux coordinates def root_fun(theta): lmbda = np.sum( - lmbda_mn(s) + lmbda_mns(s) * np.sin( xm[np.newaxis] * theta[:, np.newaxis] - xn[np.newaxis] * zeta[:, np.newaxis] ), axis=-1, + ) + np.sum( + lmbda_mnc(s) + * np.cos( + xm[np.newaxis] * theta[:, np.newaxis] + - xn[np.newaxis] * zeta[:, np.newaxis] + ), + axis=-1, ) theta_star_k = theta + lmbda # theta* = theta + lambda err = theta_star - theta_star_k # FIXME: mod by 2pi @@ -1782,6 +1817,8 @@ def compute_coord_surfaces(cls, equil, vmec_data, Nr=10, Nt=8, Nz=None, **kwargs t_nodes = t_grid.nodes t_nodes[:, 0] = t_nodes[:, 0] ** 2 + sym = "lmnc" not in vmec_data.keys() + v_nodes = cls.compute_theta_coords( vmec_data["lmns"], vmec_data["xm"], @@ -1789,29 +1826,71 @@ def compute_coord_surfaces(cls, equil, vmec_data, Nr=10, Nt=8, Nz=None, **kwargs t_nodes[:, 0], t_nodes[:, 1], t_nodes[:, 2], + lmnc=vmec_data["lmnc"] if not sym else None, ) t_nodes[:, 1] = v_nodes + if sym: + Rr_vmec, Zr_vmec = cls.vmec_interpolate( + vmec_data["rmnc"], + vmec_data["zmns"], + vmec_data["xm"], + vmec_data["xn"], + theta=r_nodes[:, 1], + phi=r_nodes[:, 2], + s=r_nodes[:, 0], + ) - Rr_vmec, Zr_vmec = cls.vmec_interpolate( - vmec_data["rmnc"], - vmec_data["zmns"], - vmec_data["xm"], - vmec_data["xn"], - theta=r_nodes[:, 1], - phi=r_nodes[:, 2], - s=r_nodes[:, 0], - ) - - Rv_vmec, Zv_vmec = cls.vmec_interpolate( - vmec_data["rmnc"], - vmec_data["zmns"], - vmec_data["xm"], - vmec_data["xn"], - theta=t_nodes[:, 1], - phi=t_nodes[:, 2], - s=t_nodes[:, 0], - ) + Rv_vmec, Zv_vmec = cls.vmec_interpolate( + vmec_data["rmnc"], + vmec_data["zmns"], + vmec_data["xm"], + vmec_data["xn"], + theta=t_nodes[:, 1], + phi=t_nodes[:, 2], + s=t_nodes[:, 0], + ) + else: + Rr_vmec = cls.vmec_interpolate( + vmec_data["rmnc"], + vmec_data["rmns"], + vmec_data["xm"], + vmec_data["xn"], + theta=r_nodes[:, 1], + phi=r_nodes[:, 2], + s=r_nodes[:, 0], + sym=False, + ) + Zr_vmec = cls.vmec_interpolate( + vmec_data["zmnc"], + vmec_data["zmns"], + vmec_data["xm"], + vmec_data["xn"], + theta=r_nodes[:, 1], + phi=r_nodes[:, 2], + s=r_nodes[:, 0], + sym=False, + ) + Rv_vmec = cls.vmec_interpolate( + vmec_data["rmnc"], + vmec_data["rmns"], + vmec_data["xm"], + vmec_data["xn"], + theta=t_nodes[:, 1], + phi=t_nodes[:, 2], + s=t_nodes[:, 0], + sym=False, + ) + Zv_vmec = cls.vmec_interpolate( + vmec_data["zmnc"], + vmec_data["zmns"], + vmec_data["xm"], + vmec_data["xn"], + theta=t_nodes[:, 1], + phi=t_nodes[:, 2], + s=t_nodes[:, 0], + sym=False, + ) coords = { "Rr_desc": Rr_desc, diff --git a/tests/conftest.py b/tests/conftest.py index 873d2c3f0a..ccab0e07a6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -335,3 +335,22 @@ def VMEC_save(SOLOVEV, tmpdir_factory): ) desc = Dataset(str(SOLOVEV["desc_nc_path"]), mode="r") return vmec, desc + + +@pytest.fixture(scope="session") +def VMEC_save_asym(tmpdir_factory): + """Save an asymmetric equilibrium in VMEC netcdf format for comparison.""" + tmpdir = tmpdir_factory.mktemp("asym_wout") + filename = tmpdir.join("wout_HELIO_asym_desc.nc") + vmec = Dataset("./tests/inputs/wout_HELIOTRON_asym_NTHETA50_NZETA100.nc", mode="r") + eq = Equilibrium.load("./tests/inputs/HELIO_asym.h5") + VMECIO.save( + eq, + filename, + surfs=vmec.variables["ns"][:], + verbose=0, + M_nyq=round(np.max(vmec.variables["xm_nyq"][:])), + N_nyq=round(np.max(vmec.variables["xn_nyq"][:]) / eq.NFP), + ) + desc = Dataset(filename, mode="r") + return vmec, desc, eq diff --git a/tests/inputs/HELIO_asym.h5 b/tests/inputs/HELIO_asym.h5 new file mode 100644 index 0000000000..c66a6cb100 Binary files /dev/null and b/tests/inputs/HELIO_asym.h5 differ diff --git a/tests/inputs/wout_HELIOTRON_asym_NTHETA50_NZETA100.nc b/tests/inputs/wout_HELIOTRON_asym_NTHETA50_NZETA100.nc new file mode 100644 index 0000000000..cc51c535a3 Binary files /dev/null and b/tests/inputs/wout_HELIOTRON_asym_NTHETA50_NZETA100.nc differ diff --git a/tests/test_vmec.py b/tests/test_vmec.py index d7ae22f2b4..0fef594b3c 100644 --- a/tests/test_vmec.py +++ b/tests/test_vmec.py @@ -368,14 +368,6 @@ def test_axis_surf_after_load(): f.close() -@pytest.mark.unit -def test_vmec_save_asym(TmpDir): - """Tests that saving a non-symmetric equilibrium runs without errors.""" - output_path = str(TmpDir.join("output.nc")) - eq = Equilibrium(L=2, M=2, N=2, NFP=3, pressure=np.array([[2, 0]]), sym=False) - VMECIO.save(eq, output_path) - - @pytest.mark.unit def test_vmec_save_kinetic(TmpDir): """Tests that saving an equilibrium with kinetic profiles runs without errors.""" @@ -874,6 +866,369 @@ def test_vmec_save_2(VMEC_save): np.testing.assert_allclose(currv_vmec, currv_desc, rtol=1e-2) +@pytest.mark.regression +@pytest.mark.slow +def test_vmec_save_asym(VMEC_save_asym): + """Tests that saving in NetCDF format agrees with VMEC.""" + vmec, desc, eq = VMEC_save_asym + # first, compare some quantities which don't require calculation + assert vmec.variables["version_"][:] == desc.variables["version_"][:] + assert vmec.variables["mgrid_mode"][:] == desc.variables["mgrid_mode"][:] + assert np.all( + np.char.compare_chararrays( + vmec.variables["mgrid_file"][:], + desc.variables["mgrid_file"][:], + "==", + False, + ) + ) + assert vmec.variables["ier_flag"][:] == desc.variables["ier_flag"][:] + assert ( + vmec.variables["lfreeb__logical__"][:] == desc.variables["lfreeb__logical__"][:] + ) + assert ( + vmec.variables["lrecon__logical__"][:] == desc.variables["lrecon__logical__"][:] + ) + assert vmec.variables["lrfp__logical__"][:] == desc.variables["lrfp__logical__"][:] + assert ( + vmec.variables["lasym__logical__"][:] == desc.variables["lasym__logical__"][:] + ) + assert vmec.variables["nfp"][:] == desc.variables["nfp"][:] + assert vmec.variables["ns"][:] == desc.variables["ns"][:] + assert vmec.variables["mpol"][:] == desc.variables["mpol"][:] + assert vmec.variables["ntor"][:] == desc.variables["ntor"][:] + assert vmec.variables["mnmax"][:] == desc.variables["mnmax"][:] + np.testing.assert_allclose(vmec.variables["xm"][:], desc.variables["xm"][:]) + np.testing.assert_allclose(vmec.variables["xn"][:], desc.variables["xn"][:]) + assert vmec.variables["mnmax_nyq"][:] == desc.variables["mnmax_nyq"][:] + np.testing.assert_allclose(vmec.variables["xm_nyq"][:], desc.variables["xm_nyq"][:]) + np.testing.assert_allclose(vmec.variables["xn_nyq"][:], desc.variables["xn_nyq"][:]) + assert vmec.variables["signgs"][:] == desc.variables["signgs"][:] + assert vmec.variables["gamma"][:] == desc.variables["gamma"][:] + assert vmec.variables["nextcur"][:] == desc.variables["nextcur"][:] + assert np.all( + np.char.compare_chararrays( + vmec.variables["pmass_type"][:], + desc.variables["pmass_type"][:], + "==", + False, + ) + ) + assert np.all( + np.char.compare_chararrays( + vmec.variables["piota_type"][:], + desc.variables["piota_type"][:], + "==", + False, + ) + ) + assert np.all( + np.char.compare_chararrays( + vmec.variables["pcurr_type"][:], + desc.variables["pcurr_type"][:], + "==", + False, + ) + ) + np.testing.assert_allclose( + vmec.variables["am"][:], desc.variables["am"][:], atol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["ai"][:], desc.variables["ai"][:], atol=1e-8 + ) + np.testing.assert_allclose( + vmec.variables["ac"][:], desc.variables["ac"][:], atol=3e-5 + ) + np.testing.assert_allclose( + vmec.variables["presf"][:], desc.variables["presf"][:], atol=2e-5 + ) + np.testing.assert_allclose(vmec.variables["pres"][:], desc.variables["pres"][:]) + np.testing.assert_allclose(vmec.variables["mass"][:], desc.variables["mass"][:]) + np.testing.assert_allclose( + vmec.variables["iotaf"][:], desc.variables["iotaf"][:], rtol=5e-4 + ) + np.testing.assert_allclose( + vmec.variables["q_factor"][:], desc.variables["q_factor"][:], rtol=5e-4 + ) + np.testing.assert_allclose( + vmec.variables["iotas"][:], desc.variables["iotas"][:], rtol=5e-4 + ) + np.testing.assert_allclose(vmec.variables["phi"][:], desc.variables["phi"][:]) + np.testing.assert_allclose(vmec.variables["phipf"][:], desc.variables["phipf"][:]) + np.testing.assert_allclose(vmec.variables["phips"][:], desc.variables["phips"][:]) + np.testing.assert_allclose( + vmec.variables["chi"][:], desc.variables["chi"][:], atol=3e-5, rtol=1e-3 + ) + np.testing.assert_allclose( + vmec.variables["chipf"][:], desc.variables["chipf"][:], atol=3e-5, rtol=1e-3 + ) + np.testing.assert_allclose( + vmec.variables["Rmajor_p"][:], desc.variables["Rmajor_p"][:] + ) + np.testing.assert_allclose( + vmec.variables["Aminor_p"][:], desc.variables["Aminor_p"][:] + ) + np.testing.assert_allclose(vmec.variables["aspect"][:], desc.variables["aspect"][:]) + np.testing.assert_allclose( + vmec.variables["volume_p"][:], desc.variables["volume_p"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["volavgB"][:], desc.variables["volavgB"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["betatotal"][:], desc.variables["betatotal"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["betapol"][:], desc.variables["betapol"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["betator"][:], desc.variables["betator"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["ctor"][:], + desc.variables["ctor"][:], + atol=1e-9, # it is a zero current solve + ) + np.testing.assert_allclose( + vmec.variables["rbtor"][:], desc.variables["rbtor"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["rbtor0"][:], desc.variables["rbtor0"][:], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["b0"][:], desc.variables["b0"][:], rtol=4e-3 + ) + np.testing.assert_allclose( + vmec.variables["buco"][20:100], desc.variables["buco"][20:100], atol=1e-15 + ) + np.testing.assert_allclose( + vmec.variables["bvco"][20:100], desc.variables["bvco"][20:100], rtol=1e-5 + ) + np.testing.assert_allclose( + vmec.variables["vp"][20:100], desc.variables["vp"][20:100], rtol=3e-4 + ) + np.testing.assert_allclose( + vmec.variables["bdotb"][20:100], desc.variables["bdotb"][20:100], rtol=3e-4 + ) + np.testing.assert_allclose( + vmec.variables["jdotb"][20:100], + desc.variables["jdotb"][20:100], + atol=4e-3, # nearly zero bc is vacuum + ) + np.testing.assert_allclose( + vmec.variables["jcuru"][20:100], desc.variables["jcuru"][20:100], atol=2 + ) + np.testing.assert_allclose( + vmec.variables["jcurv"][20:100], desc.variables["jcurv"][20:100], rtol=2 + ) + np.testing.assert_allclose( + vmec.variables["DShear"][20:100], desc.variables["DShear"][20:100], rtol=3e-2 + ) + np.testing.assert_allclose( + vmec.variables["DCurr"][20:100], + desc.variables["DCurr"][20:100], + atol=1e-4, # nearly zero bc vacuum + ) + np.testing.assert_allclose( + vmec.variables["DWell"][20:100], desc.variables["DWell"][20:100], rtol=1e-2 + ) + np.testing.assert_allclose( + vmec.variables["DGeod"][20:100], + desc.variables["DGeod"][20:100], + atol=4e-3, + rtol=1e-2, + ) + + # the Mercier stability is pretty off, + # but these are not exactly similar solutions to eachother + np.testing.assert_allclose( + vmec.variables["DMerc"][20:100], desc.variables["DMerc"][20:100], atol=4e-3 + ) + np.testing.assert_allclose( + vmec.variables["raxis_cc"][:], + desc.variables["raxis_cc"][:], + rtol=5e-5, + atol=4e-3, + ) + np.testing.assert_allclose( + vmec.variables["zaxis_cs"][:], + desc.variables["zaxis_cs"][:], + rtol=5e-5, + atol=1e-3, + ) + np.testing.assert_allclose( + vmec.variables["rmin_surf"][:], desc.variables["rmin_surf"][:], rtol=5e-3 + ) + np.testing.assert_allclose( + vmec.variables["rmax_surf"][:], desc.variables["rmax_surf"][:], rtol=5e-3 + ) + np.testing.assert_allclose( + vmec.variables["zmax_surf"][:], desc.variables["zmax_surf"][:], rtol=5e-3 + ) + np.testing.assert_allclose( + vmec.variables["beta_vol"][:], desc.variables["beta_vol"][:], rtol=5e-5 + ) + np.testing.assert_allclose( + vmec.variables["betaxis"][:], desc.variables["betaxis"][:], rtol=5e-5 + ) + # Next, calculate some quantities and compare + # the DESC wout -> DESC (should be very close) + # and the DESC wout -> VMEC wout (should be approximately close) + vol_grid = LinearGrid( + rho=np.sqrt( + abs( + vmec.variables["phi"][:].filled() + / np.max(np.abs(vmec.variables["phi"][:].filled())) + ) + )[10::10], + M=15, + N=15, + NFP=eq.NFP, + axis=False, + sym=False, + ) + bdry_grid = LinearGrid(rho=1.0, M=15, N=15, NFP=eq.NFP, axis=False, sym=False) + + def test( + nc_str, + desc_str, + negate_DESC_quant=False, + use_nyq=True, + convert_sqrt_g_or_B_rho=False, + atol_desc_desc_wout=5e-5, + rtol_desc_desc_wout=1e-5, + atol_vmec_desc_wout=1e-5, + rtol_vmec_desc_wout=1e-2, + grid=vol_grid, + ): + """Helper fxn to evaluate Fourier series from wout and compare to DESC.""" + xm = desc.variables["xm_nyq"][:] if use_nyq else desc.variables["xm"][:] + xn = desc.variables["xn_nyq"][:] if use_nyq else desc.variables["xn"][:] + + si = abs(vmec.variables["phi"][:] / np.max(np.abs(vmec.variables["phi"][:]))) + rho = grid.nodes[:, 0] + s = rho**2 + # some quantities must be negated before comparison bc + # they are negative in the wout i.e. B^theta + negate = -1 if negate_DESC_quant else 1 + + quant_from_desc_wout = VMECIO.vmec_interpolate( + desc.variables[nc_str + "c"][:], + desc.variables[nc_str + "s"][:], + xm, + xn, + theta=-grid.nodes[:, 1], # -theta bc when we save wout we reverse theta + phi=grid.nodes[:, 2], + s=s, + sym=False, + si=si, + ) + + quant_from_vmec_wout = VMECIO.vmec_interpolate( + vmec.variables[nc_str + "c"][:], + vmec.variables[nc_str + "s"][:], + xm, + xn, + # pi - theta bc VMEC, when it gets a CW angle bdry, + # changes poloidal angle to theta -> pi-theta + theta=np.pi - grid.nodes[:, 1], + phi=grid.nodes[:, 2], + s=s, + sym=False, + si=si, + ) + + data = eq.compute(["rho", "sqrt(g)", desc_str], grid=grid) + # convert sqrt(g) or B_rho->B_psi if needed + quant_desc = ( + data[desc_str] / 2 / data["rho"] + if convert_sqrt_g_or_B_rho + else data[desc_str] + ) + + # add sqrt(g) factor if currents being compared + quant_desc = ( + quant_desc * abs(data["sqrt(g)"]) / 2 / data["rho"] + if "J" in desc_str + else quant_desc + ) + + np.testing.assert_allclose( + negate * quant_desc, + quant_from_desc_wout, + atol=atol_desc_desc_wout, + rtol=rtol_desc_desc_wout, + ) + np.testing.assert_allclose( + quant_from_desc_wout, + quant_from_vmec_wout, + atol=atol_vmec_desc_wout, + rtol=rtol_vmec_desc_wout, + ) + + # R & Z & lambda + test("rmn", "R", use_nyq=False) + test("zmn", "Z", use_nyq=False, atol_vmec_desc_wout=4e-2) + + # |B| + test("bmn", "|B|", rtol_desc_desc_wout=7e-4) + + # B^zeta + test("bsupvmn", "B^zeta") # ,rtol_desc_desc_wout=6e-5) + + # B_zeta + test("bsubvmn", "B_zeta", rtol_desc_desc_wout=3e-4) + + # hard to compare to VMEC for the currents, since + # VMEC F error is worse and equilibria are not exactly similar + # just compare back to DESC + test("currumn", "J^theta", atol_vmec_desc_wout=1e4) + test("currvmn", "J^zeta", negate_DESC_quant=True, atol_vmec_desc_wout=1e5) + + # can only compare lambda, sqrt(g) B_psi B^theta and B_theta at bdry + test( + "lmn", + "lambda", + use_nyq=False, + negate_DESC_quant=True, + grid=bdry_grid, + atol_desc_desc_wout=4e-4, + atol_vmec_desc_wout=5e-2, + ) + test( + "gmn", + "sqrt(g)", + convert_sqrt_g_or_B_rho=True, + negate_DESC_quant=True, + grid=bdry_grid, + rtol_desc_desc_wout=5e-4, + rtol_vmec_desc_wout=4e-2, + ) + test( + "bsupumn", + "B^theta", + negate_DESC_quant=True, + grid=bdry_grid, + atol_vmec_desc_wout=6e-4, + ) + test( + "bsubumn", + "B_theta", + negate_DESC_quant=True, + grid=bdry_grid, + atol_desc_desc_wout=1e-4, + atol_vmec_desc_wout=4e-4, + ) + test( + "bsubsmn", + "B_rho", + grid=bdry_grid, + convert_sqrt_g_or_B_rho=True, + rtol_vmec_desc_wout=6e-2, + atol_vmec_desc_wout=9e-3, + ) + + @pytest.mark.unit @pytest.mark.mpl_image_compare(tolerance=1) def test_plot_vmec_comparison():