VMEC wout fixes

desc/magnetic_fields.py

@@ -469,7 +469,7 @@ def save_mgrid(
         nfp[:] = NFP
 
         nextcur = file.createVariable("nextcur", np.int32)
-        nextcur.long_name = "Number of coils."
+        nextcur.long_name = "Number of coils (external currents)."
         nextcur[:] = 1
 
         rmin = file.createVariable("rmin", np.float64)

desc/vmec.py

@@ -271,6 +271,7 @@ def save(cls, eq, path, surfs=128, verbose=1):  # noqa: C901 - FIXME - simplify
             np.array([" " * 100], "S" + str(file.dimensions["dim_00100"].size))
         )  # VMEC input filename: input.[input_extension]
 
+        # TODO: instead of hard-coding for fixed-boundary, allow for free-boundary?
         mgrid_mode = file.createVariable("mgrid_mode", "S1", ("dim_00001",))
         mgrid_mode[:] = stringtochar(
             np.array([""], "S" + str(file.dimensions["dim_00001"].size))
@@ -360,6 +361,10 @@ def save(cls, eq, path, surfs=128, verbose=1):  # noqa: C901 - FIXME - simplify
         gamma.long_name = "compressibility index (0 = pressure prescribed)"
         gamma[:] = 0
 
+        nextcur = file.createVariable("nextcur", np.int32)
+        nextcur.long_name = "number of coils (external currents)"
+        nextcur[:] = 0
+
         # profiles
         # TODO: add option for saving spline profiles
         power_series = stringtochar(
@@ -476,70 +481,83 @@ def save(cls, eq, path, surfs=128, verbose=1):  # noqa: C901 - FIXME - simplify
         chipf.long_name = "d(chi)/ds: poloidal flux derivative"
         chipf[:] = phipf[:] * iotaf[:]
 
-        # geometric scalars
+        # scalars computed on a Quadrature grid
+        data = eq.compute(
+            ["R0/a", "V", "<|B|>_rms", "<beta>_vol", "<beta_pol>_vol", "<beta_tor>_vol"]
+        )
+
         Rmajor_p = file.createVariable("Rmajor_p", np.float64)
         Rmajor_p.long_name = "major radius"
         Rmajor_p.units = "m"
-        Rmajor_p[:] = eq.compute("R0")["R0"]
+        Rmajor_p[:] = data["R0"]
 
         Aminor_p = file.createVariable("Aminor_p", np.float64)
         Aminor_p.long_name = "minor radius"
         Aminor_p.units = "m"
-        Aminor_p[:] = eq.compute("a")["a"]
+        Aminor_p[:] = data["a"]
 
         aspect = file.createVariable("aspect", np.float64)
         aspect.long_name = "aspect ratio = R_major / A_minor"
         aspect.units = "None"
-        aspect[:] = eq.compute("R0/a")["R0/a"]
+        aspect[:] = data["R0/a"]
 
         volume_p = file.createVariable("volume_p", np.float64)
         volume_p.long_name = "plasma volume"
         volume_p.units = "m^3"
-        volume_p[:] = eq.compute("V")["V"]
+        volume_p[:] = data["V"]
 
         volavgB = file.createVariable("volavgB", np.float64)
         volavgB.long_name = "volume average magnetic field, root mean square"
         volavgB.units = "T"
-        volavgB[:] = eq.compute("<|B|>_rms")["<|B|>_rms"]
+        volavgB[:] = data["<|B|>_rms"]
 
         betatotal = file.createVariable("betatotal", np.float64)
         betatotal.long_name = "normalized plasma pressure"
         betatotal.units = "None"
-        betatotal[:] = eq.compute("<beta>_vol")["<beta>_vol"]
+        betatotal[:] = data["<beta>_vol"]
 
         betapol = file.createVariable("betapol", np.float64)
         betapol.long_name = "normalized poloidal plasma pressure"
         betapol.units = "None"
-        betapol[:] = eq.compute("<beta_pol>_vol")["<beta_pol>_vol"]
+        betapol[:] = data["<beta_pol>_vol"]
 
         betator = file.createVariable("betator", np.float64)
         betator.long_name = "normalized toroidal plasma pressure"
         betator.units = "None"
-        betator[:] = eq.compute("<beta_tor>_vol")["<beta_tor>_vol"]
+        betator[:] = data["<beta_tor>_vol"]
+
+        # scalars computed at the last closed flux surface
+        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=[1.0], NFP=NFP)
+        data = eq.compute(["G", "I"], grid=grid)
 
         ctor = file.createVariable("ctor", np.float64)
         ctor.long_name = "total toroidal plasma current"
         ctor.units = "A"
-        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=[1.0], NFP=NFP)
-        ctor[:] = eq.compute("I", grid=grid)["I"][0] * 2 * np.pi / mu_0
+        ctor[:] = data["I"][0] * 2 * np.pi / mu_0
 
         rbtor = file.createVariable("rbtor", np.float64)
         rbtor.long_name = "<R*B_tor> on last closed flux surface"
         rbtor.units = "T*m"
