From 360ebfdf6f3da0f4c27b9e6023c9f16b1b2c6de8 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 16 Dec 2024 16:31:44 -0700 Subject: [PATCH 1/4] add wb and wp to wout --- desc/vmec.py | 29 ++++++++++++++++++++--------- tests/test_vmec.py | 4 ++++ 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/desc/vmec.py b/desc/vmec.py index 14c4cafff..5a243a350 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 = "J^2/A^2" + 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 = "J^2/A^2" + 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 ) From a326cdb8f41a5fe47ebd34dbadc17dbc9efe802c Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Mon, 16 Dec 2024 17:05:38 -0700 Subject: [PATCH 2/4] update Changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) 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 From a647c1032d43eee0036724cf3fe10cbed9191bba Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 17 Dec 2024 11:19:29 -0700 Subject: [PATCH 3/4] update Perlmutter install docs --- docs/installation.rst | 23 ++++++----------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index 9df67563b..403b1c658 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -96,17 +96,18 @@ specific JAX GPU installation instructions, as that is the main installation dif **Note that DESC does not always test on or guarantee support of the latest version of JAX (which does not have a stable 1.0 release yet), and thus older versions of GPU-accelerated versions of JAX may need to be installed, which may in turn require lower versions of JaxLib, as well as CUDA and CuDNN.** + Perlmutter (NERSC) ++++++++++++++++++++++++++++++ -These instructions were tested and confirmed to work on the Perlmutter supercomputer at NERSC on June 18, 2024. +These instructions were tested and confirmed to work on the Perlmutter supercomputer at NERSC on December 17, 2024. Set up the correct cuda environment for jax installation .. code-block:: sh - module load cudatoolkit/12.2 + module load cudatoolkit/12.4 module load cudnn/8.9.3_cuda12 - module load python + module load python/3.11 Check that you have loaded these modules @@ -118,21 +119,9 @@ Create a conda environment for DESC (`following these instructions = 1.7.0, < 2.0.0 - -to - -.. code-block:: sh - - scipy >= 1.7.0, <= 1.11.3 + pip install --upgrade "jax[cuda12]" Clone and install DESC From 22174d03f951cb400917b0f6f075ea7f0553576f Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 17 Dec 2024 15:28:25 -0700 Subject: [PATCH 4/4] change units --- desc/vmec.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/vmec.py b/desc/vmec.py index 5a243a350..c6743c8d2 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -513,12 +513,12 @@ def save(cls, eq, path, surfs=128, verbose=1, M_nyq=None, N_nyq=None): # noqa: wb = file.createVariable("wb", np.float64) wb.long_name = "plasma magnetic energy * mu_0/(4*pi^2)" - wb.units = "J^2/A^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 = "J^2/A^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