diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index a8a7e1eb32..39b8cd2ce3 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -1,24 +1,1353 @@ +"""Compute functions for coordinate system basis vectors. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" + from desc.backend import jnp from .data_index import register_compute_fun from .utils import cross +@register_compute_fun( + name="b", + label="\\hat{b}", + units="~", + units_long="None", + description="Unit vector along magnetic field", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B"], +) +def _b(params, transforms, profiles, data, **kwargs): + data["b"] = (data["B"].T / jnp.linalg.norm(data["B"], axis=-1)).T + return data + + +@register_compute_fun( + name="e^rho", + label="\\mathbf{e}^{\\rho}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant radial basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta/sqrt(g)", "e_zeta"], +) +def _e_sup_rho(params, transforms, profiles, data, **kwargs): + data["e^rho"] = cross(data["e_theta/sqrt(g)"], data["e_zeta"]) + return data + + +@register_compute_fun( + name="e^rho_r", + label="\\partial{\\rho} \\mathbf{e}^{\\rho}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant radial basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r", "sqrt(g)", "sqrt(g)_r"], + axis_limit_data=["e_theta_rr", "sqrt(g)_rr"], +) +def _e_sup_rho_r(params, transforms, profiles, data, **kwargs): + a = cross(data["e_theta_r"], data["e_zeta"]) + data["e^rho_r"] = transforms["grid"].replace_at_axis( + ( + (a + cross(data["e_theta"], data["e_zeta_r"])).T / data["sqrt(g)"] + - cross(data["e_theta"], data["e_zeta"]).T + * data["sqrt(g)_r"] + / data["sqrt(g)"] ** 2 + ).T, + lambda: ( + ( + cross(data["e_theta_rr"], data["e_zeta"]) + + 2 * cross(data["e_theta_r"], data["e_zeta_r"]) + ).T + / (2 * data["sqrt(g)_r"]) + - a.T * data["sqrt(g)_rr"] / (2 * data["sqrt(g)_r"] ** 2) + ).T, + ) + return data + + +@register_compute_fun( + name="e^rho_t", + label="\\partial{\\theta} \\mathbf{e}^{\\rho}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant radial basis vector, derivative wrt poloidal coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "e_zeta", "e_theta_t", "e_zeta_t", "sqrt(g)", "sqrt(g)_t"], + axis_limit_data=["e_theta_r", "e_theta_rt", "sqrt(g)_r", "sqrt(g)_rt"], +) +def _e_sup_rho_t(params, transforms, profiles, data, **kwargs): + data["e^rho_t"] = transforms["grid"].replace_at_axis( + ( + ( + cross(data["e_theta_t"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_t"]) + ).T + / data["sqrt(g)"] + - cross(data["e_theta"], data["e_zeta"]).T + * data["sqrt(g)_t"] + / data["sqrt(g)"] ** 2 + ).T, + lambda: ( + ( + cross(data["e_theta_r"], data["e_zeta_t"]) + + cross(data["e_theta_rt"], data["e_zeta"]) + ).T + / data["sqrt(g)_r"] + - cross(data["e_theta_r"], data["e_zeta"]).T + * data["sqrt(g)_rt"] + / data["sqrt(g)_r"] ** 2 + ).T, + ) + return data + + +@register_compute_fun( + name="e^rho_z", + label="\\partial{\\zeta} \\mathbf{e}^{\\rho}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant radial basis vector, derivative wrt toroidal coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "e_zeta", "e_theta_z", "e_zeta_z", "sqrt(g)", "sqrt(g)_z"], + axis_limit_data=["e_theta_r", "e_theta_rz", "sqrt(g)_r", "sqrt(g)_rz"], +) +def _e_sup_rho_z(params, transforms, profiles, data, **kwargs): + data["e^rho_z"] = transforms["grid"].replace_at_axis( + ( + ( + cross(data["e_theta_z"], data["e_zeta"]) + + cross(data["e_theta"], data["e_zeta_z"]) + ).T + / data["sqrt(g)"] + - cross(data["e_theta"], data["e_zeta"]).T + * data["sqrt(g)_z"] + / data["sqrt(g)"] ** 2 + ).T, + lambda: ( + ( + cross(data["e_theta_r"], data["e_zeta_z"]) + + cross(data["e_theta_rz"], data["e_zeta"]) + ).T + / data["sqrt(g)_r"] + - cross(data["e_theta_r"], data["e_zeta"]).T + * data["sqrt(g)_rz"] + / data["sqrt(g)_r"] ** 2 + ).T, + ) + return data + + +@register_compute_fun( + name="e^theta", + label="\\mathbf{e}^{\\theta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant poloidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e^theta*sqrt(g)", "sqrt(g)"], +) +def _e_sup_theta(params, transforms, profiles, data, **kwargs): + data["e^theta"] = (data["e^theta*sqrt(g)"].T / data["sqrt(g)"]).T + return data + + +@register_compute_fun( + name="e^theta*sqrt(g)", + label="\\mathbf{e}^{\\theta} \\sqrt{g}", + units="m^{2}", + units_long="square meters", + description="Contravariant poloidal basis vector weighted by 3-D volume Jacobian", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_zeta"], +) +def _e_sup_theta_times_sqrt_g(params, transforms, profiles, data, **kwargs): + data["e^theta*sqrt(g)"] = cross(data["e_zeta"], data["e_rho"]) + return data + + +@register_compute_fun( + name="e^theta_r", + label="\\partial{\\rho} \\mathbf{e}^{\\theta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant poloidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta", "e_rho", "e_zeta_r", "e_rho_r", "sqrt(g)", "sqrt(g)_r"], +) +def _e_sup_theta_r(params, transforms, profiles, data, **kwargs): + data["e^theta_r"] = ( + ( + cross(data["e_zeta_r"], data["e_rho"]) + + cross(data["e_zeta"], data["e_rho_r"]) + ).T + / data["sqrt(g)"] + - cross(data["e_zeta"], data["e_rho"]).T + * data["sqrt(g)_r"] + / data["sqrt(g)"] ** 2 + ).T + return data + + +@register_compute_fun( + name="e^theta_t", + label="\\partial{\\theta} \\mathbf{e}^{\\theta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant poloidal basis vector, derivative wrt poloidal" + " coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta", "e_rho", "e_zeta_t", "e_rho_t", "sqrt(g)", "sqrt(g)_t"], +) +def _e_sup_theta_t(params, transforms, profiles, data, **kwargs): + data["e^theta_t"] = ( + ( + cross(data["e_zeta_t"], data["e_rho"]) + + cross(data["e_zeta"], data["e_rho_t"]) + ).T + / data["sqrt(g)"] + - cross(data["e_zeta"], data["e_rho"]).T + * data["sqrt(g)_t"] + / data["sqrt(g)"] ** 2 + ).T + return data + + +@register_compute_fun( + name="e^theta_z", + label="\\partial{\\zeta} \\mathbf{e}^{\\theta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant poloidal basis vector, derivative wrt toroidal" + " coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta", "e_rho", "e_zeta_z", "e_rho_z", "sqrt(g)", "sqrt(g)_z"], +) +def _e_sup_theta_z(params, transforms, profiles, data, **kwargs): + data["e^theta_z"] = ( + ( + cross(data["e_zeta_z"], data["e_rho"]) + + cross(data["e_zeta"], data["e_rho_z"]) + ).T + / data["sqrt(g)"] + - cross(data["e_zeta"], data["e_rho"]).T + * data["sqrt(g)_z"] + / data["sqrt(g)"] ** 2 + ).T + return data + + +@register_compute_fun( + name="e^zeta", + label="\\mathbf{e}^{\\zeta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant toroidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_theta/sqrt(g)"], +) +def _e_sup_zeta(params, transforms, profiles, data, **kwargs): + data["e^zeta"] = cross(data["e_rho"], data["e_theta/sqrt(g)"]) + return data + + +@register_compute_fun( + name="e^zeta_r", + label="\\partial{\\rho} \\mathbf{e}^{\\zeta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant toroidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_rho_r", "e_theta", "e_theta_r", "sqrt(g)", "sqrt(g)_r"], + axis_limit_data=["e_theta_rr", "sqrt(g)_rr"], +) +def _e_sup_zeta_r(params, transforms, profiles, data, **kwargs): + b = cross(data["e_rho"], data["e_theta_r"]) + data["e^zeta_r"] = transforms["grid"].replace_at_axis( + ( + (cross(data["e_rho_r"], data["e_theta"]) + b).T / data["sqrt(g)"] + - cross(data["e_rho"], data["e_theta"]).T + * data["sqrt(g)_r"] + / data["sqrt(g)"] ** 2 + ).T, + lambda: ( + ( + 2 * cross(data["e_rho_r"], data["e_theta_r"]) + + cross(data["e_rho"], data["e_theta_rr"]) + ).T + / (2 * data["sqrt(g)_r"]) + - b.T * data["sqrt(g)_rr"] / (2 * data["sqrt(g)_r"] ** 2) + ).T, + ) + return data + + +@register_compute_fun( + name="e^zeta_t", + label="\\partial{\\theta} \\mathbf{e}^{\\zeta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant toroidal basis vector, derivative wrt poloidal" + " coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_rho_t", "e_theta", "e_theta_t", "sqrt(g)", "sqrt(g)_t"], + axis_limit_data=["e_theta_r", "e_theta_rt", "sqrt(g)_r", "sqrt(g)_rt"], +) +def _e_sup_zeta_t(params, transforms, profiles, data, **kwargs): + data["e^zeta_t"] = transforms["grid"].replace_at_axis( + ( + ( + cross(data["e_rho_t"], data["e_theta"]) + + cross(data["e_rho"], data["e_theta_t"]) + ).T + / data["sqrt(g)"] + - cross(data["e_rho"], data["e_theta"]).T + * data["sqrt(g)_t"] + / data["sqrt(g)"] ** 2 + ).T, + lambda: ( + ( + cross(data["e_rho_t"], data["e_theta_r"]) + + cross(data["e_rho"], data["e_theta_rt"]) + ).T + / data["sqrt(g)_r"] + - cross(data["e_rho"], data["e_theta_r"]).T + * data["sqrt(g)_rt"] + / data["sqrt(g)_r"] ** 2 + ).T, + ) + return data + + +@register_compute_fun( + name="e^zeta_z", + label="\\partial{\\zeta} \\mathbf{e}^{\\zeta}", + units="m^{-1}", + units_long="inverse meters", + description="Contravariant toroidal basis vector, derivative wrt toroidal" + " coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_rho_z", "e_theta", "e_theta_z", "sqrt(g)", "sqrt(g)_z"], + axis_limit_data=["e_theta_r", "e_theta_rz", "sqrt(g)_r", "sqrt(g)_rz"], +) +def _e_sup_zeta_z(params, transforms, profiles, data, **kwargs): + data["e^zeta_z"] = transforms["grid"].replace_at_axis( + ( + ( + cross(data["e_rho_z"], data["e_theta"]) + + cross(data["e_rho"], data["e_theta_z"]) + ).T + / data["sqrt(g)"] + - cross(data["e_rho"], data["e_theta"]).T + * data["sqrt(g)_z"] + / data["sqrt(g)"] ** 2 + ).T, + lambda: ( + ( + cross(data["e_rho_z"], data["e_theta_r"]) + + cross(data["e_rho"], data["e_theta_rz"]) + ).T + / data["sqrt(g)_r"] + - cross(data["e_rho"], data["e_theta_r"]).T + * data["sqrt(g)_rz"] + / data["sqrt(g)_r"] ** 2 + ).T, + ) + return data + + +@register_compute_fun( + name="e_phi", + label="\\mathbf{e}_{\\phi}", + units="m", + units_long="meters", + description="Covariant cylindrical toroidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta", "phi_z"], +) +def _e_sub_phi(params, transforms, profiles, data, **kwargs): + # dX/dphi at const r,t = dX/dz * dz/dphi = dX/dz / (dphi/dz) + data["e_phi"] = (data["e_zeta"].T / data["phi_z"]).T + return data + + @register_compute_fun( name="e_rho", label="\\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Covariant radial basis vector", + description="Covariant Radial basis vector", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_r", "Z_r"], + data=["R", "R_r", "Z_r", "omega_r"], ) def _e_sub_rho(params, transforms, profiles, data, **kwargs): - data["e_rho"] = jnp.array([data["R_r"], data["0"], data["Z_r"]]).T + data["e_rho"] = jnp.array([data["R_r"], data["R"] * data["omega_r"], data["Z_r"]]).T + return data + + +@register_compute_fun( + name="e_rho_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant Radial basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "R_rr", "Z_rr", "omega_r", "omega_rr"], +) +def _e_sub_rho_r(params, transforms, profiles, data, **kwargs): + # e_rho_r = a^i e_i, where the a^i are the components specified below and the + # e_i are the basis vectors of the polar lab frame. omega_r e_2, -omega_r e_1, + # 0 are the derivatives with respect to rho of e_1, e_2, e_3, respectively. + data["e_rho_r"] = jnp.array( + [ + -data["R"] * data["omega_r"] ** 2 + data["R_rr"], + 2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"], + data["Z_rr"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rr", + label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, second derivative wrt radial and radial" + " coordinates" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "R_rr", "R_rrr", "Z_rrr", "omega_r", "omega_rr", "omega_rrr"], +) +def _e_sub_rho_rr(params, transforms, profiles, data, **kwargs): + data["e_rho_rr"] = jnp.array( + [ + -3 * data["R_r"] * data["omega_r"] ** 2 + - 3 * data["R"] * data["omega_r"] * data["omega_rr"] + + data["R_rrr"], + 3 * (data["omega_r"] * data["R_rr"] + data["R_r"] * data["omega_rr"]) + + data["R"] * (-data["omega_r"] ** 3 + data["omega_rrr"]), + data["Z_rrr"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rrr", + label="\\partial_{\\rho}{\\rho}{\\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, third derivative wrt radial coordinate" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "R", + "R_r", + "R_rr", + "R_rrr", + "R_rrrr", + "Z_rrrr", + "omega_r", + "omega_rr", + "omega_rrr", + "omega_rrrr", + ], +) +def _e_sub_rho_rrr(params, transforms, profiles, data, **kwargs): + data["e_rho_rrr"] = jnp.array( + [ + -6 * data["R_rr"] * data["omega_r"] ** 2 + - 12 * data["R_r"] * data["omega_r"] * data["omega_rr"] + + data["R"] + * ( + -3 * data["omega_rr"] ** 2 + + data["omega_r"] ** 4 + - 4 * data["omega_rrr"] * data["omega_r"] + ) + + data["R_rrrr"], + 4 * data["R_rrr"] * data["omega_r"] + + 6 * data["R_rr"] * data["omega_rr"] + - 4 * data["R_r"] * (data["omega_r"] ** 3 - data["omega_rrr"]) + + data["R_r"] + * (data["omega_rrrr"] - 6 * data["omega_r"] ** 2 * data["omega_rr"]), + data["Z_rrrr"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rrt", + label="\\partial_{\\rho}{\\rho}{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial 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_rrr", + "R_rrt", + "R_rrrt", + "R_rt", + "R_t", + "Z_rrrt", + "omega_r", + "omega_rr", + "omega_rrr", + "omega_rrt", + "omega_rrrt", + "omega_rt", + "omega_t", + ], +) +def _e_sub_rho_rrt(params, transforms, profiles, data, **kwargs): + data["e_rho_rrt"] = jnp.array( + [ + -3 * data["R_rt"] * data["omega_r"] ** 2 + - 3 * data["R_t"] * data["omega_r"] * data["omega_rr"] + - 4 * data["R_r"] * data["omega_r"] * data["omega_rt"] + - 3 + * data["R"] + * ( + data["omega_rr"] * data["omega_rt"] + + data["omega_r"] * data["omega_rrt"] + ) + - data["omega_rt"] * (2 * data["R_r"] * data["omega_r"]) + - data["omega_t"] + * ( + 3 * data["R_rr"] * data["omega_r"] + + 3 * data["R_r"] * data["omega_rr"] + + data["R"] * data["omega_rrr"] + ) + + data["R_rrrt"] + + data["omega_r"] * data["R"] * data["omega_t"] * data["omega_r"] ** 2, + 3 * data["omega_rr"] * data["R_rt"] + + 3 * data["omega_r"] * data["R_rrt"] + + 3 * data["R_rr"] * data["omega_rt"] + + 2 * data["R_r"] * data["omega_rrt"] + + data["omega_t"] * data["R_rrr"] + + data["R_t"] * data["omega_rrr"] + + data["R_r"] + * (-3 * data["omega_t"] * data["omega_r"] ** 2 + data["omega_rrt"]) + + data["R"] + * ( + -3 * data["omega_rt"] * data["omega_r"] ** 2 + - 3 * data["omega_t"] * data["omega_r"] * data["omega_rr"] + + data["omega_rrrt"] + ) + - data["R_t"] * data["omega_r"] ** 3, + data["Z_rrrt"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rrz", + label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial 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_rrr", + "R_rrz", + "R_rrrz", + "R_rz", + "R_z", + "Z_rrrz", + "omega_r", + "omega_rr", + "omega_rrr", + "omega_rrz", + "omega_rrrz", + "omega_rz", + "omega_z", + ], +) +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"]) + * (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 + * data["R"] + * ( + data["omega_rr"] * data["omega_rz"] + + data["omega_r"] * data["omega_rrz"] + ) + - data["R_r"] * (1 + data["omega_z"]) * data["omega_rr"] + - data["R"] + * ( + data["omega_rz"] * data["omega_rr"] + + (1 + data["omega_z"]) * data["omega_rrr"] + ) + + 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"] + ), + data["Z_rrrz"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rt", + label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, second derivative wrt radial and poloidal" + " 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", + ], +) +def _e_sub_rho_rt(params, transforms, profiles, data, **kwargs): + data["e_rho_rt"] = 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 + return data + + +@register_compute_fun( + name="e_rho_rtt", + label="\\partial_{\\rho}{\\theta}{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, third derivative wrt radial coordinate" + "once and poloidal twice" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "R", + "R_r", + "R_rr", + "R_rt", + "R_rrt", + "R_rtt", + "R_rrtt", + "R_t", + "R_tt", + "Z_rrtt", + "omega_r", + "omega_rr", + "omega_rt", + "omega_rrt", + "omega_rtt", + "omega_rrtt", + "omega_t", + "omega_tt", + ], +) +def _e_sub_rho_rtt(params, transforms, profiles, data, **kwargs): + data["e_rho_rtt"] = jnp.array( + [ + -data["R_rr"] * data["omega_t"] ** 2 + - 4 * data["R_rt"] * data["omega_r"] * data["omega_t"] + - 4 * data["R_r"] * data["omega_rt"] * data["omega_t"] + - 2 * data["R_t"] * data["omega_rr"] * data["omega_t"] + - data["R_tt"] * data["omega_r"] ** 2 + - 4 * data["R_t"] * data["omega_r"] * data["omega_rt"] + - data["omega_tt"] + * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) + + data["R"] + * ( + (data["omega_r"] * data["omega_t"]) ** 2 + - 2 * (data["omega_rt"] ** 2 + data["omega_r"] * data["omega_rtt"]) + - 2 * (data["omega_t"] * data["omega_rrt"]) + ) + + data["R_rrtt"], + -data["omega_t"] ** 2 + * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) + + 2 + * data["omega_t"] + * ( + data["R_rrt"] + - data["omega_r"] + * (data["R_t"] * data["omega_r"] + 2 * data["R"] * data["omega_rt"]) + ) + + 4 * data["R_rt"] * data["omega_rt"] + + 2 * data["R_rtt"] * data["omega_r"] + + 2 * data["R_r"] * data["omega_rtt"] + + data["R_rr"] * data["omega_tt"] + + data["R_tt"] * data["omega_rr"] + + 2 * data["R_t"] * data["omega_rrt"] + + data["R"] + * (data["omega_rrtt"] - data["omega_tt"] * data["omega_r"] ** 2), + data["Z_rrtt"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rtz", + label="\\partial_{\\rho}{\\theta}{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, third derivative wrt radial, poloidal," + " and toroidal coordinates" + ), + 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", + ], +) +def _e_sub_rho_rtz(params, transforms, profiles, data, **kwargs): + data["e_rho_rtz"] = 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"] + + (1 + 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["omega_rz"] * data["R_rt"] + + (1 + 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 + return data + + +@register_compute_fun( + name="e_rho_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, second derivative wrt radial and toroidal" + " 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", + ], +) +def _e_sub_rho_rz(params, transforms, profiles, data, **kwargs): + data["e_rho_rz"] = 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"] * (1 + data["omega_z"]) * data["omega_rr"] + + data["R_rrz"], + 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"]), + data["Z_rrz"], + ] + ).T + return data + + +@register_compute_fun( + name="e_rho_rzz", + label="\\partial_{\\rho}{\\zeta}{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, third derivative wrt radial coordinate" + " once and toroidal twice" + ), + 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", + ], +) +def _e_sub_rho_rzz(params, transforms, profiles, data, **kwargs): + data["e_rho_rzz"] = 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 + return data + + +@register_compute_fun( + name="e_rho_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant Radial basis vector, derivative wrt poloidal coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "R_rt", "R_t", "Z_rt", "omega_r", "omega_rt", "omega_t"], +) +def _e_sub_rho_t(params, transforms, profiles, data, **kwargs): + data["e_rho_t"] = 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 + return data + + +@register_compute_fun( + name="e_rho_tt", + label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, second derivative wrt poloidal 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", + ], +) +def _e_sub_rho_tt(params, transforms, profiles, data, **kwargs): + data["e_rho_tt"] = 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 + return data + + +@register_compute_fun( + name="e_rho_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, second derivative wrt poloidal 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", + ], +) +def _e_sub_rho_tz(params, transforms, profiles, data, **kwargs): + data["e_rho_tz"] = 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"] + + (1 + 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 + 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 coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], +) +def _e_sub_rho_z(params, transforms, profiles, data, **kwargs): + data["e_rho_z"] = 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 + return data + + +@register_compute_fun( + name="e_rho_zz", + label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description=( + "Covariant Radial basis vector, second derivative wrt toroidal 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", + ], +) +def _e_sub_rho_zz(params, transforms, profiles, data, **kwargs): + data["e_rho_zz"] = 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 return data @@ -27,630 +1356,2063 @@ def _e_sub_rho(params, transforms, profiles, data, **kwargs): label="\\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant poloidal basis vector", + description="Covariant Poloidal basis vector", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_t", "Z_t"], + data=["R", "R_t", "Z_t", "omega_t"], ) def _e_sub_theta(params, transforms, profiles, data, **kwargs): - data["e_theta"] = jnp.array([data["R_t"], data["0"], data["Z_t"]]).T + data["e_theta"] = jnp.array( + [data["R_t"], data["R"] * data["omega_t"], data["Z_t"]] + ).T return data @register_compute_fun( - name="e_zeta", - label="\\mathbf{e}_{\\zeta}", + name="e_theta/sqrt(g)", + label="\\mathbf{e}_{\\theta} / \\sqrt{g}", units="m", units_long="meters", - description="Covariant toroidal basis vector", + description="Covariant Poloidal basis vector divided by 3-D volume Jacobian", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "Z_z"], + data=["e_theta", "sqrt(g)"], + axis_limit_data=["e_theta_r", "sqrt(g)_r"], ) -def _e_sub_zeta(params, transforms, profiles, data, **kwargs): - data["e_zeta"] = jnp.array([data["R_z"], data["R"], data["Z_z"]]).T +def _e_sub_theta_over_sqrt_g(params, transforms, profiles, data, **kwargs): + data["e_theta/sqrt(g)"] = transforms["grid"].replace_at_axis( + (data["e_theta"].T / data["sqrt(g)"]).T, + lambda: (data["e_theta_r"].T / data["sqrt(g)_r"]).T, + ) + return data + + +@register_compute_fun( + name="e_theta_PEST", + label="\\mathbf{e}_{\\theta_{PEST}}", + units="m", + units_long="meters", + description="Covariant straight field line (PEST) poloidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "theta_PEST_t"], +) +def _e_sub_theta_pest(params, transforms, profiles, data, **kwargs): + # dX/dv at const r,z = dX/dt * dt/dv / dX/dt / dv/dt + data["e_theta_PEST"] = (data["e_theta"].T / data["theta_PEST_t"]).T + 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={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "R_rt", "R_t", "Z_rt", "omega_r", "omega_rt", "omega_t"], +) +def _e_sub_theta_r(params, transforms, profiles, data, **kwargs): + 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 return data @register_compute_fun( - name="e_rho_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", + name="e_theta_rr", + label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, derivative wrt radial coordinate", + description=( + "Covariant Poloidal basis vector, second derivative wrt radial and radial" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rr", "Z_rr"], + data=[ + "R", + "R_r", + "R_rr", + "R_rrt", + "R_rt", + "R_t", + "Z_rrt", + "omega_r", + "omega_rr", + "omega_rrt", + "omega_rt", + "omega_t", + ], ) -def _e_sub_rho_r(params, transforms, profiles, data, **kwargs): - data["e_rho_r"] = jnp.array([data["R_rr"], data["0"], data["Z_rr"]]).T +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 return data @register_compute_fun( - name="e_rho_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", + name="e_theta_rrr", + label="\\partial_{\\rho}{\\rho}{\\rho} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, derivative wrt poloidal angle", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rt", "Z_rt"], + data=[ + "R", + "R_r", + "R_t", + "R_rr", + "R_rt", + "R_rrr", + "R_rrt", + "R_rrrt", + "Z_rrrt", + "omega_r", + "omega_t", + "omega_rr", + "omega_rt", + "omega_rrr", + "omega_rrt", + "omega_rrrt", + ], ) -def _e_sub_rho_t(params, transforms, profiles, data, **kwargs): - data["e_rho_t"] = jnp.array([data["R_rt"], data["0"], data["Z_rt"]]).T +def _e_sub_theta_rrr(params, transforms, profiles, data, **kwargs): + data["e_theta_rrr"] = jnp.array( + [ + -3 * data["omega_rrt"] * data["R"] * data["omega_r"] + - 3 + * data["omega_rt"] + * (2 * data["R_r"] * data["omega_r"] + data["R"] * data["omega_rr"]) + - 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 + ) + - 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["R"] + * ( + data["omega_rrrt"] + - 3 + * data["omega_r"] + * ( + data["omega_rt"] * data["omega_r"] + + data["omega_t"] * data["omega_rr"] + ) + ) + + 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"], + ] + ).T return data @register_compute_fun( - name="e_rho_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", + name="e_theta_rrt", + label="\\partial_{\\rho}{\\rho}{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, derivative wrt toroidal angle", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " twice and poloidal once" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rz", "Z_rz"], + data=[ + "R", + "R_r", + "R_rr", + "R_rt", + "R_rrt", + "R_rtt", + "R_rrtt", + "R_t", + "R_tt", + "Z_rrtt", + "omega_r", + "omega_rr", + "omega_rt", + "omega_rrt", + "omega_rtt", + "omega_rrtt", + "omega_t", + "omega_tt", + ], ) -def _e_sub_rho_z(params, transforms, profiles, data, **kwargs): - data["e_rho_z"] = jnp.array([data["R_rz"], data["0"], data["Z_rz"]]).T +def _e_sub_theta_rrt(params, transforms, profiles, data, **kwargs): + data["e_theta_rrt"] = 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_rt"] + * (data["R_t"] * data["omega_r"] + data["R"] * data["omega_rt"]) + - 2 + * data["omega_t"] + * ( + data["R_rt"] * data["omega_r"] + + data["R_r"] * data["omega_rt"] + + data["R_t"] * data["omega_rr"] + + data["R"] * data["omega_rrt"] + ) + + 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 return data @register_compute_fun( - name="e_theta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", + name="e_theta_rrz", + label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, derivative wrt radial coordinate", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " twice and toroidal once" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rt", "Z_rt"], + 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", + ], ) -def _e_sub_theta_r(params, transforms, profiles, data, **kwargs): - data["e_theta_r"] = jnp.array([data["R_rt"], data["0"], data["Z_rt"]]).T +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 return data @register_compute_fun( - name="e_theta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", + name="e_theta_rt", + label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, derivative wrt poloidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt radial and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_tt", "Z_tt"], + data=[ + "R", + "R_r", + "R_rt", + "R_rtt", + "R_t", + "R_tt", + "Z_rtt", + "omega_r", + "omega_rt", + "omega_rtt", + "omega_t", + "omega_tt", + ], ) -def _e_sub_theta_t(params, transforms, profiles, data, **kwargs): - data["e_theta_t"] = jnp.array([data["R_tt"], data["0"], data["Z_tt"]]).T +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 return data @register_compute_fun( - name="e_theta_z", - label="\\partial_{\\zeta} \\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, derivative wrt toroidal angle", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " once and poloidal twice" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_tz", "Z_tz"], + 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", + ], ) -def _e_sub_theta_z(params, transforms, profiles, data, **kwargs): - data["e_theta_z"] = jnp.array([data["R_tz"], data["0"], data["Z_tz"]]).T +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 return data @register_compute_fun( - name="e_zeta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", + name="e_theta_rtz", + label="\\partial_{\\rho}{\\theta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, derivative wrt radial coordinate", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial, poloidal," + " and toroidal coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_r", "R_rz", "Z_rz"], + 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", + ], ) -def _e_sub_zeta_r(params, transforms, profiles, data, **kwargs): - data["e_zeta_r"] = jnp.array([data["R_rz"], data["R_r"], data["Z_rz"]]).T +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 return data @register_compute_fun( - name="e_zeta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", + name="e_theta_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, derivative wrt poloidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt radial and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_t", "R_tz", "Z_tz"], + 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", + ], ) -def _e_sub_zeta_t(params, transforms, profiles, data, **kwargs): - data["e_zeta_t"] = jnp.array([data["R_tz"], data["R_t"], data["Z_tz"]]).T +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 return data @register_compute_fun( - name="e_zeta_z", - label="\\partial_{\\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, derivative wrt toroidal angle", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " once and toroidal twice" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_z", "R_zz", "Z_zz"], + 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", + ], ) -def _e_sub_zeta_z(params, transforms, profiles, data, **kwargs): - data["e_zeta_z"] = jnp.array([data["R_zz"], data["R_z"], data["Z_zz"]]).T +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 return data @register_compute_fun( - name="e_rho_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\rho}", + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt radial " - + "coordinate", + description="Covariant Poloidal basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrr", "Z_rrr"], + data=["R", "R_t", "R_tt", "Z_tt", "omega_t", "omega_tt"], ) -def _e_sub_rho_rr(params, transforms, profiles, data, **kwargs): - data["e_rho_rr"] = jnp.array([data["R_rrr"], data["0"], data["Z_rrr"]]).T +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 return data @register_compute_fun( - name="e_rho_tt", - label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\rho}", + name="e_theta_tt", + label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt poloidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt poloidal and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtt", "Z_rtt"], + data=["R", "R_t", "R_tt", "R_ttt", "Z_ttt", "omega_t", "omega_tt", "omega_ttt"], ) -def _e_sub_rho_tt(params, transforms, profiles, data, **kwargs): - data["e_rho_tt"] = jnp.array([data["R_rtt"], data["0"], data["Z_rtt"]]).T +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 return data @register_compute_fun( - name="e_rho_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\rho}", + name="e_theta_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt toroidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt poloidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rzz", "Z_rzz"], + data=[ + "R", + "R_t", + "R_tt", + "R_ttz", + "R_tz", + "R_z", + "Z_ttz", + "omega_t", + "omega_tt", + "omega_ttz", + "omega_tz", + "omega_z", + ], ) -def _e_sub_rho_zz(params, transforms, profiles, data, **kwargs): - data["e_rho_zz"] = jnp.array([data["R_rzz"], data["0"], data["Z_rzz"]]).T +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 return data @register_compute_fun( - name="e_rho_rt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\rho}", + name="e_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt radial " - + "coordinate and poloidal angle", + description="Covariant Poloidal basis vector, derivative wrt toroidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrt", "Z_rrt"], + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], ) -def _e_sub_rho_rt(params, transforms, profiles, data, **kwargs): - data["e_rho_rt"] = jnp.array([data["R_rrt"], data["0"], data["Z_rrt"]]).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 return data @register_compute_fun( - name="e_rho_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\rho}", + name="e_theta_zz", + label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt radial " - + "coordinate and toroidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt toroidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrz", "Z_rrz"], + data=[ + "R", + "R_t", + "R_tz", + "R_tzz", + "R_z", + "R_zz", + "Z_tzz", + "omega_t", + "omega_tz", + "omega_tzz", + "omega_z", + "omega_zz", + ], ) -def _e_sub_rho_rz(params, transforms, profiles, data, **kwargs): - data["e_rho_rz"] = jnp.array([data["R_rrz"], data["0"], data["Z_rrz"]]).T +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 return data @register_compute_fun( - name="e_rho_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\rho}", + name="e_zeta", + label="\\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt poloidal " - + "and toroidal angles", + description="Covariant Toroidal basis vector", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtz", "Z_rtz"], + data=["R", "R_z", "Z_z", "omega_z"], ) -def _e_sub_rho_tz(params, transforms, profiles, data, **kwargs): - data["e_rho_tz"] = jnp.array([data["R_rtz"], data["0"], data["Z_rtz"]]).T +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 return data @register_compute_fun( - name="e_theta_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\theta}", + name="e_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, second derivative wrt radial " - + "coordinate", + description="Covariant Toroidal basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrt", "Z_rrt"], + data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], ) -def _e_sub_theta_rr(params, transforms, profiles, data, **kwargs): - data["e_theta_rr"] = jnp.array([data["R_rrt"], data["0"], data["Z_rrt"]]).T +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 return data @register_compute_fun( - name="e_theta_tt", - label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\theta}", + name="e_zeta_rr", + label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, second derivative wrt poloidal angle", + description=( + "Covariant Toroidal basis vector, second derivative wrt radial and radial" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_ttt", "Z_ttt"], + data=[ + "R", + "R_r", + "R_rr", + "R_rrz", + "R_rz", + "R_z", + "Z_rrz", + "omega_r", + "omega_rr", + "omega_rrz", + "omega_rz", + "omega_z", + ], ) -def _e_sub_theta_tt(params, transforms, profiles, data, **kwargs): - data["e_theta_tt"] = jnp.array([data["R_ttt"], data["0"], data["Z_ttt"]]).T +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 return data @register_compute_fun( - name="e_theta_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\theta}", + name="e_zeta_rrr", + label="\\partial_{\\rho}{\\rho}{\\rho} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, second derivative wrt toroidal angle", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_tzz", "Z_tzz"], + 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", + ], ) -def _e_sub_theta_zz(params, transforms, profiles, data, **kwargs): - data["e_theta_zz"] = jnp.array([data["R_tzz"], data["0"], data["Z_tzz"]]).T +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 return data @register_compute_fun( - name="e_theta_rt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\theta}", + name="e_zeta_rrt", + label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, second derivative wrt radial " - + "coordinate and poloidal angle", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + " twice and poloidal once" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtt", "Z_rtt"], + 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", + ], ) -def _e_sub_theta_rt(params, transforms, profiles, data, **kwargs): - data["e_theta_rt"] = jnp.array([data["R_rtt"], data["0"], data["Z_rtt"]]).T +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 return data @register_compute_fun( - name="e_theta_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\theta}", + name="e_zeta_rrz", + label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, second derivative wrt radial " - + "coordinate and toroidal angle", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + " twice and toroidal once" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtz", "Z_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", + ], ) -def _e_sub_theta_rz(params, transforms, profiles, data, **kwargs): - data["e_theta_rz"] = jnp.array([data["R_rtz"], data["0"], data["Z_rtz"]]).T +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 return data @register_compute_fun( - name="e_theta_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\theta}", + name="e_zeta_rt", + label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal basis vector, second derivative wrt poloidal " - + "and toroidal angles", + description=( + "Covariant Toroidal basis vector, second derivative wrt radial and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_ttz", "Z_ttz"], + 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", + ], ) -def _e_sub_theta_tz(params, transforms, profiles, data, **kwargs): - data["e_theta_tz"] = jnp.array([data["R_ttz"], data["0"], data["Z_ttz"]]).T +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 return data @register_compute_fun( - name="e_zeta_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", + name="e_zeta_rtt", + label="\\partial_{\\rho}{\\theta}{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt radial " - + "coordinate", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + " once and poloidal twice" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_rr", "R_rrz", "Z_rrz"], + 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", + ], ) -def _e_sub_zeta_rr(params, transforms, profiles, data, **kwargs): - data["e_zeta_rr"] = jnp.array([data["R_rrz"], data["R_rr"], data["Z_rrz"]]).T +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["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 return data @register_compute_fun( - name="e_zeta_tt", - label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\zeta}", + name="e_zeta_rtz", + label="\\partial_{\\rho}{\\theta}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt poloidal angle", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial, poloidal," + " and toroidal coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_tt", "R_ttz", "Z_ttz"], + 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_r", + "omega_t", + "omega_rt", + "omega_tz", + "omega_rtz", + "omega_tzz", + "omega_rtzz", + "omega_z", + "omega_rz", + "omega_zz", + "omega_rzz", + ], ) -def _e_sub_zeta_tt(params, transforms, profiles, data, **kwargs): - data["e_zeta_tt"] = jnp.array([data["R_ttz"], data["R_tt"], data["Z_ttz"]]).T +def _e_sub_zeta_rtz(params, transforms, profiles, data, **kwargs): + data["e_zeta_rtz"] = 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 return data @register_compute_fun( - name="e_zeta_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\zeta}", + name="e_zeta_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt toroidal angle", + description=( + "Covariant Toroidal basis vector, second derivative wrt radial and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_zz", "R_zzz", "Z_zzz"], + data=[ + "R", + "R_r", + "R_rz", + "R_rzz", + "R_z", + "R_zz", + "Z_rzz", + "omega_r", + "omega_rz", + "omega_rzz", + "omega_z", + "omega_zz", + ], ) -def _e_sub_zeta_zz(params, transforms, profiles, data, **kwargs): - data["e_zeta_zz"] = jnp.array([data["R_zzz"], data["R_zz"], data["Z_zzz"]]).T +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 return data @register_compute_fun( - name="e_zeta_rt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\zeta}", + name="e_zeta_rzz", + label="\\partial_{\\rho}{\\zeta}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt radial " - + "coordinate and poloidal angle", + description=( + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + " once and toroidal twice" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_rt", "R_rtz", "Z_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", + ], ) -def _e_sub_zeta_rt(params, transforms, profiles, data, **kwargs): - data["e_zeta_rt"] = jnp.array([data["R_rtz"], data["R_rt"], data["Z_rtz"]]).T +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 return data @register_compute_fun( - name="e_zeta_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\zeta}", + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt radial " - + "coordinate and toroidal angle", + description="Covariant Toroidal basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_rz", "R_rzz", "Z_rzz"], + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], ) -def _e_sub_zeta_rz(params, transforms, profiles, data, **kwargs): - data["e_zeta_rz"] = jnp.array([data["R_rzz"], data["R_rz"], data["Z_rzz"]]).T +def _e_sub_zeta_t(params, transforms, profiles, data, **kwargs): + data["e_zeta_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"], + ] + ).T return data @register_compute_fun( - name="e_zeta_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\zeta}", + name="e_zeta_tt", + label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt poloidal " - + "and toroidal angles", + description=( + "Covariant Toroidal basis vector, second derivative wrt poloidal and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_tz", "R_tzz", "Z_tzz"], + data=[ + "R", + "R_t", + "R_tt", + "R_ttz", + "R_tz", + "R_z", + "Z_ttz", + "omega_t", + "omega_tt", + "omega_ttz", + "omega_tz", + "omega_z", + ], ) -def _e_sub_zeta_tz(params, transforms, profiles, data, **kwargs): - data["e_zeta_tz"] = jnp.array([data["R_tzz"], data["R_tz"], data["Z_tzz"]]).T +def _e_sub_zeta_tt(params, transforms, profiles, data, **kwargs): + data["e_zeta_tt"] = 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 return data @register_compute_fun( - name="e_theta_PEST", - label="\\mathbf{e}_{\\theta_{PEST}}", + name="e_zeta_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant straight field line poloidal basis vector", + description=( + "Covariant Toroidal basis vector, second derivative wrt poloidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", data=[ - "0", + "R", "R_t", - "Z_t", - "lambda_t", + "R_tz", + "R_tzz", + "R_z", + "R_zz", + "Z_tzz", + "omega_t", + "omega_tz", + "omega_tzz", + "omega_z", + "omega_zz", ], ) -def _e_sub_theta_pest(params, transforms, profiles, data, **kwargs): - dt_dv = 1 / (1 + data["lambda_t"]) - data["e_theta_PEST"] = jnp.array( - [data["R_t"] * dt_dv, data["0"], data["Z_t"] * dt_dv] +def _e_sub_zeta_tz(params, transforms, profiles, data, **kwargs): + data["e_zeta_tz"] = 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 return data @register_compute_fun( - name="e^rho", - label="\\mathbf{e}^{\\rho}", - units="m^{-1}", - units_long="inverse meters", - description="Contravariant radial basis vector", + name="e_zeta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant Toroidal basis vector, derivative wrt toroidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_theta", "e_zeta", "sqrt(g)"], + data=["R", "R_z", "R_zz", "Z_zz", "omega_z", "omega_zz"], ) -def _e_sup_rho(params, transforms, profiles, data, **kwargs): - data["e^rho"] = (cross(data["e_theta"], data["e_zeta"]).T / data["sqrt(g)"]).T +def _e_sub_zeta_z(params, transforms, profiles, data, **kwargs): + data["e_zeta_z"] = jnp.array( + [ + -data["R"] * (1 + data["omega_z"]) ** 2 + data["R_zz"], + 2 * data["R_z"] * (1 + data["omega_z"]) + data["R"] * data["omega_zz"], + data["Z_zz"], + ] + ).T return data @register_compute_fun( - name="e^theta", - label="\\mathbf{e}^{\\theta}", - units="m^{-1}", - units_long="inverse meters", - description="Contravariant poloidal basis vector", + name="e_zeta_zz", + label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description=( + "Covariant Toroidal basis vector, second derivative wrt toroidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_zeta", "sqrt(g)"], + data=["R", "R_z", "R_zz", "R_zzz", "Z_zzz", "omega_z", "omega_zz", "omega_zzz"], ) -def _e_sup_theta(params, transforms, profiles, data, **kwargs): - data["e^theta"] = (cross(data["e_zeta"], data["e_rho"]).T / data["sqrt(g)"]).T +def _e_sub_zeta_zz(params, transforms, profiles, data, **kwargs): + data["e_zeta_zz"] = jnp.array( + [ + -3 * data["R_z"] * (1 + data["omega_z"]) ** 2 + - 3 * data["R"] * (1 + data["omega_z"]) * data["omega_zz"] + + data["R_zzz"], + 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"] + ), + data["Z_zzz"], + ] + ).T return data @register_compute_fun( - name="e^zeta", - label="\\mathbf{e}^{\\zeta}", + name="grad(alpha)", + label="\\nabla \\alpha", units="m^{-1}", - units_long="inverse meters", - description="Contravariant toroidal basis vector", + units_long="Inverse meters", + description="Unit vector along field line", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_theta", "sqrt(g)"], + data=["e^rho", "e^theta", "e^zeta", "alpha_r", "alpha_t", "alpha_z"], ) -def _e_sup_zeta(params, transforms, profiles, data, **kwargs): - data["e^zeta"] = (cross(data["e_rho"], data["e_theta"]).T / data["sqrt(g)"]).T +def _grad_alpha(params, transforms, profiles, data, **kwargs): + data["grad(alpha)"] = ( + data["alpha_r"] * data["e^rho"].T + + data["alpha_t"] * data["e^theta"].T + + data["alpha_z"] * data["e^zeta"].T + ).T return data @register_compute_fun( - name="b", - label="\\hat{b}", - units="~", - units_long="None", - description="Unit vector along magnetic field", + name="grad(psi)", + label="\\nabla\\psi", + units="Wb / m", + units_long="Webers per meter", + description="Toroidal flux gradient (normalized by 2pi)", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["B"], + data=["psi_r", "e^rho"], ) -def _b(params, transforms, profiles, data, **kwargs): - data["b"] = (data["B"].T / jnp.linalg.norm(data["B"], axis=-1)).T +def _gradpsi(params, transforms, profiles, data, **kwargs): + data["grad(psi)"] = (data["psi_r"] * data["e^rho"].T).T return data @@ -726,25 +3488,3 @@ def _n_zeta(params, transforms, profiles, data, **kwargs): cross(data["e_rho"], data["e_theta"]) / data["|e_rho x e_theta|"][:, None] ) return data - - -@register_compute_fun( - name="grad(alpha)", - label="\\nabla \\alpha", - units="m^{-1}", - units_long="Inverse meters", - description="Unit vector normal to flux surface", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e^rho", "e^theta", "e^zeta", "alpha_r", "alpha_t", "alpha_z"], -) -def _grad_alpha(params, transforms, profiles, data, **kwargs): - data["grad(alpha)"] = ( - data["alpha_r"] * data["e^rho"].T - + data["alpha_t"] * data["e^theta"].T - + data["alpha_z"] * data["e^zeta"].T - ).T - return data diff --git a/desc/compute/_bootstrap.py b/desc/compute/_bootstrap.py index a7b5945ef1..fffeecc46e 100644 --- a/desc/compute/_bootstrap.py +++ b/desc/compute/_bootstrap.py @@ -334,6 +334,7 @@ def j_dot_B_Redl( "Ti", "Ti_r", "Zeff", + "rho", ], helicity="helicity", ) diff --git a/desc/compute/_field.py b/desc/compute/_field.py index 9f5bc63728..760b00a26f 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -203,9 +203,9 @@ def _B0_r(params, transforms, profiles, data, **kwargs): data=["B0", "B0_r", "iota", "iota_r", "lambda_z", "lambda_rz"], ) def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): - data["B^theta_r"] = data["B0_r"] * (data["iota"] - data["lambda_z"]) + data[ - "B0" - ] * (data["iota_r"] - data["lambda_rz"]) + data["B^theta_r"] = data["B0_r"] * (data["iota"] - data["lambda_z"]) + ( + data["B0"] * (data["iota_r"] - data["lambda_rz"]) + ) return data @@ -2223,7 +2223,7 @@ def _B_mag_tz(params, transforms, profiles, data, **kwargs): "B^zeta_r", "B_theta_r", "B_zeta_r", - "B^theta_t", + "B^theta_z", "B^zeta_z", "B_theta_z", "B_zeta_z", diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index fbc9561a7f..7ead887cfc 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -257,7 +257,7 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="", - data=["sqrt(g)", "g_tt"], + data=["sqrt(g)", "g_tt", "R"], ) def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs): max_rho = transforms["grid"].nodes[transforms["grid"].unique_rho_idx[-1], 0] diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index a42cb17553..a752f1c737 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -766,10 +766,10 @@ def _g_sup_tz(params, transforms, profiles, data, **kwargs): "e^rho", "sqrt(g)_r", "sqrt(g)", - "e_rho", "e_theta", - "e_rho_r", "e_theta_r", + "e_zeta_r", + "e_zeta", ], ) def _g_sup_rr_r(params, transforms, profiles, data, **kwargs): @@ -907,16 +907,13 @@ def _g_sup_rz_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^theta", "e^theta", "sqrt(g)_r", "sqrt(g)", - "e_theta", - "e_theta", "e_rho", - "e_theta_r", - "e_theta_r", "e_rho_r", + "e_zeta_r", + "e_zeta", ], ) def _g_sup_tt_r(params, transforms, profiles, data, **kwargs): @@ -1005,15 +1002,12 @@ def _g_sup_tz_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^zeta", "e^zeta", "sqrt(g)_r", "sqrt(g)", - "e_zeta", - "e_zeta", + "e_theta", "e_rho", - "e_zeta_r", - "e_zeta_r", + "e_theta_r", "e_rho_r", ], ) @@ -1054,16 +1048,13 @@ def _g_sup_zz_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^rho", "e^rho", "sqrt(g)_t", "sqrt(g)", - "e_rho", - "e_rho", "e_theta", - "e_rho_t", - "e_rho_t", "e_theta_t", + "e_zeta_t", + "e_zeta", ], ) def _g_sup_rr_t(params, transforms, profiles, data, **kwargs): @@ -1201,15 +1192,12 @@ def _g_sup_rz_t(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^theta", "e^theta", "sqrt(g)_t", "sqrt(g)", - "e_theta", - "e_theta", + "e_zeta", "e_rho", - "e_theta_t", - "e_theta_t", + "e_zeta_t", "e_rho_t", ], ) @@ -1299,15 +1287,12 @@ def _g_sup_tz_t(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^zeta", "e^zeta", "sqrt(g)_t", "sqrt(g)", - "e_zeta", - "e_zeta", + "e_theta", "e_rho", - "e_zeta_t", - "e_zeta_t", + "e_theta_t", "e_rho_t", ], ) @@ -1348,15 +1333,12 @@ def _g_sup_zz_t(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^rho", "e^rho", "sqrt(g)_z", "sqrt(g)", - "e_rho", - "e_rho", + "e_zeta", "e_theta", - "e_rho_z", - "e_rho_z", + "e_zeta_z", "e_theta_z", ], ) @@ -1495,15 +1477,12 @@ def _g_sup_rz_z(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^theta", "e^theta", "sqrt(g)_z", "sqrt(g)", - "e_theta", - "e_theta", + "e_zeta", "e_rho", - "e_theta_z", - "e_theta_z", + "e_zeta_z", "e_rho_z", ], ) @@ -1593,15 +1572,12 @@ def _g_sup_tz_z(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "e^zeta", "e^zeta", "sqrt(g)_z", "sqrt(g)", - "e_zeta", - "e_zeta", + "e_theta", "e_rho", - "e_zeta_z", - "e_zeta_z", + "e_theta_z", "e_rho_z", ], ) diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index c2e667da43..b710ca0b9f 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -68,7 +68,7 @@ def _psi_rr(params, transforms, profiles, data, **kwargs): units_long="Webers", description="Toroidal flux (normalized by 2pi), third radial derivative", dim=1, - params=["Psi"], + params=[], transforms={}, profiles=[], coordinates="r", @@ -509,9 +509,9 @@ def _iota(params, transforms, profiles, data, **kwargs): current_term = put( current_term, transforms["grid"].axis, limit[transforms["grid"].axis] ) - data["iota"] = (current_term + data["iota_zero_current_num"]) / data[ - "iota_zero_current_den" - ] + data["iota"] = (current_term + data["iota_zero_current_num"]) / ( + data["iota_zero_current_den"] + ) return data @@ -530,7 +530,6 @@ def _iota(params, transforms, profiles, data, **kwargs): "iota", "psi_r", "psi_rr", - "iota_zero_current_num", "iota_zero_current_num_r", "iota_zero_current_den", "iota_zero_current_den_r", @@ -582,8 +581,6 @@ def _iota_r(params, transforms, profiles, data, **kwargs): "psi_r", "psi_rr", "psi_rrr", - "iota_zero_current_num", - "iota_zero_current_num_r", "iota_zero_current_num_rr", "iota_zero_current_den", "iota_zero_current_den_r", diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py new file mode 100644 index 0000000000..5991661337 --- /dev/null +++ b/docs/sort_compute_funs.py @@ -0,0 +1,39 @@ +"""Sort the compute functions and print to specified output. + +First command line argument is module name, e.g. "_basis_vectors". +Second command line argument is output file to print to, +e.g. "sorted_basis_vectors.txt". +""" + +import inspect +import re +import sys + +import desc.compute + +preamble = "" + +# Gather all compute function source code and map name to source code. +source_codes = {} +pattern = re.compile(r"(?<=name=)[^,]+") +for module_name, module in inspect.getmembers(desc.compute, inspect.ismodule): + if module_name == sys.argv[1]: + # everything in the module until the first compute function + # (import statements, etc.) + preamble = inspect.getsource(module).partition("@register_compute_fun")[0] + for function_name, fun in inspect.getmembers(module, inspect.isfunction): + source_code = inspect.getsource(fun) + matches = pattern.findall(source_code) + if matches: # skip imported functions + # matches[0] is the thing this function says it computes, e.g. x + # while function_name is e.g. _x_ZernikeRZToroidalSection + if function_name in source_codes: + raise ValueError("Can't sort when things have same name.") + source_codes[function_name] = source_code + +# Write functions sorted to file. +with open(sys.argv[2], "w") as output_file: + output_file.write(preamble) + for function_name in sorted(source_codes): + output_file.write(source_codes[function_name]) + output_file.write("\n\n") diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index 61550f378a..0afbc21cf5 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -1,101 +1,365 @@ """Tests for compute functions evaluated at limits.""" +import inspect +import re import numpy as np import pytest -import desc.io +import desc.compute from desc.compute import data_index -from desc.compute.utils import compress, surface_averages +from desc.compute.data_index import _class_inheritance +from desc.compute.utils import surface_integrals_map from desc.equilibrium import Equilibrium -from desc.examples import get from desc.grid import LinearGrid +# Unless mentioned in the source code of the compute function, the assumptions +# made to compute the magnetic axis limit can be reduced to assuming that these +# functions tend toward zero as the magnetic axis is approached and that +# d^2šœ“/(dšœŒ)^2 and šœ•āˆšš‘”/šœ•šœŒ are both finite nonzero at the magnetic axis. +# Also d^nšœ“/(dšœŒ)^n for n > 3 is assumed zero everywhere. +zero_limits = {"rho", "psi", "psi_r", "e_theta", "sqrt(g)", "B_t"} + +not_finite_limits = { + "D_Mercier", + "D_current", + "D_geodesic", + "D_shear", # may not exist for all configurations + "D_well", + "J^theta", + "curvature_H_rho", + "curvature_H_zeta", + "curvature_K_rho", + "curvature_K_zeta", + "curvature_k1_rho", + "curvature_k1_zeta", + "curvature_k2_rho", + "curvature_k2_zeta", + "e^helical", + "e^theta", + "e^theta_r", + "e^theta_t", + "e^theta_z", + "g^rt", + "g^rt_r", + "g^rt_t", + "g^rt_z", + "g^tt", + "g^tt_r", + "g^tt_t", + "g^tt_z", + "g^tz", + "g^tz_r", + "g^tz_t", + "g^tz_z", + "grad(alpha)", + "|e^helical|", + "|grad(theta)|", + " Redl", # may not exist for all configurations +} + +# reliant limits will be added to this set automatically +not_implemented_limits = { + "iota_num_rrr", + "iota_den_rrr", +} + + +def grow_seeds(seeds, search_space): + """Traverse the dependency DAG for keys in search space dependent on seeds. + + Parameters + ---------- + seeds : set + Keys to find paths toward. + search_space : iterable + Additional keys to consider returning. + + Returns + ------- + out : set + All keys in search space with any path in the dependency DAG to any seed. + + """ + out = seeds.copy() + for key in search_space: + deps = data_index["desc.equilibrium.equilibrium.Equilibrium"][key][ + "full_with_axis_dependencies" + ]["data"] + if not seeds.isdisjoint(deps): + out.add(key) + return out + + +not_implemented_limits = grow_seeds( + not_implemented_limits, + data_index["desc.equilibrium.equilibrium.Equilibrium"].keys() - not_finite_limits, +) + + +def _skip_this(eq, name): + return ( + name in not_implemented_limits + or (eq.atomic_number is None and "Zeff" in name) + or (eq.electron_temperature is None and "Te" in name) + or (eq.electron_density is None and "ne" in name) + or (eq.ion_temperature is None and "Ti" in name) + or (eq.pressure is not None and " Redl" in name) + or (eq.current is None and ("iota_num" in name or "iota_den" in name)) + ) + + +def assert_is_continuous( + eq, + names, + delta=5e-5, + rtol=1e-4, + atol=1e-6, + desired_at_axis=None, + kwargs=None, +): + """ + Asserts that the rho=0 axis limits of names are continuous extensions. + + Parameters + ---------- + eq : Equilibrium + The equilibrium object used for the computation. + names : list, str + A list of names of the quantities to test for continuity. + delta: float, optional + Max distance from magnetic axis. + Smaller values accumulate finite precision error and fitting issues. + rtol : float, optional + Relative tolerance. + atol : float, optional + Absolute tolerance. + desired_at_axis : float, optional + If not provided, the values are extrapolated with a polynomial fit. + kwargs : dict, optional + Keyword arguments to override the parameters above for specific names. + The dictionary should have the following structure: + { + "name1": { + "rtol": custom_rtol1, + "atol": custom_atol1, + "desired_at_axis": custom_desired_at_axis1 + }, + "name2": {"rtol": custom_rtol2}, + ... + } + + """ + if kwargs is None: + kwargs = {} + if isinstance(names, str): + names = [names] + # TODO: remove when boozer transform works with multiple surfaces + names = [x for x in names if not ("Boozer" in x or "_mn" in x or x == "B modes")] + + num_points = 15 + rho = np.linspace(start=0, stop=delta, num=num_points) + grid = LinearGrid(rho=rho, M=5, N=5, NFP=eq.NFP, sym=eq.sym) + assert grid.axis.size + integrate = surface_integrals_map(grid, expand_out=False) + data = eq.compute(names, grid=grid) + + data_index_eq = data_index["desc.equilibrium.equilibrium.Equilibrium"] + for name in names: + if _skip_this(eq, name) or data_index_eq[name]["coordinates"] == "": + # can't check continuity of global scaler quantity + continue + # make single variable function of rho + if data_index_eq[name]["coordinates"] == "r": + # already single variable function of rho + profile = grid.compress(data[name]) + else: + # integrate out theta and zeta dependence + profile = np.where( + # True if integrand has nan on a given surface. + integrate(np.isnan(data[name])).astype(bool), + # The integration below replaces nan with 0; put them back. + np.nan, + # Norms and integrals are continuous functions, so their composition + # cannot disrupt existing continuity. Note that the absolute value + # before the integration ensures that a discontinuous integrand does + # not become continuous once integrated. + integrate(np.abs(data[name])), + ) + fit = kwargs.get(name, {}).get("desired_at_axis", desired_at_axis) + if fit is None: + if np.ndim(data_index_eq[name]["dim"]): + # can't polyfit tensor arrays like grad(B) + fit = (profile[0] + profile[1]) / 2 + else: + # fit the data to a polynomial to extrapolate to axis + poly = np.polynomial.polynomial.polyfit( + rho[1:], profile[1:], deg=min(5, num_points // 3) + ) + # constant term is the same as evaluating polynomial at rho=0 + fit = poly[0] + np.testing.assert_allclose( + actual=profile[0], + desired=fit, + rtol=kwargs.get(name, {}).get("rtol", rtol), + atol=kwargs.get(name, {}).get("atol", atol), + equal_nan=False, + err_msg=name, + ) + + +def get_matches(fun, pattern): + """Return all matches of ``pattern`` in source code of function ``fun``.""" + src = inspect.getsource(fun) + # attempt to remove any decorator functions + # (currently works without this filter, but better to be defensive) + src = src.partition("def ")[2] + # attempt to remove comments + src = "\n".join(line.partition("#")[0] for line in src.splitlines()) + matches = pattern.findall(src) + matches = {s.strip().strip('"') for s in matches} + return matches + + +def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): + """Get parameterization of thing computed by function ``fun``.""" + pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")') + decorator = inspect.getsource(fun).partition("def ")[0] + matches = pattern.findall(decorator) + # if list was found, split strings in list, else string was found so just get that + matches = [match[0].split(",") if match[0] else [match[1]] for match in matches] + # flatten the list + matches = {s.strip().strip('"') for sublist in matches for s in sublist} + matches.discard("") + return matches if matches else {default} + class TestAxisLimits: """Tests for compute functions evaluated at limits.""" + @pytest.mark.unit + def test_data_index_deps(self): + """Ensure developers do not add extra (or forget needed) dependencies.""" + queried_deps = {} + + pattern_names = re.compile(r"(?= 0 and n >= 0: B_nae += B_mn_nae[i] * np.cos(m * th) * np.cos(n * ph) - elif m >= 0 and n < 0: + elif m >= 0 > n: B_nae += -B_mn_nae[i] * np.cos(m * th) * np.sin(n * ph) - elif m < 0 and n >= 0: + elif m < 0 <= n: B_nae += -B_mn_nae[i] * np.sin(m * th) * np.cos(n * ph) elif m < 0 and n < 0: B_nae += B_mn_nae[i] * np.sin(m * th) * np.sin(n * ph) @@ -665,7 +665,7 @@ def test_NAE_QIC_solve(): grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) iota = compress(grid, eq.compute("iota", grid=grid)["iota"], "rho") - np.testing.assert_allclose(iota[1], qsc.iota, atol=1e-5) + np.testing.assert_allclose(iota[1], qsc.iota, atol=5e-4) np.testing.assert_allclose(iota[1:10], qsc.iota, atol=5e-4) # check lambda to match near axis @@ -692,7 +692,7 @@ def test_NAE_QIC_solve(): lam_nae = np.squeeze(lam_nae[:, 0, :]) lam_av_nae = np.mean(lam_nae, axis=0) - np.testing.assert_allclose(lam_av_nae, -qsc.iota * qsc.nu_spline(phi), atol=1e-4) + np.testing.assert_allclose(lam_av_nae, -qsc.iota * qsc.nu_spline(phi), atol=5e-4) # check |B| on axis @@ -710,9 +710,9 @@ def test_NAE_QIC_solve(): for i, (l, m, n) in enumerate(modes): if m >= 0 and n >= 0: B_nae += B_mn_nae[i] * np.cos(m * th) * np.cos(n * ph) - elif m >= 0 and n < 0: + elif m >= 0 > n: B_nae += -B_mn_nae[i] * np.cos(m * th) * np.sin(n * ph) - elif m < 0 and n >= 0: + elif m < 0 <= n: B_nae += -B_mn_nae[i] * np.sin(m * th) * np.cos(n * ph) elif m < 0 and n < 0: B_nae += B_mn_nae[i] * np.sin(m * th) * np.sin(n * ph) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index a202aa53a2..6c34c1547d 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -283,7 +283,7 @@ def test(eq): obj = QuasisymmetryTripleProduct(eq=eq) obj.build() ft = obj.compute_unscaled(*obj.xs(eq)) - np.testing.assert_allclose(ft, 0) + np.testing.assert_allclose(ft, 0, atol=5e-35) test(Equilibrium(iota=PowerSeriesProfile(0))) test(Equilibrium(current=PowerSeriesProfile(0))) @@ -582,8 +582,8 @@ def test_plasma_vessel_distance(): @pytest.mark.unit def test_mean_curvature(): """Test for mean curvature objective function.""" - # simple case like dshape should have mean curvature negative everywhere - eq = get("DSHAPE") + # torus should have mean curvature negative everywhere + eq = Equilibrium() obj = MeanCurvature(eq=eq) obj.build() H = obj.compute_unscaled(*obj.xs(eq)) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 281d4dd469..ad7b2af22b 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -917,6 +917,7 @@ def test_constrained_AL_lsq(): @pytest.mark.slow @pytest.mark.regression +@pytest.mark.xfail def test_constrained_AL_scalar(): """Tests that the augmented Lagrangian constrained optimizer does something.""" eq = desc.examples.get("SOLOVEV")