-        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=[1.0], NFP=NFP)
-        rbtor[:] = eq.compute("G", grid=grid)["G"][0]
+        rbtor[:] = data["G"][0]
+
+        # scalars computed at the magnetic axis
+        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=ONAXIS, NFP=NFP)
+        data = eq.compute(["G", "p", "R", "<|B|^2>"], grid=grid)
 
         rbtor0 = file.createVariable("rbtor0", np.float64)
         rbtor0.long_name = "<R*B_tor> on axis"
         rbtor0.units = "T*m"
-        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=ONAXIS, NFP=NFP)
-        rbtor0[:] = eq.compute("G", grid=grid)["G"][0]
+        rbtor0[:] = data["G"][0]
 
         b0 = file.createVariable("b0", np.float64)
         b0.long_name = "average B_tor on axis"
         b0.units = "T"
-        grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, rho=ONAXIS, NFP=NFP)
-        b0[:] = eq.compute("G", grid=grid)["G"][0] / eq.compute("R", grid=grid)["R"][0]
+        b0[:] = data["G"][0] / data["R"][0]
+
+        betaxis = file.createVariable("betaxis", np.float64)
+        betaxis.long_name = "2 * mu_0 * pressure / <|B|^2> on the magnetic axis"
+        betaxis.units = "None"
+        betaxis[:] = 2 * mu_0 * data["p"][0] / data["<|B|^2>"][0]
 
         # grids for computing radial profile data
         grid_full = LinearGrid(
@@ -552,19 +570,15 @@ def save(cls, eq, path, surfs=128, verbose=1):  # noqa: C901 - FIXME - simplify
             ["<|B|^2>", "<J*B>", "sqrt(g)", "J^theta*sqrt(g)", "J^zeta", "D_Mercier"],
             grid=grid_full,
         )
-        data_half = eq.compute(["I", "G"], grid=grid_half)
+        data_half = eq.compute(["I", "G", "<|B|^2>", "V_r(r)"], grid=grid_half)
 
-        bdotb = file.createVariable("bdotb", np.float64, ("radius",))
-        bdotb.long_name = "flux surface average of magnetic field squared, on full mesh"
-        bdotb.units = "T^2"
-        bdotb[:] = grid_full.compress(data_full["<|B|^2>"])
-        bdotb[0] = 0
-
-        # currents
+        # half mesh quantities
         buco = file.createVariable("buco", np.float64, ("radius",))
         buco.long_name = "Boozer toroidal current I, on half mesh"
         buco.units = "T*m"
-        buco[1:] = grid_half.compress(data_half["I"])
+        buco[1:] = -grid_half.compress(
+            data_half["I"]
+        )  # negative sign for negative Jacobian (opposite poloidal angle)
         buco[0] = 0
 
         bvco = file.createVariable("bvco", np.float64, ("radius",))
@@ -573,6 +587,27 @@ def save(cls, eq, path, surfs=128, verbose=1):  # noqa: C901 - FIXME - simplify
         bvco[1:] = grid_half.compress(data_half["G"])
         bvco[0] = 0
 
