diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index e3cd6ae457..d260ddee31 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -1728,6 +1728,7 @@ def _e_sub_rho_rrr(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.ZernikeRZToroidalSection", ], basis="basis", + aliases=["e_theta_rrr"], ) def _e_sub_rho_rrt(params, transforms, profiles, data, **kwargs): data["e_rho_rrt"] = jnp.array( @@ -1807,63 +1808,48 @@ def _e_sub_rho_rrt(params, transforms, profiles, data, **kwargs): "phi", ], basis="basis", + aliases=["e_zeta_rrr"], ) def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs): data["e_rho_rrz"] = jnp.array( [ - -2 * data["omega_rz"] * data["R_r"] * data["omega_r"] - - 2 - * (1 + data["omega_z"]) + -3 * data["R"] * data["omega_rrz"] * data["omega_r"] + - 3 + * data["omega_rz"] + * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) + - 3 + * data["omega_z"] * (data["R_rr"] * data["omega_r"] + data["R_r"] * data["omega_rr"]) - - data["R_rz"] * data["omega_r"] ** 2 - - 2 * data["R_z"] * data["omega_r"] * data["omega_rr"] - - 2 * data["R_r"] * data["omega_r"] * data["omega_rz"] - - 2 + - (1 + data["omega_z"]) * data["R"] + * (data["omega_rrr"] - data["omega_r"] ** 3) + - 3 + * data["omega_r"] * ( - data["omega_rr"] * data["omega_rz"] - + data["omega_r"] * data["omega_rrz"] + data["R_rz"] * data["omega_r"] + + data["R_z"] * data["omega_rr"] + + data["R_rr"] ) - - data["R_r"] * (1 + data["omega_z"]) * data["omega_rr"] - - data["R"] + - 3 * data["R_r"] * data["omega_rr"] + + data["R_rrrz"], + 3 + * data["R_r"] + * (data["omega_rrz"] - (1 + data["omega_z"]) * data["omega_r"] ** 2) + + 3 * data["omega_rz"] * data["R_rr"] + + data["R"] * ( - data["omega_rz"] * data["omega_rr"] - + (1 + data["omega_z"]) * data["omega_rrr"] + data["omega_rrrz"] + - 3 + * data["omega_r"] + * ( + data["omega_rz"] * data["omega_r"] + + (1 + data["omega_z"]) * data["omega_rr"] + ) ) - + data["R_rrrz"] - - data["omega_r"] - * ( - 2 * data["omega_r"] * data["R_rz"] - + 2 * data["R_r"] * data["omega_rz"] - + (1 + data["omega_z"]) * data["R_rr"] - + data["R_z"] * data["omega_rr"] - - data["R"] - * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"]) - ), - 2 * data["omega_rr"] * data["R_rz"] - + 2 * data["omega_r"] * data["R_rrz"] - + 2 * data["R_rr"] * data["omega_rz"] - + 2 * data["R_r"] * data["omega_rrz"] - + data["omega_rz"] * data["R_rr"] + (1 + data["omega_z"]) * data["R_rrr"] - + data["R_rz"] * data["omega_rr"] - + data["R_z"] * data["omega_rrr"] - - data["R_r"] - * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"]) - - data["R"] - * ( - data["omega_rz"] * data["omega_r"] ** 2 - + 2 * (1 + data["omega_z"]) * data["omega_r"] * data["omega_rr"] - - data["omega_rrrz"] - ) - + data["omega_r"] - * ( - -2 * (1 + data["omega_z"]) * data["R_r"] * data["omega_r"] - - data["R_z"] * data["omega_r"] ** 2 - - 2 * data["R"] * data["omega_r"] * data["omega_rz"] - - data["R"] * (1 + data["omega_z"]) * data["omega_rr"] - + data["R_rrz"] - ), + + 3 * data["R_rrz"] * data["omega_r"] + + 3 * data["R_rz"] * data["omega_rr"] + + data["R_z"] * (data["omega_rrr"] - data["omega_r"] ** 3), data["Z_rrrz"], ] ).T @@ -1901,7 +1887,7 @@ def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs): "omega_t", "phi", ], - aliases=["x_rrt", "x_rtr", "x_trr"], + aliases=["x_rrt", "x_rtr", "x_trr", "e_theta_rr"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.ZernikeRZToroidalSection", @@ -1969,6 +1955,7 @@ def _e_sub_rho_rt(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.ZernikeRZToroidalSection", ], basis="basis", + aliases=["e_theta_rrt"], ) def _e_sub_rho_rtt(params, transforms, profiles, data, **kwargs): data["e_rho_rtt"] = jnp.array( @@ -2055,6 +2042,7 @@ def _e_sub_rho_rtt(params, transforms, profiles, data, **kwargs): "phi", ], basis="basis", + aliases=["e_theta_rrz", "e_zeta_rrt"], ) def _e_sub_rho_rtz(params, transforms, profiles, data, **kwargs): data["e_rho_rtz"] = jnp.array( @@ -2182,6 +2170,7 @@ def _e_sub_rho_rtz(params, transforms, profiles, data, **kwargs): "phi", ], basis="basis", + aliases=["e_zeta_rr"], ) def _e_sub_rho_rz(params, transforms, profiles, data, **kwargs): data["e_rho_rz"] = jnp.array( @@ -2241,6 +2230,7 @@ def _e_sub_rho_rz(params, transforms, profiles, data, **kwargs): "phi", ], basis="basis", + aliases=["e_zeta_rrz"], ) def _e_sub_rho_rzz(params, transforms, profiles, data, **kwargs): data["e_rho_rzz"] = jnp.array( @@ -2323,8 +2313,11 @@ def _e_sub_rho_rzz(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.ZernikeRZToroidalSection", ], basis="basis", + aliases=["e_theta_r"], ) def _e_sub_rho_t(params, transforms, profiles, data, **kwargs): + # At the magnetic axis, this function returns the multivalued map whose + # image is the set { ∂ᵨ 𝐞_θ | ρ=0 } data["e_rho_t"] = jnp.array( [ -data["R"] * data["omega_t"] * data["omega_r"] + data["R_rt"], @@ -2373,6 +2366,7 @@ def _e_sub_rho_t(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.ZernikeRZToroidalSection", ], basis="basis", + aliases=["e_theta_rt"], ) def _e_sub_rho_tt(params, transforms, profiles, data, **kwargs): data["e_rho_tt"] = jnp.array( @@ -2430,6 +2424,7 @@ def _e_sub_rho_tt(params, transforms, profiles, data, **kwargs): "phi", ], basis="basis", + aliases=["e_theta_rz", "e_zeta_rt"], ) def _e_sub_rho_tz(params, transforms, profiles, data, **kwargs): data["e_rho_tz"] = jnp.array( @@ -2476,6 +2471,7 @@ def _e_sub_rho_tz(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z", "phi"], basis="basis", + aliases=["e_zeta_r"], ) def _e_sub_rho_z(params, transforms, profiles, data, **kwargs): data["e_rho_z"] = jnp.array( @@ -2522,6 +2518,7 @@ def _e_sub_rho_z(params, transforms, profiles, data, **kwargs): "phi", ], basis="basis", + aliases=["e_zeta_rz"], ) def _e_sub_rho_zz(params, transforms, profiles, data, **kwargs): data["e_rho_zz"] = jnp.array( @@ -2622,103 +2619,13 @@ def _e_sub_theta_pest(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_theta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description="Covariant Poloidal basis vector, derivative wrt radial coordinate", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_r", "R_rt", "R_t", "Z_rt", "omega_r", "omega_rt", "omega_t", "phi"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.ZernikeRZToroidalSection", - ], - basis="basis", -) -def _e_sub_theta_r(params, transforms, profiles, data, **kwargs): - # At the magnetic axis, this function returns the multivalued map whose - # image is the set { ∂ᵨ 𝐞_θ | ρ=0 } - data["e_theta_r"] = jnp.array( - [ - -data["R"] * data["omega_t"] * data["omega_r"] + data["R_rt"], - data["omega_t"] * data["R_r"] - + data["R_t"] * data["omega_r"] - + data["R"] * data["omega_rt"], - data["Z_rt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_r"] = rpz2xyz_vec(data["e_theta_r"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt radial and radial" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rrt", - "R_rt", - "R_t", - "Z_rrt", - "omega_r", - "omega_rr", - "omega_rrt", - "omega_rt", - "omega_t", - "phi", - ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.ZernikeRZToroidalSection", - ], - basis="basis", -) -def _e_sub_theta_rr(params, transforms, profiles, data, **kwargs): - data["e_theta_rr"] = jnp.array( - [ - -data["R_t"] * data["omega_r"] ** 2 - - 2 * data["R"] * data["omega_r"] * data["omega_rt"] - - data["omega_t"] - * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) - + data["R_rrt"], - 2 * data["omega_r"] * data["R_rt"] - + 2 * data["R_r"] * data["omega_rt"] - + data["omega_t"] * data["R_rr"] - + data["R_t"] * data["omega_rr"] - + data["R"] * (-data["omega_t"] * data["omega_r"] ** 2 + data["omega_rrt"]), - data["Z_rrt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rr"] = rpz2xyz_vec(data["e_theta_rr"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rrr", - label="\\partial_{\\rho \\rho \\rho} \\mathbf{e}_{\\theta}", + name="e_theta_rtt", + label="\\partial_{\\rho \\theta \\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", description=( "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " once and poloidal twice" ), dim=3, params=[], @@ -2729,19 +2636,19 @@ def _e_sub_theta_rr(params, transforms, profiles, data, **kwargs): "R", "R_r", "R_t", - "R_rr", "R_rt", - "R_rrr", - "R_rrt", - "R_rrrt", - "Z_rrrt", + "R_tt", + "R_rtt", + "R_ttt", + "R_rttt", + "Z_rttt", "omega_r", "omega_t", - "omega_rr", "omega_rt", - "omega_rrr", - "omega_rrt", - "omega_rrrt", + "omega_tt", + "omega_rtt", + "omega_ttt", + "omega_rttt", "phi", ], parameterization=[ @@ -2749,57 +2656,61 @@ def _e_sub_theta_rr(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.ZernikeRZToroidalSection", ], basis="basis", + aliases=["e_rho_ttt"], ) -def _e_sub_theta_rrr(params, transforms, profiles, data, **kwargs): - data["e_theta_rrr"] = jnp.array( +def _e_sub_theta_rtt(params, transforms, profiles, data, **kwargs): + data["e_theta_rtt"] = jnp.array( [ - -3 * data["omega_rrt"] * data["R"] * data["omega_r"] + -3 * data["R_rt"] * data["omega_t"] ** 2 - 3 - * data["omega_rt"] - * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) - - data["omega_t"] + * data["omega_t"] * ( - 3 * data["R_rr"] * data["omega_r"] - + 3 * data["R_r"] * data["omega_rr"] - + data["R"] * data["omega_rrr"] - - data["R"] * data["omega_r"] ** 3 + data["R_r"] * data["omega_tt"] + + data["R_tt"] * data["omega_r"] + + 2 * data["R_t"] * data["omega_rt"] + + data["R"] * data["omega_rt"] ) - - 3 - * data["omega_r"] - * (data["R_rt"] * data["omega_r"] + data["R_t"] * data["omega_rr"]) - + data["R_rrrt"], - 3 * data["omega_rrt"] * data["R_r"] - + 3 * data["omega_rt"] * data["R_rr"] + - data["omega_r"] + * (3 * data["R_t"] * data["omega_tt"] + data["R"] * data["omega_ttt"]) + data["R"] * ( - data["omega_rrrt"] + data["omega_r"] * data["omega_t"] ** 3 + - 3 * data["omega_rt"] * data["omega_tt"] + ) + + data["R_rttt"], + data["R_r"] * (data["omega_ttt"] - data["omega_t"] ** 3) + + data["omega_r"] + * ( + data["R_ttt"] - 3 - * data["omega_r"] - * ( - data["omega_rt"] * data["omega_r"] - + data["omega_t"] * data["omega_rr"] - ) + * data["omega_t"] + * (data["R_t"] * data["omega_t"] + data["R"] * data["omega_tt"]) ) - + data["omega_t"] * (data["R_rrr"] - 3 * data["R_r"] * data["omega_r"] ** 2) - + 3 * data["R_rrt"] * data["omega_r"] - + 3 * data["R_rt"] * data["omega_rr"] - + data["R_t"] * (data["omega_rrr"] - data["omega_r"] ** 3), - data["Z_rrrt"], + + 3 + * ( + data["R_rt"] * data["omega_tt"] + + data["R_tt"] * data["omega_rt"] + + data["R_rtt"] * data["omega_t"] + + data["R_t"] * data["omega_rtt"] + ) + + data["R"] + * (data["omega_rttt"] - 3 * data["omega_t"] ** 2 * data["omega_rt"]), + data["Z_rttt"], ] ).T if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rrr"] = rpz2xyz_vec(data["e_theta_rrr"], phi=data["phi"]) + data["e_theta_rtt"] = rpz2xyz_vec(data["e_theta_rtt"], phi=data["phi"]) return data @register_compute_fun( - name="e_theta_rrt", - label="\\partial_{\\rho \\rho \\theta} \\mathbf{e}_{\\theta}", + name="e_theta_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - " twice and poloidal once" + "Covariant Poloidal basis vector, third derivative wrt radial, poloidal," + " and toroidal coordinates" ), dim=3, params=[], @@ -2809,1456 +2720,56 @@ def _e_sub_theta_rrr(params, transforms, profiles, data, **kwargs): data=[ "R", "R_r", - "R_rr", - "R_rt", - "R_rrt", - "R_rtt", - "R_rrtt", "R_t", + "R_rt", "R_tt", - "Z_rrtt", + "R_rtt", + "R_ttz", + "R_rttz", + "R_tz", + "R_rtz", + "R_z", + "R_rz", + "Z_rttz", "omega_r", - "omega_rr", - "omega_rt", - "omega_rrt", - "omega_rtt", - "omega_rrtt", "omega_t", + "omega_rt", "omega_tt", + "omega_rtt", + "omega_ttz", + "omega_rttz", + "omega_tz", + "omega_rtz", + "omega_z", + "omega_rz", "phi", ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.ZernikeRZToroidalSection", - ], basis="basis", + aliases=["e_rho_ttz", "e_zeta_rtt"], ) -def _e_sub_theta_rrt(params, transforms, profiles, data, **kwargs): - data["e_theta_rrt"] = jnp.array( +def _e_sub_theta_rtz(params, transforms, profiles, data, **kwargs): + data["e_theta_rtz"] = jnp.array( [ - -2 * data["omega_t"] * data["omega_rt"] * data["R_r"] - - data["omega_t"] ** 2 * data["R_rr"] - - data["R_r"] * data["omega_tt"] * data["omega_r"] - - data["R"] - * ( - data["omega_rtt"] * data["omega_r"] - + data["omega_tt"] * data["omega_rr"] - ) + -2 * data["omega_rz"] * data["R_t"] * data["omega_t"] - 2 - * data["omega_rt"] - * (data["R_t"] * data["omega_r"] + data["R"] * data["omega_rt"]) + * (1 + data["omega_z"]) + * (data["R_rt"] * data["omega_t"] + data["R_t"] * data["omega_rt"]) + - data["R_rz"] * data["omega_t"] ** 2 + - 2 * data["R_z"] * data["omega_t"] * data["omega_rt"] + - 2 * data["R_r"] * data["omega_t"] * data["omega_tz"] - 2 - * data["omega_t"] + * data["R"] + * ( + data["omega_rt"] * data["omega_tz"] + + data["omega_t"] * data["omega_rtz"] + ) + - data["R_r"] * (1 + data["omega_z"]) * data["omega_tt"] + - data["R"] * ( - data["R_rt"] * data["omega_r"] - + data["R_r"] * data["omega_rt"] - + data["R_t"] * data["omega_rr"] - + data["R"] * data["omega_rrt"] + data["omega_rz"] * data["omega_tt"] + + (1 + data["omega_z"]) * data["omega_rtt"] ) - + data["R_rrtt"] - - data["omega_r"] - * ( - data["omega_tt"] * data["R_r"] - + data["R_tt"] * data["omega_r"] - + 2 * data["omega_t"] * data["R_rt"] - + 2 * data["R_t"] * data["omega_rt"] - + data["R"] - * (-data["omega_t"] ** 2 * data["omega_r"] + data["omega_rtt"]) - ), - data["omega_rtt"] * data["R_r"] - + data["omega_tt"] * data["R_rr"] - + data["R_rtt"] * data["omega_r"] - + data["R_tt"] * data["omega_rr"] - + 2 * data["omega_rt"] * data["R_rt"] - + 2 * data["omega_t"] * data["R_rrt"] - + 2 * data["R_rt"] * data["omega_rt"] - + 2 * data["R_t"] * data["omega_rrt"] - + data["R_r"] - * (-data["omega_t"] ** 2 * data["omega_r"] + data["omega_rtt"]) - + data["R"] - * ( - -2 * data["omega_t"] * data["omega_rt"] * data["omega_r"] - - data["omega_t"] ** 2 * data["omega_rr"] - + data["omega_rrtt"] - ) - + data["omega_r"] - * ( - -data["omega_t"] ** 2 * data["R_r"] - - data["R"] * data["omega_tt"] * data["omega_r"] - - 2 - * data["omega_t"] - * (data["R_t"] * data["omega_r"] + data["R"] * data["omega_rt"]) - + data["R_rtt"] - ), - data["Z_rrtt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rrt"] = rpz2xyz_vec(data["e_theta_rrt"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rrz", - label="\\partial_{\\rho \\rho \\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - " twice and toroidal once" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rt", - "R_rrt", - "R_rtz", - "R_rrtz", - "R_rz", - "R_rrz", - "R_t", - "R_tz", - "R_z", - "Z_rrtz", - "omega_r", - "omega_rr", - "omega_rt", - "omega_rrt", - "omega_rtz", - "omega_rrtz", - "omega_rz", - "omega_rrz", - "omega_t", - "omega_tz", - "omega_z", - "phi", - ], - basis="basis", -) -def _e_sub_theta_rrz(params, transforms, profiles, data, **kwargs): - data["e_theta_rrz"] = jnp.array( - [ - -data["omega_rz"] * data["R_t"] * data["omega_r"] - - (1 + data["omega_z"]) - * (data["R_rt"] * data["omega_r"] + data["R_t"] * data["omega_rr"]) - - data["R_r"] * data["omega_tz"] * data["omega_r"] - - data["R"] - * ( - data["omega_rtz"] * data["omega_r"] - + data["omega_tz"] * data["omega_rr"] - ) - - data["omega_rt"] - * ( - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"] - ) - - data["omega_t"] - * ( - data["omega_rz"] * data["R_r"] - + (1 + data["omega_z"]) * data["R_rr"] - + data["R_rz"] * data["omega_r"] - + data["R_z"] * data["omega_rr"] - + data["R_r"] * data["omega_rz"] - + data["R"] * data["omega_rrz"] - ) - - data["R_r"] * (1 + data["omega_z"]) * data["omega_rt"] - - data["R"] - * ( - data["omega_rz"] * data["omega_rt"] - + (1 + data["omega_z"]) * data["omega_rrt"] - ) - + data["R_rrtz"] - - data["omega_r"] - * ( - data["omega_tz"] * data["R_r"] - + data["R_tz"] * data["omega_r"] - + data["omega_t"] * data["R_rz"] - + data["R_t"] * data["omega_rz"] - + data["R_rt"] - + data["omega_z"] * data["R_rt"] - + data["R_z"] * data["omega_rt"] - + data["R"] - * ( - -(1 + data["omega_z"]) * data["omega_t"] * data["omega_r"] - + data["omega_rtz"] - ) - ), - data["omega_rtz"] * data["R_r"] - + data["omega_tz"] * data["R_rr"] - + data["R_rtz"] * data["omega_r"] - + data["R_tz"] * data["omega_rr"] - + data["omega_rt"] * data["R_rz"] - + data["omega_t"] * data["R_rrz"] - + data["R_rt"] * data["omega_rz"] - + data["R_t"] * data["omega_rrz"] - + data["R_rrt"] - + data["omega_rz"] * data["R_rt"] - + data["omega_z"] * data["R_rrt"] - + data["R_rz"] * data["omega_rt"] - + data["R_z"] * data["omega_rrt"] - + data["R_r"] - * ( - -(1 + data["omega_z"]) * data["omega_t"] * data["omega_r"] - + data["omega_rtz"] - ) - + data["R"] - * ( - -data["omega_rz"] * data["omega_t"] * data["omega_r"] - - (1 + data["omega_z"]) - * ( - data["omega_rt"] * data["omega_r"] - + data["omega_t"] * data["omega_rr"] - ) - + data["omega_rrtz"] - ) - + data["omega_r"] - * ( - -(1 + data["omega_z"]) * data["R_t"] * data["omega_r"] - - data["R"] * data["omega_tz"] * data["omega_r"] - - data["omega_t"] - * ( - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"] - ) - - data["R"] * (1 + data["omega_z"]) * data["omega_rt"] - + data["R_rtz"] - ), - data["Z_rrtz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rrz"] = rpz2xyz_vec(data["e_theta_rrz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt radial and poloidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rt", - "R_rtt", - "R_t", - "R_tt", - "Z_rtt", - "omega_r", - "omega_rt", - "omega_rtt", - "omega_t", - "omega_tt", - "phi", - ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.ZernikeRZToroidalSection", - ], - basis="basis", -) -def _e_sub_theta_rt(params, transforms, profiles, data, **kwargs): - data["e_theta_rt"] = jnp.array( - [ - -data["omega_t"] ** 2 * data["R_r"] - - data["R"] * data["omega_tt"] * data["omega_r"] - - 2 - * data["omega_t"] - * (data["R_t"] * data["omega_r"] + data["R"] * data["omega_rt"]) - + data["R_rtt"], - data["omega_tt"] * data["R_r"] - + data["R_tt"] * data["omega_r"] - + 2 * data["omega_t"] * data["R_rt"] - + 2 * data["R_t"] * data["omega_rt"] - + data["R"] * (-data["omega_t"] ** 2 * data["omega_r"] + data["omega_rtt"]), - data["Z_rtt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rt"] = rpz2xyz_vec(data["e_theta_rt"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rtt", - label="\\partial_{\\rho \\theta \\theta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - " once and poloidal twice" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_t", - "R_rt", - "R_tt", - "R_rtt", - "R_ttt", - "R_rttt", - "Z_rttt", - "omega_r", - "omega_t", - "omega_rt", - "omega_tt", - "omega_rtt", - "omega_ttt", - "omega_rttt", - "phi", - ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.ZernikeRZToroidalSection", - ], - basis="basis", -) -def _e_sub_theta_rtt(params, transforms, profiles, data, **kwargs): - data["e_theta_rtt"] = jnp.array( - [ - -3 * data["R_rt"] * data["omega_t"] ** 2 - - 3 - * data["omega_t"] - * ( - data["R_r"] * data["omega_tt"] - + data["R_tt"] * data["omega_r"] - + 2 * data["R_t"] * data["omega_rt"] - + data["R"] * data["omega_rt"] - ) - - data["omega_r"] - * (3 * data["R_t"] * data["omega_tt"] + data["R"] * data["omega_ttt"]) - + data["R"] - * ( - data["omega_r"] * data["omega_t"] ** 3 - - 3 * data["omega_rt"] * data["omega_tt"] - ) - + data["R_rttt"], - data["R_r"] * (data["omega_ttt"] - data["omega_t"] ** 3) - + data["omega_r"] - * ( - data["R_ttt"] - - 3 - * data["omega_t"] - * (data["R_t"] * data["omega_t"] + data["R"] * data["omega_tt"]) - ) - + 3 - * ( - data["R_rt"] * data["omega_tt"] - + data["R_tt"] * data["omega_rt"] - + data["R_rtt"] * data["omega_t"] - + data["R_t"] * data["omega_rtt"] - ) - + data["R"] - * (data["omega_rttt"] - 3 * data["omega_t"] ** 2 * data["omega_rt"]), - data["Z_rttt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rtt"] = rpz2xyz_vec(data["e_theta_rtt"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rtz", - label="\\partial_{\\rho \\theta \\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial, poloidal," - " and toroidal coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_t", - "R_rt", - "R_tt", - "R_rtt", - "R_ttz", - "R_rttz", - "R_tz", - "R_rtz", - "R_z", - "R_rz", - "Z_rttz", - "omega_r", - "omega_t", - "omega_rt", - "omega_tt", - "omega_rtt", - "omega_ttz", - "omega_rttz", - "omega_tz", - "omega_rtz", - "omega_z", - "omega_rz", - "phi", - ], - basis="basis", -) -def _e_sub_theta_rtz(params, transforms, profiles, data, **kwargs): - data["e_theta_rtz"] = jnp.array( - [ - -2 * data["omega_rz"] * data["R_t"] * data["omega_t"] - - 2 - * (1 + data["omega_z"]) - * (data["R_rt"] * data["omega_t"] + data["R_t"] * data["omega_rt"]) - - data["R_rz"] * data["omega_t"] ** 2 - - 2 * data["R_z"] * data["omega_t"] * data["omega_rt"] - - 2 * data["R_r"] * data["omega_t"] * data["omega_tz"] - - 2 - * data["R"] - * ( - data["omega_rt"] * data["omega_tz"] - + data["omega_t"] * data["omega_rtz"] - ) - - data["R_r"] * (1 + data["omega_z"]) * data["omega_tt"] - - data["R"] - * ( - data["omega_rz"] * data["omega_tt"] - + (1 + data["omega_z"]) * data["omega_rtt"] - ) - + data["R_rttz"] - - data["omega_r"] - * ( - 2 * data["omega_t"] * data["R_tz"] - + 2 * data["R_t"] * data["omega_tz"] - + (1 + data["omega_z"]) * data["R_tt"] - + data["R_z"] * data["omega_tt"] - - data["R"] - * ((1 + data["omega_z"]) * data["omega_t"] ** 2 - data["omega_ttz"]) - ), - 2 * data["omega_rt"] * data["R_tz"] - + 2 * data["omega_t"] * data["R_rtz"] - + 2 * data["R_rt"] * data["omega_tz"] - + 2 * data["R_t"] * data["omega_rtz"] - + data["omega_rz"] * data["R_tt"] - + (1 + data["omega_z"]) * data["R_rtt"] - + data["R_rz"] * data["omega_tt"] - + data["R_z"] * data["omega_rtt"] - - data["R_r"] - * ((1 + data["omega_z"]) * data["omega_t"] ** 2 - data["omega_ttz"]) - - data["R"] - * ( - data["omega_rz"] * data["omega_t"] ** 2 - + (1 + data["omega_z"]) * 2 * data["omega_t"] * data["omega_rt"] - - data["omega_rttz"] - ) - + data["omega_r"] - * ( - -2 * (1 + data["omega_z"]) * data["R_t"] * data["omega_t"] - - data["R_z"] * data["omega_t"] ** 2 - - 2 * data["R"] * data["omega_t"] * data["omega_tz"] - - data["R"] * (1 + data["omega_z"]) * data["omega_tt"] - + data["R_ttz"] - ), - data["Z_rttz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rtz"] = rpz2xyz_vec(data["e_theta_rtz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt radial and toroidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rt", - "R_rtz", - "R_rz", - "R_t", - "R_tz", - "R_z", - "Z_rtz", - "omega_r", - "omega_rt", - "omega_rtz", - "omega_rz", - "omega_t", - "omega_tz", - "omega_z", - "phi", - ], - basis="basis", -) -def _e_sub_theta_rz(params, transforms, profiles, data, **kwargs): - data["e_theta_rz"] = jnp.array( - [ - -(1 + data["omega_z"]) * data["R_t"] * data["omega_r"] - - data["R"] * data["omega_tz"] * data["omega_r"] - - data["omega_t"] - * ( - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"] - ) - - data["R"] * (1 + data["omega_z"]) * data["omega_rt"] - + data["R_rtz"], - data["omega_tz"] * data["R_r"] - + data["R_tz"] * data["omega_r"] - + data["omega_t"] * data["R_rz"] - + data["R_t"] * data["omega_rz"] - + data["R_rt"] - + data["omega_z"] * data["R_rt"] - + data["R_z"] * data["omega_rt"] - + data["R"] - * ( - -(1 + data["omega_z"]) * data["omega_t"] * data["omega_r"] - + data["omega_rtz"] - ), - data["Z_rtz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rz"] = rpz2xyz_vec(data["e_theta_rz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_rzz", - label="\\partial_{\\rho \\zeta \\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - " once and toroidal twice" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_t", - "R_rt", - "R_tz", - "R_rtz", - "R_tzz", - "R_rtzz", - "R_z", - "R_rz", - "R_zz", - "R_rzz", - "Z_rtzz", - "omega_t", - "omega_rt", - "omega_tz", - "omega_rtz", - "omega_tzz", - "omega_rtzz", - "omega_r", - "omega_z", - "omega_zz", - "omega_rz", - "omega_rzz", - "phi", - ], - basis="basis", -) -def _e_sub_theta_rzz(params, transforms, profiles, data, **kwargs): - data["e_theta_rzz"] = jnp.array( - [ - -2 * (1 + data["omega_z"]) * data["omega_rz"] * data["R_t"] - - (1 + data["omega_z"]) ** 2 * data["R_rt"] - - 2 * data["R_rz"] * (1 + data["omega_z"]) * data["omega_t"] - - 2 - * data["R_z"] - * ( - data["omega_rz"] * data["omega_t"] - + (1 + data["omega_z"]) * data["omega_rt"] - ) - - data["R_r"] * data["omega_zz"] * data["omega_t"] - - data["R"] - * ( - data["omega_rzz"] * data["omega_t"] - + data["omega_zz"] * data["omega_rt"] - ) - - 2 * data["R_r"] * (1 + data["omega_z"]) * data["omega_tz"] - - 2 - * data["R"] - * ( - data["omega_rz"] * data["omega_tz"] - + (1 + data["omega_z"]) * data["omega_rtz"] - ) - + data["R_rtzz"] - - data["omega_r"] - * ( - data["omega_zz"] * data["R_t"] - + data["R_zz"] * data["omega_t"] - + 2 * (1 + data["omega_z"]) * data["R_tz"] - + 2 * data["R_z"] * data["omega_tz"] - - data["R"] - * ((1 + data["omega_z"]) ** 2 * data["omega_t"] - data["omega_tzz"]) - ), - data["omega_rzz"] * data["R_t"] - + data["omega_zz"] * data["R_rt"] - + data["R_rzz"] * data["omega_t"] - + data["R_zz"] * data["omega_rt"] - + 2 * data["omega_rz"] * data["R_tz"] - + 2 * (1 + data["omega_z"]) * data["R_rtz"] - + 2 * data["R_rz"] * data["omega_tz"] - + 2 * data["R_z"] * data["omega_rtz"] - - data["R_r"] - * ((1 + data["omega_z"]) ** 2 * data["omega_t"] - data["omega_tzz"]) - - data["R"] - * ( - 2 * (1 + data["omega_z"]) * data["omega_rz"] * data["omega_t"] - + (1 + data["omega_z"]) ** 2 * data["omega_rt"] - - data["omega_rtzz"] - ) - + data["omega_r"] - * ( - -((1 + data["omega_z"]) ** 2) * data["R_t"] - - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_t"] - - data["R"] * data["omega_zz"] * data["omega_t"] - - 2 * data["R"] * (1 + data["omega_z"]) * data["omega_tz"] - + data["R_tzz"] - ), - data["Z_rtzz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_rzz"] = rpz2xyz_vec(data["e_theta_rzz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description="Covariant Poloidal basis vector, derivative wrt poloidal coordinate", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_t", "R_tt", "Z_tt", "omega_t", "omega_tt", "phi"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], - basis="basis", -) -def _e_sub_theta_t(params, transforms, profiles, data, **kwargs): - data["e_theta_t"] = jnp.array( - [ - -data["R"] * data["omega_t"] ** 2 + data["R_tt"], - 2 * data["R_t"] * data["omega_t"] + data["R"] * data["omega_tt"], - data["Z_tt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_t"] = rpz2xyz_vec(data["e_theta_t"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_tt", - label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt poloidal and poloidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_t", - "R_tt", - "R_ttt", - "Z_ttt", - "omega_t", - "omega_tt", - "omega_ttt", - "phi", - ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], - basis="basis", -) -def _e_sub_theta_tt(params, transforms, profiles, data, **kwargs): - data["e_theta_tt"] = jnp.array( - [ - -3 * data["R_t"] * data["omega_t"] ** 2 - - 3 * data["R"] * data["omega_t"] * data["omega_tt"] - + data["R_ttt"], - 3 * (data["omega_t"] * data["R_tt"] + data["R_t"] * data["omega_tt"]) - + data["R"] * (-data["omega_t"] ** 3 + data["omega_ttt"]), - data["Z_ttt"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_tt"] = rpz2xyz_vec(data["e_theta_tt"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt poloidal and toroidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_t", - "R_tt", - "R_ttz", - "R_tz", - "R_z", - "Z_ttz", - "omega_t", - "omega_tt", - "omega_ttz", - "omega_tz", - "omega_z", - "phi", - ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.FourierRZToroidalSurface", - ], - basis="basis", -) -def _e_sub_theta_tz(params, transforms, profiles, data, **kwargs): - data["e_theta_tz"] = jnp.array( - [ - -2 * (1 + data["omega_z"]) * data["R_t"] * data["omega_t"] - - data["R_z"] * data["omega_t"] ** 2 - - 2 * data["R"] * data["omega_t"] * data["omega_tz"] - - data["R"] * (1 + data["omega_z"]) * data["omega_tt"] - + data["R_ttz"], - 2 * data["omega_t"] * data["R_tz"] - + 2 * data["R_t"] * data["omega_tz"] - + (1 + data["omega_z"]) * data["R_tt"] - + data["R_z"] * data["omega_tt"] - - data["R"] - * ((1 + data["omega_z"]) * data["omega_t"] ** 2 - data["omega_ttz"]), - data["Z_ttz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_tz"] = rpz2xyz_vec(data["e_theta_tz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description="Covariant Poloidal basis vector, derivative wrt toroidal coordinate", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z", "phi"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.FourierRZToroidalSurface", - ], - basis="basis", -) -def _e_sub_theta_z(params, transforms, profiles, data, **kwargs): - data["e_theta_z"] = jnp.array( - [ - -data["R"] * (1 + data["omega_z"]) * data["omega_t"] + data["R_tz"], - (1 + data["omega_z"]) * data["R_t"] - + data["R_z"] * data["omega_t"] - + data["R"] * data["omega_tz"], - data["Z_tz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_z"] = rpz2xyz_vec(data["e_theta_z"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_theta_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt toroidal and toroidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_t", - "R_tz", - "R_tzz", - "R_z", - "R_zz", - "Z_tzz", - "omega_t", - "omega_tz", - "omega_tzz", - "omega_z", - "omega_zz", - "phi", - ], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.FourierRZToroidalSurface", - ], - basis="basis", -) -def _e_sub_theta_zz(params, transforms, profiles, data, **kwargs): - data["e_theta_zz"] = jnp.array( - [ - -((1 + data["omega_z"]) ** 2) * data["R_t"] - - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_t"] - - data["R"] * data["omega_zz"] * data["omega_t"] - - 2 * data["R"] * (1 + data["omega_z"]) * data["omega_tz"] - + data["R_tzz"], - data["omega_zz"] * data["R_t"] - + data["R_zz"] * data["omega_t"] - + 2 * (1 + data["omega_z"]) * data["R_tz"] - + 2 * data["R_z"] * data["omega_tz"] - - data["R"] - * ((1 + data["omega_z"]) ** 2 * data["omega_t"] - data["omega_tzz"]), - data["Z_tzz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_theta_zz"] = rpz2xyz_vec(data["e_theta_zz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta", - label="\\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant Toroidal basis vector", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_z", "Z_z", "omega_z", "phi"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.FourierRZToroidalSurface", - ], - basis="basis", -) -def _e_sub_zeta(params, transforms, profiles, data, **kwargs): - data["e_zeta"] = jnp.array( - [data["R_z"], data["R"] * (1 + data["omega_z"]), data["Z_z"]] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta"] = rpz2xyz_vec(data["e_zeta"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant Toroidal basis vector, derivative wrt radial coordinate", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z", "phi"], - basis="basis", -) -def _e_sub_zeta_r(params, transforms, profiles, data, **kwargs): - data["e_zeta_r"] = jnp.array( - [ - -data["R"] * (1 + data["omega_z"]) * data["omega_r"] + data["R_rz"], - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"], - data["Z_rz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_r"] = rpz2xyz_vec(data["e_zeta_r"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, second derivative wrt radial and radial" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rrz", - "R_rz", - "R_z", - "Z_rrz", - "omega_r", - "omega_rr", - "omega_rrz", - "omega_rz", - "omega_z", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rr(params, transforms, profiles, data, **kwargs): - data["e_zeta_rr"] = jnp.array( - [ - -2 * (1 + data["omega_z"]) * data["R_r"] * data["omega_r"] - - data["R_z"] * data["omega_r"] ** 2 - - 2 * data["R"] * data["omega_r"] * data["omega_rz"] - - data["R"] * data["omega_rr"] * (1 + data["omega_z"]) - + data["R_rrz"], - 2 * data["omega_r"] * data["R_rz"] - + 2 * data["R_r"] * data["omega_rz"] - + data["R_rr"] * (1 + data["omega_z"]) - + data["R_z"] * data["omega_rr"] - - data["R"] - * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"]), - data["Z_rrz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rr"] = rpz2xyz_vec(data["e_zeta_rr"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rrr", - label="\\partial_{\\rho \\rho \\rho} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, third derivative wrt radial coordinate" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_z", - "R_rr", - "R_rz", - "R_rrr", - "R_rrz", - "R_rrrz", - "Z_rrrz", - "omega_r", - "omega_z", - "omega_rr", - "omega_rz", - "omega_rrr", - "omega_rrz", - "omega_rrrz", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rrr(params, transforms, profiles, data, **kwargs): - data["e_zeta_rrr"] = jnp.array( - [ - -3 * data["R"] * data["omega_rrz"] * data["omega_r"] - - 3 - * data["omega_rz"] - * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) - - 3 - * data["omega_z"] - * (data["R_rr"] * data["omega_r"] + data["R_r"] * data["omega_rr"]) - - (1 + data["omega_z"]) - * data["R"] - * (data["omega_rrr"] - data["omega_r"] ** 3) - - 3 - * data["omega_r"] - * ( - data["R_rz"] * data["omega_r"] - + data["R_z"] * data["omega_rr"] - + data["R_rr"] - ) - - 3 * data["R_r"] * data["omega_rr"] - + data["R_rrrz"], - 3 - * data["R_r"] - * (data["omega_rrz"] - (1 + data["omega_z"]) * data["omega_r"] ** 2) - + 3 * data["omega_rz"] * data["R_rr"] - + data["R"] - * ( - data["omega_rrrz"] - - 3 - * data["omega_r"] - * ( - data["omega_rz"] * data["omega_r"] - + (1 + data["omega_z"]) * data["omega_rr"] - ) - ) - + (1 + data["omega_z"]) * data["R_rrr"] - + 3 * data["R_rrz"] * data["omega_r"] - + 3 * data["R_rz"] * data["omega_rr"] - + data["R_z"] * (data["omega_rrr"] - data["omega_r"] ** 3), - data["Z_rrrz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rrr"] = rpz2xyz_vec(data["e_zeta_rrr"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rrt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, third derivative wrt radial coordinate" - " twice and poloidal once" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rt", - "R_rrt", - "R_rtz", - "R_rrtz", - "R_rz", - "R_rrz", - "R_t", - "R_tz", - "R_z", - "Z_rrtz", - "omega_r", - "omega_rr", - "omega_rt", - "omega_rrt", - "omega_rtz", - "omega_rrtz", - "omega_rz", - "omega_rrz", - "omega_t", - "omega_tz", - "omega_z", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rrt(params, transforms, profiles, data, **kwargs): - data["e_zeta_rrt"] = jnp.array( - [ - -(data["omega_rz"] * data["R_t"] * data["omega_r"]) - - (1 + data["omega_z"]) - * (data["R_rt"] * data["omega_r"] + data["R_t"] * data["omega_rr"]) - - data["R_r"] * data["omega_tz"] * data["omega_r"] - - data["R"] - * ( - data["omega_rtz"] * data["omega_r"] - + data["omega_tz"] * data["omega_rr"] - ) - - data["omega_rt"] - * ( - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"] - ) - - data["omega_t"] - * ( - data["omega_rz"] * data["R_r"] - + (1 + data["omega_z"]) * data["R_rr"] - + data["R_rz"] * data["omega_r"] - + data["R_z"] * data["omega_rr"] - + data["R_r"] * data["omega_rz"] - + data["R"] * data["omega_rrz"] - ) - - data["R_r"] * (1 + data["omega_z"]) * data["omega_rt"] - - data["R"] - * ( - data["omega_rz"] * data["omega_rt"] - + (1 + data["omega_z"]) * data["omega_rrt"] - ) - + data["R_rrtz"] - - data["omega_r"] - * ( - data["omega_tz"] * data["R_r"] - + data["R_tz"] * data["omega_r"] - + data["omega_t"] * data["R_rz"] - + data["R_t"] * data["omega_rz"] - + data["R_rt"] - + data["omega_z"] * data["R_rt"] - + data["R_z"] * data["omega_rt"] - + data["R"] - * ( - -(1 + data["omega_z"]) * data["omega_t"] * data["omega_r"] - + data["omega_rtz"] - ) - ), - data["omega_rtz"] * data["R_r"] - + data["omega_tz"] * data["R_rr"] - + data["R_rtz"] * data["omega_r"] - + data["R_tz"] * data["omega_rr"] - + data["omega_rt"] * data["R_rz"] - + data["omega_t"] * data["R_rrz"] - + data["R_rt"] * data["omega_rz"] - + data["R_t"] * data["omega_rrz"] - + data["R_rrt"] - + data["omega_rz"] * data["R_rt"] - + data["omega_z"] * data["R_rrt"] - + data["R_rz"] * data["omega_rt"] - + data["R_z"] * data["omega_rrt"] - + data["R_r"] - * ( - -(1 + data["omega_z"]) * data["omega_t"] * data["omega_r"] - + data["omega_rtz"] - ) - + data["R"] - * ( - -data["omega_rz"] * data["omega_t"] * data["omega_r"] - - (1 + data["omega_z"]) - * ( - data["omega_rt"] * data["omega_r"] - + data["omega_t"] * data["omega_rr"] - ) - + data["omega_rrtz"] - ) - + data["omega_r"] - * ( - -(1 + data["omega_z"]) * data["R_t"] * data["omega_r"] - - data["R"] * data["omega_tz"] * data["omega_r"] - - data["omega_t"] - * ( - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"] - ) - - data["R"] * (1 + data["omega_z"]) * data["omega_rt"] - + data["R_rtz"] - ), - data["Z_rrtz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rrt"] = rpz2xyz_vec(data["e_zeta_rrt"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rrz", - label="\\partial_{\\rho \\rho \\zeta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, third derivative wrt radial coordinate" - " twice and toroidal once" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rz", - "R_rrz", - "R_rzz", - "R_rrzz", - "R_z", - "R_zz", - "Z_rrzz", - "omega_r", - "omega_rr", - "omega_rz", - "omega_rrz", - "omega_rzz", - "omega_rrzz", - "omega_z", - "omega_zz", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rrz(params, transforms, profiles, data, **kwargs): - data["e_zeta_rrz"] = jnp.array( - [ - -2 * ((1 + data["omega_z"]) * data["omega_rz"]) * data["R_r"] - - ((1 + data["omega_z"]) ** 2) * data["R_rr"] - - 2 * data["R_rz"] * (1 + data["omega_z"]) * data["omega_r"] - - 2 - * data["R_z"] - * ( - data["omega_rz"] * data["omega_r"] - + (1 + data["omega_z"]) * data["omega_rr"] - ) - - data["R_r"] * data["omega_zz"] * data["omega_r"] - - data["R"] - * ( - data["omega_rzz"] - * data["omega_r"] - * data["omega_zz"] - * data["omega_rr"] - ) - - 2 * data["R_r"] * (1 + data["omega_z"]) * data["omega_rz"] - - 2 - * data["R"] - * (data["omega_rz"] ** 2 + (1 + data["omega_z"]) * data["omega_rrz"]) - + data["R_rrzz"] - - data["omega_r"] - * ( - data["omega_zz"] * data["R_r"] - + data["R_zz"] * data["omega_r"] - + 2 * (1 + data["omega_z"]) * data["R_rz"] - + 2 * data["R_z"] * data["omega_rz"] - - data["R"] - * ((1 + data["omega_z"]) ** 2 * data["omega_r"] - data["omega_rzz"]) - ), - data["omega_rzz"] * data["R_r"] - + data["omega_zz"] * data["R_rr"] - + data["R_rzz"] * data["omega_r"] - + data["R_zz"] * data["omega_rr"] - + 2 * data["omega_rz"] * data["R_rz"] - + 2 * (1 + data["omega_z"]) * data["R_rrz"] - + 2 * data["R_rz"] * data["omega_rz"] - + 2 * data["R_z"] * data["omega_rrz"] - - data["R_r"] - * ((1 + data["omega_z"]) ** 2 * data["omega_r"] - data["omega_rzz"]) - - data["R"] - * ( - 2 * (1 + data["omega_z"]) * data["omega_rz"] * data["omega_r"] - + (1 + data["omega_z"]) ** 2 * data["omega_rr"] - - data["omega_rrzz"] - ) - + data["omega_r"] - * ( - -((1 + data["omega_z"]) ** 2) * data["R_r"] - - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_r"] - - data["R"] * data["omega_zz"] * data["omega_r"] - - 2 * data["R"] * (1 + data["omega_z"]) * data["omega_rz"] - + data["R_rzz"] - ), - data["Z_rrzz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rrz"] = rpz2xyz_vec(data["e_zeta_rrz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, second derivative wrt radial and poloidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rt", - "R_rtz", - "R_rz", - "R_t", - "R_tz", - "R_z", - "Z_rtz", - "omega_r", - "omega_rt", - "omega_rtz", - "omega_rz", - "omega_t", - "omega_tz", - "omega_z", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rt(params, transforms, profiles, data, **kwargs): - data["e_zeta_rt"] = jnp.array( - [ - -(1 + data["omega_z"]) * data["R_t"] * data["omega_r"] - - data["R"] * data["omega_tz"] * data["omega_r"] - - data["omega_t"] - * ( - (1 + data["omega_z"]) * data["R_r"] - + data["R_z"] * data["omega_r"] - + data["R"] * data["omega_rz"] - ) - - data["R"] * (1 + data["omega_z"]) * data["omega_rt"] - + data["R_rtz"], - data["omega_tz"] * data["R_r"] - + data["R_tz"] * data["omega_r"] - + data["omega_t"] * data["R_rz"] - + data["R_t"] * data["omega_rz"] - + data["R_rt"] - + data["omega_z"] * data["R_rt"] - + data["R_z"] * data["omega_rt"] - + data["R"] - * ( - -(1 + data["omega_z"]) * data["omega_t"] * data["omega_r"] - + data["omega_rtz"] - ), - data["Z_rtz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rt"] = rpz2xyz_vec(data["e_zeta_rt"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rtt", - label="\\partial_{\\rho \\theta \\theta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, third derivative wrt radial coordinate" - " once and poloidal twice" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_t", - "R_rt", - "R_tt", - "R_rtt", - "R_ttz", - "R_rttz", - "R_tz", - "R_rtz", - "R_z", - "R_rz", - "Z_rttz", - "omega_r", - "omega_t", - "omega_rt", - "omega_tt", - "omega_rtt", - "omega_ttz", - "omega_rttz", - "omega_tz", - "omega_rtz", - "omega_z", - "omega_rz", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rtt(params, transforms, profiles, data, **kwargs): - data["e_zeta_rtt"] = jnp.array( - [ - -2 * data["omega_rz"] * data["R_t"] * data["omega_t"] - - 2 - * (1 + data["omega_z"]) - * (data["R_rt"] * data["omega_t"] + data["R_t"] * data["omega_rt"]) - - data["R_rz"] * data["omega_t"] ** 2 - - data["R_z"] * 2 * data["omega_t"] * data["omega_rt"] - - 2 * data["R_r"] * data["omega_t"] * data["omega_tz"] - - 2 - * data["R"] - * ( - data["omega_rt"] * data["omega_tz"] - + data["omega_t"] * data["omega_rtz"] - ) - - data["R_r"] * (1 + data["omega_z"]) * data["omega_tt"] - - data["R"] - * ( - data["omega_rz"] * data["omega_tt"] - + (1 + data["omega_z"]) * data["omega_rtt"] - ) - + data["R_rttz"] + + data["R_rttz"] - data["omega_r"] * ( 2 * data["omega_t"] * data["R_tz"] @@ -4296,18 +2807,18 @@ def _e_sub_zeta_rtt(params, transforms, profiles, data, **kwargs): ] ).T if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rtt"] = rpz2xyz_vec(data["e_zeta_rtt"], phi=data["phi"]) + data["e_theta_rtz"] = rpz2xyz_vec(data["e_theta_rtz"], phi=data["phi"]) return data @register_compute_fun( - name="e_zeta_rtz", - label="\\partial_{\\rho \\theta \\zeta} \\mathbf{e}_{\\zeta}", + name="e_theta_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", description=( - "Covariant Toroidal basis vector, third derivative wrt radial, poloidal," - " and toroidal coordinates" + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " once and toroidal twice" ), dim=3, params=[], @@ -4328,23 +2839,24 @@ def _e_sub_zeta_rtt(params, transforms, profiles, data, **kwargs): "R_zz", "R_rzz", "Z_rtzz", - "omega_r", "omega_t", "omega_rt", "omega_tz", "omega_rtz", "omega_tzz", "omega_rtzz", + "omega_r", "omega_z", - "omega_rz", "omega_zz", + "omega_rz", "omega_rzz", "phi", ], basis="basis", + aliases=["e_rho_tzz", "e_zeta_rtz"], ) -def _e_sub_zeta_rtz(params, transforms, profiles, data, **kwargs): - data["e_zeta_rtz"] = jnp.array( +def _e_sub_theta_rzz(params, transforms, profiles, data, **kwargs): + data["e_theta_rzz"] = jnp.array( [ -2 * (1 + data["omega_z"]) * data["omega_rz"] * data["R_t"] - (1 + data["omega_z"]) ** 2 * data["R_rt"] @@ -4358,10 +2870,8 @@ def _e_sub_zeta_rtz(params, transforms, profiles, data, **kwargs): - data["R_r"] * data["omega_zz"] * data["omega_t"] - data["R"] * ( - data["omega_rzz"] - * data["omega_t"] - * data["omega_zz"] - * data["omega_rt"] + data["omega_rzz"] * data["omega_t"] + + data["omega_zz"] * data["omega_rt"] ) - 2 * data["R_r"] * (1 + data["omega_z"]) * data["omega_tz"] - 2 @@ -4386,217 +2896,117 @@ def _e_sub_zeta_rtz(params, transforms, profiles, data, **kwargs): + data["R_zz"] * data["omega_rt"] + 2 * data["omega_rz"] * data["R_tz"] + 2 * (1 + data["omega_z"]) * data["R_rtz"] - + 2 * data["R_rz"] * data["omega_tz"] - + 2 * data["R_z"] * data["omega_rtz"] - - data["R_r"] - * ((1 + data["omega_z"]) ** 2 * data["omega_t"] - data["omega_tzz"]) - - data["R"] - * ( - 2 * (1 + data["omega_z"]) * data["omega_rz"] * data["omega_t"] - + (1 + data["omega_z"]) ** 2 * data["omega_rt"] - - data["omega_rtzz"] - ) - + data["omega_r"] - * ( - -((1 + data["omega_z"]) ** 2) * data["R_t"] - - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_t"] - - data["R"] * data["omega_zz"] * data["omega_t"] - - 2 * data["R"] * (1 + data["omega_z"]) * data["omega_tz"] - + data["R_tzz"] - ), - data["Z_rtzz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rtz"] = rpz2xyz_vec(data["e_zeta_rtz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, second derivative wrt radial and toroidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rz", - "R_rzz", - "R_z", - "R_zz", - "Z_rzz", - "omega_r", - "omega_rz", - "omega_rzz", - "omega_z", - "omega_zz", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rz(params, transforms, profiles, data, **kwargs): - data["e_zeta_rz"] = jnp.array( - [ - -((1 + data["omega_z"]) ** 2) * data["R_r"] - - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_r"] - - data["R"] * data["omega_zz"] * data["omega_r"] - - 2 * data["R"] * (1 + data["omega_z"]) * data["omega_rz"] - + data["R_rzz"], - data["omega_zz"] * data["R_r"] - + data["R_zz"] * data["omega_r"] - + 2 * (1 + data["omega_z"]) * data["R_rz"] - + 2 * data["R_z"] * data["omega_rz"] - - data["R"] - * ((1 + data["omega_z"]) ** 2 * data["omega_r"] - data["omega_rzz"]), - data["Z_rzz"], - ] - ).T - if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rz"] = rpz2xyz_vec(data["e_zeta_rz"], phi=data["phi"]) - return data - - -@register_compute_fun( - name="e_zeta_rzz", - label="\\partial_{\\rho \\zeta \\zeta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, third derivative wrt radial coordinate" - " once and toroidal twice" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_z", - "R_rz", - "R_zz", - "R_rzz", - "R_zzz", - "R_rzzz", - "Z_rzzz", - "omega_r", - "omega_z", - "omega_rz", - "omega_zz", - "omega_rzz", - "omega_zzz", - "omega_rzzz", - "phi", - ], - basis="basis", -) -def _e_sub_zeta_rzz(params, transforms, profiles, data, **kwargs): - data["e_zeta_rzz"] = jnp.array( - [ - -3 * data["R_rz"] * (1 + data["omega_z"]) ** 2 - - 6 * data["R_z"] * (1 + data["omega_z"]) * data["omega_rz"] - - 3 * data["R_r"] * (1 + data["omega_z"]) * data["omega_zz"] - - 3 - * data["R"] - * ( - data["omega_rz"] * data["omega_zz"] - + (1 + data["omega_z"]) * data["omega_rzz"] - ) - + data["R_rzzz"] - - data["omega_r"] - * ( - 3 * (1 + data["omega_z"]) * data["R_zz"] - + 3 * data["R_z"] * data["omega_zz"] - - data["R"] - * ( - 1 - + 3 * data["omega_z"] - + 3 * data["omega_z"] ** 2 - + data["omega_z"] ** 3 - - data["omega_zzz"] - ) - ), - 3 * data["omega_rz"] * data["R_zz"] - + 3 * (1 + data["omega_z"]) * data["R_rzz"] - + 3 * data["R_rz"] * data["omega_zz"] - + 3 * data["R_z"] * data["omega_rzz"] + + 2 * data["R_rz"] * data["omega_tz"] + + 2 * data["R_z"] * data["omega_rtz"] - data["R_r"] - * ( - 1 - + 3 * data["omega_z"] - + 3 * data["omega_z"] ** 2 - + data["omega_z"] ** 3 - - data["omega_zzz"] - ) + * ((1 + data["omega_z"]) ** 2 * data["omega_t"] - data["omega_tzz"]) - data["R"] * ( - 3 * data["omega_rz"] * (1 + data["omega_z"] * (1 + data["omega_z"])) - - data["omega_rzzz"] + 2 * (1 + data["omega_z"]) * data["omega_rz"] * data["omega_t"] + + (1 + data["omega_z"]) ** 2 * data["omega_rt"] + - data["omega_rtzz"] ) + data["omega_r"] * ( - -3 * data["R_z"] * (1 + data["omega_z"]) ** 2 - - 3 * data["R"] * (1 + data["omega_z"]) * data["omega_zz"] - + data["R_zzz"] + -((1 + data["omega_z"]) ** 2) * data["R_t"] + - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_t"] + - data["R"] * data["omega_zz"] * data["omega_t"] + - 2 * data["R"] * (1 + data["omega_z"]) * data["omega_tz"] + + data["R_tzz"] ), - data["Z_rzzz"], + data["Z_rtzz"], ] ).T if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_rzz"] = rpz2xyz_vec(data["e_zeta_rzz"], phi=data["phi"]) + data["e_theta_rzz"] = rpz2xyz_vec(data["e_theta_rzz"], phi=data["phi"]) return data @register_compute_fun( - name="e_zeta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant Toroidal basis vector, derivative wrt poloidal coordinate", + description="Covariant Poloidal basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z", "phi"], + data=["R", "R_t", "R_tt", "Z_tt", "omega_t", "omega_tt", "phi"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.FourierRZToroidalSurface", + "desc.geometry.core.Surface", ], basis="basis", ) -def _e_sub_zeta_t(params, transforms, profiles, data, **kwargs): - data["e_zeta_t"] = jnp.array( +def _e_sub_theta_t(params, transforms, profiles, data, **kwargs): + data["e_theta_t"] = jnp.array( [ - -data["R"] * (1 + data["omega_z"]) * data["omega_t"] + data["R_tz"], - (1 + data["omega_z"]) * data["R_t"] - + data["R_z"] * data["omega_t"] - + data["R"] * data["omega_tz"], - data["Z_tz"], + -data["R"] * data["omega_t"] ** 2 + data["R_tt"], + 2 * data["R_t"] * data["omega_t"] + data["R"] * data["omega_tt"], + data["Z_tt"], + ] + ).T + if kwargs.get("basis", "rpz").lower() == "xyz": + data["e_theta_t"] = rpz2xyz_vec(data["e_theta_t"], phi=data["phi"]) + return data + + +@register_compute_fun( + name="e_theta_tt", + label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description=( + "Covariant Poloidal basis vector, second derivative wrt poloidal and poloidal" + " coordinates" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "R", + "R_t", + "R_tt", + "R_ttt", + "Z_ttt", + "omega_t", + "omega_tt", + "omega_ttt", + "phi", + ], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], + basis="basis", +) +def _e_sub_theta_tt(params, transforms, profiles, data, **kwargs): + data["e_theta_tt"] = jnp.array( + [ + -3 * data["R_t"] * data["omega_t"] ** 2 + - 3 * data["R"] * data["omega_t"] * data["omega_tt"] + + data["R_ttt"], + 3 * (data["omega_t"] * data["R_tt"] + data["R_t"] * data["omega_tt"]) + + data["R"] * (-data["omega_t"] ** 3 + data["omega_ttt"]), + data["Z_ttt"], ] ).T if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_t"] = rpz2xyz_vec(data["e_zeta_t"], phi=data["phi"]) + data["e_theta_tt"] = rpz2xyz_vec(data["e_theta_tt"], phi=data["phi"]) return data @register_compute_fun( - name="e_zeta_tt", - label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\zeta}", + name="e_theta_tz", + label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", description=( - "Covariant Toroidal basis vector, second derivative wrt poloidal and poloidal" + "Covariant Poloidal basis vector, second derivative wrt poloidal and toroidal" " coordinates" ), dim=3, @@ -4624,9 +3034,10 @@ def _e_sub_zeta_t(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.FourierRZToroidalSurface", ], basis="basis", + aliases=["e_zeta_tt"], ) -def _e_sub_zeta_tt(params, transforms, profiles, data, **kwargs): - data["e_zeta_tt"] = jnp.array( +def _e_sub_theta_tz(params, transforms, profiles, data, **kwargs): + data["e_theta_tz"] = jnp.array( [ -2 * (1 + data["omega_z"]) * data["R_t"] * data["omega_t"] - data["R_z"] * data["omega_t"] ** 2 @@ -4643,17 +3054,51 @@ def _e_sub_zeta_tt(params, transforms, profiles, data, **kwargs): ] ).T if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_tt"] = rpz2xyz_vec(data["e_zeta_tt"], phi=data["phi"]) + data["e_theta_tz"] = rpz2xyz_vec(data["e_theta_tz"], phi=data["phi"]) + return data + + +@register_compute_fun( + name="e_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector, derivative wrt toroidal coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z", "phi"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], + basis="basis", + aliases=["e_zeta_t"], +) +def _e_sub_theta_z(params, transforms, profiles, data, **kwargs): + data["e_theta_z"] = jnp.array( + [ + -data["R"] * (1 + data["omega_z"]) * data["omega_t"] + data["R_tz"], + (1 + data["omega_z"]) * data["R_t"] + + data["R_z"] * data["omega_t"] + + data["R"] * data["omega_tz"], + data["Z_tz"], + ] + ).T + if kwargs.get("basis", "rpz").lower() == "xyz": + data["e_theta_z"] = rpz2xyz_vec(data["e_theta_z"], phi=data["phi"]) return data @register_compute_fun( - name="e_zeta_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\zeta}", + name="e_theta_zz", + label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", description=( - "Covariant Toroidal basis vector, second derivative wrt poloidal and toroidal" + "Covariant Poloidal basis vector, second derivative wrt toroidal and toroidal" " coordinates" ), dim=3, @@ -4681,9 +3126,10 @@ def _e_sub_zeta_tt(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.FourierRZToroidalSurface", ], basis="basis", + aliases=["e_zeta_tz"], ) -def _e_sub_zeta_tz(params, transforms, profiles, data, **kwargs): - data["e_zeta_tz"] = jnp.array( +def _e_sub_theta_zz(params, transforms, profiles, data, **kwargs): + data["e_theta_zz"] = jnp.array( [ -((1 + data["omega_z"]) ** 2) * data["R_t"] - 2 * data["R_z"] * (1 + data["omega_z"]) * data["omega_t"] @@ -4700,7 +3146,126 @@ def _e_sub_zeta_tz(params, transforms, profiles, data, **kwargs): ] ).T if kwargs.get("basis", "rpz").lower() == "xyz": - data["e_zeta_tz"] = rpz2xyz_vec(data["e_zeta_tz"], phi=data["phi"]) + data["e_theta_zz"] = rpz2xyz_vec(data["e_theta_zz"], phi=data["phi"]) + return data + + +@register_compute_fun( + name="e_zeta", + label="\\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant Toroidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_z", "Z_z", "omega_z", "phi"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], + basis="basis", +) +def _e_sub_zeta(params, transforms, profiles, data, **kwargs): + data["e_zeta"] = jnp.array( + [data["R_z"], data["R"] * (1 + data["omega_z"]), data["Z_z"]] + ).T + if kwargs.get("basis", "rpz").lower() == "xyz": + data["e_zeta"] = rpz2xyz_vec(data["e_zeta"], phi=data["phi"]) + return data + + +@register_compute_fun( + name="e_zeta_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + " once and toroidal twice" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "R", + "R_r", + "R_z", + "R_rz", + "R_zz", + "R_rzz", + "R_zzz", + "R_rzzz", + "Z_rzzz", + "omega_r", + "omega_z", + "omega_rz", + "omega_zz", + "omega_rzz", + "omega_zzz", + "omega_rzzz", + "phi", + ], + basis="basis", +) +def _e_sub_zeta_rzz(params, transforms, profiles, data, **kwargs): + data["e_zeta_rzz"] = jnp.array( + [ + -3 * data["R_rz"] * (1 + data["omega_z"]) ** 2 + - 6 * data["R_z"] * (1 + data["omega_z"]) * data["omega_rz"] + - 3 * data["R_r"] * (1 + data["omega_z"]) * data["omega_zz"] + - 3 + * data["R"] + * ( + data["omega_rz"] * data["omega_zz"] + + (1 + data["omega_z"]) * data["omega_rzz"] + ) + + data["R_rzzz"] + - data["omega_r"] + * ( + 3 * (1 + data["omega_z"]) * data["R_zz"] + + 3 * data["R_z"] * data["omega_zz"] + - data["R"] + * ( + 1 + + 3 * data["omega_z"] + + 3 * data["omega_z"] ** 2 + + data["omega_z"] ** 3 + - data["omega_zzz"] + ) + ), + 3 * data["omega_rz"] * data["R_zz"] + + 3 * (1 + data["omega_z"]) * data["R_rzz"] + + 3 * data["R_rz"] * data["omega_zz"] + + 3 * data["R_z"] * data["omega_rzz"] + - data["R_r"] + * ( + 1 + + 3 * data["omega_z"] + + 3 * data["omega_z"] ** 2 + + data["omega_z"] ** 3 + - data["omega_zzz"] + ) + - data["R"] + * ( + 3 * data["omega_rz"] * (1 + data["omega_z"] * (1 + data["omega_z"])) + - data["omega_rzzz"] + ) + + data["omega_r"] + * ( + -3 * data["R_z"] * (1 + data["omega_z"]) ** 2 + - 3 * data["R"] * (1 + data["omega_z"]) * data["omega_zz"] + + data["R_zzz"] + ), + data["Z_rzzz"], + ] + ).T + if kwargs.get("basis", "rpz").lower() == "xyz": + data["e_zeta_rzz"] = rpz2xyz_vec(data["e_zeta_rzz"], phi=data["phi"]) return data diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 270c6f6612..d9d736aa95 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -164,7 +164,7 @@ def _V_rrr_of_r(params, transforms, profiles, data, **kwargs): def _A_of_z(params, transforms, profiles, data, **kwargs): data["A(z)"] = surface_integrals( transforms["grid"], - jnp.abs(data["|e_rho x e_theta|"]), + data["|e_rho x e_theta|"], surface_label="zeta", expand_out=True, ) diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 1ccb989da1..2368deab44 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -307,6 +307,7 @@ def _e_rho_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa data=[], parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + aliases=["e_theta_r"], ) def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -328,8 +329,12 @@ def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.surface.FourierRZToroidalSurface", + parameterization=[ + "desc.geometry.surface.FourierRZToroidalSurface", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + aliases=["e_zeta_r"], ) def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -337,29 +342,6 @@ def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar return data -@register_compute_fun( - name="e_theta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description="Covariant poloidal basis vector, derivative wrt radial coordinate", - dim=3, - params=[], - transforms={ - "grid": [], - }, - profiles=[], - coordinates="tz", - data=[], - parameterization="desc.geometry.surface.FourierRZToroidalSurface", - basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", -) -def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["e_theta_r"] = coords - return data - - @register_compute_fun( name="e_theta_rr", label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}", @@ -377,6 +359,7 @@ def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kw data=[], parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + aliases=["e_rho_rt"], ) def _e_theta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -384,29 +367,6 @@ def _e_theta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **k return data -@register_compute_fun( - name="e_zeta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant toroidal basis vector, derivative wrt radial coordinate", - dim=3, - params=[], - transforms={ - "grid": [], - }, - profiles=[], - coordinates="tz", - data=[], - parameterization="desc.geometry.surface.FourierRZToroidalSurface", - basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", -) -def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["e_zeta_r"] = coords - return data - - @register_compute_fun( name="e_zeta_rr", label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", @@ -422,8 +382,12 @@ def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.surface.FourierRZToroidalSurface", + parameterization=[ + "desc.geometry.surface.FourierRZToroidalSurface", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + aliases=["e_rho_rz"], ) def _e_zeta_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -649,29 +613,6 @@ def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwarg return data -@register_compute_fun( - name="e_rho_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", - units="m", - units_long="meters", - description="Covariant radial basis vector, derivative wrt toroidal angle", - dim=3, - params=[], - transforms={ - "grid": [], - }, - profiles=[], - coordinates="rt", - data=[], - parameterization="desc.geometry.surface.ZernikeRZToroidalSection", - basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", -) -def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["e_rho_z"] = coords - return data - - @register_compute_fun( name="e_theta_z", label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", @@ -688,6 +629,7 @@ def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar data=[], parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + aliases=["e_zeta_t"], ) def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -695,76 +637,6 @@ def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw return data -@register_compute_fun( - name="e_zeta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant toroidal basis vector, derivative wrt radial coordinate", - dim=3, - params=[], - transforms={ - "grid": [], - }, - profiles=[], - coordinates="rt", - data=[], - parameterization="desc.geometry.surface.ZernikeRZToroidalSection", - basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", -) -def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["e_zeta_r"] = coords - return data - - -@register_compute_fun( - name="e_zeta_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant toroidal basis vector," - " second derivative wrt radial coordinate", - dim=3, - params=[], - transforms={ - "grid": [], - }, - profiles=[], - coordinates="rt", - data=[], - parameterization="desc.geometry.surface.ZernikeRZToroidalSection", - basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", -) -def _e_zeta_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["e_zeta_rr"] = coords - return data - - -@register_compute_fun( - name="e_zeta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant toroidal basis vector, derivative wrt poloidal angle", - dim=3, - params=[], - transforms={ - "grid": [], - }, - profiles=[], - coordinates="rt", - data=[], - parameterization="desc.geometry.surface.ZernikeRZToroidalSection", - basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", -) -def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["e_zeta_t"] = coords - return data - - @register_compute_fun( name="e_zeta_z", label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index a6b2a82967..eebaadfca0 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -8,8 +8,7 @@ def find_permutations(primary, separator="_"): """Finds permutations of quantity names for aliases.""" - split_name = primary.split(separator) - primary_permutation = split_name[-1] + prefix, primary_permutation = primary.rsplit(separator, 1) primary_permutation = deque(primary_permutation) new_permutations = [] @@ -18,10 +17,7 @@ def find_permutations(primary, separator="_"): new_permutations.append(list(primary_permutation)) # join new permutation to form alias keys - aliases = [ - "".join(split_name[:-1]) + separator + "".join(perm) - for perm in new_permutations - ] + aliases = [prefix + separator + "".join(perm) for perm in new_permutations] aliases = np.unique(aliases) aliases = np.delete(aliases, np.where(aliases == primary)) @@ -63,7 +59,7 @@ def register_compute_fun( profiles, coordinates, data, - aliases=[], + aliases=None, parameterization="desc.equilibrium.equilibrium.Equilibrium", grid_type=None, axis_limit_data=None, @@ -98,7 +94,7 @@ def register_compute_fun( a flux function, etc. data : list of str Names of other items in the data index needed to compute qty. - aliases : list + aliases : list of str Aliases of `name`. Will be stored in the data dictionary as a copy of `name`s data. parameterization : str or list of str @@ -114,6 +110,8 @@ def register_compute_fun( Should only list *direct* dependencies. The full dependencies will be built recursively at runtime using each quantity's direct dependencies. """ + if aliases is None: + aliases = [] if not isinstance(parameterization, (tuple, list)): parameterization = [parameterization] @@ -127,9 +125,16 @@ def register_compute_fun( } for kw in kwargs: allowed_kwargs.add(kw) - permutable_names = ["R_", "Z_", "phi_", "lambda_", "omega_"] - if not aliases and "".join(name.split("_")[:-1]) + "_" in permutable_names: - aliases = find_permutations(name) + splits = name.rsplit("_", 1) + if ( + len(splits) > 1 + # Only look for permutations of partial derivatives of same coordinate system. + and {"r", "t", "z"}.issuperset(splits[-1]) + ): + aliases_temp = np.append(np.array(aliases), find_permutations(name)) + for alias in aliases: + aliases_temp = np.append(aliases_temp, find_permutations(alias)) + aliases = np.unique(aliases_temp) def _decorator(func): d = { diff --git a/tests/inputs/master_compute_data.pkl b/tests/inputs/master_compute_data.pkl index 3b3da05f77..836a76b534 100644 Binary files a/tests/inputs/master_compute_data.pkl and b/tests/inputs/master_compute_data.pkl differ diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 53728cb180..52fab38dbe 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -68,8 +68,9 @@ def test_aliases(): # manual case primary_data = eq.compute("e_rho_rt") - alias_data = eq.compute("x_rrt") + alias_data = eq.compute(["x_rrt", "e_theta_rr"]) np.testing.assert_allclose(primary_data["e_rho_rt"], alias_data["x_rrt"]) + np.testing.assert_allclose(primary_data["e_rho_rt"], alias_data["e_theta_rr"]) @pytest.mark.unit