diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 414c014182..70bc1fa6e5 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -12,7 +12,7 @@ from desc.backend import jnp from .data_index import register_compute_fun -from .utils import cross, safediv +from .utils import cross, dot, safediv @register_compute_fun( @@ -34,7 +34,7 @@ def _b(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e^rho", + name="e^rho", # ∇ρ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates. label="\\mathbf{e}^{\\rho}", units="m^{-1}", units_long="inverse meters", @@ -927,7 +927,7 @@ def _e_sup_theta_zz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e^zeta", + name="e^zeta", # ∇ζ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates. label="\\mathbf{e}^{\\zeta}", units="m^{-1}", units_long="inverse meters", @@ -3396,7 +3396,7 @@ def _n_zeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["e_theta", "e_phi", "phi_t"], + data=["e_theta", "e_phi|r,t", "phi_t"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.FourierRZToroidalSurface", @@ -3404,5 +3404,235 @@ def _n_zeta(params, transforms, profiles, data, **kwargs): ], ) def _e_sub_theta_rp(params, transforms, profiles, data, **kwargs): - data["e_theta|r,p"] = data["e_theta"] - (data["e_phi"].T * data["phi_t"]).T + data["e_theta|r,p"] = data["e_theta"] - (data["e_phi|r,t"].T * data["phi_t"]).T + return data + + +@register_compute_fun( + name="e_rho|a,z", + label="\\mathbf{e}_{\\rho} |_{\\alpha, \\zeta}", + units="m", + units_long="meters", + description="Tangent vector along radial field line label", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_alpha", "alpha_r"], +) +def _e_rho_az(params, transforms, profiles, data, **kwargs): + # constant α and ζ + data["e_rho|a,z"] = ( + data["e_rho"] - data["e_alpha"] * data["alpha_r"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="e_alpha", + label="\\mathbf{e}_{\\alpha}", + units="m", + units_long="meters", + description="Tangent vector along poloidal field line label", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "alpha_t"], +) +def _e_alpha(params, transforms, profiles, data, **kwargs): + # constant ρ and ζ + data["e_alpha"] = data["e_theta"] / data["alpha_t"][:, jnp.newaxis] + return data + + +@register_compute_fun( + name="e_alpha_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\alpha}", + units="m", + units_long="meters", + description="Tangent vector along poloidal field line label, derivative wrt" + " DESC poloidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "alpha_t", "e_theta_t", "alpha_tt"], +) +def _e_alpha_t(params, transforms, profiles, data, **kwargs): + data["e_alpha_t"] = ( + data["e_theta_t"] / data["alpha_t"][:, jnp.newaxis] + - data["e_theta"] * (data["alpha_tt"] / data["alpha_t"] ** 2)[:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="e_alpha_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\alpha}", + units="m", + units_long="meters", + description="Tangent vector along poloidal field line label, " + "derivative wrt DESC toroidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "alpha_t", "e_theta_z", "alpha_tz"], +) +def _e_alpha_z(params, transforms, profiles, data, **kwargs): + data["e_alpha_z"] = ( + data["e_theta_z"] / data["alpha_t"][:, jnp.newaxis] + - data["e_theta"] * (data["alpha_tz"] / data["alpha_t"] ** 2)[:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="e_zeta|r,a", # Same as B/(B⋅∇ζ). + label="\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha} " + "= \\frac{\\mathbf{B}}{\\mathbf{B} \\cdot \\nabla \\zeta}", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta", "e_alpha", "alpha_z"], +) +def _e_zeta_ra(params, transforms, profiles, data, **kwargs): + data["e_zeta|r,a"] = ( + data["e_zeta"] - data["e_alpha"] * data["alpha_z"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_t", + label="\\partial_{\\theta} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha})", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line, " + "derivative wrt DESC poloidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta_t", "e_alpha", "alpha_z", "e_alpha_t", "alpha_tz"], +) +def _e_zeta_ra_t(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_t"] = data["e_zeta_t"] - ( + data["e_alpha_t"] * data["alpha_z"][:, jnp.newaxis] + + data["e_alpha"] * data["alpha_tz"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_a", + label="\\partial_{\\alpha} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha})", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line, derivative " + "wrt field line poloidal label", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["(e_zeta|r,a)_t", "alpha_t"], +) +def _e_zeta_ra_a(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_a"] = data["(e_zeta|r,a)_t"] / data["alpha_t"][:, jnp.newaxis] + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_z", + label="\\partial_{\\zeta} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha})", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line, " + "derivative wrt DESC toroidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta_z", "e_alpha", "alpha_z", "e_alpha_z", "alpha_zz"], +) +def _e_zeta_ra_z(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_z"] = data["e_zeta_z"] - ( + data["e_alpha_z"] * data["alpha_z"][:, jnp.newaxis] + + data["e_alpha"] * data["alpha_zz"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_z|r,a", + label="\\partial_{\\zeta} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha}) " + "|_{\\rho, \\alpha}", + units="m", + units_long="meters", + description="Curvature vector along field line", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["(e_zeta|r,a)_z", "(e_zeta|r,a)_a", "alpha_z"], +) +def _e_zeta_z_ra(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_z|r,a"] = ( + data["(e_zeta|r,a)_z"] + - data["(e_zeta|r,a)_a"] * data["alpha_z"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="|e_zeta|r,a|", # Often written as |dℓ/dζ| = |B/(B⋅∇ζ)|. + label="|\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha}|" + " = \\frac{|\\mathbf{B}|}{|\\mathbf{B} \\cdot \\nabla \\zeta|}", + units="m", + units_long="meters", + description="Differential length along field line", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta|r,a"], +) +def _d_ell_d_zeta(params, transforms, profiles, data, **kwargs): + data["|e_zeta|r,a|"] = jnp.linalg.norm(data["e_zeta|r,a"], axis=-1) + return data + + +@register_compute_fun( + name="|e_zeta|r,a|_z|r,a", + label="\\partial_{\\zeta} |\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha}| " + "|_{\\rho, \\alpha}", + units="m", + units_long="meters", + description="Differential length along field line, derivative along field line", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|e_zeta|r,a|", "(e_zeta|r,a)_z|r,a", "e_zeta|r,a"], +) +def _d_ell_d_zeta_z(params, transforms, profiles, data, **kwargs): + data["|e_zeta|r,a|_z|r,a"] = ( + dot(data["(e_zeta|r,a)_z|r,a"], data["e_zeta|r,a"]) / data["|e_zeta|r,a|"] + ) return data diff --git a/desc/compute/_core.py b/desc/compute/_core.py index a157d91412..2947e244a2 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -1518,7 +1518,43 @@ def _alpha_r(params, transforms, profiles, data, **kwargs): data=["theta_PEST_t", "phi_t", "iota"], ) def _alpha_t(params, transforms, profiles, data, **kwargs): - data["alpha_t"] = data["theta_PEST_t"] + data["iota"] * data["phi_t"] + data["alpha_t"] = data["theta_PEST_t"] - data["iota"] * data["phi_t"] + return data + + +@register_compute_fun( + name="alpha_tt", + label="\\partial_{\\theta \\theta} \\alpha", + units="~", + units_long="None", + description="Field line label, second derivative wrt poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_tt", "phi_tt", "iota"], +) +def _alpha_tt(params, transforms, profiles, data, **kwargs): + data["alpha_tt"] = data["theta_PEST_tt"] - data["iota"] * data["phi_tt"] + return data + + +@register_compute_fun( + name="alpha_tz", + label="\\partial_{\\theta \\zeta} \\alpha", + units="~", + units_long="None", + description="Field line label, derivative wrt poloidal and toroidal coordinates", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_tz", "phi_tz", "iota"], +) +def _alpha_tz(params, transforms, profiles, data, **kwargs): + data["alpha_tz"] = data["theta_PEST_tz"] - data["iota"] * data["phi_tz"] return data @@ -1540,6 +1576,24 @@ def _alpha_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="alpha_zz", + label="\\partial_{\\zeta \\zeta} \\alpha", + units="~", + units_long="None", + description="Field line label, second derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_zz", "phi_zz", "iota"], +) +def _alpha_zz(params, transforms, profiles, data, **kwargs): + data["alpha_zz"] = data["theta_PEST_zz"] - data["iota"] * data["phi_zz"] + return data + + @register_compute_fun( name="lambda", label="\\lambda", @@ -2983,6 +3037,44 @@ def _theta_PEST_t(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="theta_PEST_tt", + label="\\partial_{\\theta \\theta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, second " + "derivative wrt poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_tt"], +) +def _theta_PEST_tt(params, transforms, profiles, data, **kwargs): + data["theta_PEST_tt"] = data["lambda_tt"] + return data + + +@register_compute_fun( + name="theta_PEST_tz", + label="\\partial_{\\theta \\zeta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "poloidal and toroidal coordinates", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_tz"], +) +def _theta_PEST_tz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_tz"] = data["lambda_tz"] + return data + + @register_compute_fun( name="theta_PEST_z", label="\\partial_{\\zeta} \\vartheta", @@ -3002,6 +3094,25 @@ def _theta_PEST_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="theta_PEST_zz", + label="\\partial_{\\zeta \\zeta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, second " + "derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_zz"], +) +def _theta_PEST_zz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_zz"] = data["lambda_zz"] + return data + + @register_compute_fun( name="theta_r", label="\\partial_{\\rho} \\theta", diff --git a/desc/compute/_field.py b/desc/compute/_field.py index c3c95d86a7..eef1354e3d 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -483,6 +483,49 @@ def _B_sup_zeta_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="B^zeta_a", + label="\\partial_{\\alpha} B^{\\zeta}", + units="T \\cdot m^{-1}", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, derivative wrt field" + " line poloidal label" + ), + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B^zeta_t", "alpha_t"], +) +def _B_sup_zeta_a(params, transforms, profiles, data, **kwargs): + # constant ρ and ζ + data["B^zeta_a"] = data["B^zeta_t"] / data["alpha_t"] + return data + + +@register_compute_fun( + name="B^zeta_z|r,a", + label="\\partial_{\\zeta} B^{\\zeta} |_{\\rho, \\alpha}", + units="T \\cdot m^{-1}", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, derivative along field" + " line" + ), + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B^zeta_z", "B^zeta_a", "alpha_z"], +) +def _B_sup_zeta_z_ra(params, transforms, profiles, data, **kwargs): + data["B^zeta_z|r,a"] = data["B^zeta_z"] - data["B^zeta_a"] * data["alpha_z"] + return data + + @register_compute_fun( name="B_z", label="\\partial_{\\zeta} \\mathbf{B}", @@ -2295,6 +2338,43 @@ def _B_mag_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="|B|_a", + label="\\partial_{\\alpha} (|\\mathbf{B}|) |_{\\rho, \\zeta}", + units="T", + units_long="Tesla", + description="Magnitude of magnetic field, derivative wrt field line angle", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|B|_t", "alpha_t"], +) +def _B_mag_alpha(params, transforms, profiles, data, **kwargs): + # constant ρ and ζ + data["|B|_a"] = data["|B|_t"] / data["alpha_t"] + return data + + +@register_compute_fun( + name="|B|_z|r,a", + label="\\partial_{\\zeta} (|\\mathbf{B}|) |_{\\rho, \\alpha}", + units="T", + units_long="Tesla", + description="Magnitude of magnetic field, derivative along field line", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|B|_z", "|B|_a", "alpha_z"], +) +def _B_mag_z_constant_rho_alpha(params, transforms, profiles, data, **kwargs): + data["|B|_z|r,a"] = data["|B|_z"] - data["|B|_a"] * data["alpha_z"] + return data + + @register_compute_fun( name="|B|_rr", label="\\partial_{\\rho\\rho} |\\mathbf{B}|", @@ -3322,7 +3402,7 @@ def _kappa(params, transforms, profiles, data, **kwargs): label="\\kappa_n", units="m^{-1}", units_long="Inverse meters", - description="Normal curvature vector of magnetic field lines", + description="Normal curvature of magnetic field lines", dim=1, params=[], transforms={}, @@ -3340,7 +3420,7 @@ def _kappa_n(params, transforms, profiles, data, **kwargs): label="\\kappa_g", units="m^{-1}", units_long="Inverse meters", - description="Geodesic curvature vector of magnetic field lines", + description="Geodesic curvature of magnetic field lines", dim=1, params=[], transforms={}, diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index 3a842e9378..754116a98d 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -14,7 +14,7 @@ from desc.backend import jnp from .data_index import register_compute_fun -from .utils import cross, dot, safediv, safenorm +from .utils import cross, dot, safediv, safenorm, surface_averages @register_compute_fun( @@ -59,6 +59,27 @@ def _sqrtg_pest(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="sqrt(g)_Clebsch", + label="\\sqrt{g}_{\\text{Clebsch}}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of Clebsch field line coordinate system (ρ,α,ζ)" + " where ζ is the DESC toroidal coordinate.", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["sqrt(g)", "alpha_t"], +) +def _sqrtg_clebsch(params, transforms, profiles, data, **kwargs): + # Same as dot(data["e_rho|a,z"], cross(data["e_alpha"], data["e_zeta|r,a"])), but + # more efficient as it avoids computing radial derivative of alpha and hence iota. + data["sqrt(g)_Clebsch"] = data["sqrt(g)"] / data["alpha_t"] + return data + + @register_compute_fun( name="|e_theta x e_zeta|", label="|\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|", @@ -229,6 +250,27 @@ def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="|e_rho x e_alpha|", + label="|\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\alpha}|", + units="m^{2}", + units_long="square meters", + description="2D Jacobian determinant for constant zeta surface in Clebsch " + "field line coordinates (ρ,α,ζ) where ζ is the DESC toroidal coordinate.", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|e_rho x e_theta|", "alpha_t"], +) +def _e_rho_x_e_alpha(params, transforms, profiles, data, **kwargs): + # Same as safenorm(cross(data["e_rho|a,z"], data["e_alpha"]), axis=-1), but more + # efficient as it avoids computing radial derivative of iota and stream functions. + data["|e_rho x e_alpha|"] = data["|e_rho x e_theta|"] / jnp.abs(data["alpha_t"]) + return data + + @register_compute_fun( name="|e_rho x e_theta|_r", label="\\partial_{\\rho} |\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\theta}|", @@ -1331,7 +1373,7 @@ def _g_sup_zz(params, transforms, profiles, data, **kwargs): label="g^{\\rho\\theta}", units="m^{-2}", units_long="inverse square meters", - description="Radial/Poloidal element of contravariant metric tensor", + description="Radial/Poloidal (ρ, θ) element of contravariant metric tensor", dim=1, params=[], transforms={}, @@ -1759,6 +1801,32 @@ def _gradrho(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="<|grad(rho)|>", # same as S(r) / V_r(r) + label="\\langle \\vert \\nabla \\rho \\vert \\rangle|", + units="m^{-1}", + units_long="inverse meters", + description="Magnitude of contravariant radial basis vector, flux surface average", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["|grad(rho)|", "sqrt(g)"], + axis_limit_data=["sqrt(g)_r"], + resolution_requirement="tz", +) +def _gradrho_norm_fsa(params, transforms, profiles, data, **kwargs): + data["<|grad(rho)|>"] = surface_averages( + transforms["grid"], + data["|grad(rho)|"], + sqrt_g=transforms["grid"].replace_at_axis( + data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True + ), + ) + return data + + @register_compute_fun( name="|grad(psi)|", label="|\\nabla\\psi|", @@ -1847,12 +1915,12 @@ def _gradzeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["|B|", "b", "grad(alpha)", "grad(|B|)"], + data=["|B|^2", "b", "grad(alpha)", "grad(|B|)"], ) def _gbdrift(params, transforms, profiles, data, **kwargs): data["gbdrift"] = ( 1 - / data["|B|"] ** 2 + / data["|B|^2"] * dot(data["b"], cross(data["grad(|B|)"], data["grad(alpha)"])) ) return data @@ -1874,11 +1942,11 @@ def _gbdrift(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["p_r", "psi_r", "|B|", "gbdrift"], + data=["p_r", "psi_r", "|B|^2", "gbdrift"], ) def _cvdrift(params, transforms, profiles, data, **kwargs): dp_dpsi = mu_0 * data["p_r"] / data["psi_r"] - data["cvdrift"] = 1 / data["|B|"] ** 2 * dp_dpsi + data["gbdrift"] + data["cvdrift"] = 1 / data["|B|^2"] * dp_dpsi + data["gbdrift"] return data @@ -1892,16 +1960,16 @@ def _cvdrift(params, transforms, profiles, data, **kwargs): units="1/(T-m^{2})", units_long="inverse Tesla meters^2", description="Radial component of the geometric part of the curvature drift" - + " used for local stability analyses, Gamma_c, epsilon_eff etc.", + + " used for local stability analyses, gyrokinetics, Gamma_c.", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["|B|", "b", "e^rho", "grad(|B|)"], + data=["|B|^2", "b", "e^rho", "grad(|B|)"], ) def _cvdrift0(params, transforms, profiles, data, **kwargs): data["cvdrift0"] = ( - 1 / data["|B|"] ** 2 * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"]))) + 1 / data["|B|^2"] * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"]))) ) return data diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index bfd408487c..cf0a32714e 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -105,7 +105,11 @@ def register_compute_fun( # noqa: C901 or `'desc.equilibrium.Equilibrium'`. resolution_requirement : str Resolution requirements in coordinates. I.e. "r" expects radial resolution - in the grid, "rtz" expects a grid with radial, poloidal, and toroidal resolution. + in the grid. Likewise, "rtz" is shorthand for "rho, theta, zeta" and indicates + the computation expects a grid with radial, poloidal, and toroidal resolution. + If the computation simply performs pointwise operations, instead of a + reduction (such as integration) over a coordinate, then an empty string may + be used to indicate no requirements. source_grid_requirement : dict Attributes of the source grid that the compute function requires. Also assumes dependencies were computed on such a grid. diff --git a/docs/adding_compute_funs.rst b/docs/adding_compute_funs.rst index 66e3c2c047..a1dcad35b3 100644 --- a/docs/adding_compute_funs.rst +++ b/docs/adding_compute_funs.rst @@ -27,6 +27,7 @@ The full code is below: coordinates="rtz", data=["sqrt(g)", "B_zeta_t", "B_theta_z"], axis_limit_data=["sqrt(g)_r", "B_zeta_rt", "B_theta_rz"], + resolution_requirement="", parameterization="desc.equilibrium.equilibrium.Equilibrium", ) def _J_sup_rho(params, transforms, profiles, data, **kwargs): @@ -96,6 +97,12 @@ metadata about the quantity. The necessary fields are detailed below: the quantities mentioned above to evaluate the magnetic axis limit. These dependencies are specified in ``axis_limit_data``. The dependencies specified in this list are marked to be computed only when there is a node at the magnetic axis. +* ``resolution_requirement``: Resolution requirements in coordinates. + I.e. "r" expects radial resolution in the grid. Likewise, "rtz" is shorthand for + "rho, theta, zeta" and indicates the computation expects a grid with radial, + poloidal, and toroidal resolution. If the computation simply performs + pointwise operations, instead of a reduction (such as integration) over a + coordinate, then an empty string may be used to indicate no requirements. * ``parameterization``: what sorts of DESC objects is this function for. Most functions will just be for ``Equilibrium``, but some methods may also be for ``desc.geometry.core.Curve``, or specific types eg ``desc.geometry.curve.FourierRZCurve``. If a quantity is computed differently diff --git a/docs/write_variables.py b/docs/write_variables.py index eec46136f1..13b15b9061 100644 --- a/docs/write_variables.py +++ b/docs/write_variables.py @@ -78,6 +78,12 @@ def write_csv(parameterization): The keyword argument ``basis='xyz'`` can be used to convert the variables into Cartesian coordinates :math:`(X,Y,Z)`. ``basis`` must be one of ``{'rpz', 'xyz'}``. +Our convention to denote partial derivatives is an underscore followed by the first +letter of the coordinate that the partial derivative is taken with respect to. Unless +otherwise specified or implied by the variable name, these partial derivatives are +those of the DESC :math:`\rho, \theta, \zeta` coordinate system. For example, ``|B|_z`` +is :math:`(\partial \vert B \vert / \partial\zeta)|_{\rho, \theta}`. + """ block = """ diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index f230d7a816..c787da5046 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/inputs/master_compute_data_xyz.pkl b/tests/inputs/master_compute_data_xyz.pkl index 0de07ede29..33247963ca 100644 Binary files a/tests/inputs/master_compute_data_xyz.pkl and b/tests/inputs/master_compute_data_xyz.pkl differ diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 21c768e215..bf1e910232 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -10,6 +10,7 @@ from desc.coils import FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, SplineXYZCoil from desc.compute import data_index, rpz2xyz_vec +from desc.compute.utils import dot from desc.equilibrium import Equilibrium from desc.examples import get from desc.geometry import ( @@ -1052,8 +1053,8 @@ def test_magnetic_pressure_gradient(DummyStellarator): num_rho = 110 grid = LinearGrid(NFP=eq.NFP, rho=num_rho) drho = grid.nodes[1, 0] - data = eq.compute(["|B|", "grad(|B|^2)_rho"], grid=grid) - B2_r = np.convolve(data["|B|"] ** 2, FD_COEF_1_4, "same") / drho + data = eq.compute(["|B|^2", "grad(|B|^2)_rho"], grid=grid) + B2_r = np.convolve(data["|B|^2"], FD_COEF_1_4, "same") / drho np.testing.assert_allclose( data["grad(|B|^2)_rho"][3:-2], B2_r[3:-2], @@ -1065,8 +1066,8 @@ def test_magnetic_pressure_gradient(DummyStellarator): num_theta = 90 grid = LinearGrid(NFP=eq.NFP, theta=num_theta) dtheta = grid.nodes[1, 1] - data = eq.compute(["|B|", "grad(|B|^2)_theta"], grid=grid) - B2_t = np.convolve(data["|B|"] ** 2, FD_COEF_1_4, "same") / dtheta + data = eq.compute(["|B|^2", "grad(|B|^2)_theta"], grid=grid) + B2_t = np.convolve(data["|B|^2"], FD_COEF_1_4, "same") / dtheta np.testing.assert_allclose( data["grad(|B|^2)_theta"][2:-2], B2_t[2:-2], @@ -1078,8 +1079,8 @@ def test_magnetic_pressure_gradient(DummyStellarator): num_zeta = 90 grid = LinearGrid(NFP=eq.NFP, zeta=num_zeta) dzeta = grid.nodes[1, 2] - data = eq.compute(["|B|", "grad(|B|^2)_zeta"], grid=grid) - B2_z = np.convolve(data["|B|"] ** 2, FD_COEF_1_4, "same") / dzeta + data = eq.compute(["|B|^2", "grad(|B|^2)_zeta"], grid=grid) + B2_z = np.convolve(data["|B|^2"], FD_COEF_1_4, "same") / dzeta np.testing.assert_allclose( data["grad(|B|^2)_zeta"][2:-2], B2_z[2:-2], @@ -1738,3 +1739,35 @@ def test_surface_equilibrium_geometry(): rtol=3e-13, atol=1e-13, ) + + +@pytest.mark.unit +def test_parallel_grad(): + """Test geometric and physical methods of computing parallel gradients agree.""" + eq = get("W7-X") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(2, 2, 2, 4, 4, 4) + data = eq.compute( + [ + "e_zeta|r,a", + "B", + "B^zeta", + "|B|_z|r,a", + "grad(|B|)", + "|e_zeta|r,a|_z|r,a", + "B^zeta_z|r,a", + "|B|", + ], + ) + np.testing.assert_allclose(data["e_zeta|r,a"], (data["B"].T / data["B^zeta"]).T) + np.testing.assert_allclose( + data["|B|_z|r,a"], dot(data["grad(|B|)"], data["e_zeta|r,a"]) + ) + np.testing.assert_allclose( + data["|e_zeta|r,a|_z|r,a"], + data["|B|_z|r,a"] / np.abs(data["B^zeta"]) + - data["|B|"] + * data["B^zeta_z|r,a"] + * np.sign(data["B^zeta"]) + / data["B^zeta"] ** 2, + )