Add documentation for compute function keyword arguments (#987)

Resolves #707

desc/compute/__init__.py

@@ -26,6 +26,8 @@
 # just need to import all the submodules here to register everything in the
 # data_index
 
+from desc.utils import flatten_list
+
 from . import (
     _basis_vectors,
     _bootstrap,
@@ -40,7 +42,7 @@
     _stability,
     _surface,
 )
-from .data_index import data_index
+from .data_index import all_kwargs, allowed_kwargs, data_index
 from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec
 from .utils import (
     compute,

desc/compute/_bootstrap.py

@@ -31,7 +31,8 @@
     coordinates="r",
     data=["sqrt(g)", "V_r(r)", "|B|", "<|B|^2>", "max_tz |B|"],
     axis_limit_data=["sqrt(g)_r", "V_rr(r)"],
-    n_gauss="n_gauss",
+    n_gauss="int: Number of quadrature points to use for estimating trapped fraction. "
+    + "Default 20.",
 )
 def _trapped_fraction(params, transforms, profiles, data, **kwargs):
     """Evaluate the effective trapped particle fraction.
@@ -333,7 +334,7 @@ def compute_J_dot_B_Redl(geom_data, profile_data, helicity_N=None):
         "Zeff",
         "rho",
     ],
-    helicity="helicity",
+    helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)",
 )
 def _J_dot_B_Redl(params, transforms, profiles, data, **kwargs):
     """Compute the bootstrap current 〈𝐉 ⋅ 𝐁〉.
@@ -395,7 +396,8 @@ def _J_dot_B_Redl(params, transforms, profiles, data, **kwargs):
     profiles=["current"],
     coordinates="r",
     data=["rho", "psi_r", "p_r", "current", "<|B|^2>", "<J*B> Redl"],
-    degree="degree",
+    degree="int: Degree of polynomial used for fitting current profile. "
+    + "Default grid.num_rho-1",
 )
 def _current_Redl(params, transforms, profiles, data, **kwargs):
     """Compute the current profile consistent with the Redl bootstrap current.

desc/compute/_curve.py

@@ -58,7 +58,7 @@ def _ds(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["x"],
     parameterization="desc.geometry.core.Curve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _X_curve(params, transforms, profiles, data, **kwargs):
     coords = data["x"]
@@ -82,7 +82,7 @@ def _X_curve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["x"],
     parameterization="desc.geometry.core.Curve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _Y_Curve(params, transforms, profiles, data, **kwargs):
     coords = data["x"]
@@ -106,7 +106,7 @@ def _Y_Curve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["x"],
     parameterization="desc.geometry.core.Curve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _R_Curve(params, transforms, profiles, data, **kwargs):
     coords = data["x"]
@@ -130,7 +130,7 @@ def _R_Curve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["x"],
     parameterization="desc.geometry.core.Curve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _phi_Curve(params, transforms, profiles, data, **kwargs):
     coords = data["x"]
@@ -175,7 +175,7 @@ def _Z_Curve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     # create planar curve at Z==0
@@ -212,7 +212,7 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     r = transforms["r"].transform(params["r_n"], dz=0)
@@ -258,7 +258,7 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     r = transforms["r"].transform(params["r_n"], dz=0)
@@ -309,7 +309,7 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.FourierPlanarCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     r = transforms["r"].transform(params["r_n"], dz=0)
@@ -369,7 +369,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierRZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     R = transforms["R"].transform(params["R_n"], dz=0)
@@ -404,7 +404,7 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierRZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     R0 = transforms["R"].transform(params["R_n"], dz=0)
@@ -438,7 +438,7 @@ def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierRZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     R0 = transforms["R"].transform(params["R_n"], dz=0)
@@ -476,7 +476,7 @@ def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierRZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     R0 = transforms["R"].transform(params["R_n"], dz=0)
@@ -514,7 +514,7 @@ def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     X = transforms["X"].transform(params["X_n"], dz=0)
@@ -547,7 +547,7 @@ def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     dX = transforms["X"].transform(params["X_n"], dz=1)
@@ -582,7 +582,7 @@ def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     d2X = transforms["X"].transform(params["X_n"], dz=2)
@@ -617,7 +617,7 @@ def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=[],
     parameterization="desc.geometry.curve.FourierXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     d3X = transforms["X"].transform(params["X_n"], dz=3)
@@ -650,7 +650,7 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.SplineXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     xq = data["s"]
@@ -705,7 +705,7 @@ def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.SplineXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     xq = data["s"]
@@ -791,7 +791,7 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.SplineXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     xq = data["s"]
@@ -876,7 +876,7 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     coordinates="s",
     data=["s"],
     parameterization="desc.geometry.curve.SplineXYZCurve",
-    basis="basis",
+    basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
 )
 def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
     xq = data["s"]

desc/compute/_equil.py

@@ -564,7 +564,7 @@ def _e_sup_helical(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="e^helical*sqrt(g)",
-    label=" \\sqrt{g}(B^{\\theta} \\nabla \\zeta - B^{\\zeta} \\nabla \\theta)",
+    label="\\sqrt{g}(B^{\\theta} \\nabla \\zeta - B^{\\zeta} \\nabla \\theta)",
     units="T \\cdot m^{2}",
     units_long="Tesla * square meter",
     description="Helical basis vector weighted by 3-D volume Jacobian",
@@ -717,7 +717,7 @@ def _W_Btor(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="",
     data=["p", "sqrt(g)"],
-    gamma="gamma",
+    gamma="float: Adiabatic index. Default 0",
 )
 def _W_p(params, transforms, profiles, data, **kwargs):
     data["W_p"] = jnp.sum(data["p"] * data["sqrt(g)"] * transforms["grid"].weights) / (

desc/compute/_field.py

@@ -108,7 +108,7 @@ def _B_sup_zeta(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="B",
-    label="B",
+    label="\\mathbf{B}",
     units="T",
     units_long="Tesla",
     description="Magnetic field",
@@ -3412,7 +3412,7 @@ def _L_grad_B(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="K_vc",
-    label="\\mathbf{K}_{VC} = \\mathbf{n} \\times \\mathbf{B}",
+    label="\\mathbf{K}_{VC} = \\frac{1}{\\mu_0}\\mathbf{n} \\times \\mathbf{B}",
     units="A \\cdot m^{-1}",
     units_long="Amps / meter",
     description="Virtual casing sheet current",

desc/compute/_metric.py

@@ -1832,7 +1832,7 @@ def _gradzeta(params, transforms, profiles, data, **kwargs):
     # Exact definition of the magnetic drifts taken from
     # eqn. 48 of Introduction to Quasisymmetry by Landreman
     # https://tinyurl.com/54udvaa4
-    label="\\mathrm{gradB-drift} = 1/B^{2} * (\\mathbf{b}\\times\\nabla (B)) \\cdot"
+    label="\\mathrm{gbdrift} = 1/B^{2} (\\mathbf{b}\\times\\nabla B) \\cdot"
     + "\\nabla \\alpha",
     units="1/(T-m^{2})",
     units_long="inverse Tesla meters^2",
@@ -1859,7 +1859,7 @@ def _gbdrift(params, transforms, profiles, data, **kwargs):
     # Exact definition of the magnetic drifts taken from
     # eqn. 48 of Introduction to Quasisymmetry by Landreman
     # https://tinyurl.com/54udvaa4
-    label="\\mathrm{curvature-drift} = 1/B^{3} * (\\mathbf{b}\\times\\nabla(p + B^2/2))"
+    label="\\mathrm{cvdrift} = 1/B^{3} (\\mathbf{b}\\times\\nabla(p + B^2/2))"
     + "\\cdot \\nabla \\alpha",
     units="1/(T-m^{2})",
     units_long="inverse Tesla meters^2",
@@ -1883,8 +1883,8 @@ def _cvdrift(params, transforms, profiles, data, **kwargs):
     # Exact definition of the magnetic drifts taken from
     # eqn. 48 of Introduction to Quasisymmetry by Landreman
     # https://tinyurl.com/54udvaa4
-    label="\\mathrm{curvature-drift-1} = 1/B^{2} * (\\mathbf{b}\\times\\nabla(B))"
-    + "\\cdot \\nabla rho",
+    label="\\mathrm{cvdrift0} = 1/B^{2} (\\mathbf{b}\\times\\nabla B)"
+    + "\\cdot \\nabla \\rho",
     units="1/(T-m^{2})",
     units_long="inverse Tesla meters^2",
     description="Radial component of the geometric part of the curvature drift"

desc/compute/_omnigenity.py

@@ -31,6 +31,8 @@
     profiles=[],
     coordinates="rtz",
     data=["B_theta"],
+    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_theta_mn(params, transforms, profiles, data, **kwargs):
     data["B_theta_mn"] = transforms["B"].fit(data["B_theta"])
@@ -50,6 +52,8 @@ def _B_theta_mn(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["B_zeta"],
+    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_zeta_mn(params, transforms, profiles, data, **kwargs):
     data["B_zeta_mn"] = transforms["B"].fit(data["B_zeta"])
@@ -69,6 +73,8 @@ def _B_zeta_mn(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["B_theta_mn", "B_zeta_mn"],
+    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 _w_mn(params, transforms, profiles, data, **kwargs):
     w_mn = jnp.zeros((transforms["w"].basis.num_modes,))
@@ -103,6 +109,8 @@ def _w_mn(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["w_Boozer_mn"],
+    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 _w(params, transforms, profiles, data, **kwargs):
     data["w_Boozer"] = transforms["w"].transform(data["w_Boozer_mn"])
@@ -122,6 +130,8 @@ def _w(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["w_Boozer_mn"],
+    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 _w_t(params, transforms, profiles, data, **kwargs):
     data["w_Boozer_t"] = transforms["w"].transform(data["w_Boozer_mn"], dt=1)
@@ -141,6 +151,8 @@ def _w_t(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["w_Boozer_mn"],
+    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 _w_z(params, transforms, profiles, data, **kwargs):
     data["w_Boozer_z"] = transforms["w"].transform(data["w_Boozer_mn"], dz=1)
@@ -272,6 +284,8 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["sqrt(g)_B", "|B|", "rho", "theta_B", "zeta_B"],
+    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
@@ -296,6 +310,8 @@ def _B_mn(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=[],
+    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_modes(params, transforms, profiles, data, **kwargs):
     data["B modes"] = transforms["B"].basis.modes
@@ -304,8 +320,8 @@ def _B_modes(params, transforms, profiles, data, **kwargs):
 
 @register_compute_fun(
     name="f_C",
-    label="(M \\iota - N) (\\mathbf{B} \\times \\nabla \\psi) \\cdot \\nabla B"
-    + " - (M G + N I) \\mathbf{B} \\cdot \\nabla B",
+    label="[(M \\iota - N) (\\mathbf{B} \\times \\nabla \\psi)"
+    + " - (M G + N I) \\mathbf{B}] \\cdot \\nabla B",
     units="T^{3}",
     units_long="Tesla cubed",
     description="Two-term quasisymmetry metric",
@@ -315,7 +331,7 @@ def _B_modes(params, transforms, profiles, data, **kwargs):
     profiles=[],
     coordinates="rtz",
     data=["iota", "B0", "B_theta", "B_zeta", "|B|_t", "|B|_z", "G", "I", "B*grad(|B|)"],
-    helicity="helicity",
+    helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)",
 )
 def _f_C(params, transforms, profiles, data, **kwargs):
     M, N = kwargs.get("helicity", (1, 0))
@@ -418,8 +434,8 @@ def _omni_angle(params, transforms, profiles, data, **kwargs):
     data=["alpha", "h"],
     aliases=["zeta_B"],
     parameterization="desc.magnetic_fields._core.OmnigenousField",
-    helicity="helicity",
-    iota="iota",
+    helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)",
+    iota="float: Value of rotational transform on the Omnigenous surface. Default 1.0",
 )
 def _omni_map(params, transforms, profiles, data, **kwargs):
     M, N = kwargs.get("helicity", (1, 0))