Add Ability to Write Boozmn.nc style output files

CHANGELOG.md

@@ -3,6 +3,9 @@ Changelog
 
 - 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/_qs.py

@@ -165,6 +165,31 @@ def _nu(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="nu_mn",
+    label="\\nu_{mn} = (\\zeta_{B} - \\zeta)_{mn}",
+    units="rad",
+    units_long="radians",
+    description="Boozer harmonics of Boozer toroidal stream function",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "nu",
+        "jac_Boozer_DESC",
+        "Boozer transform matrix",
+    ],
+)
+def _nu_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform matrix"]
+    data["nu_mn"] = (norm @ (data["jac_Boozer_DESC"] * data["nu"])) / transforms[
+        "B"
+    ].grid.num_nodes
+    return data
+
+
 @register_compute_fun(
     name="nu_t",
     label="\\partial_{\\theta} \\nu",
@@ -240,25 +265,70 @@ def _zeta_B(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="sqrt(g)_B",
-    label="\\sqrt{g}_{B}",
+    name="jac_Boozer_DESC",
+    label="\\frac{\\partial(\\theta_B,\\zeta_B)}{\\theta_{DESC},\\zeta_{DESC}}",
     units="~",
     units_long="None",
-    description="Jacobian determinant of Boozer coordinates",
+    description="Jacobian determinant from Boozer coordinates to DESC coordinates.",
     dim=1,
     params=[],
     transforms={},
     profiles=[],
     coordinates="rtz",
     data=["lambda_t", "lambda_z", "nu_t", "nu_z", "iota"],
 )
-def _sqrtg_B(params, transforms, profiles, data, **kwargs):
-    data["sqrt(g)_B"] = (1 + data["lambda_t"]) * (1 + data["nu_z"]) + (
+def _jac_Boozer_DESC(params, transforms, profiles, data, **kwargs):
+    data["jac_Boozer_DESC"] = (1 + data["lambda_t"]) * (1 + data["nu_z"]) + (
         data["iota"] - data["lambda_z"]
     ) * data["nu_t"]
     return data
 
 
+@register_compute_fun(
+    name="sqrt(g)_B",
+    label="\\sqrt{g}_B",
+    units="~",
+    units_long="None",
+    description="Jacobian determinant from (rho, theta_B, zeta_B)"
+    " Boozer coordinates to (R,phi,Z) lab frame.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["jac_Boozer_DESC", "sqrt(g)"],
+)
+def _sqrtg_B(params, transforms, profiles, data, **kwargs):
+    data["sqrt(g)_B"] = data["sqrt(g)"] / data["jac_Boozer_DESC"]
+    return data
+
+
+@register_compute_fun(
+    name="sqrt(g)_B_mn",
+    label="\\sqrt{g}_{B,mn}",
+    units="~",
+    units_long="None",
+    description="Boozer harmonics of Jacobian determinant from (rho, theta_B, zeta_B)"
+    " Boozer coordinates to (R,phi,Z) lab frame.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "sqrt(g)_B",
+        "jac_Boozer_DESC",
+        "Boozer transform matrix",
+    ],
+)
+def _sqrtg_B_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform matrix"]
+    data["sqrt(g)_B_mn"] = (
+        norm @ (data["jac_Boozer_DESC"] * data["sqrt(g)_B"])
+    ) / transforms["B"].grid.num_nodes
+    return data
+
+
 @register_compute_fun(
     name="|B|_mn",
     label="B_{mn}^{\\mathrm{Boozer}}",
@@ -270,16 +340,67 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs):
     transforms={"B": [[0, 0, 0]]},
     profiles=[],
     coordinates="rtz",
-    data=["sqrt(g)_B", "|B|", "rho", "theta_B", "zeta_B"],
+    data=[
+        "jac_Boozer_DESC",
+        "|B|",
+        "Boozer transform matrix",
+    ],
 )
 def _B_mn(params, transforms, profiles, data, **kwargs):
