From 4fed74dc8de4354e932aa2d3d7f89ce2dca2359b Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 21 Aug 2024 15:07:13 -0400 Subject: [PATCH] address comments --- CHANGELOG.md | 6 +++--- desc/compute/_omnigenity.py | 35 ++++++++++++++++++----------------- desc/vmec_utils.py | 5 +++-- tests/test_vmec.py | 1 + 4 files changed, 25 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fba9b55dd4..28c8126356 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ 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. +- Adds ``desc.vmec_utils.make_boozmn_output `` for writing boozmn.nc style output files +for compatibility with other codes which expect such files from the Booz_Xform code. +- Renames compute quantity ``sqrt(g)_B`` to ``sqrt(g)_Boozer_DESC`` to more accurately reflect what the quantiy is (the jacobian from (rho,theta_B,zeta_B) to (rho,theta,zeta)), and adds a new function to compute ``sqrt(g)_B`` which is the jacobian from (rho,theta_B,zeta_B) to (R,phi,Z). v0.12.1 ------- @@ -201,9 +204,6 @@ v0.10.3 use forward or reverse mode automatic differentiation. - Adds ``desc.compat.rescale`` for rescaling equilibria to a specified size and field strength. -- Adds ``desc.vmec_utils.make_boozmn_output `` for writing boozmn.nc style output files -for compatibility with other codes which expect such files from the Booz_Xform code. -- Renames compute quantity ``sqrt(g)_B`` to ``jac_Boozer_DESC`` to more accurately reflect what the quantiy is (the jacobian from (rho,theta_B,zeta_B) to (rho,theta,zeta)), and adds a new function to compute ``sqrt(g)_B`` which is the jacobian from (rho,theta_B,zeta_B) to (R,phi,Z) - Adds new keyword ``surface_fixed`` to ``PlasmaVesselDistance`` objective which says whether or not the surface comparing the distance from the plasma to is fixed or not. If True, then the surface coordinates can be precomputed, saving on computation during diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index e74d63e1e1..a38ee3eb9a 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -202,10 +202,10 @@ def _nu(params, transforms, profiles, data, **kwargs): ], ) def _nu_mn(params, transforms, profiles, data, **kwargs): - norm = data["Boozer transform matrix"] - data["nu_mn"] = (norm @ (data["sqrt(g)_Boozer_DESC"] * data["nu"])) / transforms[ - "B" - ].grid.num_nodes + transform_matrix = data["Boozer transform matrix"] + data["nu_mn"] = ( + transform_matrix @ (data["sqrt(g)_Boozer_DESC"] * data["nu"]) + ) / transforms["B"].grid.num_nodes return data @@ -297,7 +297,7 @@ def _zeta_B(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["lambda_t", "lambda_z", "nu_t", "nu_z", "iota"], ) -def _sqrt_g__Boozer_DESC(params, transforms, profiles, data, **kwargs): +def _sqrt_g_Boozer_DESC(params, transforms, profiles, data, **kwargs): data["sqrt(g)_Boozer_DESC"] = (1 + data["lambda_t"]) * (1 + data["nu_z"]) + ( data["iota"] - data["lambda_z"] ) * data["nu_t"] @@ -342,9 +342,9 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs): ], ) def _sqrtg_B_mn(params, transforms, profiles, data, **kwargs): - norm = data["Boozer transform matrix"] + transform_matrix = data["Boozer transform matrix"] data["sqrt(g)_B_mn"] = ( - norm @ (data["sqrt(g)_Boozer_DESC"] * data["sqrt(g)_B"]) + transform_matrix @ (data["sqrt(g)_Boozer_DESC"] * data["sqrt(g)_B"]) ) / transforms["B"].grid.num_nodes return data @@ -367,11 +367,12 @@ def _sqrtg_B_mn(params, transforms, profiles, data, **kwargs): ], M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", + aliases=["|B|_mn"], ) def _B_mn(params, transforms, profiles, data, **kwargs): - norm = data["Boozer transform matrix"] + transform_matrix = data["Boozer transform matrix"] data["|B|_mn_B"] = ( - norm @ (data["sqrt(g)_Boozer_DESC"] * data["|B|"]) + transform_matrix @ (data["sqrt(g)_Boozer_DESC"] * data["|B|"]) ) / transforms["B"].grid.num_nodes return data @@ -394,10 +395,10 @@ def _B_mn(params, transforms, profiles, data, **kwargs): ], ) def _R_mn(params, transforms, profiles, data, **kwargs): - norm = data["Boozer transform matrix"] - data["R_mn_B"] = (norm @ (data["sqrt(g)_Boozer_DESC"] * data["R"])) / transforms[ - "B" - ].grid.num_nodes + transform_matrix = data["Boozer transform matrix"] + data["R_mn_B"] = ( + transform_matrix @ (data["sqrt(g)_Boozer_DESC"] * data["R"]) + ) / transforms["B"].grid.num_nodes return data @@ -419,10 +420,10 @@ def _R_mn(params, transforms, profiles, data, **kwargs): ], ) def _Z_mn(params, transforms, profiles, data, **kwargs): - norm = data["Boozer transform matrix"] - data["Z_mn_B"] = (norm @ (data["sqrt(g)_Boozer_DESC"] * data["Z"])) / transforms[ - "B" - ].grid.num_nodes + transform_matrix = data["Boozer transform matrix"] + data["Z_mn_B"] = ( + transform_matrix @ (data["sqrt(g)_Boozer_DESC"] * data["Z"]) + ) / transforms["B"].grid.num_nodes return data diff --git a/desc/vmec_utils.py b/desc/vmec_utils.py index 320fe4706d..fbd7677c9f 100644 --- a/desc/vmec_utils.py +++ b/desc/vmec_utils.py @@ -540,6 +540,9 @@ def make_boozmn_output( # noqa: 16 fxn too complex r_full = np.sqrt(s_full) r_half = np.sqrt(s_half) + # rho=1.0 here is a dummy value to ensure we just have one surface, + # as we are only using the DoubleFourierSeries transform which has + # no rho dependence grid = LinearGrid(M=2 * M_booz, N=2 * N_booz, NFP=eq.NFP, rho=1.0, sym=False) transforms = get_transforms( @@ -790,8 +793,6 @@ def compute_data_sin_sym(grid, data): lasym.long_name = "0 if the configuration is stellarator-symmetric, 1 if not" lasym[:] = not eq.sym - # FIXME: ns? - ns = file.createVariable("ns_b", np.int32) ns.long_name = "Number of radial surfaces at which data is outputted minus 1" ns[:] = surfs diff --git a/tests/test_vmec.py b/tests/test_vmec.py index 5f3b68310d..463dc97c7e 100644 --- a/tests/test_vmec.py +++ b/tests/test_vmec.py @@ -1384,6 +1384,7 @@ def test_make_boozmn_output_against_hidden_symmetries_booz_xform(TmpDir): "rmin_b", "pres_b", "phip_b", + "ns_b", ]: np.testing.assert_allclose( file.variables[key], file_cpp.variables[key], err_msg=key