Add Ability to Write Boozmn.nc style output files

CHANGELOG.md

@@ -122,6 +122,9 @@ 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

@@ -179,6 +179,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",
+        "sqrt(g)_Boozer_DESC",
+        "Boozer transform matrix",
+    ],
+)
+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
+    return data
+
+
 @register_compute_fun(
     name="nu_t",
     label="\\partial_{\\theta} \\nu",
@@ -254,27 +279,73 @@ def _zeta_B(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="sqrt(g)_B",
-    label="\\sqrt{g}_{B}",
+    name="sqrt(g)_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 (rho, theta_B, zeta_B)"
+    " to DESC coordinates (rho,theta,zeta).",
     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 _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"]
     return data
 
 
 @register_compute_fun(
-    name="|B|_mn",
+    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=["sqrt(g)_Boozer_DESC", "sqrt(g)"],
+)
+def _sqrtg_B(params, transforms, profiles, data, **kwargs):
+    data["sqrt(g)_B"] = data["sqrt(g)"] / data["sqrt(g)_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",
+        "sqrt(g)_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["sqrt(g)_Boozer_DESC"] * data["sqrt(g)_B"])
+    ) / transforms["B"].grid.num_nodes
+    return data
+
+
+@register_compute_fun(
+    name="|B|_mn_B",
     label="B_{mn}^{\\mathrm{Boozer}}",
     units="T",
     units_long="Tesla",
@@ -284,18 +355,69 @@ 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=[
+        "sqrt(g)_Boozer_DESC",
+        "|B|",
+        "Boozer transform matrix",
+    ],
     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",
 )
 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_B"] = (
+        norm @ (data["sqrt(g)_Boozer_DESC"] * data["|B|"])
+    ) / transforms["B"].grid.num_nodes
+    return data
+
+
+@register_compute_fun(
+    name="R_mn_B",
+    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=[
+        "sqrt(g)_Boozer_DESC",
+        "R",
+        "Boozer transform matrix",
+    ],
+)
+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
+    return data
+
+
+@register_compute_fun(
+    name="Z_mn_B",
+    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=[
+        "sqrt(g)_Boozer_DESC",
+        "Z",
+        "Boozer transform matrix",
+    ],
+)
+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
     return data
 
 
@@ -319,6 +441,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)"

desc/compute/utils.py

@@ -364,11 +364,14 @@ def get_transforms(keys, obj, grid, jitable=False, **kwargs):
         elif c == "B":  # used for Boozer transform
             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=False,
@@ -382,7 +385,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=False,

desc/objectives/_omnigenity.py

@@ -144,7 +144,7 @@ def build(self, use_jit=True, verbose=1):
             "resolution for surface averages",
         )
 
-        self._data_keys = ["|B|_mn"]
+        self._data_keys = ["|B|_mn_B"]
 
         timer = Timer()
         if verbose > 0:
@@ -210,7 +210,7 @@ def compute(self, params, constants=None):
             transforms=constants["transforms"],
             profiles=constants["profiles"],
         )
-        B_mn = constants["matrix"] @ data["|B|_mn"]
+        B_mn = constants["matrix"] @ data["|B|_mn_B"]
         return B_mn[constants["idx"]]
 
     @property
@@ -753,7 +753,7 @@ def build(self, use_jit=True, verbose=1):
             field_grid = self._field_grid
 
         self._dim_f = field_grid.num_nodes
-        self._eq_data_keys = ["|B|_mn"]
+        self._eq_data_keys = ["|B|_mn_B"]
         self._field_data_keys = ["|B|", "theta_B", "zeta_B"]
 
         errorif(
@@ -925,7 +925,7 @@ def compute(self, params_1=None, params_2=None, constants=None):
             )
         ).T
         B_eta_alpha = jnp.matmul(
-            constants["eq_transforms"]["B"].basis.evaluate(nodes), eq_data["|B|_mn"]
+            constants["eq_transforms"]["B"].basis.evaluate(nodes), eq_data["|B|_mn_B"]
         )
         omnigenity_error = B_eta_alpha - field_data["|B|"]
         weights = (self.eta_weight + 1) / 2 + (self.eta_weight - 1) / 2 * jnp.cos(

desc/plotting.py

@@ -2425,7 +2425,7 @@ def plot_boozer_modes(  # noqa: C901
     ylabel_fontsize = kwargs.pop("ylabel_fontsize", None)
 
     basis = get_transforms(
-        "|B|_mn", obj=eq, grid=Grid(np.array([])), M_booz=M_booz, N_booz=N_booz
+        "|B|_mn_B", obj=eq, grid=Grid(np.array([])), M_booz=M_booz, N_booz=N_booz
     )["B"].basis
     if helicity:
         matrix, modes, symidx = ptolemy_linear_transform(
@@ -2437,12 +2437,12 @@ def plot_boozer_modes(  # noqa: C901
     for i, r in enumerate(rho):
         grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=np.array(r))
         transforms = get_transforms(
-            "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz
+            "|B|_mn_B", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz
         )
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            data = eq.compute("|B|_mn", grid=grid, transforms=transforms)
-        b_mn = np.atleast_2d(matrix @ data["|B|_mn"])
+            data = eq.compute("|B|_mn_B", grid=grid, transforms=transforms)
+        b_mn = np.atleast_2d(matrix @ data["|B|_mn_B"])
         B_mn = np.vstack((B_mn, b_mn)) if B_mn.size else b_mn
 
     zidx = np.where((modes[:, 1:] == np.array([[0, 0]])).all(axis=1))[0]
@@ -2493,7 +2493,7 @@ def plot_boozer_modes(  # noqa: C901
             )
 
     plot_data = {
-        "|B|_mn": B_mn,
+        "|B|_mn_B": B_mn,
         "B modes": modes,
         "rho": rho,
     }
@@ -2638,12 +2638,12 @@ def plot_boozer_surface(
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             data = thing.compute(
-                "|B|_mn", grid=grid_compute, M_booz=M_booz, N_booz=N_booz
+                "|B|_mn_B", grid=grid_compute, M_booz=M_booz, N_booz=N_booz
             )
         B_transform = get_transforms(
-            "|B|_mn", obj=thing, grid=grid_plot, M_booz=M_booz, N_booz=N_booz
+            "|B|_mn_B", obj=thing, grid=grid_plot, M_booz=M_booz, N_booz=N_booz
         )["B"]
-        B = B_transform.transform(data["|B|_mn"]).reshape(
+        B = B_transform.transform(data["|B|_mn_B"]).reshape(
             (grid_plot.num_theta, grid_plot.num_zeta), order="F"
         )
         theta_B = (
@@ -2845,7 +2845,7 @@ def plot_qs_error(  # noqa: 16 fxn too complex
         grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=np.array(r))
         if fB:
             transforms = get_transforms(
-                "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz
+                "|B|_mn_B", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz
             )
             if i == 0:  # only need to do this once for the first rho surface
                 matrix, modes, idx = ptolemy_linear_transform(
@@ -2856,9 +2856,9 @@ def plot_qs_error(  # noqa: 16 fxn too complex
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore")
                 data = eq.compute(
-                    ["|B|_mn", "B modes"], grid=grid, transforms=transforms
+                    ["|B|_mn_B", "B modes"], grid=grid, transforms=transforms
                 )
-            B_mn = matrix @ data["|B|_mn"]
+            B_mn = matrix @ data["|B|_mn_B"]
             f_b = np.sqrt(np.sum(B_mn[idx] ** 2)) / np.sqrt(np.sum(B_mn**2))
             f_B = np.append(f_B, f_b)
         if fC: