address comments

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

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
 
 

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

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