Add Ability to Write Boozmn.nc style output files

CHANGELOG.md

@@ -21,6 +21,10 @@ New Features
 - Adds batching feature to singular integrals.
 - ``desc.objectives.CoilSetMinDistance`` and ``desc.objectives.PlasmaCoilSetMinDistance`` now include the option to use a softmin which can give smoother gradients. They also both now have a ``dist_chunk_size`` option to break up the distance calculation into smaller pieces to save memory
 - Adds a new function ``desc.coils.initialize_helical_coils`` for creating an initial guess for stage 2 helical coil optimization.
+- 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)_Boozer`` which is the jacobian from (rho,theta_B,zeta_B) to (R,phi,Z).
+- Allows specification of Nyquist spectrum maximum modenumbers when using ``VMECIO.save`` to save a DESC .h5 file as a VMEC-format wout file
 
 Bug Fixes
 

CONTRIBUTING.rst

@@ -169,6 +169,9 @@ If the 2nd case is the reason, then you must update the ``master_compute_data_rp
 
 -  First, run the test with ``pytest tests -k test_compute_everything`` and inspect the compute quantities whose values are in error, to ensure that only the quantities you expect to be different are shown (and that the new values are indeed the correct ones, you should have a test elsewhere for that though).
 -  If the values are as expected and only the expected compute quantities are different, then replace the block
+-  Then, if you have re-defined a variable, consider adding a ``DeprecationWarning`` to it to allow a grace period so old code
+   does not break, but is informed of the change. Add the old name an alias in the ``aliases`` part of the ``register_compute_function`` decorator,
+   and add to the ``deprecated_names`` dictionary in ``desc/compute/data_index.py`` the deprecated name as a key and the corresponding value as the new name for that variable.
 
 .. code-block:: python
 

desc/compute/_omnigenity.py

@@ -231,6 +231,62 @@ def _nu(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="nu_B_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={"B": [[0, 0, 0]], "grid": []},
+    profiles=[],
+    coordinates="rtz",
+    data=[
+        "sqrt(g)_Boozer_DESC",
+        "nu",
+        "rho",
+        "theta_B",
+        "zeta_B",
+        "Boozer transform modes norm",
+    ],
+    resolution_requirement="tz",
+    grid_requirement={"is_meshgrid": True},
+    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 _nu_B_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform modes norm"]
+    grid = transforms["grid"]
+
+    def fun(rho, theta_B, zeta_B, sqrtg_B_desc, quant):
+        # this fits Boozer modes on a single surface
+        nodes = jnp.array([rho, theta_B, zeta_B]).T
+        quant_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 @ (sqrtg_B_desc * quant))
+            / transforms["B"].grid.num_nodes
+        )
+        return quant_mn
+
+    def reshape(x):
+        return grid.meshgrid_reshape(x, "rtz").reshape((grid.num_rho, -1))
+
+    rho, theta_B, zeta_B, sqrtg_B_desc, nu = map(
+        reshape,
+        (
+            data["rho"],
+            data["theta_B"],
+            data["zeta_B"],
+            data["sqrt(g)_Boozer_DESC"],
+            data["nu"],
+        ),
+    )
+    nu_B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, nu)
+    data["nu_B_mn"] = nu_B_mn.flatten()
+    return data
+
+
 @register_compute_fun(
     name="nu_t",
     label="\\partial_{\\theta} \\nu",
@@ -306,20 +362,22 @@ 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 from Boozer to DESC 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=["theta_PEST_t", "theta_PEST_z", "phi_t", "phi_z", "nu_t", "nu_z", "iota"],
+    aliases=["sqrt(g)_B"],
 )
-def _sqrtg_B(params, transforms, profiles, data, **kwargs):
-    data["sqrt(g)_B"] = (
+def _sqrt_g_Boozer_DESC(params, transforms, profiles, data, **kwargs):
+    data["sqrt(g)_Boozer_DESC"] = (
         data["theta_PEST_t"] * (data["phi_z"] + data["nu_z"])
         - data["theta_PEST_z"] * (data["phi_t"] + data["nu_t"])
         + data["iota"] * (data["nu_t"] * data["phi_z"] - data["nu_z"] * data["phi_t"])
@@ -328,7 +386,83 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="|B|_mn",
+    name="sqrt(g)_Boozer",
+    label="\\sqrt{g}_Boozer",
+    units="m^{3}",
+    units_long="cubic meters",
+    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)_Boozer"] = data["sqrt(g)"] / data["sqrt(g)_Boozer_DESC"]
+    return data
+
+
+@register_compute_fun(
+    name="sqrt(g)_Boozer_mn",
+    label="\\sqrt{g}_{B,mn}",
+    units="m^{3}",
+    units_long="cubic meters",
+    description="Boozer harmonics of Jacobian determinant from (rho, theta_B, zeta_B)"
+    " Boozer coordinates to (R,phi,Z) lab frame.",
+    dim=1,
+    params=[],
+    transforms={"B": [[0, 0, 0]], "grid": []},
+    profiles=[],
+    coordinates="rtz",
+    resolution_requirement="tz",
+    grid_requirement={"is_meshgrid": True},
+    data=[
+        "sqrt(g)_Boozer",
+        "sqrt(g)_Boozer_DESC",
+        "rho",
+        "theta_B",
+        "zeta_B",
+        "Boozer transform modes norm",
+    ],
+    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 _sqrtg_Boozer_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform modes norm"]
+    grid = transforms["grid"]
+
+    def fun(rho, theta_B, zeta_B, sqrtg_B_desc, quant):
+        # this fits Boozer modes on a single surface
+        nodes = jnp.array([rho, theta_B, zeta_B]).T
+        quant_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 @ (sqrtg_B_desc * quant))
+            / transforms["B"].grid.num_nodes
+        )
+        return quant_mn
+
+    def reshape(x):
+        return grid.meshgrid_reshape(x, "rtz").reshape((grid.num_rho, -1))
+
+    rho, theta_B, zeta_B, sqrtg_B_desc, sqrtg_B = map(
+        reshape,
+        (
+            data["rho"],
+            data["theta_B"],
+            data["zeta_B"],
+            data["sqrt(g)_Boozer_DESC"],
+            data["sqrt(g)_Boozer"],
+        ),
+    )
+    sqrtg_B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, sqrtg_B)
+    data["sqrt(g)_Boozer_mn"] = sqrtg_B_mn.flatten()
+    return data
+
+
+@register_compute_fun(
+    name="|B|_mn_B",
     label="B_{mn}^{\\mathrm{Boozer}}",
     units="T",
     units_long="Tesla",
@@ -338,35 +472,161 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs):
     transforms={"B": [[0, 0, 0]], "grid": []},
     profiles=[],
     coordinates="rtz",
-    data=["sqrt(g)_B", "|B|", "rho", "theta_B", "zeta_B"],
+    data=[
+        "sqrt(g)_Boozer_DESC",
+        "|B|",
+        "rho",
+        "theta_B",
+        "zeta_B",
+        "Boozer transform modes norm",
+    ],
     resolution_requirement="tz",
     grid_requirement={"is_meshgrid": True},
     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 = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1))
+    norm = data["Boozer transform modes norm"]
     grid = transforms["grid"]
 
-    def fun(rho, theta_B, zeta_B, sqrtg_B, B):
+    def fun(rho, theta_B, zeta_B, sqrtg_B_desc, quant):
         # this fits Boozer modes on a single surface
         nodes = jnp.array([rho, theta_B, zeta_B]).T
         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 @ (sqrtg_B * B))
+            * (transforms["B"].basis.evaluate(nodes).T @ (sqrtg_B_desc * quant))
             / transforms["B"].grid.num_nodes
         )
         return B_mn
 
     def reshape(x):
         return grid.meshgrid_reshape(x, "rtz").reshape((grid.num_rho, -1))
 
-    rho, theta_B, zeta_B, sqrtg_B, B = map(
+    rho, theta_B, zeta_B, sqrtg_B_desc, B = map(
         reshape,
-        (data["rho"], data["theta_B"], data["zeta_B"], data["sqrt(g)_B"], data["|B|"]),
+        (
+            data["rho"],
+            data["theta_B"],
+            data["zeta_B"],
+            data["sqrt(g)_Boozer_DESC"],
+            data["|B|"],
+        ),
+    )
+    B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, B)
+    data["|B|_mn_B"] = B_mn.flatten()
+    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]], "grid": []},
+    profiles=[],
+    coordinates="rtz",
+    resolution_requirement="tz",
+    grid_requirement={"is_meshgrid": True},
+    data=[
+        "R",
+        "sqrt(g)_Boozer_DESC",
+        "rho",
+        "theta_B",
+        "zeta_B",
+        "Boozer transform modes norm",
+    ],
+    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 _R_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform modes norm"]
+    grid = transforms["grid"]
+
+    def fun(rho, theta_B, zeta_B, sqrtg_B_desc, quant):
+        # this fits Boozer modes on a single surface
+        nodes = jnp.array([rho, theta_B, zeta_B]).T
+        quant_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 @ (sqrtg_B_desc * quant))
+            / transforms["B"].grid.num_nodes
+        )
+        return quant_mn
+
+    def reshape(x):
+        return grid.meshgrid_reshape(x, "rtz").reshape((grid.num_rho, -1))
+
+    rho, theta_B, zeta_B, sqrtg_B_desc, R = map(
+        reshape,
+        (
+            data["rho"],
+            data["theta_B"],
+            data["zeta_B"],
+            data["sqrt(g)_Boozer_DESC"],
+            data["R"],
+        ),
+    )
+    R_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, R)
+    data["R_mn_B"] = R_mn.flatten()
+    return data
+
+
+@register_compute_fun(
+    name="Z_mn_B",
+    label="Z_{mn}^{\\mathrm{Boozer}}",
+    units="m",
+    units_long="meters",
+    description="Boozer harmonics of vertical coordinate of a flux surface",
+    dim=1,
+    params=[],
+    transforms={"B": [[0, 0, 0]], "grid": []},
+    profiles=[],
+    coordinates="rtz",
+    resolution_requirement="tz",
+    grid_requirement={"is_meshgrid": True},
+    data=[
+        "Z",
+        "sqrt(g)_Boozer_DESC",
+        "rho",
+        "theta_B",
+        "zeta_B",
+        "Boozer transform modes norm",
+    ],
+    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 _Z_mn(params, transforms, profiles, data, **kwargs):
+    norm = data["Boozer transform modes norm"]
+    grid = transforms["grid"]
+
+    def fun(rho, theta_B, zeta_B, sqrtg_B_desc, quant):
+        # this fits Boozer modes on a single surface
+        nodes = jnp.array([rho, theta_B, zeta_B]).T
+        quant_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 @ (sqrtg_B_desc * quant))
+            / transforms["B"].grid.num_nodes
+        )
+        return quant_mn
+
+    def reshape(x):
+        return grid.meshgrid_reshape(x, "rtz").reshape((grid.num_rho, -1))
+
+    rho, theta_B, zeta_B, sqrtg_B_desc, Z = map(
+        reshape,
+        (
+            data["rho"],
+            data["theta_B"],
+            data["zeta_B"],
+            data["sqrt(g)_Boozer_DESC"],
+            data["Z"],
+        ),
     )
-    B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B, B)
-    data["|B|_mn"] = B_mn.flatten()
+    Z_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, Z)
+    data["Z_mn_B"] = Z_mn.flatten()
     return data
 
 
@@ -390,6 +650,28 @@ 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. This norm is used as a"
+    "weight when performing the integral of the Boozer transform to get the "
+    "correct Boozer Fourier amplitudes.",
+    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="f_C",
     label="[(M \\iota - N) (\\mathbf{B} \\times \\nabla \\psi)"