+        beta_vol = file.createVariable("beta_vol", np.float64, ("radius",))
+        beta_vol.long_name = "2 * mu_0 * pressure / <|B|^2>, on half mesh"
+        beta_vol.units = "None"
+        beta_vol[1:] = 2 * mu_0 * p_half / grid_half.compress(data_half["<|B|^2>"])
+        beta_vol[0] = 0.0
+
+        vp = file.createVariable("vp", np.float64, ("radius",))
+        vp.long_name = "dV/ds normalized by 4*pi^2, on half mesh"
+        vp.units = "m^3"
+        vp[1:] = grid_half.compress(data_half["V_r(r)"]) / (
+            8 * np.pi**2 * grid_half.compress(data_half["rho"])
+        )
+        vp[0] = 0
+
+        # full mesh quantities
+        bdotb = file.createVariable("bdotb", np.float64, ("radius",))
+        bdotb.long_name = "flux surface average of magnetic field squared, on full mesh"
+        bdotb.units = "T^2"
+        bdotb[:] = grid_full.compress(data_full["<|B|^2>"])
+        bdotb[0] = 0
+
         jdotb = file.createVariable("jdotb", np.float64, ("radius",))
         jdotb.long_name = "flux surface average of J*B, on full mesh"
         jdotb.units = "N/m^3"
@@ -582,12 +617,12 @@ def save(cls, eq, path, surfs=128, verbose=1):  # noqa: C901 - FIXME - simplify
         jcuru = file.createVariable("jcuru", np.float64, ("radius",))
         jcuru.long_name = "flux surface average of sqrt(g)*J^theta, on full mesh"
         jcuru.units = "A/m^3"
-        jcuru[:] = surface_averages(
+        jcuru[:] = -surface_averages(
             grid_full,
             data_full["J^theta*sqrt(g)"] / (2 * data_full["rho"]),
             sqrt_g=data_full["sqrt(g)"],
             expand_out=False,
-        )
+        )  # negative sign for negative Jacobian (opposite poloidal angle)
         jcuru[0] = 0
 
         jcurv = file.createVariable("jcurv", np.float64, ("radius",))
@@ -1008,7 +1043,7 @@ def fullfit(x):
         bsubsmns[0, :] = (  # linear extrapolation for coefficient at the magnetic axis
             s[1, :] - (s[2, :] - s[1, :]) / (s_full[2] - s_full[1]) * s_full[1]
         )
-        # Todo: evaluate current at rho=0 nodes instead of extrapolation
+        # TODO: evaluate current at rho=0 nodes instead of extrapolation
         if not eq.sym:
             bsubsmnc[:, :] = c
             bsubsmnc[0, :] = (
@@ -1144,7 +1179,7 @@ def fullfit(x):
         currumnc[0, :] = (  # linear extrapolation for coefficient at the magnetic axis
             s[1, :] - (c[2, :] - c[1, :]) / (s_full[2] - s_full[1]) * s_full[1]
         )
-        # Todo: evaluate current at rho=0 nodes instead of extrapolation
+        # TODO: evaluate current at rho=0 nodes instead of extrapolation
         if not eq.sym:
             currumns[:, :] = s
             currumns[0, :] = (
@@ -1198,7 +1233,7 @@ def fullfit(x):
         currvmnc[0, :] = -(  # linear extrapolation for coefficient at the magnetic axis
             s[1, :] - (c[2, :] - c[1, :]) / (s_full[2] - s_full[1]) * s_full[1]
         )
-        # Todo: evaluate current at rho=0 nodes instead of extrapolation
+        # TODO: evaluate current at rho=0 nodes instead of extrapolation
         if not eq.sym:
             currvmns[:, :] = -s
             currumns[0, :] = -(
@@ -1208,21 +1243,6 @@ def fullfit(x):
         if verbose > 1:
             timer.disp("J^zeta*sqrt(g)")
 
-        # beta_vol = 2 μ₀ p / ⟨|B|^2⟩
-        beta_vol = file.createVariable("beta_vol", np.float64, ("radius",))
-        beta_vol[0] = 0.0
-        beta_vol[1:] = 2 * mu_0 * p_half / half_grid.compress(data_half_grid["<|B|^2>"])
-        beta_vol.long_name = "2 * mu_0 * pressure / <|B|^2>, on half mesh"
-        beta_vol.units = "None"
-
-        # betaxis = beta_vol at the magnetic axis
-        betaxis = file.createVariable("betaxis", np.float64)
-        betaxis[:] = (
-            2 * mu_0 * p_full[0] / full_grid.compress(data_full_grid["<|B|^2>"])[0]
-        )
-        betaxis.long_name = "2 * mu_0 * pressure / <|B|^2> on the magnetic axis"
-        betaxis.units = "None"
-
         # TODO: these output quantities need to be added
         bdotgradv = file.createVariable("bdotgradv", np.float64, ("radius",))
         bdotgradv[:] = np.zeros((file.dimensions["radius"].size,))
@@ -1251,7 +1271,8 @@ def fullfit(x):
         am_aux_s[:] = -np.ones((file.dimensions["ndfmax"].size,))
 
         extcur = file.createVariable("extcur", np.float64)
-        extcur[:] = 0.0
+        extcur.long_name = "external current?"
+        extcur[:] = np.nan  # VMEC gives nothing for fixed-boundary solutions
 
         fsql = file.createVariable("fsql", np.float64)
         fsql[:] = 1e-16
@@ -1271,9 +1292,6 @@ def fullfit(x):
         itfsq = file.createVariable("itfsq", np.int32)
         itfsq[:] = 1
 
-        nextcur = file.createVariable("nextcur", np.int32)
-        nextcur[:] = 0
-
         niter = file.createVariable("niter", np.int32)
         niter[:] = 1
 
@@ -1283,10 +1301,6 @@ def fullfit(x):
         specw = file.createVariable("specw", np.float64, ("radius",))
         specw[:] = np.zeros((file.dimensions["radius"].size,))
 
-        # this is not the same as DESC's "V(r)"
-        vp = file.createVariable("vp", np.float64, ("radius",))
-        vp[:] = np.zeros((file.dimensions["radius"].size,))
-
         # this is not the same as DESC's "W_B"
         wb = file.createVariable("wb", np.float64)
         wb[:] = 0.0

tests/test_vmec.py

@@ -430,6 +430,7 @@ def test_vmec_save_1(VMEC_save):
     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"][:],
@@ -515,34 +516,25 @@ def test_vmec_save_1(VMEC_save):
         vmec.variables["b0"][:], desc.variables["b0"][:], rtol=5e-5
     )
     np.testing.assert_allclose(
-        np.abs(vmec.variables["bdotb"][20:100]),
-        np.abs(desc.variables["bdotb"][20:100]),
-        rtol=1e-6,
+        vmec.variables["buco"][20:100], desc.variables["buco"][20:100], rtol=1e-5
     )
     np.testing.assert_allclose(
-        np.abs(vmec.variables["buco"][20:100]),
-        np.abs(desc.variables["buco"][20:100]),
-        rtol=1e-5,
+        vmec.variables["bvco"][20:100], desc.variables["bvco"][20:100], rtol=1e-5
     )
     np.testing.assert_allclose(
-        np.abs(vmec.variables["bvco"][20:100]),
-        np.abs(desc.variables["bvco"][20:100]),
-        rtol=1e-5,
+        vmec.variables["vp"][20:100], desc.variables["vp"][20:100], rtol=1e-6
     )
     np.testing.assert_allclose(
-        np.abs(vmec.variables["jdotb"][20:100]),
-        np.abs(desc.variables["jdotb"][20:100]),
-        rtol=1e-5,
+        vmec.variables["bdotb"][20:100], desc.variables["bdotb"][20:100], rtol=1e-6
     )
     np.testing.assert_allclose(
-        np.abs(vmec.variables["jcuru"][20:100]),
-        np.abs(desc.variables["jcuru"][20:100]),
-        rtol=1e-2,
+        vmec.variables["jdotb"][20:100], desc.variables["jdotb"][20:100], rtol=1e-5
     )
     np.testing.assert_allclose(
-        np.abs(vmec.variables["jcurv"][20:100]),
-        np.abs(desc.variables["jcurv"][20:100]),
-        rtol=3e-2,
+        vmec.variables["jcuru"][20:100], desc.variables["jcuru"][20:100], rtol=1e-2
+    )
+    np.testing.assert_allclose(
+        vmec.variables["jcurv"][20:100], desc.variables["jcurv"][20:100], rtol=3e-2
     )
     np.testing.assert_allclose(
         vmec.variables["DShear"][20:100], desc.variables["DShear"][20:100], rtol=1e-2