-    nodes = jnp.array([data["rho"], data["theta_B"], data["zeta_B"]]).T
-    norm = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1))
-    data["|B|_mn"] = (
-        norm  # 1 if m=n=0, 2 if m=0 or n=0, 4 if m!=0 and n!=0
-        * (transforms["B"].basis.evaluate(nodes).T @ (data["sqrt(g)_B"] * data["|B|"]))
-        / transforms["B"].grid.num_nodes
-    )
+    norm = data["Boozer transform matrix"]
+    data["|B|_mn"] = (norm @ (data["jac_Boozer_DESC"] * data["|B|"])) / transforms[
+        "B"
+    ].grid.num_nodes
+    return data
+
+
+@register_compute_fun(
+    name="R_mn",
+    label="R_{mn}^{\\mathrm{Boozer}}",
+    units="m",
+    units_long="meters",
+    description="Boozer harmonics of radial toroidal coordinate of a flux surface",
+    dim=1,
+    params=[],
+    transforms={"B": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "jac_Boozer_DESC",
+        "R",
+        "Boozer transform matrix",
+    ],
+)
+def _R_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform matrix"]
+    data["R_mn"] = (norm @ (data["jac_Boozer_DESC"] * data["R"])) / transforms[
+        "B"
+    ].grid.num_nodes
+    return data
+
+
+@register_compute_fun(
+    name="Z_mn",
+    label="Z_{mn}^{\\mathrm{Boozer}}",
+    units="m",
+    units_long="meters",
+    description="Boozer harmonics of vertical toroidal coordinate of a flux surface",
+    dim=1,
+    params=[],
+    transforms={"B": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "jac_Boozer_DESC",
+        "Z",
+        "Boozer transform matrix",
+    ],
+)
+def _Z_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform matrix"]
+    data["Z_mn"] = (norm @ (data["jac_Boozer_DESC"] * data["Z"])) / transforms[
+        "B"
+    ].grid.num_nodes
     return data
 
 
@@ -301,6 +422,47 @@ def _B_modes(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="Boozer transform modes norm",
+    label="",
+    units="~",
+    units_long="None",
+    description="Inner product norm for boozer modes basis.",
+    dim=1,
+    params=[],
+    transforms={"B": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=[],
+)
+def _boozer_modes_norm(params, transforms, profiles, data, **kwargs):
+    # norm is 1 if m=n=0, 2 if m=0 or n=0, 4 if m!=0 and n!=0
+    norm = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1))
+    data["Boozer transform modes norm"] = norm
+    return data
+
+
+@register_compute_fun(
+    name="Boozer transform matrix",
+    label="",
+    units="~",
+    units_long="None",
+    description="Matrix which performs boozer transformation.",
+    dim=1,
+    params=[],
+    transforms={"B": [[0, 0, 0]]},
+    profiles=[],
+    coordinates="rtz",
+    data=["rho", "theta_B", "zeta_B", "Boozer transform modes norm"],
+)
+def _boozer_transform_matrix(params, transforms, profiles, data, **kwargs):
+    nodes = jnp.array([data["rho"], data["theta_B"], data["zeta_B"]]).T
+    data["Boozer transform matrix"] = (
+        data["Boozer transform modes norm"] * transforms["B"].basis.evaluate(nodes)
+    ).T
+    return data
+
+
 @register_compute_fun(
     name="f_C",
     label="(M \\iota - N) (\\mathbf{B} \\times \\nabla \\psi) \\cdot \\nabla B"

desc/compute/utils.py

@@ -342,11 +342,14 @@ def get_transforms(keys, obj, grid, jitable=False, **kwargs):
         elif c == "B":
             transforms["B"] = Transform(
                 grid,
-                DoubleFourierSeries(
-                    M=kwargs.get("M_booz", 2 * obj.M),
-                    N=kwargs.get("N_booz", 2 * obj.N),
-                    NFP=obj.NFP,
-                    sym=obj.R_basis.sym,
+                kwargs.get(
+                    "B_basis",
+                    DoubleFourierSeries(
+                        M=kwargs.get("M_booz", 2 * obj.M),
+                        N=kwargs.get("N_booz", 2 * obj.N),
+                        NFP=obj.NFP,
+                        sym=kwargs.get("sym", obj.R_basis.sym),
+                    ),
                 ),
                 derivs=derivs["B"],
                 build=True,
@@ -360,7 +363,7 @@ def get_transforms(keys, obj, grid, jitable=False, **kwargs):
                     M=kwargs.get("M_booz", 2 * obj.M),
                     N=kwargs.get("N_booz", 2 * obj.N),
                     NFP=obj.NFP,
-                    sym=obj.Z_basis.sym,
+                    sym=kwargs.get("sym", obj.Z_basis.sym),
                 ),
                 derivs=derivs["w"],
                 build=True,