diff --git a/CHANGELOG.md b/CHANGELOG.md index 348e970cd..d39af8d57 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ New Feature - Adds an option ``scaled_termination`` (defaults to True) to all of the desc optimizers to measure the norms for ``xtol`` and ``gtol`` in the scaled norm provided by ``x_scale`` (which defaults to using an adaptive scaling based on the Jacobian or Hessian). This should make things more robust when optimizing parameters with widely different magnitudes. The old behavior can be recovered by passing ``options={"scaled_termination": False}``. - ``desc.objectives.Omnigenity`` is now vectorized and able to optimize multiple surfaces at the same time. Previously it was required to use a different objective for each surface. - Adds a new objective ``desc.objectives.MirrorRatio`` for targeting a particular mirror ratio on each flux surface, for either an ``Equilibrium`` or ``OmnigenousField``. +- Adds the output quantities ``wb`` and ``wp`` to ``VMECIO.save``. Bug Fixes diff --git a/desc/vmec.py b/desc/vmec.py index 14c4cafff..c6743c8d2 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -289,7 +289,16 @@ def save(cls, eq, path, surfs=128, verbose=1, M_nyq=None, N_nyq=None): # noqa: grid_full = LinearGrid(M=M_nyq, N=N_nyq, NFP=NFP, rho=r_full) data_quad = eq.compute( - ["R0/a", "V", "<|B|>_rms", "_vol", "_vol", "_vol"] + [ + "R0/a", + "V", + "W_B", + "W_p", + "<|B|>_rms", + "_vol", + "_vol", + "_vol", + ] ) data_axis = eq.compute(["G", "p", "R", "<|B|^2>", "<|B|>"], grid=grid_axis) data_lcfs = eq.compute(["G", "I", "R", "Z"], grid=grid_lcfs) @@ -502,6 +511,16 @@ def save(cls, eq, path, surfs=128, verbose=1, M_nyq=None, N_nyq=None): # noqa: betator.units = "None" betator[:] = data_quad["_vol"] + wb = file.createVariable("wb", np.float64) + wb.long_name = "plasma magnetic energy * mu_0/(4*pi^2)" + wb.units = "T^2*m^3" + wb[:] = data_quad["W_B"] * mu_0 / (4 * np.pi**2) + + wp = file.createVariable("wp", np.float64) + wp.long_name = "plasma thermodynamic energy * mu_0/(4*pi^2)" + wp.units = "T^2*m^3" + wp[:] = np.abs(data_quad["W_p"]) * mu_0 / (4 * np.pi**2) + # scalars computed at the magnetic axis rbtor0 = file.createVariable("rbtor0", np.float64) @@ -1338,16 +1357,8 @@ 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 "W_B" - wb = file.createVariable("wb", np.float64) - wb[:] = 0.0 - wdot = file.createVariable("wdot", np.float64, ("time",)) wdot[:] = np.zeros((file.dimensions["time"].size,)) - - # this is not the same as DESC's "W_p" - wp = file.createVariable("wp", np.float64) - wp[:] = 0.0 """ file.close() diff --git a/tests/test_vmec.py b/tests/test_vmec.py index 91b34667a..2edca5c28 100644 --- a/tests/test_vmec.py +++ b/tests/test_vmec.py @@ -496,6 +496,10 @@ def test_vmec_save_1(VMEC_save): np.testing.assert_allclose( vmec.variables["betator"][:], desc.variables["betator"][:], rtol=1e-5 ) + np.testing.assert_allclose(vmec.variables["wb"][:], desc.variables["wb"][:]) + np.testing.assert_allclose( + vmec.variables["wp"][:], desc.variables["wp"][:], rtol=1e-6 + ) np.testing.assert_allclose( vmec.variables["ctor"][:], desc.variables["ctor"][:], rtol=1e-5 )