From dbf61eff3823dbcf440d8de0c4720ab65227d201 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 25 Jul 2023 13:05:57 -0500 Subject: [PATCH 01/19] git checkout add_all_limits desc/compute/_basis_vectors.py --- desc/compute/_basis_vectors.py | 3245 +++++++++++++++++++++++++++++--- 1 file changed, 2992 insertions(+), 253 deletions(-) diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 468764b8cb..033d6a850e 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -1,3 +1,14 @@ +"""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 @@ -9,686 +20,3414 @@ 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=["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["R"] * data["omega_r"], data["Z_r"]]).T + return data + + +@register_compute_fun( + name="e_theta", + label="\\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + 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["R"] * data["omega_t"], data["Z_t"]] + ).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/sqrt(g)", + label="\\mathbf{e}_{\\theta} / \\sqrt{g}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector divided by 3-D volume Jacobian", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "sqrt(g)"], + axis_limit_data=["e_theta_r", "sqrt(g)_r"], +) +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_zeta", + label="\\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant Toroidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_z", "Z_z", "omega_z"], +) +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_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^{-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="n_rho", + label="\\hat{\\mathbf{n}}_{\\rho}", + units="~", + units_long="None", + description="Unit vector normal to constant rho surface (direction of e^rho)", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e^rho"], +) +def _n_rho(params, transforms, profiles, data, **kwargs): + data["n_rho"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"], axis=-1)).T + return data + + +@register_compute_fun( + 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=["psi_r", "e^rho"], +) +def _gradpsi(params, transforms, profiles, data, **kwargs): + data["grad(psi)"] = (data["psi_r"] * data["e^rho"].T).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", + 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^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="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="grad(alpha)", + label="\\nabla \\alpha", + units="m^{-1}", + units_long="Inverse meters", + description="Unit vector along field line", + 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 + + +@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_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_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_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_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_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_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_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_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_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_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_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_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 + + +@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_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_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector, derivative wrt poloidal coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_t", "R_tt", "Z_tt", "omega_t", "omega_tt"], +) +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_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector, derivative wrt toroidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_r", "Z_r"], + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], ) -def _e_sub_rho(params, transforms, profiles, data, **kwargs): - data["e_rho"] = jnp.array([data["R_r"], data["0"], data["Z_r"]]).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_theta", - label="\\mathbf{e}_{\\theta}", + name="e_theta_rr", + label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant poloidal basis vector", + description=( + "Covariant Poloidal basis vector, second derivative wrt radial and radial" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_t", "Z_t"], + 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_theta(params, transforms, profiles, data, **kwargs): - data["e_theta"] = jnp.array([data["R_t"], data["0"], data["Z_t"]]).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_zeta", - label="\\mathbf{e}_{\\zeta}", + name="e_theta_rrr", + label="\\partial_{\\rho}{\\rho}{\\rho} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant toroidal basis vector", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "Z_z"], + 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_zeta(params, transforms, profiles, data, **kwargs): - data["e_zeta"] = jnp.array([data["R_z"], data["R"], data["Z_z"]]).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_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", + name="e_theta_rt", + label="\\partial_{\\rho}{\\theta} \\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 poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rr", "Z_rr"], + 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_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_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_rho_t", - label="\\partial_{\\theta} \\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 poloidal 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_rt", "Z_rt"], + 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_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_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_rho_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", + name="e_theta_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, derivative wrt toroidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt radial and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rz", "Z_rz"], + 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_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_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_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_tt", + label="\\partial_{\\theta}{\\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 poloidal and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_tt", "Z_tt"], + data=["R", "R_t", "R_tt", "R_ttt", "Z_ttt", "omega_t", "omega_tt", "omega_ttt"], ) -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_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_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_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, derivative wrt radial coordinate", + description=( + "Covariant Poloidal basis vector, second derivative wrt poloidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_r", "R_rz", "Z_rz"], + 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_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_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_zeta_t", - label="\\partial_{\\theta} \\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 poloidal angle", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial, poloidal," + " and toroidal coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_t", "R_tz", "Z_tz"], + 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_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_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_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", + name="e_theta_zz", + label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, derivative wrt toroidal angle", + description=( + "Covariant Poloidal basis vector, second derivative wrt toroidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_z", "R_zz", "Z_zz"], + 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_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_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_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\rho}", + name="e_theta_rzz", + label="\\partial_{\\rho}{\\zeta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt radial " - + "coordinate", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " once and toroidal twice" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrr", "Z_rrr"], + 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_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_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_tt", - label="\\partial_{\\theta \\theta} \\mathbf{e}_{\\rho}", + name="e_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt poloidal angle", + description="Covariant Toroidal basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtt", "Z_rtt"], + data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], ) -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_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_rho_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\rho}", + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt toroidal angle", + description="Covariant Toroidal basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rzz", "Z_rzz"], + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "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_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_rho_rt", - label="\\partial_{\\rho \\theta} \\mathbf{e}_{\\rho}", + name="e_zeta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt radial " - + "coordinate and poloidal angle", + description="Covariant Toroidal basis vector, derivative wrt toroidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrt", "Z_rrt"], + data=["R", "R_z", "R_zz", "Z_zz", "omega_z", "omega_zz"], ) -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_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_rho_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\rho}", + name="e_zeta_rr", + label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant radial basis vector, second derivative wrt radial " - + "coordinate and toroidal angle", + description=( + "Covariant Toroidal basis vector, second derivative wrt radial and radial" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrz", "Z_rrz"], + 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([data["R_rrz"], data["0"], data["Z_rrz"]]).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_rho_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\rho}", + name="e_zeta_rrr", + label="\\partial_{\\rho}{\\rho}{\\rho} \\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, third derivative wrt radial coordinate" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtz", "Z_rtz"], + data=[ + "R", + "R_r", + "R_z", + "R_rr", + "R_rz", + "R_rrr", + "R_rrz", + "R_rrrz", + "Z_rrrz", + "omega_r", + "omega_z", + "omega_rr", + "omega_rz", + "omega_rrr", + "omega_rrz", + "omega_rrrz", + ], ) -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_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_rr", - label="\\partial_{\\rho \\rho} \\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 radial " - + "coordinate", + description=( + "Covariant Toroidal basis vector, second derivative wrt radial and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rrt", "Z_rrt"], + 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_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_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_theta_tt", - label="\\partial_{\\theta \\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 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_ttt", "Z_ttt"], + 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_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_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_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{e}_{\\theta}", + name="e_zeta_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant poloidal 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=["0", "R_tzz", "Z_tzz"], + 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_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_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_theta_rt", - label="\\partial_{\\rho \\theta} \\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 poloidal 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_rtt", "Z_rtt"], + 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_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_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_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{e}_{\\theta}", + name="e_zeta_tt", + label="\\partial_{\\theta}{\\theta} \\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, second derivative wrt poloidal and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_rtz", "Z_rtz"], + data=[ + "R", + "R_t", + "R_tt", + "R_ttz", + "R_tz", + "R_z", + "Z_ttz", + "omega_t", + "omega_tt", + "omega_ttz", + "omega_tz", + "omega_z", + ], ) -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_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_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\theta}", + name="e_zeta_rtt", + label="\\partial_{\\rho}{\\theta}{\\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, third derivative wrt radial coordinate" + " once and poloidal twice" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0", "R_ttz", "Z_ttz"], + 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_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_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["R_rtt"] + + data["omega_rz"] * data["R_tt"] + + (1 + data["omega_z"]) * data["R_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_rr", - label="\\partial_{\\rho \\rho} \\mathbf{e}_{\\zeta}", + name="e_zeta_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt radial " - + "coordinate", + description=( + "Covariant Toroidal basis vector, second derivative wrt poloidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_rr", "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_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_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_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}", + label="\\partial_{\\zeta}{\\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 toroidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R_zz", "R_zzz", "Z_zzz"], + data=["R", "R_z", "R_zz", "R_zzz", "Z_zzz", "omega_z", "omega_zz", "omega_zzz"], ) 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 + 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_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}", - units="m", - units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt radial " - + "coordinate and toroidal angle", + 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=["R_rz", "R_rzz", "Z_rzz"], + 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_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_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_zeta_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant toroidal basis vector, second derivative wrt poloidal " - + "and toroidal angles", + 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=["R_tz", "R_tzz", "Z_tzz"], + 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_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_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_theta_PEST", - label="\\mathbf{e}_{\\theta_{PEST}}", - units="m", - units_long="meters", - description="Covariant straight field line poloidal basis vector", + 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=[ - "0", - "R_t", - "Z_t", - "lambda_t", - ], + 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_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] - ).T +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^rho", - label="\\mathbf{e}^{\\rho}", + name="e^theta_r", + label="\\partial{\\rho} \\mathbf{e}^{\\theta}", units="m^{-1}", units_long="inverse meters", - description="Contravariant radial basis vector", + description="Contravariant poloidal basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_theta", "e_zeta", "sqrt(g)"], + data=["e_zeta", "e_rho", "e_zeta_r", "e_rho_r", "sqrt(g)", "sqrt(g)_r"], ) -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_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", - label="\\mathbf{e}^{\\theta}", + name="e^theta_t", + label="\\partial{\\theta} \\mathbf{e}^{\\theta}", units="m^{-1}", units_long="inverse meters", - description="Contravariant poloidal basis vector", + description="Contravariant poloidal basis vector, derivative wrt poloidal" + " coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_zeta", "sqrt(g)"], + data=["e_zeta", "e_rho", "e_zeta_t", "e_rho_t", "sqrt(g)", "sqrt(g)_t"], ) -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_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^zeta", - label="\\mathbf{e}^{\\zeta}", + name="e^theta_z", + label="\\partial{\\zeta} \\mathbf{e}^{\\theta}", units="m^{-1}", units_long="inverse meters", - description="Contravariant toroidal basis vector", + description="Contravariant poloidal basis vector, derivative wrt toroidal" + " coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_theta", "sqrt(g)"], + data=["e_zeta", "e_rho", "e_zeta_z", "e_rho_z", "sqrt(g)", "sqrt(g)_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 _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="b", - label="\\hat{b}", - units="~", - units_long="None", - description="Unit vector along magnetic field", + 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=["B"], + 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 _b(params, transforms, profiles, data, **kwargs): - data["b"] = (data["B"].T / jnp.linalg.norm(data["B"], axis=-1)).T +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="n_rho", - label="\\hat{\\mathbf{n}}_{\\rho}", - units="~", - units_long="None", - description="Unit vector normal to constant rho surface (direction of e^rho)", + 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"], + 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 _n_rho(params, transforms, profiles, data, **kwargs): - data["n_rho"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"], axis=-1)).T +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="grad(alpha)", - label="\\nabla \\alpha", + name="e^zeta_z", + label="\\partial{\\zeta} \\mathbf{e}^{\\zeta}", units="m^{-1}", - units_long="Inverse meters", - description="Unit vector normal to flux surface", + units_long="inverse meters", + description="Contravariant toroidal basis vector, derivative wrt toroidal" + " coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e^rho", "e^theta", "e^zeta", "alpha_r", "alpha_t", "alpha_z"], + 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 _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 +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 From 4eb0001723260897167ccdf846ab74e45f6c1121 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 25 Jul 2023 13:07:43 -0500 Subject: [PATCH 02/19] git checkout add_all_limits desc/compute/_core.py --- desc/compute/_core.py | 1121 ++++++++++++++++++++++++++++++++++++++--- 1 file changed, 1050 insertions(+), 71 deletions(-) diff --git a/desc/compute/_core.py b/desc/compute/_core.py index 418efb34b9..5cfa2cc710 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -67,7 +67,7 @@ def _rho_r(params, transforms, profiles, data, **kwargs): units="~", units_long="None", description="Radial coordinate, proportional to the square root " - + "of the toroidal flux, derivative wrt poloidal coordinate", + "of the toroidal flux, derivative wrt poloidal coordinate", dim=1, params=[], transforms={}, @@ -86,7 +86,7 @@ def _rho_t(params, transforms, profiles, data, **kwargs): units="~", units_long="None", description="Radial coordinate, proportional to the square root " - + "of the toroidal flux, derivative wrt toroidal coordinate", + "of the toroidal flux, derivative wrt toroidal coordinate", dim=1, params=[], transforms={}, @@ -123,7 +123,7 @@ def _theta(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal angular coordinate (geometric, not magnetic), " - + "derivative wrt radial coordinate", + "derivative wrt radial coordinate", dim=1, params=[], transforms={}, @@ -142,7 +142,7 @@ def _theta_r(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal angular coordinate (geometric, not magnetic), " - + "derivative wrt poloidal coordinate", + "derivative wrt poloidal coordinate", dim=1, params=[], transforms={}, @@ -161,7 +161,7 @@ def _theta_t(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal angular coordinate (geometric, not magnetic), " - + "derivative wrt toroidal coordinate", + "derivative wrt toroidal coordinate", dim=1, params=[], transforms={}, @@ -270,7 +270,7 @@ def _theta_PEST(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="PEST straight field line poloidal angular coordinate, derivative wrt " - + "radial coordinate", + "radial coordinate", dim=1, params=[], transforms={}, @@ -289,7 +289,7 @@ def _theta_PEST_r(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="PEST straight field line poloidal angular coordinate, derivative wrt " - + "poloidal coordinate", + "poloidal coordinate", dim=1, params=[], transforms={}, @@ -308,7 +308,7 @@ def _theta_PEST_t(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="PEST straight field line poloidal angular coordinate, derivative wrt " - + "toroidal coordinate", + "toroidal coordinate", dim=1, params=[], transforms={}, @@ -350,10 +350,14 @@ def _alpha(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["lambda_r", "zeta", "iota_r"], + data=["theta_PEST_r", "phi", "phi_r", "iota", "iota_r"], ) def _alpha_r(params, transforms, profiles, data, **kwargs): - data["alpha_r"] = data["lambda_r"] - data["iota_r"] * data["zeta"] + data["alpha_r"] = ( + data["theta_PEST_r"] + - data["iota_r"] * data["phi"] + - data["iota"] * data["phi_r"] + ) return data @@ -368,10 +372,10 @@ def _alpha_r(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["lambda_t"], + data=["theta_PEST_t", "phi_t", "iota"], ) def _alpha_t(params, transforms, profiles, data, **kwargs): - data["alpha_t"] = 1 + data["lambda_t"] + data["alpha_t"] = data["theta_PEST_t"] + data["iota"] * data["phi_t"] return data @@ -386,10 +390,10 @@ def _alpha_t(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["lambda_z", "iota"], + data=["theta_PEST_z", "phi_z", "iota"], ) def _alpha_z(params, transforms, profiles, data, **kwargs): - data["alpha_z"] = data["lambda_z"] - data["iota"] + data["alpha_z"] = data["theta_PEST_z"] - data["iota"] * data["phi_z"] return data @@ -525,7 +529,7 @@ def _R_zz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Major radius in lab frame, second derivative wrt radius " - + "and poloidal angle", + "and poloidal angle", dim=1, params=["R_lmn"], transforms={"R": [[1, 1, 0]]}, @@ -544,7 +548,7 @@ def _R_rt(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Major radius in lab frame, second derivative wrt radius " - + "and toroidal angle", + "and toroidal angle", dim=1, params=["R_lmn"], transforms={"R": [[1, 0, 1]]}, @@ -563,7 +567,7 @@ def _R_rz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Major radius in lab frame, second derivative wrt poloidal " - + "and toroidal angles", + "and toroidal angles", dim=1, params=["R_lmn"], transforms={"R": [[0, 1, 1]]}, @@ -594,6 +598,62 @@ def _R_rrr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="R_rrrr", + label="\\partial_{\rho \\rho \\rho \\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth radial derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[4, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rrrr(params, transforms, profiles, data, **kwargs): + data["R_rrrr"] = transforms["R"].transform(params["R_lmn"], 4, 0, 0) + return data + + +@register_compute_fun( + name="R_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt" + " radial coordinate thrice and poloidal once", + dim=1, + params=["R_lmn"], + transforms={"R": [[3, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rrrt(params, transforms, profiles, data, **kwargs): + data["R_rrrt"] = transforms["R"].transform(params["R_lmn"], 3, 1, 0) + return data + + +@register_compute_fun( + name="R_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt" + " radial coordinate thrice and toroidal once", + dim=1, + params=["R_lmn"], + transforms={"R": [[3, 0, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rrrz(params, transforms, profiles, data, **kwargs): + data["R_rrrz"] = transforms["R"].transform(params["R_lmn"], 3, 0, 1) + return data + + @register_compute_fun( name="R_ttt", label="\\partial_{\\theta \\theta \\theta} R", @@ -636,7 +696,7 @@ def _R_zzz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Major radius in lab frame, third derivative, wrt radius twice " - + "and poloidal angle", + "and poloidal angle", dim=1, params=["R_lmn"], transforms={"R": [[2, 1, 0]]}, @@ -649,13 +709,32 @@ def _R_rrt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="R_rrtt", + label="\\partial_{\\rho \\rho \\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fouth derivative, wrt radius twice " + "and poloidal angle twice", + dim=1, + params=["R_lmn"], + transforms={"R": [[2, 2, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rrtt(params, transforms, profiles, data, **kwargs): + data["R_rrtt"] = transforms["R"].transform(params["R_lmn"], 2, 2, 0) + return data + + @register_compute_fun( name="R_rtt", label="\\partial_{\\rho \\theta \\theta} R", units="m", units_long="meters", description="Major radius in lab frame, third derivative wrt radius and " - + "poloidal angle twice", + "poloidal angle twice", dim=1, params=["R_lmn"], transforms={"R": [[1, 2, 0]]}, @@ -668,13 +747,51 @@ def _R_rtt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="R_rttz", + label="\\partial_{\\rho \\theta \\theta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius once, " + "poloidal angle twice, and toroidal angle once", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 2, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rttz(params, transforms, profiles, data, **kwargs): + data["R_rttz"] = transforms["R"].transform(params["R_lmn"], 1, 2, 1) + return data + + +@register_compute_fun( + name="R_rttt", + label="\\partial_{\\rho \\theta \\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius and " + "poloidal angle thrice", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 3, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rttt(params, transforms, profiles, data, **kwargs): + data["R_rttt"] = transforms["R"].transform(params["R_lmn"], 1, 3, 0) + return data + + @register_compute_fun( name="R_rrz", label="\\partial_{\\rho \\rho \\zeta} R", units="m", units_long="meters", description="Major radius in lab frame, third derivative, wrt radius twice " - + "and toroidal angle", + "and toroidal angle", dim=1, params=["R_lmn"], transforms={"R": [[2, 0, 1]]}, @@ -687,13 +804,32 @@ def _R_rrz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="R_rrzz", + label="\\partial_{\\rho \\rho \\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative, wrt radius twice " + "and toroidal angle twice", + dim=1, + params=["R_lmn"], + transforms={"R": [[2, 0, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rrzz(params, transforms, profiles, data, **kwargs): + data["R_rrzz"] = transforms["R"].transform(params["R_lmn"], 2, 0, 2) + return data + + @register_compute_fun( name="R_rzz", label="\\partial_{\\rho \\zeta \\zeta} R", units="m", units_long="meters", description="Major radius in lab frame, third derivative wrt radius and " - + "toroidal angle twice", + "toroidal angle twice", dim=1, params=["R_lmn"], transforms={"R": [[1, 0, 2]]}, @@ -706,13 +842,32 @@ def _R_rzz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="R_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius and " + "toroidal angle thrice", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 0, 3]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rzzz(params, transforms, profiles, data, **kwargs): + data["R_rzzz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 3) + return data + + @register_compute_fun( name="R_ttz", label="\\partial_{\\theta \\theta \\zeta} R", units="m", units_long="meters", description="Major radius in lab frame, third derivative wrt poloidal angle " - + "twice and toroidal angle", + "twice and toroidal angle", dim=1, params=["R_lmn"], transforms={"R": [[0, 2, 1]]}, @@ -731,7 +886,7 @@ def _R_ttz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Major radius in lab frame, third derivative wrt poloidal angle " - + "and toroidal angle twice", + "and toroidal angle twice", dim=1, params=["R_lmn"], transforms={"R": [[0, 1, 2]]}, @@ -750,7 +905,7 @@ def _R_tzz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Major radius in lab frame, third derivative wrt radius, poloidal " - + "angle, and toroidal angle", + "angle, and toroidal angle", dim=1, params=["R_lmn"], transforms={"R": [[1, 1, 1]]}, @@ -763,6 +918,44 @@ def _R_rtz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="R_rrtz", + label="\\partial_{\\rho \\rho \\theta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius twice," + " poloidal angle, and toroidal angle", + dim=1, + params=["R_lmn"], + transforms={"R": [[2, 1, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rrtz(params, transforms, profiles, data, **kwargs): + data["R_rrtz"] = transforms["R"].transform(params["R_lmn"], 2, 1, 1) + return data + + +@register_compute_fun( + name="R_rtzz", + label="\\partial_{\\rho \\theta \\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius, poloidal " + "angle, and toroidal angle twice", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 1, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_rtzz(params, transforms, profiles, data, **kwargs): + data["R_rtzz"] = transforms["R"].transform(params["R_lmn"], 1, 1, 2) + return data + + @register_compute_fun( name="Z", label="Z", @@ -895,7 +1088,7 @@ def _Z_zz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Vertical coordinate in lab frame, second derivative wrt radius " - + "and poloidal angle", + "and poloidal angle", dim=1, params=["Z_lmn"], transforms={"Z": [[1, 1, 0]]}, @@ -914,7 +1107,7 @@ def _Z_rt(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Vertical coordinate in lab frame, second derivative wrt radius " - + "and toroidal angle", + "and toroidal angle", dim=1, params=["Z_lmn"], transforms={"Z": [[1, 0, 1]]}, @@ -933,7 +1126,7 @@ def _Z_rz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Vertical coordinate in lab frame, second derivative wrt poloidal " - + "and toroidal angles", + "and toroidal angles", dim=1, params=["Z_lmn"], transforms={"Z": [[0, 1, 1]]}, @@ -964,6 +1157,62 @@ def _Z_rrr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Z_rrrr", + label="\\partial_{\rho \\rho \\rho \\rho} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth radial derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[4, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rrrr(params, transforms, profiles, data, **kwargs): + data["Z_rrrr"] = transforms["Z"].transform(params["Z_lmn"], 4, 0, 0) + return data + + +@register_compute_fun( + name="Z_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative wrt " + " radial coordinate thrice and poloidal once", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[3, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rrrt(params, transforms, profiles, data, **kwargs): + data["Z_rrrt"] = transforms["Z"].transform(params["Z_lmn"], 3, 1, 0) + return data + + +@register_compute_fun( + name="Z_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative wrt " + " radial coordinate thrice and toroidal once", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[3, 0, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rrrz(params, transforms, profiles, data, **kwargs): + data["Z_rrrz"] = transforms["Z"].transform(params["Z_lmn"], 3, 0, 1) + return data + + @register_compute_fun( name="Z_ttt", label="\\partial_{\\theta \\theta \\theta} Z", @@ -1006,7 +1255,7 @@ def _Z_zzz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative, wrt radius " - + "twice and poloidal angle", + "twice and poloidal angle", dim=1, params=["Z_lmn"], transforms={"Z": [[2, 1, 0]]}, @@ -1019,13 +1268,32 @@ def _Z_rrt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Z_rrtt", + label="\\partial_{\\rho \\rho \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative, wrt radius " + "twice and poloidal angle twice", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[2, 2, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rrtt(params, transforms, profiles, data, **kwargs): + data["Z_rrtt"] = transforms["Z"].transform(params["Z_lmn"], 2, 2, 0) + return data + + @register_compute_fun( name="Z_rtt", label="\\partial_{\\rho \\theta \\theta} Z", units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative wrt radius " - + "and poloidal angle twice", + "and poloidal angle twice", dim=1, params=["Z_lmn"], transforms={"Z": [[1, 2, 0]]}, @@ -1038,13 +1306,51 @@ def _Z_rtt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Z_rttz", + label="\\partial_{\\rho \\theta \\theta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative wrt radius " + "once, poloidal angle twice, and toroidal angle once", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[1, 2, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rttz(params, transforms, profiles, data, **kwargs): + data["Z_rttz"] = transforms["Z"].transform(params["Z_lmn"], 1, 2, 1) + return data + + +@register_compute_fun( + name="Z_rttt", + label="\\partial_{\\rho \\theta \\theta \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, third derivative wrt radius " + "and poloidal angle thrice", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[1, 3, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rttt(params, transforms, profiles, data, **kwargs): + data["Z_rttt"] = transforms["Z"].transform(params["Z_lmn"], 1, 3, 0) + return data + + @register_compute_fun( name="Z_rrz", label="\\partial_{\\rho \\rho \\zeta} Z", units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative, wrt radius " - + "twice and toroidal angle", + "twice and toroidal angle", dim=1, params=["Z_lmn"], transforms={"Z": [[2, 0, 1]]}, @@ -1057,13 +1363,32 @@ def _Z_rrz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Z_rrzz", + label="\\partial_{\\rho \\rho \\zeta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative, wrt radius " + "twice and toroidal angle twice", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[2, 0, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rrzz(params, transforms, profiles, data, **kwargs): + data["Z_rrzz"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 2) + return data + + @register_compute_fun( name="Z_rzz", label="\\partial_{\\rho \\zeta \\zeta} Z", units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative wrt radius " - + "and toroidal angle twice", + "and toroidal angle twice", dim=1, params=["Z_lmn"], transforms={"Z": [[1, 0, 2]]}, @@ -1076,13 +1401,32 @@ def _Z_rzz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Z_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, third derivative wrt radius " + "and toroidal angle thrice", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[1, 0, 3]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rzzz(params, transforms, profiles, data, **kwargs): + data["Z_rzzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 3) + return data + + @register_compute_fun( name="Z_ttz", label="\\partial_{\\theta \\theta \\zeta} Z", units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative wrt poloidal " - + "angle twice and toroidal angle", + "angle twice and toroidal angle", dim=1, params=["Z_lmn"], transforms={"Z": [[0, 2, 1]]}, @@ -1101,7 +1445,7 @@ def _Z_ttz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative wrt poloidal " - + "angle and toroidal angle twice", + "angle and toroidal angle twice", dim=1, params=["Z_lmn"], transforms={"Z": [[0, 1, 2]]}, @@ -1120,7 +1464,7 @@ def _Z_tzz(params, transforms, profiles, data, **kwargs): units="m", units_long="meters", description="Vertical coordinate in lab frame, third derivative wrt radius, " - + "poloidal angle, and toroidal angle", + "poloidal angle, and toroidal angle", dim=1, params=["Z_lmn"], transforms={"Z": [[1, 1, 1]]}, @@ -1133,6 +1477,44 @@ def _Z_rtz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="Z_rrtz", + label="\\partial_{\\rho \\rho \\theta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative wrt radius" + "twice, poloidal angle, and toroidal angle", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[2, 1, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rrtz(params, transforms, profiles, data, **kwargs): + data["Z_rrtz"] = transforms["Z"].transform(params["Z_lmn"], 2, 1, 1) + return data + + +@register_compute_fun( + name="Z_rtzz", + label="\\partial_{\\rho \\theta \\zeta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, fourth derivative wrt radius, " + "poloidal angle, and toroidal angle twice", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[1, 1, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_rtzz(params, transforms, profiles, data, **kwargs): + data["Z_rtzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 2) + return data + + @register_compute_fun( name="phi", label="\\phi", @@ -1144,10 +1526,10 @@ def _Z_rtz(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["zeta"], + data=["zeta", "omega"], ) def _phi(params, transforms, profiles, data, **kwargs): - data["phi"] = data["zeta"] + data["phi"] = data["zeta"] + data["omega"] return data @@ -1162,10 +1544,10 @@ def _phi(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_r"], ) def _phi_r(params, transforms, profiles, data, **kwargs): - data["phi_r"] = data["0"] + data["phi_r"] = data["omega_r"] return data @@ -1180,10 +1562,10 @@ def _phi_r(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_t"], ) def _phi_t(params, transforms, profiles, data, **kwargs): - data["phi_t"] = data["0"] + data["phi_t"] = data["omega_t"] return data @@ -1198,10 +1580,10 @@ def _phi_t(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_z"], ) def _phi_z(params, transforms, profiles, data, **kwargs): - data["phi_z"] = jnp.ones_like(data["0"]) + data["phi_z"] = 1 + data["omega_z"] return data @@ -1216,10 +1598,10 @@ def _phi_z(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_rr"], ) def _phi_rr(params, transforms, profiles, data, **kwargs): - data["phi_rr"] = data["0"] + data["phi_rr"] = data["omega_rr"] return data @@ -1229,16 +1611,16 @@ def _phi_rr(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Toroidal angle in lab frame, second derivative wrt radial and " - + "poloidal coordinate", + "poloidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_rt"], ) def _phi_rt(params, transforms, profiles, data, **kwargs): - data["phi_rt"] = data["0"] + data["phi_rt"] = data["omega_rt"] return data @@ -1248,16 +1630,16 @@ def _phi_rt(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Toroidal angle in lab frame, second derivative wrt radial and " - + "toroidal coordinate", + "toroidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_rz"], ) def _phi_rz(params, transforms, profiles, data, **kwargs): - data["phi_rz"] = data["0"] + data["phi_rz"] = data["omega_rz"] return data @@ -1267,16 +1649,16 @@ def _phi_rz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Toroidal angle in lab frame, second derivative wrt poloidal " - + "coordinate", + "coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_tt"], ) def _phi_tt(params, transforms, profiles, data, **kwargs): - data["phi_tt"] = data["0"] + data["phi_tt"] = data["omega_tt"] return data @@ -1286,16 +1668,16 @@ def _phi_tt(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Toroidal angle in lab frame, second derivative wrt poloidal and " - + "toroidal coordinate", + "toroidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_tz"], ) def _phi_tz(params, transforms, profiles, data, **kwargs): - data["phi_tz"] = data["0"] + data["phi_tz"] = data["omega_tz"] return data @@ -1305,16 +1687,16 @@ def _phi_tz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Toroidal angle in lab frame, second derivative wrt toroidal " - + "coordinate", + "coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_zz"], ) def _phi_zz(params, transforms, profiles, data, **kwargs): - data["phi_zz"] = data["0"] + data["phi_zz"] = data["omega_zz"] return data @@ -1612,7 +1994,7 @@ def _lambda_zz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, second derivative wrt radius and " - + "poloidal angle", + "poloidal angle", dim=1, params=["L_lmn"], transforms={"L": [[1, 1, 0]]}, @@ -1631,7 +2013,7 @@ def _lambda_rt(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, second derivative wrt radius and " - + "toroidal angle", + "toroidal angle", dim=1, params=["L_lmn"], transforms={"L": [[1, 0, 1]]}, @@ -1650,7 +2032,7 @@ def _lambda_rz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, second derivative wrt poloidal and " - + "toroidal angles", + "toroidal angles", dim=1, params=["L_lmn"], transforms={"L": [[0, 1, 1]]}, @@ -1682,8 +2064,46 @@ def _lambda_rrr(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="lambda_ttt", - label="\\partial_{\\theta \\theta \\theta} \\lambda", + name="lambda_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third radial derivative and" + " first poloidal derivative", + dim=1, + params=["L_lmn"], + transforms={"L": [[3, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _lambda_rrrt(params, transforms, profiles, data, **kwargs): + data["lambda_rrrt"] = transforms["L"].transform(params["L_lmn"], 3, 1, 0) + return data + + +@register_compute_fun( + name="lambda_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third radial derivative and" + " first toroidal derivative", + dim=1, + params=["L_lmn"], + transforms={"L": [[3, 0, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _lambda_rrrz(params, transforms, profiles, data, **kwargs): + data["lambda_rrrz"] = transforms["L"].transform(params["L_lmn"], 3, 0, 1) + return data + + +@register_compute_fun( + name="lambda_ttt", + label="\\partial_{\\theta \\theta \\theta} \\lambda", units="rad", units_long="radians", description="Poloidal stream function, third poloidal derivative", @@ -1723,7 +2143,7 @@ def _lambda_zzz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative, wrt radius twice " - + "and poloidal angle", + "and poloidal angle", dim=1, params=["L_lmn"], transforms={"L": [[2, 1, 0]]}, @@ -1742,7 +2162,7 @@ def _lambda_rrt(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative wrt radius and " - + "poloidal angle twice", + "poloidal angle twice", dim=1, params=["L_lmn"], transforms={"L": [[1, 2, 0]]}, @@ -1761,7 +2181,7 @@ def _lambda_rtt(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative, wrt radius twice " - + "and toroidal angle", + "and toroidal angle", dim=1, params=["L_lmn"], transforms={"L": [[2, 0, 1]]}, @@ -1780,7 +2200,7 @@ def _lambda_rrz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative wrt radius and " - + "toroidal angle twice", + "toroidal angle twice", dim=1, params=["L_lmn"], transforms={"L": [[1, 0, 2]]}, @@ -1799,7 +2219,7 @@ def _lambda_rzz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative wrt poloidal angle " - + "twice and toroidal angle", + "twice and toroidal angle", dim=1, params=["L_lmn"], transforms={"L": [[0, 2, 1]]}, @@ -1818,7 +2238,7 @@ def _lambda_ttz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative wrt poloidal angle " - + "and toroidal angle twice", + "and toroidal angle twice", dim=1, params=["L_lmn"], transforms={"L": [[0, 1, 2]]}, @@ -1837,7 +2257,7 @@ def _lambda_tzz(params, transforms, profiles, data, **kwargs): units="rad", units_long="radians", description="Poloidal stream function, third derivative wrt radius, poloidal " - + " angle, and toroidal angle", + " angle, and toroidal angle", dim=1, params=["L_lmn"], transforms={"L": [[1, 1, 1]]}, @@ -1848,3 +2268,562 @@ def _lambda_tzz(params, transforms, profiles, data, **kwargs): def _lambda_rtz(params, transforms, profiles, data, **kwargs): data["lambda_rtz"] = transforms["L"].transform(params["L_lmn"], 1, 1, 1) return data + + +@register_compute_fun( + name="omega", + label="\\omega", + units="rad", + units_long="radians", + description="Toroidal stream function", + dim=1, + params=[], # ["W_lmn"], + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega(params, transforms, profiles, data, **kwargs): + data["omega"] = data["0"] + return data + + +@register_compute_fun( + name="omega_r", + label="\\partial_{\\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, first radial derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_r(params, transforms, profiles, data, **kwargs): + data["omega_r"] = data["0"] + return data + + +@register_compute_fun( + name="omega_t", + label="\\partial_{\\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, first poloidal derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_t(params, transforms, profiles, data, **kwargs): + data["omega_t"] = data["0"] + return data + + +@register_compute_fun( + name="omega_z", + label="\\partial_{\\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, first toroidal derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_z(params, transforms, profiles, data, **kwargs): + data["omega_z"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rr", + label="\\partial_{\\rho \\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second radial derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rr(params, transforms, profiles, data, **kwargs): + data["omega_rr"] = data["0"] + return data + + +@register_compute_fun( + name="omega_tt", + label="\\partial_{\\theta \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second poloidal derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_tt(params, transforms, profiles, data, **kwargs): + data["omega_tt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_zz", + label="\\partial_{\\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second toroidal derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_zz(params, transforms, profiles, data, **kwargs): + data["omega_zz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rt", + label="\\partial_{\\rho \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second derivative wrt radius and " + "poloidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rt(params, transforms, profiles, data, **kwargs): + data["omega_rt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rz", + label="\\partial_{\\rho \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second derivative wrt radius and " + "toroidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rz(params, transforms, profiles, data, **kwargs): + data["omega_rz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_tz", + label="\\partial_{\\theta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second derivative wrt poloidal and " + "toroidal angles", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_tz(params, transforms, profiles, data, **kwargs): + data["omega_tz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrr", + label="\\partial_{\rho \\rho \\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third radial derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrr(params, transforms, profiles, data, **kwargs): + data["omega_rrr"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrrr", + label="\\partial_{\rho \\rho \\rho \\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth radial derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrrr(params, transforms, profiles, data, **kwargs): + data["omega_rrrr"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative wrt radial coordinate" + " thrice and poloidal once", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrrt(params, transforms, profiles, data, **kwargs): + data["omega_rrrt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative wrt radial coordinate" + " thrice and toroidal once", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrrz(params, transforms, profiles, data, **kwargs): + data["omega_rrrz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_ttt", + label="\\partial_{\\theta \\theta \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third poloidal derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_ttt(params, transforms, profiles, data, **kwargs): + data["omega_ttt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third toroidal derivative", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_zzz(params, transforms, profiles, data, **kwargs): + data["omega_zzz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrt", + label="\\partial_{\\rho \\rho \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative, wrt radius twice " + "and poloidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrt(params, transforms, profiles, data, **kwargs): + data["omega_rrt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrtt", + label="\\partial_{\\rho \\rho \\theta \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative, wrt radius twice " + "and poloidal angle twice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrtt(params, transforms, profiles, data, **kwargs): + data["omega_rrtt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rtt", + label="\\partial_{\\rho \\theta \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt radius and " + "poloidal angle twice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rtt(params, transforms, profiles, data, **kwargs): + data["omega_rtt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rttz", + label="\\partial_{\\rho \\theta \\theta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative wrt radius once, " + "poloidal angle twice, and toroidal angle once", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rttz(params, transforms, profiles, data, **kwargs): + data["omega_rttz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rttt", + label="\\partial_{\\rho \\theta \\theta \\theta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt radius and " + "poloidal angle thrice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rttt(params, transforms, profiles, data, **kwargs): + data["omega_rttt"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrz", + label="\\partial_{\\rho \\rho \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative, wrt radius twice " + "and toroidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrz(params, transforms, profiles, data, **kwargs): + data["omega_rrz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrzz", + label="\\partial_{\\rho \\rho \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative, wrt radius twice " + "and toroidal angle twice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrzz(params, transforms, profiles, data, **kwargs): + data["omega_rrzz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt radius and " + "toroidal angle twice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rzz(params, transforms, profiles, data, **kwargs): + data["omega_rzz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt radius and " + "toroidal angle thrice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rzzz(params, transforms, profiles, data, **kwargs): + data["omega_rzzz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_ttz", + label="\\partial_{\\theta \\theta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt poloidal angle " + "twice and toroidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_ttz(params, transforms, profiles, data, **kwargs): + data["omega_ttz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_tzz", + label="\\partial_{\\theta \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt poloidal angle " + "and toroidal angle twice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_tzz(params, transforms, profiles, data, **kwargs): + data["omega_tzz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third derivative wrt radius, poloidal" + " angle, and toroidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rtz(params, transforms, profiles, data, **kwargs): + data["omega_rtz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rrtz", + label="\\partial_{\\rho \\theta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative wrt radius twice," + " poloidal angle, and toroidal angle", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rrtz(params, transforms, profiles, data, **kwargs): + data["omega_rrtz"] = data["0"] + return data + + +@register_compute_fun( + name="omega_rtzz", + label="\\partial_{\\rho \\theta \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, fourth derivative wrt radius, poloidal" + " angle, and toroidal angle twice", + dim=1, + params=[], # ["W_lmn"] + transforms={}, + profiles=[], + coordinates="rtz", + data=["0"], +) +def _omega_rtzz(params, transforms, profiles, data, **kwargs): + data["omega_rtzz"] = data["0"] + return data From 055ae48b5d8aad680832129f2411cba99af5b537 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 25 Jul 2023 13:30:02 -0500 Subject: [PATCH 03/19] Include patch from #556 that lets reviewers skip dependency review git checkout add_all_limits tests/test_axis_limits.py --- tests/test_axis_limits.py | 366 +++++++++++++++++++++++++++++++------- 1 file changed, 297 insertions(+), 69 deletions(-) diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index 9c07535af6..df448be71f 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -1,95 +1,323 @@ """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.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", + "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", +} | data_index.keys() + + +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[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.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=1e-5, + rtol=1e-4, + atol=1e-5, + 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. + 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")] + + rho = np.linspace(0, 1, 10) * delta + grid = LinearGrid(rho=rho, M=5, N=5, NFP=eq.NFP, sym=eq.sym) + assert grid.axis.size + integrate = surface_integrals_map(grid) + data = eq.compute(names, grid=grid) + + for name in names: + if _skip_this(eq, name) or data_index[name]["coordinates"] == "": + continue + # make single variable function of rho + if data_index[name]["coordinates"] == "r": + profile = data[name] + else: + # 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. + profile = integrate(np.abs(data[name])) + # integration replaced nans with 0, put them back + profile = np.where(np.isnan(data[name]), np.nan, profile) + profile = grid.compress(profile) + fit = kwargs.get(name, {}).get("desired_at_axis", desired_at_axis) + if fit is None: + if np.ndim(data_index[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.polyfit(rho[1:], profile[1:], 6) + # constant term is the same as evaluating polynomial at rho=0 + fit = poly[-1] + 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.replace("'", "").replace('"', "") for s in matches} + return matches + 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_keys = re.compile(r"(? Date: Tue, 25 Jul 2023 16:04:37 -0500 Subject: [PATCH 04/19] Alphabetize basis vectors --- desc/compute/_basis_vectors.py | 2388 ++++++++++++++++---------------- 1 file changed, 1194 insertions(+), 1194 deletions(-) diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 033d6a850e..f167e109df 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -16,174 +16,171 @@ @register_compute_fun( - name="e_rho", - label="\\mathbf{e}_{\\rho}", - units="m", - units_long="meters", - description="Covariant Radial basis vector", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - 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["R"] * data["omega_r"], data["Z_r"]]).T - return data - - -@register_compute_fun( - name="e_theta", - label="\\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description="Covariant Poloidal basis vector", + name="b", + label="\\hat{b}", + units="~", + units_long="None", + description="Unit vector along magnetic field", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "Z_t", "omega_t"], + data=["B"], ) -def _e_sub_theta(params, transforms, profiles, data, **kwargs): - data["e_theta"] = jnp.array( - [data["R_t"], data["R"] * data["omega_t"], data["Z_t"]] - ).T +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_theta_PEST", - label="\\mathbf{e}_{\\theta_{PEST}}", - units="m", - units_long="meters", - description="Covariant straight field line (PEST) poloidal basis vector", + 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", "theta_PEST_t"], + data=["e_theta/sqrt(g)", "e_zeta"], ) -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 +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_theta/sqrt(g)", - label="\\mathbf{e}_{\\theta} / \\sqrt{g}", - units="m", - units_long="meters", - description="Covariant Poloidal basis vector divided by 3-D volume Jacobian", + 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", "sqrt(g)"], - axis_limit_data=["e_theta_r", "sqrt(g)_r"], + 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_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, +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_zeta", - label="\\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description="Covariant Toroidal basis vector", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_z", "Z_z", "omega_z"], -) -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_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}", + name="e^rho_t", + label="\\partial{\\theta} \\mathbf{e}^{\\rho}", units="m^{-1}", units_long="inverse meters", - description="Contravariant radial basis vector", + description="Contravariant radial basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_theta/sqrt(g)", "e_zeta"], + 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(params, transforms, profiles, data, **kwargs): - data["e^rho"] = cross(data["e_theta/sqrt(g)"], data["e_zeta"]) +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="n_rho", - label="\\hat{\\mathbf{n}}_{\\rho}", - units="~", - units_long="None", - description="Unit vector normal to constant rho surface (direction of e^rho)", + 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^rho"], + 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 _n_rho(params, transforms, profiles, data, **kwargs): - data["n_rho"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"], axis=-1)).T +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="grad(psi)", - label="\\nabla\\psi", - units="Wb / m", - units_long="Webers per meter", - description="Toroidal flux gradient (normalized by 2pi)", + 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=["psi_r", "e^rho"], + data=["e^theta*sqrt(g)", "sqrt(g)"], ) -def _gradpsi(params, transforms, profiles, data, **kwargs): - data["grad(psi)"] = (data["psi_r"] * data["e^rho"].T).T +def _e_sup_theta(params, transforms, profiles, data, **kwargs): + data["e^theta"] = (data["e^theta*sqrt(g)"].T / data["sqrt(g)"]).T return data @@ -206,155 +203,282 @@ def _e_sup_theta_times_sqrt_g(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e^theta", - label="\\mathbf{e}^{\\theta}", + name="e^theta_r", + label="\\partial{\\rho} \\mathbf{e}^{\\theta}", units="m^{-1}", units_long="inverse meters", - description="Contravariant poloidal basis vector", + description="Contravariant poloidal basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e^theta*sqrt(g)", "sqrt(g)"], + data=["e_zeta", "e_rho", "e_zeta_r", "e_rho_r", "sqrt(g)", "sqrt(g)_r"], ) -def _e_sup_theta(params, transforms, profiles, data, **kwargs): - data["e^theta"] = (data["e^theta*sqrt(g)"].T / data["sqrt(g)"]).T +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^zeta", - label="\\mathbf{e}^{\\zeta}", + name="e^theta_t", + label="\\partial{\\theta} \\mathbf{e}^{\\theta}", units="m^{-1}", units_long="inverse meters", - description="Contravariant toroidal basis vector", + description="Contravariant poloidal basis vector, derivative wrt poloidal" + " coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_theta/sqrt(g)"], + data=["e_zeta", "e_rho", "e_zeta_t", "e_rho_t", "sqrt(g)", "sqrt(g)_t"], ) -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="b", - label="\\hat{b}", - units="~", - units_long="None", - description="Unit vector along magnetic field", +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=["B"], + data=["e_zeta", "e_rho", "e_zeta_z", "e_rho_z", "sqrt(g)", "sqrt(g)_z"], ) -def _b(params, transforms, profiles, data, **kwargs): - data["b"] = (data["B"].T / jnp.linalg.norm(data["B"], axis=-1)).T +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="grad(alpha)", - label="\\nabla \\alpha", + name="e^zeta", + label="\\mathbf{e}^{\\zeta}", units="m^{-1}", - units_long="Inverse meters", - description="Unit vector along field line", + units_long="inverse meters", + description="Contravariant toroidal basis vector", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e^rho", "e^theta", "e^zeta", "alpha_r", "alpha_t", "alpha_z"], + data=["e_rho", "e_theta/sqrt(g)"], ) -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 +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_rho_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", + 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 Radial basis vector, derivative wrt radial coordinate", + description="Covariant cylindrical toroidal basis vector", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_r", "R_rr", "Z_rr", "omega_r", "omega_rr"], + data=["e_zeta", "phi_z"], ) -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 +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_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", + name="e_rho", + label="\\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Covariant Radial basis vector, derivative wrt poloidal coordinate", + description="Covariant Radial basis vector", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_r", "R_rt", "R_t", "Z_rt", "omega_r", "omega_rt", "omega_t"], + data=["R", "R_r", "Z_r", "omega_r"], ) -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 +def _e_sub_rho(params, transforms, profiles, data, **kwargs): + 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_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", + name="e_rho_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Covariant Radial basis vector, derivative wrt toroidal coordinate", + description="Covariant Radial basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], + data=["R", "R_r", "R_rr", "Z_rr", "omega_r", "omega_rr"], ) -def _e_sub_rho_z(params, transforms, profiles, data, **kwargs): - data["e_rho_z"] = jnp.array( +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"] * (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"], + -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 @@ -440,13 +564,13 @@ def _e_sub_rho_rrr(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_rho_rt", - label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\rho}", + name="e_rho_rrt", + label="\\partial_{\\rho}{\\rho}{\\theta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", description=( - "Covariant Radial basis vector, second derivative wrt radial and poloidal" - " coordinates" + "Covariant Radial basis vector, third derivative wrt radial coordinate" + " twice and poloidal once" ), dim=3, params=[], @@ -457,71 +581,23 @@ def _e_sub_rho_rrr(params, transforms, profiles, data, **kwargs): "R", "R_r", "R_rr", + "R_rrr", "R_rrt", + "R_rrrt", "R_rt", "R_t", - "Z_rrt", + "Z_rrrt", "omega_r", "omega_rr", + "omega_rrr", "omega_rrt", + "omega_rrrt", "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_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( +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"] @@ -562,55 +638,6 @@ def _e_sub_rho_rrt(params, transforms, profiles, data, **kwargs): 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_rrz", label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\rho}", @@ -707,12 +734,12 @@ def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_rho_tt", - label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\rho}", + name="e_rho_rt", + label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", description=( - "Covariant Radial basis vector, second derivative wrt poloidal and poloidal" + "Covariant Radial basis vector, second derivative wrt radial and poloidal" " coordinates" ), dim=3, @@ -723,33 +750,32 @@ def _e_sub_rho_rrz(params, transforms, profiles, data, **kwargs): data=[ "R", "R_r", + "R_rr", + "R_rrt", "R_rt", - "R_rtt", "R_t", - "R_tt", - "Z_rtt", + "Z_rrt", "omega_r", + "omega_rr", + "omega_rrt", "omega_rt", - "omega_rtt", "omega_t", - "omega_tt", ], ) -def _e_sub_rho_tt(params, transforms, profiles, data, **kwargs): - data["e_rho_tt"] = jnp.array( +def _e_sub_rho_rt(params, transforms, profiles, data, **kwargs): + data["e_rho_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"], + -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 @@ -831,69 +857,6 @@ def _e_sub_rho_rtt(params, transforms, profiles, data, **kwargs): 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_rtz", label="\\partial_{\\rho}{\\theta}{\\zeta} \\mathbf{e}_{\\rho}", @@ -1030,12 +993,12 @@ def _e_sub_rho_rtz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_rho_zz", - label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\rho}", + name="e_rho_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", description=( - "Covariant Radial basis vector, second derivative wrt toroidal and toroidal" + "Covariant Radial basis vector, second derivative wrt radial and toroidal" " coordinates" ), dim=3, @@ -1046,33 +1009,33 @@ def _e_sub_rho_rtz(params, transforms, profiles, data, **kwargs): data=[ "R", "R_r", + "R_rr", + "R_rrz", "R_rz", - "R_rzz", "R_z", - "R_zz", - "Z_rzz", + "Z_rrz", "omega_r", + "omega_rr", + "omega_rrz", "omega_rz", - "omega_rzz", "omega_z", - "omega_zz", ], ) -def _e_sub_rho_zz(params, transforms, profiles, data, **kwargs): - data["e_rho_zz"] = jnp.array( +def _e_sub_rho_rz(params, transforms, profiles, data, **kwargs): + data["e_rho_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"] + -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"]) ** 2 * data["omega_r"] - data["omega_rzz"]), - data["Z_rzz"], + * ((1 + data["omega_z"]) * data["omega_r"] ** 2 - data["omega_rrz"]), + data["Z_rrz"], ] ).T return data @@ -1176,11 +1139,11 @@ def _e_sub_rho_rzz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_theta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", + name="e_rho_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Covariant Poloidal basis vector, derivative wrt radial coordinate", + description="Covariant Radial basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, @@ -1188,9 +1151,9 @@ def _e_sub_rho_rzz(params, transforms, profiles, data, **kwargs): 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( - [ +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"] @@ -1202,62 +1165,61 @@ def _e_sub_theta_r(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_theta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description="Covariant Poloidal basis vector, derivative wrt poloidal coordinate", - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["R", "R_t", "R_tt", "Z_tt", "omega_t", "omega_tt"], -) -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_theta_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + name="e_rho_tt", + label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Covariant Poloidal basis vector, derivative wrt toroidal coordinate", + description=( + "Covariant Radial basis vector, second derivative wrt poloidal and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], + 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_z(params, transforms, profiles, data, **kwargs): - data["e_theta_z"] = jnp.array( +def _e_sub_rho_tt(params, transforms, profiles, data, **kwargs): + data["e_rho_tt"] = jnp.array( [ - -data["R"] * (1 + data["omega_z"]) * data["omega_t"] + data["R_tz"], - (1 + data["omega_z"]) * data["R_t"] - + data["R_z"] * data["omega_t"] - + data["R"] * data["omega_tz"], - data["Z_tz"], + -data["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_rr", - label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\theta}", + name="e_rho_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", description=( - "Covariant Poloidal basis vector, second derivative wrt radial and radial" + "Covariant Radial basis vector, second derivative wrt poloidal and toroidal" " coordinates" ), dim=3, @@ -1268,116 +1230,85 @@ def _e_sub_theta_z(params, transforms, profiles, data, **kwargs): data=[ "R", "R_r", - "R_rr", - "R_rrt", "R_rt", + "R_rtz", + "R_rz", "R_t", - "Z_rrt", + "R_tz", + "R_z", + "Z_rtz", "omega_r", - "omega_rr", - "omega_rrt", "omega_rt", + "omega_rtz", + "omega_rz", "omega_t", + "omega_tz", + "omega_z", ], ) -def _e_sub_theta_rr(params, transforms, profiles, data, **kwargs): - data["e_theta_rr"] = jnp.array( +def _e_sub_rho_tz(params, transforms, profiles, data, **kwargs): + data["e_rho_tz"] = jnp.array( [ - -data["R_t"] * data["omega_r"] ** 2 - - 2 * data["R"] * data["omega_r"] * data["omega_rt"] + -((1 + data["omega_z"]) * data["R_t"] * data["omega_r"]) + - data["R"] * data["omega_tz"] * data["omega_r"] - 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"], + * ( + (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_theta_rrr", - label="\\partial_{\\rho}{\\rho}{\\rho} \\mathbf{e}_{\\theta}", + name="e_rho_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - ), + description="Covariant Radial basis vector, derivative wrt toroidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - 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", - ], + data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], ) -def _e_sub_theta_rrr(params, transforms, profiles, data, **kwargs): - data["e_theta_rrr"] = jnp.array( +def _e_sub_rho_z(params, transforms, profiles, data, **kwargs): + data["e_rho_z"] = 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"], + -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_rt", - label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\theta}", + name="e_rho_zz", + label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", description=( - "Covariant Poloidal basis vector, second derivative wrt radial and poloidal" + "Covariant Radial basis vector, second derivative wrt toroidal and toroidal" " coordinates" ), dim=3, @@ -1388,47 +1319,254 @@ def _e_sub_theta_rrr(params, transforms, profiles, data, **kwargs): data=[ "R", "R_r", - "R_rt", - "R_rtt", - "R_t", - "R_tt", - "Z_rtt", + "R_rz", + "R_rzz", + "R_z", + "R_zz", + "Z_rzz", "omega_r", - "omega_rt", - "omega_rtt", - "omega_t", - "omega_tt", + "omega_rz", + "omega_rzz", + "omega_z", + "omega_zz", ], ) -def _e_sub_theta_rt(params, transforms, profiles, data, **kwargs): - data["e_theta_rt"] = jnp.array( +def _e_sub_rho_zz(params, transforms, profiles, data, **kwargs): + data["e_rho_zz"] = 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"], + -((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_theta_rrt", - label="\\partial_{\\rho}{\\rho}{\\theta} \\mathbf{e}_{\\theta}", + name="e_theta", + label="\\mathbf{e}_{\\theta}", units="m", units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - " twice and poloidal once" - ), + description="Covariant Poloidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + 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["R"] * data["omega_t"], data["Z_t"]] + ).T + return data + + +@register_compute_fun( + name="e_theta/sqrt(g)", + label="\\mathbf{e}_{\\theta} / \\sqrt{g}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector divided by 3-D volume Jacobian", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "sqrt(g)"], + axis_limit_data=["e_theta_r", "sqrt(g)_r"], +) +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_theta_rr", + label="\\partial_{\\rho}{\\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description=( + "Covariant Poloidal basis vector, second derivative wrt radial and radial" + " coordinates" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "R", + "R_r", + "R_rr", + "R_rrt", + "R_rt", + "R_t", + "Z_rrt", + "omega_r", + "omega_rr", + "omega_rrt", + "omega_rt", + "omega_t", + ], +) +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_theta_rrr", + label="\\partial_{\\rho}{\\rho}{\\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + 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_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_theta_rrt", + label="\\partial_{\\rho}{\\rho}{\\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description=( + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " twice and poloidal once" + ), dim=3, params=[], transforms={}, @@ -1519,13 +1657,13 @@ def _e_sub_theta_rrt(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_theta_rz", - label="\\partial_{\\rho}{\\zeta} \\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, second derivative wrt radial and toroidal" - " coordinates" + "Covariant Poloidal basis vector, third derivative wrt radial coordinate" + " twice and toroidal once" ), dim=3, params=[], @@ -1535,113 +1673,49 @@ def _e_sub_theta_rrt(params, transforms, profiles, data, **kwargs): 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_rtz", + "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_rz(params, transforms, profiles, data, **kwargs): - data["e_theta_rz"] = jnp.array( +def _e_sub_theta_rrz(params, transforms, profiles, data, **kwargs): + data["e_theta_rrz"] = jnp.array( [ - -(1 + data["omega_z"]) * data["R_t"] * data["omega_r"] - - data["R"] * data["omega_tz"] * data["omega_r"] - - data["omega_t"] + -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["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_theta_rrz", - label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, third derivative wrt radial coordinate" - " twice and toroidal once" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rt", - "R_rrt", - "R_rtz", - "R_rrtz", - "R_rz", - "R_rrz", - "R_t", - "R_tz", - "R_z", - "Z_rrtz", - "omega_r", - "omega_rr", - "omega_rt", - "omega_rrt", - "omega_rtz", - "omega_rrtz", - "omega_rz", - "omega_rrz", - "omega_t", - "omega_tz", - "omega_z", - ], -) -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_t"] * ( data["omega_rz"] * data["R_r"] + (1 + data["omega_z"]) * data["R_rr"] @@ -1720,12 +1794,12 @@ def _e_sub_theta_rrz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_theta_tt", - label="\\partial_{\\theta}{\\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, second derivative wrt poloidal and poloidal" + "Covariant Poloidal basis vector, second derivative wrt radial and poloidal" " coordinates" ), dim=3, @@ -1733,17 +1807,36 @@ def _e_sub_theta_rrz(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "R_tt", "R_ttt", "Z_ttt", "omega_t", "omega_tt", "omega_ttt"], + 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_tt(params, transforms, profiles, data, **kwargs): - data["e_theta_tt"] = jnp.array( +def _e_sub_theta_rt(params, transforms, profiles, data, **kwargs): + data["e_theta_rt"] = 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"], + -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 @@ -1825,55 +1918,6 @@ def _e_sub_theta_rtt(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="e_theta_tz", - label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\theta}", - units="m", - units_long="meters", - description=( - "Covariant Poloidal basis vector, second derivative wrt poloidal and toroidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_t", - "R_tt", - "R_ttz", - "R_tz", - "R_z", - "Z_ttz", - "omega_t", - "omega_tt", - "omega_ttz", - "omega_tz", - "omega_z", - ], -) -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_theta_rtz", label="\\partial_{\\rho}{\\theta}{\\zeta} \\mathbf{e}_{\\theta}", @@ -1978,12 +2022,12 @@ def _e_sub_theta_rtz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_theta_zz", - label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\theta}", + name="e_theta_rz", + label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", description=( - "Covariant Poloidal basis vector, second derivative wrt toroidal and toroidal" + "Covariant Poloidal basis vector, second derivative wrt radial and toroidal" " coordinates" ), dim=3, @@ -1993,34 +2037,49 @@ def _e_sub_theta_rtz(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=[ "R", + "R_r", + "R_rt", + "R_rtz", + "R_rz", "R_t", "R_tz", - "R_tzz", "R_z", - "R_zz", - "Z_tzz", + "Z_rtz", + "omega_r", + "omega_rt", + "omega_rtz", + "omega_rz", "omega_t", "omega_tz", - "omega_tzz", "omega_z", - "omega_zz", ], ) -def _e_sub_theta_zz(params, transforms, profiles, data, **kwargs): - data["e_theta_zz"] = jnp.array( +def _e_sub_theta_rz(params, transforms, profiles, data, **kwargs): + data["e_theta_rz"] = 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"], + -(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 @@ -2133,46 +2192,123 @@ def _e_sub_theta_rzz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_zeta_r", - label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant Toroidal basis vector, derivative wrt radial coordinate", + description="Covariant Poloidal basis vector, derivative wrt poloidal coordinate", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], + data=["R", "R_t", "R_tt", "Z_tt", "omega_t", "omega_tt"], ) -def _e_sub_zeta_r(params, transforms, profiles, data, **kwargs): - data["e_zeta_r"] = jnp.array( +def _e_sub_theta_t(params, transforms, profiles, data, **kwargs): + data["e_theta_t"] = jnp.array( [ - -data["R"] * (1 + data["omega_z"]) * data["omega_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"], + -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_zeta_t", - label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", + name="e_theta_tt", + label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant Toroidal basis vector, derivative wrt poloidal coordinate", + description=( + "Covariant Poloidal basis vector, second derivative wrt poloidal and poloidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], + data=["R", "R_t", "R_tt", "R_ttt", "Z_ttt", "omega_t", "omega_tt", "omega_ttt"], ) -def _e_sub_zeta_t(params, transforms, profiles, data, **kwargs): - data["e_zeta_t"] = jnp.array( +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_theta_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description=( + "Covariant Poloidal basis vector, second derivative wrt poloidal and toroidal" + " coordinates" + ), + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=[ + "R", + "R_t", + "R_tt", + "R_ttz", + "R_tz", + "R_z", + "Z_ttz", + "omega_t", + "omega_tt", + "omega_ttz", + "omega_tz", + "omega_z", + ], +) +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_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant Poloidal basis vector, derivative wrt toroidal coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], +) +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"] @@ -2185,24 +2321,95 @@ def _e_sub_zeta_t(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_zeta_z", - label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", + name="e_theta_zz", + label="\\partial_{\\zeta}{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Covariant Toroidal basis vector, derivative wrt toroidal coordinate", + description=( + "Covariant Poloidal basis vector, second derivative wrt toroidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "R_zz", "Z_zz", "omega_z", "omega_zz"], + 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_zeta_z(params, transforms, profiles, data, **kwargs): - data["e_zeta_z"] = jnp.array( +def _e_sub_theta_zz(params, transforms, profiles, data, **kwargs): + data["e_theta_zz"] = 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"], + -((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_zeta", + label="\\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant Toroidal basis vector", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_z", "Z_z", "omega_z"], +) +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_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant Toroidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "R_rz", "R_z", "Z_rz", "omega_r", "omega_rz", "omega_z"], +) +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 @@ -2335,70 +2542,6 @@ def _e_sub_zeta_rrr(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="e_zeta_rt", - label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, second derivative wrt radial and poloidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rt", - "R_rtz", - "R_rz", - "R_t", - "R_tz", - "R_z", - "Z_rtz", - "omega_r", - "omega_rt", - "omega_rtz", - "omega_rz", - "omega_t", - "omega_tz", - "omega_z", - ], -) -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_rrt", label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\zeta}", @@ -2537,13 +2680,13 @@ def _e_sub_zeta_rrt(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_zeta_rz", - label="\\partial_{\\rho}{\\zeta} \\mathbf{e}_{\\zeta}", + name="e_zeta_rrz", + label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", description=( - "Covariant Toroidal basis vector, second derivative wrt radial and toroidal" - " coordinates" + "Covariant Toroidal basis vector, third derivative wrt radial coordinate" + " twice and toroidal once" ), dim=3, params=[], @@ -2553,75 +2696,26 @@ def _e_sub_zeta_rrt(params, transforms, profiles, data, **kwargs): data=[ "R", "R_r", + "R_rr", "R_rz", + "R_rrz", "R_rzz", + "R_rrzz", "R_z", "R_zz", - "Z_rzz", + "Z_rrzz", "omega_r", + "omega_rr", "omega_rz", + "omega_rrz", "omega_rzz", + "omega_rrzz", "omega_z", "omega_zz", ], ) -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_rrz", - label="\\partial_{\\rho}{\\rho}{\\zeta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, third derivative wrt radial coordinate" - " twice and toroidal once" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_r", - "R_rr", - "R_rz", - "R_rrz", - "R_rzz", - "R_rrzz", - "R_z", - "R_zz", - "Z_rrzz", - "omega_r", - "omega_rr", - "omega_rz", - "omega_rrz", - "omega_rzz", - "omega_rrzz", - "omega_z", - "omega_zz", - ], -) -def _e_sub_zeta_rrz(params, transforms, profiles, data, **kwargs): - data["e_zeta_rrz"] = jnp.array( +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"] @@ -2685,12 +2779,12 @@ def _e_sub_zeta_rrz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e_zeta_tt", - label="\\partial_{\\theta}{\\theta} \\mathbf{e}_{\\zeta}", + name="e_zeta_rt", + label="\\partial_{\\rho}{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", description=( - "Covariant Toroidal basis vector, second derivative wrt poloidal and poloidal" + "Covariant Toroidal basis vector, second derivative wrt radial and poloidal" " coordinates" ), dim=3, @@ -2700,34 +2794,49 @@ def _e_sub_zeta_rrz(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=[ "R", + "R_r", + "R_rt", + "R_rtz", + "R_rz", "R_t", - "R_tt", - "R_ttz", "R_tz", "R_z", - "Z_ttz", + "Z_rtz", + "omega_r", + "omega_rt", + "omega_rtz", + "omega_rz", "omega_t", - "omega_tt", - "omega_ttz", "omega_tz", "omega_z", ], ) -def _e_sub_zeta_tt(params, transforms, profiles, data, **kwargs): - data["e_zeta_tt"] = jnp.array( +def _e_sub_zeta_rt(params, transforms, profiles, data, **kwargs): + data["e_zeta_rt"] = 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"], + -(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 @@ -2835,55 +2944,6 @@ def _e_sub_zeta_rtt(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="e_zeta_tz", - label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\zeta}", - units="m", - units_long="meters", - description=( - "Covariant Toroidal basis vector, second derivative wrt poloidal and toroidal" - " coordinates" - ), - dim=3, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=[ - "R", - "R_t", - "R_tz", - "R_tzz", - "R_z", - "R_zz", - "Z_tzz", - "omega_t", - "omega_tz", - "omega_tzz", - "omega_z", - "omega_zz", - ], -) -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_zeta_rtz", label="\\partial_{\\rho}{\\theta}{\\zeta} \\mathbf{e}_{\\zeta}", @@ -2993,12 +3053,12 @@ def _e_sub_zeta_rtz(params, transforms, profiles, data, **kwargs): @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 and toroidal" + "Covariant Toroidal basis vector, second derivative wrt radial and toroidal" " coordinates" ), dim=3, @@ -3006,25 +3066,36 @@ def _e_sub_zeta_rtz(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "R_zz", "R_zzz", "Z_zzz", "omega_z", "omega_zz", "omega_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( +def _e_sub_zeta_rz(params, transforms, profiles, data, **kwargs): + data["e_zeta_rz"] = 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"] + -((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 - + 3 * data["omega_z"] - + 3 * data["omega_z"] ** 2 - + data["omega_z"] ** 3 - - data["omega_zzz"] - ), - data["Z_zzz"], + * ((1 + data["omega_z"]) ** 2 * data["omega_r"] - data["omega_rzz"]), + data["Z_rzz"], ] ).T return data @@ -3119,315 +3190,244 @@ def _e_sub_zeta_rzz(params, transforms, profiles, data, **kwargs): @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", + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant Toroidal basis vector, derivative wrt poloidal 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"], + data=["R", "R_t", "R_tz", "R_z", "Z_tz", "omega_t", "omega_tz", "omega_z"], ) -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, - ) +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^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", + 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 poloidal" + " coordinates" + ), 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"], + 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_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, - ) +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_r", - label="\\partial{\\rho} \\mathbf{e}^{\\theta}", - units="m^{-1}", - units_long="inverse meters", - description="Contravariant poloidal basis vector, derivative wrt radial coordinate", + name="e_zeta_tz", + label="\\partial_{\\theta}{\\zeta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description=( + "Covariant Toroidal basis vector, second derivative wrt poloidal and toroidal" + " coordinates" + ), dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_zeta", "e_rho", "e_zeta_r", "e_rho_r", "sqrt(g)", "sqrt(g)_r"], + 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_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 +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^theta_t", - label="\\partial{\\theta} \\mathbf{e}^{\\theta}", - units="m^{-1}", - units_long="inverse meters", - description="Contravariant poloidal basis vector, derivative wrt poloidal" - " coordinate", + 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_zeta", "e_rho", "e_zeta_t", "e_rho_t", "sqrt(g)", "sqrt(g)_t"], + data=["R", "R_z", "R_zz", "Z_zz", "omega_z", "omega_zz"], ) -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 +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_z", - label="\\partial{\\zeta} \\mathbf{e}^{\\theta}", - units="m^{-1}", - units_long="inverse meters", - description="Contravariant poloidal basis vector, derivative wrt toroidal" - " coordinate", + 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_zeta", "e_rho", "e_zeta_z", "e_rho_z", "sqrt(g)", "sqrt(g)_z"], + data=["R", "R_z", "R_zz", "R_zzz", "Z_zzz", "omega_z", "omega_zz", "omega_zzz"], ) -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 +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_r", - label="\\partial{\\rho} \\mathbf{e}^{\\zeta}", + name="grad(alpha)", + label="\\nabla \\alpha", units="m^{-1}", - units_long="inverse meters", - description="Contravariant toroidal basis vector, derivative wrt radial coordinate", + units_long="Inverse meters", + description="Unit vector along field line", 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"], + data=["e^rho", "e^theta", "e^zeta", "alpha_r", "alpha_t", "alpha_z"], ) -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, - ) +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="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", + 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=["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"], + data=["psi_r", "e^rho"], ) -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, - ) +def _gradpsi(params, transforms, profiles, data, **kwargs): + data["grad(psi)"] = (data["psi_r"] * data["e^rho"].T).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", + name="n_rho", + label="\\hat{\\mathbf{n}}_{\\rho}", + units="~", + units_long="None", + description="Unit vector normal to constant rho surface (direction of e^rho)", 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"], + data=["e^rho"], ) -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, - ) +def _n_rho(params, transforms, profiles, data, **kwargs): + data["n_rho"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"], axis=-1)).T return data From 311ff0a185aea6b316dcf48b5f2a002790d8cdc6 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 25 Jul 2023 16:17:36 -0500 Subject: [PATCH 05/19] Add script to sort compute functions --- docs/sort_compute_funs.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 docs/sort_compute_funs.py diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py new file mode 100644 index 0000000000..79e6cf66af --- /dev/null +++ b/docs/sort_compute_funs.py @@ -0,0 +1,38 @@ +"""Sort the compute functions and print to specified output. + +First command line argument is module name, e.g. "_basis_vectors". +Second command line output is file to print to, +e.g. "sorted_basis_vectors.txt". +""" + +import inspect +import re +import sys + +import desc.compute + + +def get_matches(fun, pattern): + """Return all matches of ``pattern`` in source code of function ``fun``.""" + src = inspect.getsource(fun) + matches = pattern.findall(src) + matches = [s.replace("'", "").replace('"', "") for s in matches] + return matches + + +# Gather all compute function source code and map to quantity name. +src_codes = {} +pattern_name = re.compile(r"(?<=name=)[^,]+") +for module_name, module in inspect.getmembers(desc.compute, inspect.ismodule): + if module_name == sys.argv[1]: + for _, fun in inspect.getmembers(module, inspect.isfunction): + # quantities that this function computes + name = get_matches(fun, pattern_name) + if len(name) > 0: # skip imported functions + src_codes[name[0]] = inspect.getsource(fun) + +# Write compute functions sorted by name to file. +with open(sys.argv[2], "w") as output_file: + for name in sorted(src_codes): + output_file.write(src_codes[name]) + output_file.write("\n\n") From fe97a88720075073d8cf049d0b56dbab77c6a905 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 25 Jul 2023 16:36:40 -0500 Subject: [PATCH 06/19] simplify matching pattern --- docs/sort_compute_funs.py | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py index 79e6cf66af..b8274a85c7 100644 --- a/docs/sort_compute_funs.py +++ b/docs/sort_compute_funs.py @@ -16,7 +16,6 @@ def get_matches(fun, pattern): """Return all matches of ``pattern`` in source code of function ``fun``.""" src = inspect.getsource(fun) matches = pattern.findall(src) - matches = [s.replace("'", "").replace('"', "") for s in matches] return matches From 42777b12b0d63ae967b5c9357deb21f639b79782 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 25 Jul 2023 16:41:55 -0500 Subject: [PATCH 07/19] Sort _core using script --- desc/compute/_core.py | 2566 ++++++++++++++++++++--------------------- 1 file changed, 1283 insertions(+), 1283 deletions(-) diff --git a/desc/compute/_core.py b/desc/compute/_core.py index 5cfa2cc710..a370fdb4e4 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -24,466 +24,414 @@ def _0(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="rho", - label="\\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - + "of the toroidal flux", + name="R", + label="R", + units="m", + units_long="meters", + description="Major radius in lab frame", dim=1, - params=[], - transforms={"grid": []}, + params=["R_lmn"], + transforms={"R": [[0, 0, 0]]}, profiles=[], - coordinates="r", + coordinates="rtz", data=[], ) -def _rho(params, transforms, profiles, data, **kwargs): - data["rho"] = transforms["grid"].nodes[:, 0] - return data - - -@register_compute_fun( - name="rho_r", - label="\\partial_{\\rho} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - + "of the toroidal flux, derivative wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], -) -def _rho_r(params, transforms, profiles, data, **kwargs): - data["rho_r"] = jnp.ones_like(data["0"]) - return data - - -@register_compute_fun( - name="rho_t", - label="\\partial_{\\theta} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - "of the toroidal flux, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], -) -def _rho_t(params, transforms, profiles, data, **kwargs): - data["rho_t"] = data["0"] - return data - - -@register_compute_fun( - name="rho_z", - label="\\partial_{\\zeta} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - "of the toroidal flux, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], -) -def _rho_z(params, transforms, profiles, data, **kwargs): - data["rho_z"] = data["0"] +def _R(params, transforms, profiles, data, **kwargs): + data["R"] = transforms["R"].transform(params["R_lmn"], 0, 0, 0) return data @register_compute_fun( - name="theta", - label="\\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic)", + name="R_r", + label="\\partial_{\\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, first radial derivative", dim=1, - params=[], - transforms={"grid": []}, + params=["R_lmn"], + transforms={"R": [[1, 0, 0]]}, profiles=[], - coordinates="t", + coordinates="rtz", data=[], ) -def _theta(params, transforms, profiles, data, **kwargs): - data["theta"] = transforms["grid"].nodes[:, 1] +def _R_r(params, transforms, profiles, data, **kwargs): + data["R_r"] = transforms["R"].transform(params["R_lmn"], 1, 0, 0) return data @register_compute_fun( - name="theta_r", - label="\\partial_{\\rho} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt radial coordinate", + name="R_rr", + label="\\partial_{\\rho \\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second radial derivative", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[2, 0, 0]]}, profiles=[], - coordinates="t", - data=["0"], + coordinates="rtz", + data=[], ) -def _theta_r(params, transforms, profiles, data, **kwargs): - data["theta_r"] = data["0"] +def _R_rr(params, transforms, profiles, data, **kwargs): + data["R_rr"] = transforms["R"].transform(params["R_lmn"], 2, 0, 0) return data @register_compute_fun( - name="theta_t", - label="\\partial_{\\theta} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt poloidal coordinate", + name="R_rrr", + label="\\partial_{\rho \\rho \\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third radial derivative", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[3, 0, 0]]}, profiles=[], - coordinates="t", - data=["0"], + coordinates="rtz", + data=[], ) -def _theta_t(params, transforms, profiles, data, **kwargs): - data["theta_t"] = jnp.ones_like(data["0"]) +def _R_rrr(params, transforms, profiles, data, **kwargs): + data["R_rrr"] = transforms["R"].transform(params["R_lmn"], 3, 0, 0) return data @register_compute_fun( - name="theta_z", - label="\\partial_{\\zeta} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt toroidal coordinate", + name="R_rrrr", + label="\\partial_{\rho \\rho \\rho \\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth radial derivative", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[4, 0, 0]]}, profiles=[], - coordinates="t", - data=["0"], + coordinates="rtz", + data=[], ) -def _theta_z(params, transforms, profiles, data, **kwargs): - data["theta_z"] = data["0"] +def _R_rrrr(params, transforms, profiles, data, **kwargs): + data["R_rrrr"] = transforms["R"].transform(params["R_lmn"], 4, 0, 0) return data @register_compute_fun( - name="zeta", - label="\\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate", + name="R_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt" + " radial coordinate thrice and poloidal once", dim=1, - params=[], - transforms={"grid": []}, + params=["R_lmn"], + transforms={"R": [[3, 1, 0]]}, profiles=[], - coordinates="z", + coordinates="rtz", data=[], ) -def _zeta(params, transforms, profiles, data, **kwargs): - data["zeta"] = transforms["grid"].nodes[:, 2] +def _R_rrrt(params, transforms, profiles, data, **kwargs): + data["R_rrrt"] = transforms["R"].transform(params["R_lmn"], 3, 1, 0) return data @register_compute_fun( - name="zeta_r", - label="\\partial_{\\rho} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate derivative, wrt radial coordinate", + name="R_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt" + " radial coordinate thrice and toroidal once", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[3, 0, 1]]}, profiles=[], - coordinates="z", - data=["0"], + coordinates="rtz", + data=[], ) -def _zeta_r(params, transforms, profiles, data, **kwargs): - data["zeta_r"] = data["0"] +def _R_rrrz(params, transforms, profiles, data, **kwargs): + data["R_rrrz"] = transforms["R"].transform(params["R_lmn"], 3, 0, 1) return data @register_compute_fun( - name="zeta_t", - label="\\partial_{\\theta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt poloidal coordinate", + name="R_rrt", + label="\\partial_{\\rho \\rho \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third derivative, wrt radius twice " + "and poloidal angle", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[2, 1, 0]]}, profiles=[], - coordinates="z", - data=["0"], + coordinates="rtz", + data=[], ) -def _zeta_t(params, transforms, profiles, data, **kwargs): - data["zeta_t"] = data["0"] +def _R_rrt(params, transforms, profiles, data, **kwargs): + data["R_rrt"] = transforms["R"].transform(params["R_lmn"], 2, 1, 0) return data @register_compute_fun( - name="zeta_z", - label="\\partial_{\\zeta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt toroidal coordinate", + name="R_rrtt", + label="\\partial_{\\rho \\rho \\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fouth derivative, wrt radius twice " + "and poloidal angle twice", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[2, 2, 0]]}, profiles=[], - coordinates="z", - data=["0"], + coordinates="rtz", + data=[], ) -def _zeta_z(params, transforms, profiles, data, **kwargs): - data["zeta_z"] = jnp.ones_like(data["0"]) +def _R_rrtt(params, transforms, profiles, data, **kwargs): + data["R_rrtt"] = transforms["R"].transform(params["R_lmn"], 2, 2, 0) return data @register_compute_fun( - name="theta_PEST", - label="\\vartheta", - units="rad", - units_long="radians", - description="PEST straight field line poloidal angular coordinate", + name="R_rrtz", + label="\\partial_{\\rho \\rho \\theta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius twice," + " poloidal angle, and toroidal angle", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[2, 1, 1]]}, profiles=[], coordinates="rtz", - data=["theta", "lambda"], + data=[], ) -def _theta_PEST(params, transforms, profiles, data, **kwargs): - data["theta_PEST"] = (data["theta"] + data["lambda"]) % (2 * jnp.pi) +def _R_rrtz(params, transforms, profiles, data, **kwargs): + data["R_rrtz"] = transforms["R"].transform(params["R_lmn"], 2, 1, 1) return data @register_compute_fun( - name="theta_PEST_r", - label="\\partial_{\\rho} \\vartheta", - units="rad", - units_long="radians", - description="PEST straight field line poloidal angular coordinate, derivative wrt " - "radial coordinate", + name="R_rrz", + label="\\partial_{\\rho \\rho \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third derivative, wrt radius twice " + "and toroidal angle", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[2, 0, 1]]}, profiles=[], coordinates="rtz", - data=["lambda_r"], + data=[], ) -def _theta_PEST_r(params, transforms, profiles, data, **kwargs): - data["theta_PEST_r"] = data["lambda_r"] +def _R_rrz(params, transforms, profiles, data, **kwargs): + data["R_rrz"] = transforms["R"].transform(params["R_lmn"], 2, 0, 1) return data @register_compute_fun( - name="theta_PEST_t", - label="\\partial_{\\theta} \\vartheta", - units="rad", - units_long="radians", - description="PEST straight field line poloidal angular coordinate, derivative wrt " - "poloidal coordinate", + name="R_rrzz", + label="\\partial_{\\rho \\rho \\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative, wrt radius twice " + "and toroidal angle twice", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[2, 0, 2]]}, profiles=[], coordinates="rtz", - data=["lambda_t"], + data=[], ) -def _theta_PEST_t(params, transforms, profiles, data, **kwargs): - data["theta_PEST_t"] = 1 + data["lambda_t"] +def _R_rrzz(params, transforms, profiles, data, **kwargs): + data["R_rrzz"] = transforms["R"].transform(params["R_lmn"], 2, 0, 2) return data @register_compute_fun( - name="theta_PEST_z", - label="\\partial_{\\zeta} \\vartheta", - units="rad", - units_long="radians", - description="PEST straight field line poloidal angular coordinate, derivative wrt " - "toroidal coordinate", + name="R_rt", + label="\\partial_{\\rho \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second derivative wrt radius " + "and poloidal angle", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[1, 1, 0]]}, profiles=[], coordinates="rtz", - data=["lambda_z"], + data=[], ) -def _theta_PEST_z(params, transforms, profiles, data, **kwargs): - data["theta_PEST_z"] = data["lambda_z"] +def _R_rt(params, transforms, profiles, data, **kwargs): + data["R_rt"] = transforms["R"].transform(params["R_lmn"], 1, 1, 0) return data @register_compute_fun( - name="alpha", - label="\\alpha", - units="~", - units_long="None", - description="Field line label, defined on [0, 2pi)", + name="R_rtt", + label="\\partial_{\\rho \\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third derivative wrt radius and " + "poloidal angle twice", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[1, 2, 0]]}, profiles=[], coordinates="rtz", - data=["theta_PEST", "zeta", "iota"], + data=[], ) -def _alpha(params, transforms, profiles, data, **kwargs): - data["alpha"] = (data["theta_PEST"] - data["iota"] * data["zeta"]) % (2 * jnp.pi) +def _R_rtt(params, transforms, profiles, data, **kwargs): + data["R_rtt"] = transforms["R"].transform(params["R_lmn"], 1, 2, 0) return data @register_compute_fun( - name="alpha_r", - label="\\partial_\\rho \\alpha", - units="~", - units_long="None", - description="Field line label, derivative wrt radial coordinate", + name="R_rttt", + label="\\partial_{\\rho \\theta \\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius and " + "poloidal angle thrice", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[1, 3, 0]]}, profiles=[], coordinates="rtz", - data=["theta_PEST_r", "phi", "phi_r", "iota", "iota_r"], + data=[], ) -def _alpha_r(params, transforms, profiles, data, **kwargs): - data["alpha_r"] = ( - data["theta_PEST_r"] - - data["iota_r"] * data["phi"] - - data["iota"] * data["phi_r"] - ) +def _R_rttt(params, transforms, profiles, data, **kwargs): + data["R_rttt"] = transforms["R"].transform(params["R_lmn"], 1, 3, 0) return data @register_compute_fun( - name="alpha_t", - label="\\partial_\\theta \\alpha", - units="~", - units_long="None", - description="Field line label, derivative wrt poloidal coordinate", + name="R_rttz", + label="\\partial_{\\rho \\theta \\theta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, fourth derivative wrt radius once, " + "poloidal angle twice, and toroidal angle once", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[1, 2, 1]]}, profiles=[], coordinates="rtz", - data=["theta_PEST_t", "phi_t", "iota"], + data=[], ) -def _alpha_t(params, transforms, profiles, data, **kwargs): - data["alpha_t"] = data["theta_PEST_t"] + data["iota"] * data["phi_t"] +def _R_rttz(params, transforms, profiles, data, **kwargs): + data["R_rttz"] = transforms["R"].transform(params["R_lmn"], 1, 2, 1) return data @register_compute_fun( - name="alpha_z", - label="\\partial_\\zeta \\alpha", - units="~", - units_long="None", - description="Field line label, derivative wrt toroidal coordinate", + name="R_rtz", + label="\\partial_{\\rho \\theta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third derivative wrt radius, poloidal " + "angle, and toroidal angle", dim=1, - params=[], - transforms={}, + params=["R_lmn"], + transforms={"R": [[1, 1, 1]]}, profiles=[], coordinates="rtz", - data=["theta_PEST_z", "phi_z", "iota"], + data=[], ) -def _alpha_z(params, transforms, profiles, data, **kwargs): - data["alpha_z"] = data["theta_PEST_z"] - data["iota"] * data["phi_z"] +def _R_rtz(params, transforms, profiles, data, **kwargs): + data["R_rtz"] = transforms["R"].transform(params["R_lmn"], 1, 1, 1) return data @register_compute_fun( - name="R", - label="R", + name="R_rtzz", + label="\\partial_{\\rho \\theta \\zeta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame", + description="Major radius in lab frame, fourth derivative wrt radius, poloidal " + "angle, and toroidal angle twice", dim=1, params=["R_lmn"], - transforms={"R": [[0, 0, 0]]}, + transforms={"R": [[1, 1, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _R(params, transforms, profiles, data, **kwargs): - data["R"] = transforms["R"].transform(params["R_lmn"], 0, 0, 0) +def _R_rtzz(params, transforms, profiles, data, **kwargs): + data["R_rtzz"] = transforms["R"].transform(params["R_lmn"], 1, 1, 2) return data @register_compute_fun( - name="R_r", - label="\\partial_{\\rho} R", + name="R_rz", + label="\\partial_{\\rho \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, first radial derivative", + description="Major radius in lab frame, second derivative wrt radius " + "and toroidal angle", dim=1, params=["R_lmn"], - transforms={"R": [[1, 0, 0]]}, + transforms={"R": [[1, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_r(params, transforms, profiles, data, **kwargs): - data["R_r"] = transforms["R"].transform(params["R_lmn"], 1, 0, 0) +def _R_rz(params, transforms, profiles, data, **kwargs): + data["R_rz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 1) return data @register_compute_fun( - name="R_t", - label="\\partial_{\\theta} R", + name="R_rzz", + label="\\partial_{\\rho \\zeta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, first poloidal derivative", + description="Major radius in lab frame, third derivative wrt radius and " + "toroidal angle twice", dim=1, params=["R_lmn"], - transforms={"R": [[0, 1, 0]]}, + transforms={"R": [[1, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_t(params, transforms, profiles, data, **kwargs): - data["R_t"] = transforms["R"].transform(params["R_lmn"], 0, 1, 0) +def _R_rzz(params, transforms, profiles, data, **kwargs): + data["R_rzz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 2) return data @register_compute_fun( - name="R_z", - label="\\partial_{\\zeta} R", + name="R_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, first toroidal derivative", + description="Major radius in lab frame, fourth derivative wrt radius and " + "toroidal angle thrice", dim=1, params=["R_lmn"], - transforms={"R": [[0, 0, 1]]}, + transforms={"R": [[1, 0, 3]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_z(params, transforms, profiles, data, **kwargs): - data["R_z"] = transforms["R"].transform(params["R_lmn"], 0, 0, 1) +def _R_rzzz(params, transforms, profiles, data, **kwargs): + data["R_rzzz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 3) return data @register_compute_fun( - name="R_rr", - label="\\partial_{\\rho \\rho} R", + name="R_t", + label="\\partial_{\\theta} R", units="m", units_long="meters", - description="Major radius in lab frame, second radial derivative", + description="Major radius in lab frame, first poloidal derivative", dim=1, params=["R_lmn"], - transforms={"R": [[2, 0, 0]]}, + transforms={"R": [[0, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rr(params, transforms, profiles, data, **kwargs): - data["R_rr"] = transforms["R"].transform(params["R_lmn"], 2, 0, 0) +def _R_t(params, transforms, profiles, data, **kwargs): + data["R_t"] = transforms["R"].transform(params["R_lmn"], 0, 1, 0) return data @@ -506,598 +454,630 @@ def _R_tt(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="R_zz", - label="\\partial_{\\zeta \\zeta} R", + name="R_ttt", + label="\\partial_{\\theta \\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third poloidal derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 3, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_ttt(params, transforms, profiles, data, **kwargs): + data["R_ttt"] = transforms["R"].transform(params["R_lmn"], 0, 3, 0) + return data + + +@register_compute_fun( + name="R_ttz", + label="\\partial_{\\theta \\theta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, second toroidal derivative", + description="Major radius in lab frame, third derivative wrt poloidal angle " + "twice and toroidal angle", dim=1, params=["R_lmn"], - transforms={"R": [[0, 0, 2]]}, + transforms={"R": [[0, 2, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_zz(params, transforms, profiles, data, **kwargs): - data["R_zz"] = transforms["R"].transform(params["R_lmn"], 0, 0, 2) +def _R_ttz(params, transforms, profiles, data, **kwargs): + data["R_ttz"] = transforms["R"].transform(params["R_lmn"], 0, 2, 1) return data @register_compute_fun( - name="R_rt", - label="\\partial_{\\rho \\theta} R", + name="R_tz", + label="\\partial_{\\theta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, second derivative wrt radius " - "and poloidal angle", + description="Major radius in lab frame, second derivative wrt poloidal " + "and toroidal angles", dim=1, params=["R_lmn"], - transforms={"R": [[1, 1, 0]]}, + transforms={"R": [[0, 1, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rt(params, transforms, profiles, data, **kwargs): - data["R_rt"] = transforms["R"].transform(params["R_lmn"], 1, 1, 0) +def _R_tz(params, transforms, profiles, data, **kwargs): + data["R_tz"] = transforms["R"].transform(params["R_lmn"], 0, 1, 1) return data @register_compute_fun( - name="R_rz", - label="\\partial_{\\rho \\zeta} R", + name="R_tzz", + label="\\partial_{\\theta \\zeta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, second derivative wrt radius " - "and toroidal angle", + description="Major radius in lab frame, third derivative wrt poloidal angle " + "and toroidal angle twice", dim=1, params=["R_lmn"], - transforms={"R": [[1, 0, 1]]}, + transforms={"R": [[0, 1, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rz(params, transforms, profiles, data, **kwargs): - data["R_rz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 1) +def _R_tzz(params, transforms, profiles, data, **kwargs): + data["R_tzz"] = transforms["R"].transform(params["R_lmn"], 0, 1, 2) return data @register_compute_fun( - name="R_tz", - label="\\partial_{\\theta \\zeta} R", + name="R_z", + label="\\partial_{\\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, second derivative wrt poloidal " - "and toroidal angles", + description="Major radius in lab frame, first toroidal derivative", dim=1, params=["R_lmn"], - transforms={"R": [[0, 1, 1]]}, + transforms={"R": [[0, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_tz(params, transforms, profiles, data, **kwargs): - data["R_tz"] = transforms["R"].transform(params["R_lmn"], 0, 1, 1) +def _R_z(params, transforms, profiles, data, **kwargs): + data["R_z"] = transforms["R"].transform(params["R_lmn"], 0, 0, 1) return data @register_compute_fun( - name="R_rrr", - label="\\partial_{\rho \\rho \\rho} R", + name="R_zz", + label="\\partial_{\\zeta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, third radial derivative", + description="Major radius in lab frame, second toroidal derivative", dim=1, params=["R_lmn"], - transforms={"R": [[3, 0, 0]]}, + transforms={"R": [[0, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rrr(params, transforms, profiles, data, **kwargs): - data["R_rrr"] = transforms["R"].transform(params["R_lmn"], 3, 0, 0) +def _R_zz(params, transforms, profiles, data, **kwargs): + data["R_zz"] = transforms["R"].transform(params["R_lmn"], 0, 0, 2) return data @register_compute_fun( - name="R_rrrr", - label="\\partial_{\rho \\rho \\rho \\rho} R", + name="R_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} R", units="m", units_long="meters", - description="Major radius in lab frame, fourth radial derivative", + description="Major radius in lab frame, third toroidal derivative", dim=1, params=["R_lmn"], - transforms={"R": [[4, 0, 0]]}, + transforms={"R": [[0, 0, 3]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rrrr(params, transforms, profiles, data, **kwargs): - data["R_rrrr"] = transforms["R"].transform(params["R_lmn"], 4, 0, 0) +def _R_zzz(params, transforms, profiles, data, **kwargs): + data["R_zzz"] = transforms["R"].transform(params["R_lmn"], 0, 0, 3) return data @register_compute_fun( - name="R_rrrt", - label="\\partial_{\rho \\rho \\rho \\theta} R", + name="X", + label="X = R \\cos{\\phi}", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt" - " radial coordinate thrice and poloidal once", + description="Cartesian X coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[3, 1, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "phi"], ) -def _R_rrrt(params, transforms, profiles, data, **kwargs): - data["R_rrrt"] = transforms["R"].transform(params["R_lmn"], 3, 1, 0) +def _X(params, transforms, profiles, data, **kwargs): + data["X"] = data["R"] * jnp.cos(data["phi"]) return data @register_compute_fun( - name="R_rrrz", - label="\\partial_{\rho \\rho \\rho \\zeta} R", + name="X_r", + label="\\partial_{\\rho} X", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt" - " radial coordinate thrice and toroidal once", + description="Cartesian X coordinate, derivative wrt radial coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[3, 0, 1]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "R_r", "phi", "phi_r"], ) -def _R_rrrz(params, transforms, profiles, data, **kwargs): - data["R_rrrz"] = transforms["R"].transform(params["R_lmn"], 3, 0, 1) +def _X_r(params, transforms, profiles, data, **kwargs): + data["X_r"] = ( + data["R_r"] * jnp.cos(data["phi"]) + - data["R"] * jnp.sin(data["phi"]) * data["phi_r"] + ) return data @register_compute_fun( - name="R_ttt", - label="\\partial_{\\theta \\theta \\theta} R", + name="X_t", + label="\\partial_{\\theta} X", units="m", units_long="meters", - description="Major radius in lab frame, third poloidal derivative", + description="Cartesian X coordinate, derivative wrt poloidal coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 3, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "R_t", "phi", "phi_t"], ) -def _R_ttt(params, transforms, profiles, data, **kwargs): - data["R_ttt"] = transforms["R"].transform(params["R_lmn"], 0, 3, 0) +def _X_t(params, transforms, profiles, data, **kwargs): + data["X_t"] = ( + data["R_t"] * jnp.cos(data["phi"]) + - data["R"] * jnp.sin(data["phi"]) * data["phi_t"] + ) return data @register_compute_fun( - name="R_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} R", + name="X_z", + label="\\partial_{\\zeta} X", units="m", units_long="meters", - description="Major radius in lab frame, third toroidal derivative", + description="Cartesian X coordinate, derivative wrt toroidal coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 0, 3]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "R_z", "phi", "phi_z"], ) -def _R_zzz(params, transforms, profiles, data, **kwargs): - data["R_zzz"] = transforms["R"].transform(params["R_lmn"], 0, 0, 3) +def _X_z(params, transforms, profiles, data, **kwargs): + data["X_z"] = ( + data["R_z"] * jnp.cos(data["phi"]) + - data["R"] * jnp.sin(data["phi"]) * data["phi_z"] + ) return data @register_compute_fun( - name="R_rrt", - label="\\partial_{\\rho \\rho \\theta} R", + name="Y", + label="Y = R \\sin{\\phi}", units="m", units_long="meters", - description="Major radius in lab frame, third derivative, wrt radius twice " - "and poloidal angle", + description="Cartesian Y coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 1, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "phi"], ) -def _R_rrt(params, transforms, profiles, data, **kwargs): - data["R_rrt"] = transforms["R"].transform(params["R_lmn"], 2, 1, 0) +def _Y(params, transforms, profiles, data, **kwargs): + data["Y"] = data["R"] * jnp.sin(data["phi"]) return data @register_compute_fun( - name="R_rrtt", - label="\\partial_{\\rho \\rho \\theta \\theta} R", + name="Y_r", + label="\\partial_{\\rho} Y", units="m", units_long="meters", - description="Major radius in lab frame, fouth derivative, wrt radius twice " - "and poloidal angle twice", + description="Cartesian Y coordinate, derivative wrt radial coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 2, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "R_r", "phi", "phi_r"], ) -def _R_rrtt(params, transforms, profiles, data, **kwargs): - data["R_rrtt"] = transforms["R"].transform(params["R_lmn"], 2, 2, 0) +def _Y_r(params, transforms, profiles, data, **kwargs): + data["Y_r"] = ( + data["R_r"] * jnp.sin(data["phi"]) + + data["R"] * jnp.cos(data["phi"]) * data["phi_r"] + ) return data @register_compute_fun( - name="R_rtt", - label="\\partial_{\\rho \\theta \\theta} R", + name="Y_t", + label="\\partial_{\\theta} Y", units="m", units_long="meters", - description="Major radius in lab frame, third derivative wrt radius and " - "poloidal angle twice", + description="Cartesian Y coordinate, derivative wrt poloidal coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 2, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "R_t", "phi", "phi_t"], ) -def _R_rtt(params, transforms, profiles, data, **kwargs): - data["R_rtt"] = transforms["R"].transform(params["R_lmn"], 1, 2, 0) +def _Y_t(params, transforms, profiles, data, **kwargs): + data["Y_t"] = ( + data["R_t"] * jnp.sin(data["phi"]) + + data["R"] * jnp.cos(data["phi"]) * data["phi_t"] + ) return data @register_compute_fun( - name="R_rttz", - label="\\partial_{\\rho \\theta \\theta \\zeta} R", + name="Y_z", + label="\\partial_{\\zeta} Y", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt radius once, " - "poloidal angle twice, and toroidal angle once", + description="Cartesian Y coordinate, derivative wrt toroidal coordinate", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 2, 1]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["R", "R_z", "phi", "phi_z"], ) -def _R_rttz(params, transforms, profiles, data, **kwargs): - data["R_rttz"] = transforms["R"].transform(params["R_lmn"], 1, 2, 1) +def _Y_z(params, transforms, profiles, data, **kwargs): + data["Y_z"] = ( + data["R_z"] * jnp.sin(data["phi"]) + + data["R"] * jnp.cos(data["phi"]) * data["phi_z"] + ) return data @register_compute_fun( - name="R_rttt", - label="\\partial_{\\rho \\theta \\theta \\theta} R", + name="Z", + label="Z", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt radius and " - "poloidal angle thrice", + description="Vertical coordinate in lab frame", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 3, 0]]}, + params=["Z_lmn"], + transforms={"Z": [[0, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rttt(params, transforms, profiles, data, **kwargs): - data["R_rttt"] = transforms["R"].transform(params["R_lmn"], 1, 3, 0) +def _Z(params, transforms, profiles, data, **kwargs): + data["Z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 0) return data @register_compute_fun( - name="R_rrz", - label="\\partial_{\\rho \\rho \\zeta} R", + name="Z_r", + label="\\partial_{\\rho} Z", units="m", units_long="meters", - description="Major radius in lab frame, third derivative, wrt radius twice " - "and toroidal angle", + description="Vertical coordinate in lab frame, first radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 0, 1]]}, + params=["Z_lmn"], + transforms={"Z": [[1, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rrz(params, transforms, profiles, data, **kwargs): - data["R_rrz"] = transforms["R"].transform(params["R_lmn"], 2, 0, 1) +def _Z_r(params, transforms, profiles, data, **kwargs): + data["Z_r"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 0) return data @register_compute_fun( - name="R_rrzz", - label="\\partial_{\\rho \\rho \\zeta \\zeta} R", + name="Z_rr", + label="\\partial_{\\rho \\rho} Z", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative, wrt radius twice " - "and toroidal angle twice", + description="Vertical coordinate in lab frame, second radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 0, 2]]}, + params=["Z_lmn"], + transforms={"Z": [[2, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rrzz(params, transforms, profiles, data, **kwargs): - data["R_rrzz"] = transforms["R"].transform(params["R_lmn"], 2, 0, 2) +def _Z_rr(params, transforms, profiles, data, **kwargs): + data["Z_rr"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 0) return data @register_compute_fun( - name="R_rzz", - label="\\partial_{\\rho \\zeta \\zeta} R", + name="Z_rrr", + label="\\partial_{\rho \\rho \\rho} Z", units="m", units_long="meters", - description="Major radius in lab frame, third derivative wrt radius and " - "toroidal angle twice", + description="Vertical coordinate in lab frame, third radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 0, 2]]}, + params=["Z_lmn"], + transforms={"Z": [[3, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rzz(params, transforms, profiles, data, **kwargs): - data["R_rzz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 2) +def _Z_rrr(params, transforms, profiles, data, **kwargs): + data["Z_rrr"] = transforms["Z"].transform(params["Z_lmn"], 3, 0, 0) return data @register_compute_fun( - name="R_rzzz", - label="\\partial_{\\rho \\zeta \\zeta \\zeta} R", + name="Z_rrrr", + label="\\partial_{\rho \\rho \\rho \\rho} Z", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt radius and " - "toroidal angle thrice", + description="Vertical coordinate in lab frame, fourth radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 0, 3]]}, + params=["Z_lmn"], + transforms={"Z": [[4, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rzzz(params, transforms, profiles, data, **kwargs): - data["R_rzzz"] = transforms["R"].transform(params["R_lmn"], 1, 0, 3) +def _Z_rrrr(params, transforms, profiles, data, **kwargs): + data["Z_rrrr"] = transforms["Z"].transform(params["Z_lmn"], 4, 0, 0) return data @register_compute_fun( - name="R_ttz", - label="\\partial_{\\theta \\theta \\zeta} R", + name="Z_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} Z", units="m", units_long="meters", - description="Major radius in lab frame, third derivative wrt poloidal angle " - "twice and toroidal angle", + description="Vertical coordinate in lab frame, fourth derivative wrt " + " radial coordinate thrice and poloidal once", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 2, 1]]}, + params=["Z_lmn"], + transforms={"Z": [[3, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_ttz(params, transforms, profiles, data, **kwargs): - data["R_ttz"] = transforms["R"].transform(params["R_lmn"], 0, 2, 1) +def _Z_rrrt(params, transforms, profiles, data, **kwargs): + data["Z_rrrt"] = transforms["Z"].transform(params["Z_lmn"], 3, 1, 0) return data @register_compute_fun( - name="R_tzz", - label="\\partial_{\\theta \\zeta \\zeta} R", + name="Z_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} Z", units="m", units_long="meters", - description="Major radius in lab frame, third derivative wrt poloidal angle " - "and toroidal angle twice", + description="Vertical coordinate in lab frame, fourth derivative wrt " + " radial coordinate thrice and toroidal once", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 1, 2]]}, + params=["Z_lmn"], + transforms={"Z": [[3, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_tzz(params, transforms, profiles, data, **kwargs): - data["R_tzz"] = transforms["R"].transform(params["R_lmn"], 0, 1, 2) +def _Z_rrrz(params, transforms, profiles, data, **kwargs): + data["Z_rrrz"] = transforms["Z"].transform(params["Z_lmn"], 3, 0, 1) return data @register_compute_fun( - name="R_rtz", - label="\\partial_{\\rho \\theta \\zeta} R", + name="Z_rrt", + label="\\partial_{\\rho \\rho \\theta} Z", units="m", units_long="meters", - description="Major radius in lab frame, third derivative wrt radius, poloidal " - "angle, and toroidal angle", + description="Vertical coordinate in lab frame, third derivative, wrt radius " + "twice and poloidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 1, 1]]}, + params=["Z_lmn"], + transforms={"Z": [[2, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rtz(params, transforms, profiles, data, **kwargs): - data["R_rtz"] = transforms["R"].transform(params["R_lmn"], 1, 1, 1) +def _Z_rrt(params, transforms, profiles, data, **kwargs): + data["Z_rrt"] = transforms["Z"].transform(params["Z_lmn"], 2, 1, 0) return data @register_compute_fun( - name="R_rrtz", - label="\\partial_{\\rho \\rho \\theta \\zeta} R", + name="Z_rrtt", + label="\\partial_{\\rho \\rho \\theta} Z", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt radius twice," - " poloidal angle, and toroidal angle", + description="Vertical coordinate in lab frame, fourth derivative, wrt radius " + "twice and poloidal angle twice", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 1, 1]]}, + params=["Z_lmn"], + transforms={"Z": [[2, 2, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rrtz(params, transforms, profiles, data, **kwargs): - data["R_rrtz"] = transforms["R"].transform(params["R_lmn"], 2, 1, 1) +def _Z_rrtt(params, transforms, profiles, data, **kwargs): + data["Z_rrtt"] = transforms["Z"].transform(params["Z_lmn"], 2, 2, 0) return data @register_compute_fun( - name="R_rtzz", - label="\\partial_{\\rho \\theta \\zeta \\zeta} R", + name="Z_rrtz", + label="\\partial_{\\rho \\rho \\theta \\zeta} Z", units="m", units_long="meters", - description="Major radius in lab frame, fourth derivative wrt radius, poloidal " - "angle, and toroidal angle twice", + description="Vertical coordinate in lab frame, fourth derivative wrt radius" + "twice, poloidal angle, and toroidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 1, 2]]}, + params=["Z_lmn"], + transforms={"Z": [[2, 1, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rtzz(params, transforms, profiles, data, **kwargs): - data["R_rtzz"] = transforms["R"].transform(params["R_lmn"], 1, 1, 2) +def _Z_rrtz(params, transforms, profiles, data, **kwargs): + data["Z_rrtz"] = transforms["Z"].transform(params["Z_lmn"], 2, 1, 1) return data @register_compute_fun( - name="Z", - label="Z", + name="Z_rrz", + label="\\partial_{\\rho \\rho \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame", + description="Vertical coordinate in lab frame, third derivative, wrt radius " + "twice and toroidal angle", dim=1, params=["Z_lmn"], - transforms={"Z": [[0, 0, 0]]}, + transforms={"Z": [[2, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z(params, transforms, profiles, data, **kwargs): - data["Z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 0) +def _Z_rrz(params, transforms, profiles, data, **kwargs): + data["Z_rrz"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 1) return data @register_compute_fun( - name="Z_r", - label="\\partial_{\\rho} Z", + name="Z_rrzz", + label="\\partial_{\\rho \\rho \\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, first radial derivative", + description="Vertical coordinate in lab frame, fourth derivative, wrt radius " + "twice and toroidal angle twice", dim=1, params=["Z_lmn"], - transforms={"Z": [[1, 0, 0]]}, + transforms={"Z": [[2, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_r(params, transforms, profiles, data, **kwargs): - data["Z_r"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 0) +def _Z_rrzz(params, transforms, profiles, data, **kwargs): + data["Z_rrzz"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 2) return data @register_compute_fun( - name="Z_t", - label="\\partial_{\\theta} Z", + name="Z_rt", + label="\\partial_{\\rho \\theta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, first poloidal derivative", + description="Vertical coordinate in lab frame, second derivative wrt radius " + "and poloidal angle", dim=1, params=["Z_lmn"], - transforms={"Z": [[0, 1, 0]]}, + transforms={"Z": [[1, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_t(params, transforms, profiles, data, **kwargs): - data["Z_t"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 0) +def _Z_rt(params, transforms, profiles, data, **kwargs): + data["Z_rt"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 0) return data @register_compute_fun( - name="Z_z", - label="\\partial_{\\zeta} Z", + name="Z_rtt", + label="\\partial_{\\rho \\theta \\theta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, first toroidal derivative", + description="Vertical coordinate in lab frame, third derivative wrt radius " + "and poloidal angle twice", dim=1, params=["Z_lmn"], - transforms={"Z": [[0, 0, 1]]}, + transforms={"Z": [[1, 2, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_z(params, transforms, profiles, data, **kwargs): - data["Z_z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 1) +def _Z_rtt(params, transforms, profiles, data, **kwargs): + data["Z_rtt"] = transforms["Z"].transform(params["Z_lmn"], 1, 2, 0) return data @register_compute_fun( - name="Z_rr", - label="\\partial_{\\rho \\rho} Z", + name="Z_rttt", + label="\\partial_{\\rho \\theta \\theta \\theta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, second radial derivative", + description="Vertical coordinate in lab frame, third derivative wrt radius " + "and poloidal angle thrice", dim=1, params=["Z_lmn"], - transforms={"Z": [[2, 0, 0]]}, + transforms={"Z": [[1, 3, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rr(params, transforms, profiles, data, **kwargs): - data["Z_rr"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 0) +def _Z_rttt(params, transforms, profiles, data, **kwargs): + data["Z_rttt"] = transforms["Z"].transform(params["Z_lmn"], 1, 3, 0) return data @register_compute_fun( - name="Z_tt", - label="\\partial_{\\theta \\theta} Z", + name="Z_rttz", + label="\\partial_{\\rho \\theta \\theta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, second poloidal derivative", + description="Vertical coordinate in lab frame, fourth derivative wrt radius " + "once, poloidal angle twice, and toroidal angle once", dim=1, params=["Z_lmn"], - transforms={"Z": [[0, 2, 0]]}, + transforms={"Z": [[1, 2, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_tt(params, transforms, profiles, data, **kwargs): - data["Z_tt"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 0) +def _Z_rttz(params, transforms, profiles, data, **kwargs): + data["Z_rttz"] = transforms["Z"].transform(params["Z_lmn"], 1, 2, 1) return data @register_compute_fun( - name="Z_zz", - label="\\partial_{\\zeta \\zeta} Z", + name="Z_rtz", + label="\\partial_{\\rho \\theta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, second toroidal derivative", + description="Vertical coordinate in lab frame, third derivative wrt radius, " + "poloidal angle, and toroidal angle", dim=1, params=["Z_lmn"], - transforms={"Z": [[0, 0, 2]]}, + transforms={"Z": [[1, 1, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_zz(params, transforms, profiles, data, **kwargs): - data["Z_zz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 2) +def _Z_rtz(params, transforms, profiles, data, **kwargs): + data["Z_rtz"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 1) return data @register_compute_fun( - name="Z_rt", - label="\\partial_{\\rho \\theta} Z", + name="Z_rtzz", + label="\\partial_{\\rho \\theta \\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, second derivative wrt radius " - "and poloidal angle", + description="Vertical coordinate in lab frame, fourth derivative wrt radius, " + "poloidal angle, and toroidal angle twice", dim=1, params=["Z_lmn"], - transforms={"Z": [[1, 1, 0]]}, + transforms={"Z": [[1, 1, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rt(params, transforms, profiles, data, **kwargs): - data["Z_rt"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 0) +def _Z_rtzz(params, transforms, profiles, data, **kwargs): + data["Z_rtzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 2) return data @@ -1121,95 +1101,76 @@ def _Z_rz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="Z_tz", - label="\\partial_{\\theta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, second derivative wrt poloidal " - "and toroidal angles", - dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 1, 1]]}, - profiles=[], - coordinates="rtz", - data=[], -) -def _Z_tz(params, transforms, profiles, data, **kwargs): - data["Z_tz"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 1) - return data - - -@register_compute_fun( - name="Z_rrr", - label="\\partial_{\rho \\rho \\rho} Z", + name="Z_rzz", + label="\\partial_{\\rho \\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, third radial derivative", + description="Vertical coordinate in lab frame, third derivative wrt radius " + "and toroidal angle twice", dim=1, params=["Z_lmn"], - transforms={"Z": [[3, 0, 0]]}, + transforms={"Z": [[1, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrr(params, transforms, profiles, data, **kwargs): - data["Z_rrr"] = transforms["Z"].transform(params["Z_lmn"], 3, 0, 0) +def _Z_rzz(params, transforms, profiles, data, **kwargs): + data["Z_rzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 2) return data @register_compute_fun( - name="Z_rrrr", - label="\\partial_{\rho \\rho \\rho \\rho} Z", + name="Z_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, fourth radial derivative", + description="Vertical coordinate in lab frame, third derivative wrt radius " + "and toroidal angle thrice", dim=1, params=["Z_lmn"], - transforms={"Z": [[4, 0, 0]]}, + transforms={"Z": [[1, 0, 3]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrrr(params, transforms, profiles, data, **kwargs): - data["Z_rrrr"] = transforms["Z"].transform(params["Z_lmn"], 4, 0, 0) +def _Z_rzzz(params, transforms, profiles, data, **kwargs): + data["Z_rzzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 3) return data @register_compute_fun( - name="Z_rrrt", - label="\\partial_{\rho \\rho \\rho \\theta} Z", + name="Z_t", + label="\\partial_{\\theta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative wrt " - " radial coordinate thrice and poloidal once", + description="Vertical coordinate in lab frame, first poloidal derivative", dim=1, params=["Z_lmn"], - transforms={"Z": [[3, 1, 0]]}, + transforms={"Z": [[0, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrrt(params, transforms, profiles, data, **kwargs): - data["Z_rrrt"] = transforms["Z"].transform(params["Z_lmn"], 3, 1, 0) +def _Z_t(params, transforms, profiles, data, **kwargs): + data["Z_t"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 0) return data @register_compute_fun( - name="Z_rrrz", - label="\\partial_{\rho \\rho \\rho \\zeta} Z", + name="Z_tt", + label="\\partial_{\\theta \\theta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative wrt " - " radial coordinate thrice and toroidal once", + description="Vertical coordinate in lab frame, second poloidal derivative", dim=1, params=["Z_lmn"], - transforms={"Z": [[3, 0, 1]]}, + transforms={"Z": [[0, 2, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrrz(params, transforms, profiles, data, **kwargs): - data["Z_rrrz"] = transforms["Z"].transform(params["Z_lmn"], 3, 0, 1) +def _Z_tt(params, transforms, profiles, data, **kwargs): + data["Z_tt"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 0) return data @@ -1232,1068 +1193,1074 @@ def _Z_ttt(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="Z_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} Z", + name="Z_ttz", + label="\\partial_{\\theta \\theta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, third toroidal derivative", + description="Vertical coordinate in lab frame, third derivative wrt poloidal " + "angle twice and toroidal angle", dim=1, params=["Z_lmn"], - transforms={"Z": [[0, 0, 3]]}, + transforms={"Z": [[0, 2, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_zzz(params, transforms, profiles, data, **kwargs): - data["Z_zzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 3) +def _Z_ttz(params, transforms, profiles, data, **kwargs): + data["Z_ttz"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 1) return data @register_compute_fun( - name="Z_rrt", - label="\\partial_{\\rho \\rho \\theta} Z", + name="Z_tz", + label="\\partial_{\\theta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, third derivative, wrt radius " - "twice and poloidal angle", + description="Vertical coordinate in lab frame, second derivative wrt poloidal " + "and toroidal angles", dim=1, params=["Z_lmn"], - transforms={"Z": [[2, 1, 0]]}, + transforms={"Z": [[0, 1, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrt(params, transforms, profiles, data, **kwargs): - data["Z_rrt"] = transforms["Z"].transform(params["Z_lmn"], 2, 1, 0) +def _Z_tz(params, transforms, profiles, data, **kwargs): + data["Z_tz"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 1) return data @register_compute_fun( - name="Z_rrtt", - label="\\partial_{\\rho \\rho \\theta} Z", + name="Z_tzz", + label="\\partial_{\\theta \\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative, wrt radius " - "twice and poloidal angle twice", + description="Vertical coordinate in lab frame, third derivative wrt poloidal " + "angle and toroidal angle twice", dim=1, params=["Z_lmn"], - transforms={"Z": [[2, 2, 0]]}, + transforms={"Z": [[0, 1, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrtt(params, transforms, profiles, data, **kwargs): - data["Z_rrtt"] = transforms["Z"].transform(params["Z_lmn"], 2, 2, 0) +def _Z_tzz(params, transforms, profiles, data, **kwargs): + data["Z_tzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 2) return data @register_compute_fun( - name="Z_rtt", - label="\\partial_{\\rho \\theta \\theta} Z", + name="Z_z", + label="\\partial_{\\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt radius " - "and poloidal angle twice", + description="Vertical coordinate in lab frame, first toroidal derivative", dim=1, params=["Z_lmn"], - transforms={"Z": [[1, 2, 0]]}, + transforms={"Z": [[0, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rtt(params, transforms, profiles, data, **kwargs): - data["Z_rtt"] = transforms["Z"].transform(params["Z_lmn"], 1, 2, 0) +def _Z_z(params, transforms, profiles, data, **kwargs): + data["Z_z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 1) return data @register_compute_fun( - name="Z_rttz", - label="\\partial_{\\rho \\theta \\theta \\zeta} Z", + name="Z_zz", + label="\\partial_{\\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative wrt radius " - "once, poloidal angle twice, and toroidal angle once", + description="Vertical coordinate in lab frame, second toroidal derivative", dim=1, params=["Z_lmn"], - transforms={"Z": [[1, 2, 1]]}, + transforms={"Z": [[0, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rttz(params, transforms, profiles, data, **kwargs): - data["Z_rttz"] = transforms["Z"].transform(params["Z_lmn"], 1, 2, 1) +def _Z_zz(params, transforms, profiles, data, **kwargs): + data["Z_zz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 2) return data @register_compute_fun( - name="Z_rttt", - label="\\partial_{\\rho \\theta \\theta \\theta} Z", + name="Z_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} Z", units="m", units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt radius " - "and poloidal angle thrice", + description="Vertical coordinate in lab frame, third toroidal derivative", dim=1, params=["Z_lmn"], - transforms={"Z": [[1, 3, 0]]}, + transforms={"Z": [[0, 0, 3]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rttt(params, transforms, profiles, data, **kwargs): - data["Z_rttt"] = transforms["Z"].transform(params["Z_lmn"], 1, 3, 0) +def _Z_zzz(params, transforms, profiles, data, **kwargs): + data["Z_zzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 3) return data @register_compute_fun( - name="Z_rrz", - label="\\partial_{\\rho \\rho \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative, wrt radius " - "twice and toroidal angle", + name="alpha", + label="\\alpha", + units="~", + units_long="None", + description="Field line label, defined on [0, 2pi)", dim=1, - params=["Z_lmn"], - transforms={"Z": [[2, 0, 1]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["theta_PEST", "zeta", "iota"], ) -def _Z_rrz(params, transforms, profiles, data, **kwargs): - data["Z_rrz"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 1) +def _alpha(params, transforms, profiles, data, **kwargs): + data["alpha"] = (data["theta_PEST"] - data["iota"] * data["zeta"]) % (2 * jnp.pi) return data @register_compute_fun( - name="Z_rrzz", - label="\\partial_{\\rho \\rho \\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative, wrt radius " - "twice and toroidal angle twice", + name="alpha_r", + label="\\partial_\\rho \\alpha", + units="~", + units_long="None", + description="Field line label, derivative wrt radial coordinate", dim=1, - params=["Z_lmn"], - transforms={"Z": [[2, 0, 2]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["theta_PEST_r", "phi", "phi_r", "iota", "iota_r"], ) -def _Z_rrzz(params, transforms, profiles, data, **kwargs): - data["Z_rrzz"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 2) +def _alpha_r(params, transforms, profiles, data, **kwargs): + data["alpha_r"] = ( + data["theta_PEST_r"] + - data["iota_r"] * data["phi"] + - data["iota"] * data["phi_r"] + ) return data @register_compute_fun( - name="Z_rzz", - label="\\partial_{\\rho \\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt radius " - "and toroidal angle twice", + name="alpha_t", + label="\\partial_\\theta \\alpha", + units="~", + units_long="None", + description="Field line label, derivative wrt poloidal coordinate", dim=1, - params=["Z_lmn"], - transforms={"Z": [[1, 0, 2]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["theta_PEST_t", "phi_t", "iota"], ) -def _Z_rzz(params, transforms, profiles, data, **kwargs): - data["Z_rzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 2) +def _alpha_t(params, transforms, profiles, data, **kwargs): + data["alpha_t"] = data["theta_PEST_t"] + data["iota"] * data["phi_t"] return data @register_compute_fun( - name="Z_rzzz", - label="\\partial_{\\rho \\zeta \\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt radius " - "and toroidal angle thrice", + name="alpha_z", + label="\\partial_\\zeta \\alpha", + units="~", + units_long="None", + description="Field line label, derivative wrt toroidal coordinate", dim=1, - params=["Z_lmn"], - transforms={"Z": [[1, 0, 3]]}, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_z", "phi_z", "iota"], +) +def _alpha_z(params, transforms, profiles, data, **kwargs): + data["alpha_z"] = data["theta_PEST_z"] - data["iota"] * data["phi_z"] + return data + + +@register_compute_fun( + name="lambda", + label="\\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function", + dim=1, + params=["L_lmn"], + transforms={"L": [[0, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rzzz(params, transforms, profiles, data, **kwargs): - data["Z_rzzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 3) +def _lambda(params, transforms, profiles, data, **kwargs): + data["lambda"] = transforms["L"].transform(params["L_lmn"], 0, 0, 0) return data @register_compute_fun( - name="Z_ttz", - label="\\partial_{\\theta \\theta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt poloidal " - "angle twice and toroidal angle", + name="lambda_r", + label="\\partial_{\\rho} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, first radial derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 2, 1]]}, + params=["L_lmn"], + transforms={"L": [[1, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_ttz(params, transforms, profiles, data, **kwargs): - data["Z_ttz"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 1) +def _lambda_r(params, transforms, profiles, data, **kwargs): + data["lambda_r"] = transforms["L"].transform(params["L_lmn"], 1, 0, 0) return data @register_compute_fun( - name="Z_tzz", - label="\\partial_{\\theta \\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt poloidal " - "angle and toroidal angle twice", + name="lambda_rr", + label="\\partial_{\\rho \\rho} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second radial derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 1, 2]]}, + params=["L_lmn"], + transforms={"L": [[2, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_tzz(params, transforms, profiles, data, **kwargs): - data["Z_tzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 2) +def _lambda_rr(params, transforms, profiles, data, **kwargs): + data["lambda_rr"] = transforms["L"].transform(params["L_lmn"], 2, 0, 0) return data @register_compute_fun( - name="Z_rtz", - label="\\partial_{\\rho \\theta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative wrt radius, " - "poloidal angle, and toroidal angle", + name="lambda_rrr", + label="\\partial_{\rho \\rho \\rho} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third radial derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[1, 1, 1]]}, + params=["L_lmn"], + transforms={"L": [[3, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rtz(params, transforms, profiles, data, **kwargs): - data["Z_rtz"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 1) +def _lambda_rrr(params, transforms, profiles, data, **kwargs): + data["lambda_rrr"] = transforms["L"].transform(params["L_lmn"], 3, 0, 0) return data @register_compute_fun( - name="Z_rrtz", - label="\\partial_{\\rho \\rho \\theta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative wrt radius" - "twice, poloidal angle, and toroidal angle", + name="lambda_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third radial derivative and" + " first poloidal derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[2, 1, 1]]}, + params=["L_lmn"], + transforms={"L": [[3, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rrtz(params, transforms, profiles, data, **kwargs): - data["Z_rrtz"] = transforms["Z"].transform(params["Z_lmn"], 2, 1, 1) +def _lambda_rrrt(params, transforms, profiles, data, **kwargs): + data["lambda_rrrt"] = transforms["L"].transform(params["L_lmn"], 3, 1, 0) return data @register_compute_fun( - name="Z_rtzz", - label="\\partial_{\\rho \\theta \\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, fourth derivative wrt radius, " - "poloidal angle, and toroidal angle twice", + name="lambda_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third radial derivative and" + " first toroidal derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[1, 1, 2]]}, + params=["L_lmn"], + transforms={"L": [[3, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z_rtzz(params, transforms, profiles, data, **kwargs): - data["Z_rtzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 2) +def _lambda_rrrz(params, transforms, profiles, data, **kwargs): + data["lambda_rrrz"] = transforms["L"].transform(params["L_lmn"], 3, 0, 1) return data @register_compute_fun( - name="phi", - label="\\phi", + name="lambda_rrt", + label="\\partial_{\\rho \\rho \\theta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame", + description="Poloidal stream function, third derivative, wrt radius twice " + "and poloidal angle", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[2, 1, 0]]}, profiles=[], coordinates="rtz", - data=["zeta", "omega"], + data=[], ) -def _phi(params, transforms, profiles, data, **kwargs): - data["phi"] = data["zeta"] + data["omega"] +def _lambda_rrt(params, transforms, profiles, data, **kwargs): + data["lambda_rrt"] = transforms["L"].transform(params["L_lmn"], 2, 1, 0) return data @register_compute_fun( - name="phi_r", - label="\\partial_{\\rho} \\phi", + name="lambda_rrz", + label="\\partial_{\\rho \\rho \\zeta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, derivative wrt radial coordinate", + description="Poloidal stream function, third derivative, wrt radius twice " + "and toroidal angle", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[2, 0, 1]]}, profiles=[], coordinates="rtz", - data=["omega_r"], + data=[], ) -def _phi_r(params, transforms, profiles, data, **kwargs): - data["phi_r"] = data["omega_r"] +def _lambda_rrz(params, transforms, profiles, data, **kwargs): + data["lambda_rrz"] = transforms["L"].transform(params["L_lmn"], 2, 0, 1) return data @register_compute_fun( - name="phi_t", - label="\\partial_{\\theta} \\phi", + name="lambda_rt", + label="\\partial_{\\rho \\theta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, derivative wrt poloidal coordinate", + description="Poloidal stream function, second derivative wrt radius and " + "poloidal angle", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[1, 1, 0]]}, profiles=[], coordinates="rtz", - data=["omega_t"], + data=[], ) -def _phi_t(params, transforms, profiles, data, **kwargs): - data["phi_t"] = data["omega_t"] +def _lambda_rt(params, transforms, profiles, data, **kwargs): + data["lambda_rt"] = transforms["L"].transform(params["L_lmn"], 1, 1, 0) return data @register_compute_fun( - name="phi_z", - label="\\partial_{\\zeta} \\phi", + name="lambda_rtt", + label="\\partial_{\\rho \\theta \\theta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, derivative wrt toroidal coordinate", + description="Poloidal stream function, third derivative wrt radius and " + "poloidal angle twice", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[1, 2, 0]]}, profiles=[], coordinates="rtz", - data=["omega_z"], + data=[], ) -def _phi_z(params, transforms, profiles, data, **kwargs): - data["phi_z"] = 1 + data["omega_z"] +def _lambda_rtt(params, transforms, profiles, data, **kwargs): + data["lambda_rtt"] = transforms["L"].transform(params["L_lmn"], 1, 2, 0) return data @register_compute_fun( - name="phi_rr", - label="\\partial_{\\rho \\rho} \\phi", + name="lambda_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt radial coordinate", + description="Poloidal stream function, third derivative wrt radius, poloidal " + " angle, and toroidal angle", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[1, 1, 1]]}, profiles=[], coordinates="rtz", - data=["omega_rr"], + data=[], ) -def _phi_rr(params, transforms, profiles, data, **kwargs): - data["phi_rr"] = data["omega_rr"] +def _lambda_rtz(params, transforms, profiles, data, **kwargs): + data["lambda_rtz"] = transforms["L"].transform(params["L_lmn"], 1, 1, 1) return data @register_compute_fun( - name="phi_rt", - label="\\partial_{\\rho \\theta} \\phi", + name="lambda_rz", + label="\\partial_{\\rho \\zeta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt radial and " - "poloidal coordinate", + description="Poloidal stream function, second derivative wrt radius and " + "toroidal angle", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[1, 0, 1]]}, profiles=[], coordinates="rtz", - data=["omega_rt"], + data=[], ) -def _phi_rt(params, transforms, profiles, data, **kwargs): - data["phi_rt"] = data["omega_rt"] +def _lambda_rz(params, transforms, profiles, data, **kwargs): + data["lambda_rz"] = transforms["L"].transform(params["L_lmn"], 1, 0, 1) return data @register_compute_fun( - name="phi_rz", - label="\\partial_{\\rho \\zeta} \\phi", + name="lambda_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt radial and " - "toroidal coordinate", + description="Poloidal stream function, third derivative wrt radius and " + "toroidal angle twice", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[1, 0, 2]]}, profiles=[], coordinates="rtz", - data=["omega_rz"], + data=[], ) -def _phi_rz(params, transforms, profiles, data, **kwargs): - data["phi_rz"] = data["omega_rz"] +def _lambda_rzz(params, transforms, profiles, data, **kwargs): + data["lambda_rzz"] = transforms["L"].transform(params["L_lmn"], 1, 0, 2) return data @register_compute_fun( - name="phi_tt", - label="\\partial_{\\theta \\theta} \\phi", + name="lambda_t", + label="\\partial_{\\theta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt poloidal " - "coordinate", + description="Poloidal stream function, first poloidal derivative", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 1, 0]]}, profiles=[], coordinates="rtz", - data=["omega_tt"], + data=[], ) -def _phi_tt(params, transforms, profiles, data, **kwargs): - data["phi_tt"] = data["omega_tt"] +def _lambda_t(params, transforms, profiles, data, **kwargs): + data["lambda_t"] = transforms["L"].transform(params["L_lmn"], 0, 1, 0) return data @register_compute_fun( - name="phi_tz", - label="\\partial_{\\theta \\zeta} \\phi", + name="lambda_tt", + label="\\partial_{\\theta \\theta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt poloidal and " - "toroidal coordinate", + description="Poloidal stream function, second poloidal derivative", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 2, 0]]}, profiles=[], coordinates="rtz", - data=["omega_tz"], + data=[], ) -def _phi_tz(params, transforms, profiles, data, **kwargs): - data["phi_tz"] = data["omega_tz"] +def _lambda_tt(params, transforms, profiles, data, **kwargs): + data["lambda_tt"] = transforms["L"].transform(params["L_lmn"], 0, 2, 0) return data @register_compute_fun( - name="phi_zz", - label="\\partial_{\\zeta \\zeta} \\phi", + name="lambda_ttt", + label="\\partial_{\\theta \\theta \\theta} \\lambda", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt toroidal " - "coordinate", + description="Poloidal stream function, third poloidal derivative", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 3, 0]]}, profiles=[], coordinates="rtz", - data=["omega_zz"], + data=[], ) -def _phi_zz(params, transforms, profiles, data, **kwargs): - data["phi_zz"] = data["omega_zz"] +def _lambda_ttt(params, transforms, profiles, data, **kwargs): + data["lambda_ttt"] = transforms["L"].transform(params["L_lmn"], 0, 3, 0) return data @register_compute_fun( - name="X", - label="X = R \\cos{\\phi}", - units="m", - units_long="meters", - description="Cartesian X coordinate", + name="lambda_ttz", + label="\\partial_{\\theta \\theta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative wrt poloidal angle " + "twice and toroidal angle", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 2, 1]]}, profiles=[], coordinates="rtz", - data=["R", "phi"], + data=[], ) -def _X(params, transforms, profiles, data, **kwargs): - data["X"] = data["R"] * jnp.cos(data["phi"]) +def _lambda_ttz(params, transforms, profiles, data, **kwargs): + data["lambda_ttz"] = transforms["L"].transform(params["L_lmn"], 0, 2, 1) return data @register_compute_fun( - name="X_r", - label="\\partial_{\\rho} X", - units="m", - units_long="meters", - description="Cartesian X coordinate, derivative wrt radial coordinate", + name="lambda_tz", + label="\\partial_{\\theta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second derivative wrt poloidal and " + "toroidal angles", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 1, 1]]}, profiles=[], coordinates="rtz", - data=["R", "R_r", "phi", "phi_r"], + data=[], ) -def _X_r(params, transforms, profiles, data, **kwargs): - data["X_r"] = ( - data["R_r"] * jnp.cos(data["phi"]) - - data["R"] * jnp.sin(data["phi"]) * data["phi_r"] - ) +def _lambda_tz(params, transforms, profiles, data, **kwargs): + data["lambda_tz"] = transforms["L"].transform(params["L_lmn"], 0, 1, 1) return data @register_compute_fun( - name="X_t", - label="\\partial_{\\theta} X", - units="m", - units_long="meters", - description="Cartesian X coordinate, derivative wrt poloidal coordinate", + name="lambda_tzz", + label="\\partial_{\\theta \\zeta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative wrt poloidal angle " + "and toroidal angle twice", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 1, 2]]}, profiles=[], coordinates="rtz", - data=["R", "R_t", "phi", "phi_t"], + data=[], ) -def _X_t(params, transforms, profiles, data, **kwargs): - data["X_t"] = ( - data["R_t"] * jnp.cos(data["phi"]) - - data["R"] * jnp.sin(data["phi"]) * data["phi_t"] - ) +def _lambda_tzz(params, transforms, profiles, data, **kwargs): + data["lambda_tzz"] = transforms["L"].transform(params["L_lmn"], 0, 1, 2) return data @register_compute_fun( - name="X_z", - label="\\partial_{\\zeta} X", - units="m", - units_long="meters", - description="Cartesian X coordinate, derivative wrt toroidal coordinate", + name="lambda_z", + label="\\partial_{\\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, first toroidal derivative", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 0, 1]]}, profiles=[], coordinates="rtz", - data=["R", "R_z", "phi", "phi_z"], + data=[], ) -def _X_z(params, transforms, profiles, data, **kwargs): - data["X_z"] = ( - data["R_z"] * jnp.cos(data["phi"]) - - data["R"] * jnp.sin(data["phi"]) * data["phi_z"] - ) +def _lambda_z(params, transforms, profiles, data, **kwargs): + data["lambda_z"] = transforms["L"].transform(params["L_lmn"], 0, 0, 1) return data @register_compute_fun( - name="Y", - label="Y = R \\sin{\\phi}", - units="m", - units_long="meters", - description="Cartesian Y coordinate", + name="lambda_zz", + label="\\partial_{\\zeta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second toroidal derivative", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 0, 2]]}, profiles=[], coordinates="rtz", - data=["R", "phi"], + data=[], ) -def _Y(params, transforms, profiles, data, **kwargs): - data["Y"] = data["R"] * jnp.sin(data["phi"]) +def _lambda_zz(params, transforms, profiles, data, **kwargs): + data["lambda_zz"] = transforms["L"].transform(params["L_lmn"], 0, 0, 2) return data @register_compute_fun( - name="Y_r", - label="\\partial_{\\rho} Y", - units="m", - units_long="meters", - description="Cartesian Y coordinate, derivative wrt radial coordinate", + name="lambda_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third toroidal derivative", dim=1, - params=[], - transforms={}, + params=["L_lmn"], + transforms={"L": [[0, 0, 3]]}, profiles=[], coordinates="rtz", - data=["R", "R_r", "phi", "phi_r"], + data=[], ) -def _Y_r(params, transforms, profiles, data, **kwargs): - data["Y_r"] = ( - data["R_r"] * jnp.sin(data["phi"]) - + data["R"] * jnp.cos(data["phi"]) * data["phi_r"] - ) +def _lambda_zzz(params, transforms, profiles, data, **kwargs): + data["lambda_zzz"] = transforms["L"].transform(params["L_lmn"], 0, 0, 3) return data @register_compute_fun( - name="Y_t", - label="\\partial_{\\theta} Y", - units="m", - units_long="meters", - description="Cartesian Y coordinate, derivative wrt poloidal coordinate", + name="omega", + label="\\omega", + units="rad", + units_long="radians", + description="Toroidal stream function", dim=1, - params=[], + params=[], # ["W_lmn"], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "phi", "phi_t"], + data=["0"], ) -def _Y_t(params, transforms, profiles, data, **kwargs): - data["Y_t"] = ( - data["R_t"] * jnp.sin(data["phi"]) - + data["R"] * jnp.cos(data["phi"]) * data["phi_t"] - ) +def _omega(params, transforms, profiles, data, **kwargs): + data["omega"] = data["0"] return data @register_compute_fun( - name="Y_z", - label="\\partial_{\\zeta} Y", - units="m", - units_long="meters", - description="Cartesian Y coordinate, derivative wrt toroidal coordinate", + name="omega_r", + label="\\partial_{\\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, first radial derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "phi", "phi_z"], + data=["0"], ) -def _Y_z(params, transforms, profiles, data, **kwargs): - data["Y_z"] = ( - data["R_z"] * jnp.sin(data["phi"]) - + data["R"] * jnp.cos(data["phi"]) * data["phi_z"] - ) +def _omega_r(params, transforms, profiles, data, **kwargs): + data["omega_r"] = data["0"] return data @register_compute_fun( - name="lambda", - label="\\lambda", + name="omega_rr", + label="\\partial_{\\rho \\rho} \\omega", units="rad", units_long="radians", - description="Poloidal stream function", + description="Toroidal stream function, second radial derivative", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda(params, transforms, profiles, data, **kwargs): - data["lambda"] = transforms["L"].transform(params["L_lmn"], 0, 0, 0) +def _omega_rr(params, transforms, profiles, data, **kwargs): + data["omega_rr"] = data["0"] return data @register_compute_fun( - name="lambda_r", - label="\\partial_{\\rho} \\lambda", + name="omega_rrr", + label="\\partial_{\rho \\rho \\rho} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, first radial derivative", + description="Toroidal stream function, third radial derivative", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 0, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_r(params, transforms, profiles, data, **kwargs): - data["lambda_r"] = transforms["L"].transform(params["L_lmn"], 1, 0, 0) +def _omega_rrr(params, transforms, profiles, data, **kwargs): + data["omega_rrr"] = data["0"] return data @register_compute_fun( - name="lambda_t", - label="\\partial_{\\theta} \\lambda", + name="omega_rrrr", + label="\\partial_{\rho \\rho \\rho \\rho} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, first poloidal derivative", + description="Toroidal stream function, fourth radial derivative", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 1, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_t(params, transforms, profiles, data, **kwargs): - data["lambda_t"] = transforms["L"].transform(params["L_lmn"], 0, 1, 0) +def _omega_rrrr(params, transforms, profiles, data, **kwargs): + data["omega_rrrr"] = data["0"] return data @register_compute_fun( - name="lambda_z", - label="\\partial_{\\zeta} \\lambda", + name="omega_rrrt", + label="\\partial_{\rho \\rho \\rho \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, first toroidal derivative", + description="Toroidal stream function, fourth derivative wrt radial coordinate" + " thrice and poloidal once", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_z(params, transforms, profiles, data, **kwargs): - data["lambda_z"] = transforms["L"].transform(params["L_lmn"], 0, 0, 1) +def _omega_rrrt(params, transforms, profiles, data, **kwargs): + data["omega_rrrt"] = data["0"] return data @register_compute_fun( - name="lambda_rr", - label="\\partial_{\\rho \\rho} \\lambda", + name="omega_rrrz", + label="\\partial_{\rho \\rho \\rho \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, second radial derivative", + description="Toroidal stream function, fourth derivative wrt radial coordinate" + " thrice and toroidal once", dim=1, - params=["L_lmn"], - transforms={"L": [[2, 0, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rr(params, transforms, profiles, data, **kwargs): - data["lambda_rr"] = transforms["L"].transform(params["L_lmn"], 2, 0, 0) +def _omega_rrrz(params, transforms, profiles, data, **kwargs): + data["omega_rrrz"] = data["0"] return data @register_compute_fun( - name="lambda_tt", - label="\\partial_{\\theta \\theta} \\lambda", + name="omega_rrt", + label="\\partial_{\\rho \\rho \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, second poloidal derivative", + description="Toroidal stream function, third derivative, wrt radius twice " + "and poloidal angle", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 2, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_tt(params, transforms, profiles, data, **kwargs): - data["lambda_tt"] = transforms["L"].transform(params["L_lmn"], 0, 2, 0) +def _omega_rrt(params, transforms, profiles, data, **kwargs): + data["omega_rrt"] = data["0"] return data @register_compute_fun( - name="lambda_zz", - label="\\partial_{\\zeta \\zeta} \\lambda", + name="omega_rrtt", + label="\\partial_{\\rho \\rho \\theta \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, second toroidal derivative", + description="Toroidal stream function, fourth derivative, wrt radius twice " + "and poloidal angle twice", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 2]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_zz(params, transforms, profiles, data, **kwargs): - data["lambda_zz"] = transforms["L"].transform(params["L_lmn"], 0, 0, 2) +def _omega_rrtt(params, transforms, profiles, data, **kwargs): + data["omega_rrtt"] = data["0"] return data @register_compute_fun( - name="lambda_rt", - label="\\partial_{\\rho \\theta} \\lambda", + name="omega_rrtz", + label="\\partial_{\\rho \\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, second derivative wrt radius and " - "poloidal angle", + description="Toroidal stream function, fourth derivative wrt radius twice," + " poloidal angle, and toroidal angle", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 1, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rt(params, transforms, profiles, data, **kwargs): - data["lambda_rt"] = transforms["L"].transform(params["L_lmn"], 1, 1, 0) +def _omega_rrtz(params, transforms, profiles, data, **kwargs): + data["omega_rrtz"] = data["0"] return data @register_compute_fun( - name="lambda_rz", - label="\\partial_{\\rho \\zeta} \\lambda", + name="omega_rrz", + label="\\partial_{\\rho \\rho \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, second derivative wrt radius and " - "toroidal angle", + description="Toroidal stream function, third derivative, wrt radius twice " + "and toroidal angle", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rz(params, transforms, profiles, data, **kwargs): - data["lambda_rz"] = transforms["L"].transform(params["L_lmn"], 1, 0, 1) +def _omega_rrz(params, transforms, profiles, data, **kwargs): + data["omega_rrz"] = data["0"] return data @register_compute_fun( - name="lambda_tz", - label="\\partial_{\\theta \\zeta} \\lambda", + name="omega_rrzz", + label="\\partial_{\\rho \\rho \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, second derivative wrt poloidal and " - "toroidal angles", + description="Toroidal stream function, fourth derivative, wrt radius twice " + "and toroidal angle twice", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 1, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_tz(params, transforms, profiles, data, **kwargs): - data["lambda_tz"] = transforms["L"].transform(params["L_lmn"], 0, 1, 1) +def _omega_rrzz(params, transforms, profiles, data, **kwargs): + data["omega_rrzz"] = data["0"] return data @register_compute_fun( - name="lambda_rrr", - label="\\partial_{\rho \\rho \\rho} \\lambda", + name="omega_rt", + label="\\partial_{\\rho \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third radial derivative", + description="Toroidal stream function, second derivative wrt radius and " + "poloidal angle", dim=1, - params=["L_lmn"], - transforms={"L": [[3, 0, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rrr(params, transforms, profiles, data, **kwargs): - data["lambda_rrr"] = transforms["L"].transform(params["L_lmn"], 3, 0, 0) +def _omega_rt(params, transforms, profiles, data, **kwargs): + data["omega_rt"] = data["0"] return data @register_compute_fun( - name="lambda_rrrt", - label="\\partial_{\rho \\rho \\rho \\theta} \\lambda", + name="omega_rtt", + label="\\partial_{\\rho \\theta \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third radial derivative and" - " first poloidal derivative", + description="Toroidal stream function, third derivative wrt radius and " + "poloidal angle twice", dim=1, - params=["L_lmn"], - transforms={"L": [[3, 1, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rrrt(params, transforms, profiles, data, **kwargs): - data["lambda_rrrt"] = transforms["L"].transform(params["L_lmn"], 3, 1, 0) +def _omega_rtt(params, transforms, profiles, data, **kwargs): + data["omega_rtt"] = data["0"] return data @register_compute_fun( - name="lambda_rrrz", - label="\\partial_{\rho \\rho \\rho \\zeta} \\lambda", + name="omega_rttt", + label="\\partial_{\\rho \\theta \\theta \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third radial derivative and" - " first toroidal derivative", + description="Toroidal stream function, third derivative wrt radius and " + "poloidal angle thrice", dim=1, - params=["L_lmn"], - transforms={"L": [[3, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rrrz(params, transforms, profiles, data, **kwargs): - data["lambda_rrrz"] = transforms["L"].transform(params["L_lmn"], 3, 0, 1) +def _omega_rttt(params, transforms, profiles, data, **kwargs): + data["omega_rttt"] = data["0"] return data @register_compute_fun( - name="lambda_ttt", - label="\\partial_{\\theta \\theta \\theta} \\lambda", + name="omega_rttz", + label="\\partial_{\\rho \\theta \\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third poloidal derivative", + description="Toroidal stream function, fourth derivative wrt radius once, " + "poloidal angle twice, and toroidal angle once", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 3, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_ttt(params, transforms, profiles, data, **kwargs): - data["lambda_ttt"] = transforms["L"].transform(params["L_lmn"], 0, 3, 0) +def _omega_rttz(params, transforms, profiles, data, **kwargs): + data["omega_rttz"] = data["0"] return data @register_compute_fun( - name="lambda_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} \\lambda", + name="omega_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third toroidal derivative", + description="Toroidal stream function, third derivative wrt radius, poloidal" + " angle, and toroidal angle", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 3]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_zzz(params, transforms, profiles, data, **kwargs): - data["lambda_zzz"] = transforms["L"].transform(params["L_lmn"], 0, 0, 3) +def _omega_rtz(params, transforms, profiles, data, **kwargs): + data["omega_rtz"] = data["0"] return data @register_compute_fun( - name="lambda_rrt", - label="\\partial_{\\rho \\rho \\theta} \\lambda", + name="omega_rtzz", + label="\\partial_{\\rho \\theta \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative, wrt radius twice " - "and poloidal angle", + description="Toroidal stream function, fourth derivative wrt radius, poloidal" + " angle, and toroidal angle twice", dim=1, - params=["L_lmn"], - transforms={"L": [[2, 1, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rrt(params, transforms, profiles, data, **kwargs): - data["lambda_rrt"] = transforms["L"].transform(params["L_lmn"], 2, 1, 0) +def _omega_rtzz(params, transforms, profiles, data, **kwargs): + data["omega_rtzz"] = data["0"] return data @register_compute_fun( - name="lambda_rtt", - label="\\partial_{\\rho \\theta \\theta} \\lambda", + name="omega_rz", + label="\\partial_{\\rho \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt radius and " - "poloidal angle twice", + description="Toroidal stream function, second derivative wrt radius and " + "toroidal angle", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 2, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rtt(params, transforms, profiles, data, **kwargs): - data["lambda_rtt"] = transforms["L"].transform(params["L_lmn"], 1, 2, 0) +def _omega_rz(params, transforms, profiles, data, **kwargs): + data["omega_rz"] = data["0"] return data @register_compute_fun( - name="lambda_rrz", - label="\\partial_{\\rho \\rho \\zeta} \\lambda", + name="omega_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative, wrt radius twice " - "and toroidal angle", + description="Toroidal stream function, third derivative wrt radius and " + "toroidal angle twice", dim=1, - params=["L_lmn"], - transforms={"L": [[2, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rrz(params, transforms, profiles, data, **kwargs): - data["lambda_rrz"] = transforms["L"].transform(params["L_lmn"], 2, 0, 1) +def _omega_rzz(params, transforms, profiles, data, **kwargs): + data["omega_rzz"] = data["0"] return data @register_compute_fun( - name="lambda_rzz", - label="\\partial_{\\rho \\zeta \\zeta} \\lambda", + name="omega_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt radius and " - "toroidal angle twice", + description="Toroidal stream function, third derivative wrt radius and " + "toroidal angle thrice", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 0, 2]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rzz(params, transforms, profiles, data, **kwargs): - data["lambda_rzz"] = transforms["L"].transform(params["L_lmn"], 1, 0, 2) +def _omega_rzzz(params, transforms, profiles, data, **kwargs): + data["omega_rzzz"] = data["0"] return data @register_compute_fun( - name="lambda_ttz", - label="\\partial_{\\theta \\theta \\zeta} \\lambda", + name="omega_t", + label="\\partial_{\\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt poloidal angle " - "twice and toroidal angle", + description="Toroidal stream function, first poloidal derivative", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 2, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_ttz(params, transforms, profiles, data, **kwargs): - data["lambda_ttz"] = transforms["L"].transform(params["L_lmn"], 0, 2, 1) +def _omega_t(params, transforms, profiles, data, **kwargs): + data["omega_t"] = data["0"] return data @register_compute_fun( - name="lambda_tzz", - label="\\partial_{\\theta \\zeta \\zeta} \\lambda", + name="omega_tt", + label="\\partial_{\\theta \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt poloidal angle " - "and toroidal angle twice", + description="Toroidal stream function, second poloidal derivative", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 1, 2]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_tzz(params, transforms, profiles, data, **kwargs): - data["lambda_tzz"] = transforms["L"].transform(params["L_lmn"], 0, 1, 2) +def _omega_tt(params, transforms, profiles, data, **kwargs): + data["omega_tt"] = data["0"] return data @register_compute_fun( - name="lambda_rtz", - label="\\partial_{\\rho \\theta \\zeta} \\lambda", + name="omega_ttt", + label="\\partial_{\\theta \\theta \\theta} \\omega", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt radius, poloidal " - " angle, and toroidal angle", + description="Toroidal stream function, third poloidal derivative", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 1, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _lambda_rtz(params, transforms, profiles, data, **kwargs): - data["lambda_rtz"] = transforms["L"].transform(params["L_lmn"], 1, 1, 1) +def _omega_ttt(params, transforms, profiles, data, **kwargs): + data["omega_ttt"] = data["0"] return data @register_compute_fun( - name="omega", - label="\\omega", + name="omega_ttz", + label="\\partial_{\\theta \\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal stream function", + description="Toroidal stream function, third derivative wrt poloidal angle " + "twice and toroidal angle", dim=1, - params=[], # ["W_lmn"], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _omega(params, transforms, profiles, data, **kwargs): - data["omega"] = data["0"] +def _omega_ttz(params, transforms, profiles, data, **kwargs): + data["omega_ttz"] = data["0"] return data @register_compute_fun( - name="omega_r", - label="\\partial_{\\rho} \\omega", + name="omega_tz", + label="\\partial_{\\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal stream function, first radial derivative", + description="Toroidal stream function, second derivative wrt poloidal and " + "toroidal angles", dim=1, params=[], # ["W_lmn"] transforms={}, @@ -2301,17 +2268,18 @@ def _omega(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["0"], ) -def _omega_r(params, transforms, profiles, data, **kwargs): - data["omega_r"] = data["0"] +def _omega_tz(params, transforms, profiles, data, **kwargs): + data["omega_tz"] = data["0"] return data @register_compute_fun( - name="omega_t", - label="\\partial_{\\theta} \\omega", + name="omega_tzz", + label="\\partial_{\\theta \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal stream function, first poloidal derivative", + description="Toroidal stream function, third derivative wrt poloidal angle " + "and toroidal angle twice", dim=1, params=[], # ["W_lmn"] transforms={}, @@ -2319,8 +2287,8 @@ def _omega_r(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["0"], ) -def _omega_t(params, transforms, profiles, data, **kwargs): - data["omega_t"] = data["0"] +def _omega_tzz(params, transforms, profiles, data, **kwargs): + data["omega_tzz"] = data["0"] return data @@ -2343,11 +2311,11 @@ def _omega_z(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="omega_rr", - label="\\partial_{\\rho \\rho} \\omega", + name="omega_zz", + label="\\partial_{\\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal stream function, second radial derivative", + description="Toroidal stream function, second toroidal derivative", dim=1, params=[], # ["W_lmn"] transforms={}, @@ -2355,17 +2323,17 @@ def _omega_z(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["0"], ) -def _omega_rr(params, transforms, profiles, data, **kwargs): - data["omega_rr"] = data["0"] +def _omega_zz(params, transforms, profiles, data, **kwargs): + data["omega_zz"] = data["0"] return data @register_compute_fun( - name="omega_tt", - label="\\partial_{\\theta \\theta} \\omega", + name="omega_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal stream function, second poloidal derivative", + description="Toroidal stream function, third toroidal derivative", dim=1, params=[], # ["W_lmn"] transforms={}, @@ -2373,457 +2341,489 @@ def _omega_rr(params, transforms, profiles, data, **kwargs): coordinates="rtz", data=["0"], ) -def _omega_tt(params, transforms, profiles, data, **kwargs): - data["omega_tt"] = data["0"] +def _omega_zzz(params, transforms, profiles, data, **kwargs): + data["omega_zzz"] = data["0"] return data @register_compute_fun( - name="omega_zz", - label="\\partial_{\\zeta \\zeta} \\omega", + name="phi", + label="\\phi", units="rad", units_long="radians", - description="Toroidal stream function, second toroidal derivative", + description="Toroidal angle in lab frame", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["zeta", "omega"], ) -def _omega_zz(params, transforms, profiles, data, **kwargs): - data["omega_zz"] = data["0"] +def _phi(params, transforms, profiles, data, **kwargs): + data["phi"] = data["zeta"] + data["omega"] return data @register_compute_fun( - name="omega_rt", - label="\\partial_{\\rho \\theta} \\omega", + name="phi_r", + label="\\partial_{\\rho} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, second derivative wrt radius and " - "poloidal angle", + description="Toroidal angle in lab frame, derivative wrt radial coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_r"], ) -def _omega_rt(params, transforms, profiles, data, **kwargs): - data["omega_rt"] = data["0"] +def _phi_r(params, transforms, profiles, data, **kwargs): + data["phi_r"] = data["omega_r"] return data @register_compute_fun( - name="omega_rz", - label="\\partial_{\\rho \\zeta} \\omega", + name="phi_rr", + label="\\partial_{\\rho \\rho} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, second derivative wrt radius and " - "toroidal angle", + description="Toroidal angle in lab frame, second derivative wrt radial coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_rr"], ) -def _omega_rz(params, transforms, profiles, data, **kwargs): - data["omega_rz"] = data["0"] +def _phi_rr(params, transforms, profiles, data, **kwargs): + data["phi_rr"] = data["omega_rr"] return data @register_compute_fun( - name="omega_tz", - label="\\partial_{\\theta \\zeta} \\omega", + name="phi_rt", + label="\\partial_{\\rho \\theta} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, second derivative wrt poloidal and " - "toroidal angles", + description="Toroidal angle in lab frame, second derivative wrt radial and " + "poloidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_rt"], ) -def _omega_tz(params, transforms, profiles, data, **kwargs): - data["omega_tz"] = data["0"] +def _phi_rt(params, transforms, profiles, data, **kwargs): + data["phi_rt"] = data["omega_rt"] return data @register_compute_fun( - name="omega_rrr", - label="\\partial_{\rho \\rho \\rho} \\omega", + name="phi_rz", + label="\\partial_{\\rho \\zeta} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, third radial derivative", + description="Toroidal angle in lab frame, second derivative wrt radial and " + "toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_rz"], ) -def _omega_rrr(params, transforms, profiles, data, **kwargs): - data["omega_rrr"] = data["0"] +def _phi_rz(params, transforms, profiles, data, **kwargs): + data["phi_rz"] = data["omega_rz"] return data @register_compute_fun( - name="omega_rrrr", - label="\\partial_{\rho \\rho \\rho \\rho} \\omega", + name="phi_t", + label="\\partial_{\\theta} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, fourth radial derivative", + description="Toroidal angle in lab frame, derivative wrt poloidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_t"], ) -def _omega_rrrr(params, transforms, profiles, data, **kwargs): - data["omega_rrrr"] = data["0"] +def _phi_t(params, transforms, profiles, data, **kwargs): + data["phi_t"] = data["omega_t"] return data @register_compute_fun( - name="omega_rrrt", - label="\\partial_{\rho \\rho \\rho \\theta} \\omega", + name="phi_tt", + label="\\partial_{\\theta \\theta} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, fourth derivative wrt radial coordinate" - " thrice and poloidal once", + description="Toroidal angle in lab frame, second derivative wrt poloidal " + "coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_tt"], ) -def _omega_rrrt(params, transforms, profiles, data, **kwargs): - data["omega_rrrt"] = data["0"] +def _phi_tt(params, transforms, profiles, data, **kwargs): + data["phi_tt"] = data["omega_tt"] return data @register_compute_fun( - name="omega_rrrz", - label="\\partial_{\rho \\rho \\rho \\zeta} \\omega", + name="phi_tz", + label="\\partial_{\\theta \\zeta} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, fourth derivative wrt radial coordinate" - " thrice and toroidal once", + description="Toroidal angle in lab frame, second derivative wrt poloidal and " + "toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_tz"], ) -def _omega_rrrz(params, transforms, profiles, data, **kwargs): - data["omega_rrrz"] = data["0"] +def _phi_tz(params, transforms, profiles, data, **kwargs): + data["phi_tz"] = data["omega_tz"] return data @register_compute_fun( - name="omega_ttt", - label="\\partial_{\\theta \\theta \\theta} \\omega", + name="phi_z", + label="\\partial_{\\zeta} \\phi", units="rad", units_long="radians", - description="Toroidal stream function, third poloidal derivative", + description="Toroidal angle in lab frame, derivative wrt toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["omega_z"], ) -def _omega_ttt(params, transforms, profiles, data, **kwargs): - data["omega_ttt"] = data["0"] +def _phi_z(params, transforms, profiles, data, **kwargs): + data["phi_z"] = 1 + data["omega_z"] + return data + + +@register_compute_fun( + name="phi_zz", + label="\\partial_{\\zeta \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, second derivative wrt toroidal " + "coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_zz"], +) +def _phi_zz(params, transforms, profiles, data, **kwargs): + data["phi_zz"] = data["omega_zz"] + return data + + +@register_compute_fun( + name="rho", + label="\\rho", + units="~", + units_long="None", + description="Radial coordinate, proportional to the square root " + + "of the toroidal flux", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=[], +) +def _rho(params, transforms, profiles, data, **kwargs): + data["rho"] = transforms["grid"].nodes[:, 0] return data @register_compute_fun( - name="omega_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} \\omega", - units="rad", - units_long="radians", - description="Toroidal stream function, third toroidal derivative", + name="rho_r", + label="\\partial_{\\rho} \\rho", + units="~", + units_long="None", + description="Radial coordinate, proportional to the square root " + + "of the toroidal flux, derivative wrt radial coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="r", data=["0"], ) -def _omega_zzz(params, transforms, profiles, data, **kwargs): - data["omega_zzz"] = data["0"] +def _rho_r(params, transforms, profiles, data, **kwargs): + data["rho_r"] = jnp.ones_like(data["0"]) return data @register_compute_fun( - name="omega_rrt", - label="\\partial_{\\rho \\rho \\theta} \\omega", - units="rad", - units_long="radians", - description="Toroidal stream function, third derivative, wrt radius twice " - "and poloidal angle", + name="rho_t", + label="\\partial_{\\theta} \\rho", + units="~", + units_long="None", + description="Radial coordinate, proportional to the square root " + "of the toroidal flux, derivative wrt poloidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="r", data=["0"], ) -def _omega_rrt(params, transforms, profiles, data, **kwargs): - data["omega_rrt"] = data["0"] +def _rho_t(params, transforms, profiles, data, **kwargs): + data["rho_t"] = data["0"] return data @register_compute_fun( - name="omega_rrtt", - label="\\partial_{\\rho \\rho \\theta \\theta} \\omega", - units="rad", - units_long="radians", - description="Toroidal stream function, fourth derivative, wrt radius twice " - "and poloidal angle twice", + name="rho_z", + label="\\partial_{\\zeta} \\rho", + units="~", + units_long="None", + description="Radial coordinate, proportional to the square root " + "of the toroidal flux, derivative wrt toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="r", data=["0"], ) -def _omega_rrtt(params, transforms, profiles, data, **kwargs): - data["omega_rrtt"] = data["0"] +def _rho_z(params, transforms, profiles, data, **kwargs): + data["rho_z"] = data["0"] return data @register_compute_fun( - name="omega_rtt", - label="\\partial_{\\rho \\theta \\theta} \\omega", + name="theta", + label="\\theta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt radius and " - "poloidal angle twice", + description="Poloidal angular coordinate (geometric, not magnetic)", dim=1, - params=[], # ["W_lmn"] - transforms={}, + params=[], + transforms={"grid": []}, profiles=[], - coordinates="rtz", - data=["0"], + coordinates="t", + data=[], ) -def _omega_rtt(params, transforms, profiles, data, **kwargs): - data["omega_rtt"] = data["0"] +def _theta(params, transforms, profiles, data, **kwargs): + data["theta"] = transforms["grid"].nodes[:, 1] return data @register_compute_fun( - name="omega_rttz", - label="\\partial_{\\rho \\theta \\theta \\zeta} \\omega", + name="theta_PEST", + label="\\vartheta", units="rad", units_long="radians", - description="Toroidal stream function, fourth derivative wrt radius once, " - "poloidal angle twice, and toroidal angle once", + description="PEST straight field line poloidal angular coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["theta", "lambda"], ) -def _omega_rttz(params, transforms, profiles, data, **kwargs): - data["omega_rttz"] = data["0"] +def _theta_PEST(params, transforms, profiles, data, **kwargs): + data["theta_PEST"] = (data["theta"] + data["lambda"]) % (2 * jnp.pi) return data @register_compute_fun( - name="omega_rttt", - label="\\partial_{\\rho \\theta \\theta \\theta} \\omega", + name="theta_PEST_r", + label="\\partial_{\\rho} \\vartheta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt radius and " - "poloidal angle thrice", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "radial coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["lambda_r"], ) -def _omega_rttt(params, transforms, profiles, data, **kwargs): - data["omega_rttt"] = data["0"] +def _theta_PEST_r(params, transforms, profiles, data, **kwargs): + data["theta_PEST_r"] = data["lambda_r"] return data @register_compute_fun( - name="omega_rrz", - label="\\partial_{\\rho \\rho \\zeta} \\omega", + name="theta_PEST_t", + label="\\partial_{\\theta} \\vartheta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative, wrt radius twice " - "and toroidal angle", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "poloidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["lambda_t"], ) -def _omega_rrz(params, transforms, profiles, data, **kwargs): - data["omega_rrz"] = data["0"] +def _theta_PEST_t(params, transforms, profiles, data, **kwargs): + data["theta_PEST_t"] = 1 + data["lambda_t"] return data @register_compute_fun( - name="omega_rrzz", - label="\\partial_{\\rho \\rho \\zeta \\zeta} \\omega", + name="theta_PEST_z", + label="\\partial_{\\zeta} \\vartheta", units="rad", units_long="radians", - description="Toroidal stream function, fourth derivative, wrt radius twice " - "and toroidal angle twice", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], coordinates="rtz", - data=["0"], + data=["lambda_z"], ) -def _omega_rrzz(params, transforms, profiles, data, **kwargs): - data["omega_rrzz"] = data["0"] +def _theta_PEST_z(params, transforms, profiles, data, **kwargs): + data["theta_PEST_z"] = data["lambda_z"] return data @register_compute_fun( - name="omega_rzz", - label="\\partial_{\\rho \\zeta \\zeta} \\omega", + name="theta_r", + label="\\partial_{\\rho} \\theta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt radius and " - "toroidal angle twice", + description="Poloidal angular coordinate (geometric, not magnetic), " + "derivative wrt radial coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="t", data=["0"], ) -def _omega_rzz(params, transforms, profiles, data, **kwargs): - data["omega_rzz"] = data["0"] +def _theta_r(params, transforms, profiles, data, **kwargs): + data["theta_r"] = data["0"] return data @register_compute_fun( - name="omega_rzzz", - label="\\partial_{\\rho \\zeta \\zeta \\zeta} \\omega", + name="theta_t", + label="\\partial_{\\theta} \\theta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt radius and " - "toroidal angle thrice", + description="Poloidal angular coordinate (geometric, not magnetic), " + "derivative wrt poloidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="t", data=["0"], ) -def _omega_rzzz(params, transforms, profiles, data, **kwargs): - data["omega_rzzz"] = data["0"] +def _theta_t(params, transforms, profiles, data, **kwargs): + data["theta_t"] = jnp.ones_like(data["0"]) return data @register_compute_fun( - name="omega_ttz", - label="\\partial_{\\theta \\theta \\zeta} \\omega", + name="theta_z", + label="\\partial_{\\zeta} \\theta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt poloidal angle " - "twice and toroidal angle", + description="Poloidal angular coordinate (geometric, not magnetic), " + "derivative wrt toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="t", data=["0"], ) -def _omega_ttz(params, transforms, profiles, data, **kwargs): - data["omega_ttz"] = data["0"] +def _theta_z(params, transforms, profiles, data, **kwargs): + data["theta_z"] = data["0"] return data @register_compute_fun( - name="omega_tzz", - label="\\partial_{\\theta \\zeta \\zeta} \\omega", + name="zeta", + label="\\zeta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt poloidal angle " - "and toroidal angle twice", + description="Toroidal angular coordinate", dim=1, - params=[], # ["W_lmn"] - transforms={}, + params=[], + transforms={"grid": []}, profiles=[], - coordinates="rtz", - data=["0"], + coordinates="z", + data=[], ) -def _omega_tzz(params, transforms, profiles, data, **kwargs): - data["omega_tzz"] = data["0"] +def _zeta(params, transforms, profiles, data, **kwargs): + data["zeta"] = transforms["grid"].nodes[:, 2] return data @register_compute_fun( - name="omega_rtz", - label="\\partial_{\\rho \\theta \\zeta} \\omega", + name="zeta_r", + label="\\partial_{\\rho} \\zeta", units="rad", units_long="radians", - description="Toroidal stream function, third derivative wrt radius, poloidal" - " angle, and toroidal angle", + description="Toroidal angular coordinate derivative, wrt radial coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="z", data=["0"], ) -def _omega_rtz(params, transforms, profiles, data, **kwargs): - data["omega_rtz"] = data["0"] +def _zeta_r(params, transforms, profiles, data, **kwargs): + data["zeta_r"] = data["0"] return data @register_compute_fun( - name="omega_rrtz", - label="\\partial_{\\rho \\theta \\zeta} \\omega", + name="zeta_t", + label="\\partial_{\\theta} \\zeta", units="rad", units_long="radians", - description="Toroidal stream function, fourth derivative wrt radius twice," - " poloidal angle, and toroidal angle", + description="Toroidal angular coordinate, derivative wrt poloidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="z", data=["0"], ) -def _omega_rrtz(params, transforms, profiles, data, **kwargs): - data["omega_rrtz"] = data["0"] +def _zeta_t(params, transforms, profiles, data, **kwargs): + data["zeta_t"] = data["0"] return data @register_compute_fun( - name="omega_rtzz", - label="\\partial_{\\rho \\theta \\zeta \\zeta} \\omega", + name="zeta_z", + label="\\partial_{\\zeta} \\zeta", units="rad", units_long="radians", - description="Toroidal stream function, fourth derivative wrt radius, poloidal" - " angle, and toroidal angle twice", + description="Toroidal angular coordinate, derivative wrt toroidal coordinate", dim=1, - params=[], # ["W_lmn"] + params=[], transforms={}, profiles=[], - coordinates="rtz", + coordinates="z", data=["0"], ) -def _omega_rtzz(params, transforms, profiles, data, **kwargs): - data["omega_rtzz"] = data["0"] +def _zeta_z(params, transforms, profiles, data, **kwargs): + data["zeta_z"] = jnp.ones_like(data["0"]) return data From 0ca2226cd271fc14427e86acca04cf44cc23dcc8 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Mon, 31 Jul 2023 14:44:01 -0500 Subject: [PATCH 08/19] Merge branch 'master' into basis_vectors and ... Remove unneeded dependencies so that tests/test_axis_limits.py passes. Other changes were formatting changes forced by flake8 precommit. --- desc/compute/__init__.py | 58 ++++---- desc/compute/_basis_vectors.py | 12 -- desc/compute/_bootstrap.py | 25 ++-- desc/compute/_field.py | 8 +- desc/compute/_geometry.py | 2 +- desc/compute/_metric.py | 60 +++------ desc/compute/_profiles.py | 11 +- desc/compute/data_index.py | 61 ++++++++- desc/compute/utils.py | 232 +++++++++++++++++++------------- desc/equilibrium/coords.py | 5 +- desc/equilibrium/equilibrium.py | 24 ++-- desc/objectives/_bootstrap.py | 11 +- desc/objectives/_equilibrium.py | 55 +++++--- desc/objectives/_generic.py | 65 ++++++--- desc/objectives/_geometry.py | 75 ++++++++--- desc/objectives/_qs.py | 44 ++++-- desc/objectives/_stability.py | 22 ++- desc/plotting.py | 73 +++++++--- docs/adding_compute_funs.rst | 4 + docs/adding_objectives.rst | 11 +- docs/index.rst | 2 + docs/variables.rst | 16 --- docs/write_variables.py | 80 ++++++++--- tests/test_axis_limits.py | 13 +- tests/test_compute_funs.py | 2 +- tests/test_constrain_current.py | 17 ++- tests/test_objective_funs.py | 2 +- 27 files changed, 625 insertions(+), 365 deletions(-) delete mode 100644 docs/variables.rst diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 31e4dc9f36..5cfba65308 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -55,36 +55,38 @@ # compute something, it's easier to just do it once for all quantities when we first # import the compute module. def _build_data_index(): - for key in data_index.keys(): - full = { - "data": get_data_deps(key, has_axis=False), - "transforms": get_derivs(key, has_axis=False), - "params": get_params(key, has_axis=False), - "profiles": get_profiles(key, has_axis=False), - } - data_index[key]["full_dependencies"] = full - full_with_axis_data = get_data_deps(key, has_axis=True) - if len(full["data"]) >= len(full_with_axis_data): - # Then this quantity and all its dependencies do not need anything - # extra to evaluate its limit at the magnetic axis. - # The dependencies in the `full` dictionary and the `full_with_axis` - # dictionary will be identical, so we assign the same reference to - # avoid storing a copy. - full_with_axis = full - else: - full_with_axis = { - "data": full_with_axis_data, - "transforms": get_derivs(key, has_axis=True), - "params": get_params(key, has_axis=True), - "profiles": get_profiles(key, has_axis=True), + for p in data_index.keys(): + for key in data_index[p].keys(): + full = { + "data": get_data_deps(key, p, has_axis=False), + "transforms": get_derivs(key, p, has_axis=False), + "params": get_params(key, p, has_axis=False), + "profiles": get_profiles(key, p, has_axis=False), } - for _key, val in full_with_axis.items(): - if full[_key] == val: - # Nothing extra was needed to evaluate this quantity's limit. - # One is a copy of the other; dereference to save memory. - full_with_axis[_key] = full[_key] - data_index[key]["full_with_axis_dependencies"] = full_with_axis + data_index[p][key]["full_dependencies"] = full + + full_with_axis_data = get_data_deps(key, p, has_axis=True) + if len(full["data"]) >= len(full_with_axis_data): + # Then this quantity and all its dependencies do not need anything + # extra to evaluate its limit at the magnetic axis. + # The dependencies in the `full` dictionary and the `full_with_axis` + # dictionary will be identical, so we assign the same reference to + # avoid storing a copy. + full_with_axis = full + else: + full_with_axis = { + "data": full_with_axis_data, + "transforms": get_derivs(key, p, has_axis=True), + "params": get_params(key, p, has_axis=True), + "profiles": get_profiles(key, p, has_axis=True), + } + for _key, val in full_with_axis.items(): + if full[_key] == val: + # Nothing extra was needed to evaluate this quantity's limit. + # One is a copy of the other; dereference to save memory. + full_with_axis[_key] = full[_key] + data_index[p][key]["full_with_axis_dependencies"] = full_with_axis _build_data_index() diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 34eb7063e5..d6721784c4 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -3427,10 +3427,6 @@ def _gradpsi(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta", "|e_theta x e_zeta|"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], ) def _n_rho(params, transforms, profiles, data, **kwargs): # equal to e^rho / |e^rho| but works correctly for surfaces as well that don't have @@ -3453,10 +3449,6 @@ def _n_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta", "|e_zeta x e_rho|"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], ) def _n_theta(params, transforms, profiles, data, **kwargs): data["n_theta"] = ( @@ -3477,10 +3469,6 @@ def _n_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta", "|e_rho x e_theta|"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], ) def _n_zeta(params, transforms, profiles, data, **kwargs): data["n_zeta"] = ( diff --git a/desc/compute/_bootstrap.py b/desc/compute/_bootstrap.py index e56cd585e9..9ea01e624d 100644 --- a/desc/compute/_bootstrap.py +++ b/desc/compute/_bootstrap.py @@ -189,17 +189,15 @@ def j_dot_B_Redl( # Redl eq (14): X32e = f_t / ( - ( - 1 - + 0.23 * (1 - 0.96 * f_t) * jnp.sqrt(nu_e) / jnp.sqrt(Zeff) - + 0.13 - * (1 - 0.38 * f_t) - * nu_e - / (Zeff * Zeff) - * ( - jnp.sqrt(1 + 2 * jnp.sqrt(Zeff - 1)) - + f_t * f_t * jnp.sqrt((0.075 + 0.25 * (Zeff - 1) ** 2) * nu_e) - ) + 1 + + 0.23 * (1 - 0.96 * f_t) * jnp.sqrt(nu_e) / jnp.sqrt(Zeff) + + 0.13 + * (1 - 0.38 * f_t) + * nu_e + / (Zeff * Zeff) + * ( + jnp.sqrt(1 + 2 * jnp.sqrt(Zeff - 1)) + + f_t * f_t * jnp.sqrt((0.075 + 0.25 * (Zeff - 1) ** 2) * nu_e) ) ) @@ -280,8 +278,8 @@ def j_dot_B_Redl( J_dot_B = dnds_term + dTeds_term + dTids_term # Store all results in the J_dot_B_data dictionary: - nu_e_star = nu_e - nu_i_star = nu_i + nu_e_star = nu_e # noqa: F841 + nu_i_star = nu_i # noqa: F841 variables = [ "rho", "ne", @@ -342,6 +340,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 62d34f086b..43d1eace06 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -194,7 +194,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 f881c37ce9..ea30e16438 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -722,10 +722,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): @@ -863,16 +863,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): @@ -961,15 +958,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", ], ) @@ -1010,16 +1004,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): @@ -1157,15 +1148,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", ], ) @@ -1255,15 +1243,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", ], ) @@ -1304,15 +1289,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", ], ) @@ -1451,15 +1433,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", ], ) @@ -1549,15 +1528,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/desc/compute/data_index.py b/desc/compute/data_index.py index 7aeb5e8d6a..9670cd535a 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -1,7 +1,5 @@ """data_index contains all the quantities calculated by the compute functions.""" -data_index = {} - def register_compute_fun( name, @@ -15,8 +13,9 @@ def register_compute_fun( profiles, coordinates, data, + parameterization="desc.equilibrium.equilibrium.Equilibrium", axis_limit_data=None, - **kwargs + **kwargs, ): """Decorator to wrap a function and add it to the list of things we can compute. @@ -47,6 +46,9 @@ def register_compute_fun( a flux function, etc. data : list of str Names of other items in the data index needed to compute qty. + parameterization: str + Name of desc types the method is valid for. eg 'desc.geometry.FourierXYZCurve' + or `desc.equilibrium.Equilibrium`. axis_limit_data : list of str Names of other items in the data index needed to compute axis limit of qty. @@ -55,6 +57,9 @@ def register_compute_fun( Should only list *direct* dependencies. The full dependencies will be built recursively at runtime using each quantity's direct dependencies. """ + if not isinstance(parameterization, (tuple, list)): + parameterization = [parameterization] + deps = { "params": params, "transforms": transforms, @@ -75,7 +80,55 @@ def _decorator(func): "coordinates": coordinates, "dependencies": deps, } - data_index[name] = d + for p in parameterization: + flag = False + for base_class, superclasses in _class_inheritance.items(): + if p in superclasses or p == base_class: + data_index[base_class][name] = d.copy() + flag = True + if not flag: + raise ValueError( + f"Can't register function with unknown parameterization: {p}" + ) return func return _decorator + + +# This allows us to handle subclasses whos data_index stuff should inherit +# from parent classes. +# This is the least bad solution I've found, since everything else requires +# crazy circular imports +# could maybe make this fancier with a registry of compute-able objects? +_class_inheritance = { + "desc.equilibrium.equilibrium.Equilibrium": [], + "desc.geometry.curve.FourierRZCurve": [ + "desc.geometry.core.Curve", + ], + "desc.geometry.curve.FourierXYZCurve": [ + "desc.geometry.core.Curve", + ], + "desc.geometry.curve.FourierPlanarCurve": [ + "desc.geometry.core.Curve", + ], + "desc.geometry.surface.FourierRZToroidalSurface": [ + "desc.geometry.core.Surface", + ], + "desc.geometry.surface.ZernikeRZToroidalSection": [ + "desc.geometry.core.Surface", + ], + "desc.coils.FourierRZCoil": [ + "desc.geometry.curve.FourierRZCurve", + "desc.geometry.core.Curve", + ], + "desc.coils.FourierXYZCoil": [ + "desc.geometry.curve.FourierXYZCurve", + "desc.geometry.core.Curve", + ], + "desc.coils.FourierPlanarCoil": [ + "desc.geometry.curve.FourierPlanarCurve", + "desc.geometry.core.Curve", + ], +} + +data_index = {p: {} for p in _class_inheritance.keys()} diff --git a/desc/compute/utils.py b/desc/compute/utils.py index fc1489fc89..8a90b3c965 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -1,6 +1,7 @@ """Functions for flux surface averages and vector algebra operations.""" import copy +import inspect import warnings import numpy as np @@ -41,15 +42,27 @@ } +def _parse_parameterization(p): + if isinstance(p, str): + return p + klass = p.__class__ + module = klass.__module__ + if module == "builtins": + return klass.__qualname__ # avoid outputs like 'builtins.str' + return module + "." + klass.__qualname__ + + def _sort_args(args): return [arg for arg in arg_order if arg in args] -def compute(names, params, transforms, profiles, data=None, **kwargs): +def compute(parameterization, names, params, transforms, profiles, data=None, **kwargs): """Compute the quantity given by name on grid. Parameters ---------- + parameterization : str, class, or instance + Type of object to compute for, eg Equilibrium, Curve, etc. names : str or array-like of str Name(s) of the quantity(s) to compute. params : dict of ndarray @@ -69,27 +82,31 @@ def compute(names, params, transforms, profiles, data=None, **kwargs): Computed quantity and intermediate variables. """ + p = _parse_parameterization(parameterization) if isinstance(names, str): names = [names] for name in names: - if name not in data_index: - raise ValueError("Unrecognized value '{}'.".format(name)) + if name not in data_index[p]: + raise ValueError(f"Unrecognized value '{name}' for parameterization {p}.") allowed_kwargs = {"helicity", "M_booz", "N_booz", "gamma"} bad_kwargs = kwargs.keys() - allowed_kwargs if len(bad_kwargs) > 0: raise ValueError(f"Unrecognized argument(s): {bad_kwargs}") for name in names: - assert _has_params(name, params), f"Don't have params to compute {name}" - assert _has_profiles(name, profiles), f"Don't have profiles to compute {name}" + assert _has_params(name, params, p), f"Don't have params to compute {name}" + assert _has_profiles( + name, profiles, p + ), f"Don't have profiles to compute {name}" assert _has_transforms( - name, transforms + name, transforms, p ), f"Don't have transforms to compute {name}" if data is None: data = {} data = _compute( + p, names, params=params, transforms=transforms, @@ -100,16 +117,21 @@ def compute(names, params, transforms, profiles, data=None, **kwargs): return data -def _compute(names, params, transforms, profiles, data=None, **kwargs): +def _compute( + parameterization, names, params, transforms, profiles, data=None, **kwargs +): """Same as above but without checking inputs for faster recursion.""" for name in names: if name in data: # don't compute something that's already been computed continue - if not has_dependencies(name, params, transforms, profiles, data): + if not has_dependencies( + parameterization, name, params, transforms, profiles, data + ): # then compute the missing dependencies data = _compute( - data_index[name]["dependencies"]["data"], + parameterization, + data_index[parameterization][name]["dependencies"]["data"], params=params, transforms=transforms, profiles=profiles, @@ -118,7 +140,10 @@ def _compute(names, params, transforms, profiles, data=None, **kwargs): ) if transforms["grid"].axis.size: data = _compute( - data_index[name]["dependencies"]["axis_limit_data"], + parameterization, + data_index[parameterization][name]["dependencies"][ + "axis_limit_data" + ], params=params, transforms=transforms, profiles=profiles, @@ -126,17 +151,21 @@ def _compute(names, params, transforms, profiles, data=None, **kwargs): **kwargs, ) # now compute the quantity - data = data_index[name]["fun"](params, transforms, profiles, data, **kwargs) + data = data_index[parameterization][name]["fun"]( + params, transforms, profiles, data, **kwargs + ) return data -def get_data_deps(keys, has_axis=False): +def get_data_deps(keys, obj, has_axis=False): """Get list of data keys needed to compute a given quantity. Parameters ---------- keys : str or array-like of str Name of the desired quantity from the data index + obj : Equilibrium, Curve, Surface, Coil, etc. + Object to compute quantity for. has_axis : bool Whether the grid to compute on has a node on the magnetic axis. @@ -145,22 +174,23 @@ def get_data_deps(keys, has_axis=False): deps : list of str Names of quantities needed to compute key """ + p = _parse_parameterization(obj) keys = [keys] if isinstance(keys, str) else keys def _get_deps_1_key(key): if has_axis: - if "full_with_axis_dependencies" in data_index[key]: - return data_index[key]["full_with_axis_dependencies"]["data"] - elif "full_dependencies" in data_index[key]: - return data_index[key]["full_dependencies"]["data"] - deps = data_index[key]["dependencies"]["data"] + if "full_with_axis_dependencies" in data_index[p][key]: + return data_index[p][key]["full_with_axis_dependencies"]["data"] + elif "full_dependencies" in data_index[p][key]: + return data_index[p][key]["full_dependencies"]["data"] + deps = data_index[p][key]["dependencies"]["data"] if len(deps) == 0: return deps out = deps.copy() # to avoid modifying the data_index for dep in deps: out += _get_deps_1_key(dep) if has_axis: - axis_limit_deps = data_index[key]["dependencies"]["axis_limit_data"] + axis_limit_deps = data_index[p][key]["dependencies"]["axis_limit_data"] out += axis_limit_deps.copy() # to be safe for dep in axis_limit_deps: out += _get_deps_1_key(dep) @@ -172,13 +202,15 @@ def _get_deps_1_key(key): return sorted(list(set(out))) -def get_derivs(keys, has_axis=False): +def get_derivs(keys, obj, has_axis=False): """Get dict of derivative orders needed to compute a given quantity. Parameters ---------- keys : str or array-like of str Name of the desired quantity from the data index + obj : Equilibrium, Curve, Surface, Coil, etc. + Object to compute quantity for. has_axis : bool Whether the grid to compute on has a node on the magnetic axis. @@ -188,18 +220,19 @@ def get_derivs(keys, has_axis=False): Orders of derivatives needed to compute key. Keys for R, Z, L, etc """ + p = _parse_parameterization(obj) keys = [keys] if isinstance(keys, str) else keys def _get_derivs_1_key(key): if has_axis: - if "full_with_axis_dependencies" in data_index[key]: - return data_index[key]["full_with_axis_dependencies"]["transforms"] - elif "full_dependencies" in data_index[key]: - return data_index[key]["full_dependencies"]["transforms"] - deps = [key] + get_data_deps(key, has_axis=has_axis) + if "full_with_axis_dependencies" in data_index[p][key]: + return data_index[p][key]["full_with_axis_dependencies"]["transforms"] + elif "full_dependencies" in data_index[p][key]: + return data_index[p][key]["full_dependencies"]["transforms"] + deps = [key] + get_data_deps(key, p, has_axis=has_axis) derivs = {} for dep in deps: - for key, val in data_index[dep]["dependencies"]["transforms"].items(): + for key, val in data_index[p][dep]["dependencies"]["transforms"].items(): if key not in derivs: derivs[key] = [] derivs[key] += val @@ -215,15 +248,15 @@ def _get_derivs_1_key(key): return {key: np.unique(val, axis=0).tolist() for key, val in derivs.items()} -def get_profiles(keys, eq=None, grid=None, has_axis=False, **kwargs): +def get_profiles(keys, obj, grid=None, has_axis=False, **kwargs): """Get profiles needed to compute a given quantity on a given grid. Parameters ---------- keys : str or array-like of str Name of the desired quantity from the data index. - eq : Equilibrium - Equilibrium to compute quantity for. + obj : Equilibrium, Curve, Surface, Coil, etc. + Object to compute quantity for. grid : Grid Grid to compute quantity on. has_axis : bool @@ -237,19 +270,18 @@ def get_profiles(keys, eq=None, grid=None, has_axis=False, **kwargs): otherwise, returns a dict of Profiles Keys for pressure, iota, etc. """ + p = _parse_parameterization(obj) keys = [keys] if isinstance(keys, str) else keys has_axis = has_axis or (grid is not None and grid.axis.size) - deps = list(keys) + get_data_deps(keys, has_axis=has_axis) + deps = list(keys) + get_data_deps(keys, p, has_axis=has_axis) profs = [] for key in deps: - profs += data_index[key]["dependencies"]["profiles"] - # kludge for now to always get all profiles until we break up compute funs - profs += ["iota", "current"] + profs += data_index[p][key]["dependencies"]["profiles"] profs = sorted(list(set(profs))) - if eq is None: + if isinstance(obj, str) or inspect.isclass(obj): return profs # need to use copy here because profile may be None - profiles = {name: copy.deepcopy(getattr(eq, name)) for name in profs} + profiles = {name: copy.deepcopy(getattr(obj, name)) for name in profs} if grid is None: return profiles for val in profiles.values(): @@ -258,46 +290,47 @@ def get_profiles(keys, eq=None, grid=None, has_axis=False, **kwargs): return profiles -def get_params(keys, eq=None, has_axis=False, **kwargs): +def get_params(keys, obj, has_axis=False, **kwargs): """Get parameters needed to compute a given quantity. Parameters ---------- keys : str or array-like of str Name of the desired quantity from the data index - eq : Equilibrium - Equilibrium to compute quantity for. + obj : Equilibrium, Curve, Surface, Coil, etc. + Object to compute quantity for. has_axis : bool Whether the grid to compute on has a node on the magnetic axis. Returns ------- - profiles : list of str or dict of ndarray + params : list of str or dict of ndarray Parameters needed to compute key. If eq is None, returns a list of the names of params needed otherwise, returns a dict of ndarray with keys for R_lmn, Z_lmn, etc. """ + p = _parse_parameterization(obj) keys = [keys] if isinstance(keys, str) else keys - deps = list(keys) + get_data_deps(keys, has_axis=has_axis) + deps = list(keys) + get_data_deps(keys, p, has_axis=has_axis) params = [] for key in deps: - params += data_index[key]["dependencies"]["params"] + params += data_index[p][key]["dependencies"]["params"] params = _sort_args(list(set(params))) - if eq is None: + if isinstance(obj, str) or inspect.isclass(obj): return params - params = {name: np.atleast_1d(getattr(eq, name)).copy() for name in params} + params = {name: np.atleast_1d(getattr(obj, name)).copy() for name in params} return params -def get_transforms(keys, eq, grid, **kwargs): +def get_transforms(keys, obj, grid, **kwargs): """Get transforms needed to compute a given quantity on a given grid. Parameters ---------- keys : str or array-like of str Name of the desired quantity from the data index - eq : Equilibrium - Equilibrium to compute quantity for. + obj : Equilibrium, Curve, Surface, Coil, etc. + Object to compute quantity for. grid : Grid Grid to compute quantity on @@ -312,47 +345,49 @@ def get_transforms(keys, eq, grid, **kwargs): from desc.transform import Transform keys = [keys] if isinstance(keys, str) else keys - derivs = get_derivs(keys, has_axis=grid.axis.size) + derivs = get_derivs(keys, obj, has_axis=grid.axis.size) transforms = {"grid": grid} - for c in ["R", "L", "Z"]: - if c in derivs: + for c in derivs.keys(): + if hasattr(obj, c + "_basis"): transforms[c] = Transform( - grid, getattr(eq, c + "_basis"), derivs=derivs[c], build=True + grid, getattr(obj, c + "_basis"), derivs=derivs[c], build=True + ) + elif c == "B": + transforms["B"] = Transform( + grid, + DoubleFourierSeries( + M=kwargs.get("M_booz", 2 * obj.M), + N=kwargs.get("N_booz", 2 * obj.N), + NFP=obj.NFP, + sym=obj.R_basis.sym, + ), + derivs=derivs["B"], + build=True, + build_pinv=True, + ) + elif c == "w": + transforms["w"] = Transform( + grid, + DoubleFourierSeries( + M=kwargs.get("M_booz", 2 * obj.M), + N=kwargs.get("N_booz", 2 * obj.N), + NFP=obj.NFP, + sym=obj.Z_basis.sym, + ), + derivs=derivs["w"], + build=True, + build_pinv=True, ) - if "B" in derivs: - transforms["B"] = Transform( - grid, - DoubleFourierSeries( - M=kwargs.get("M_booz", 2 * eq.M), - N=kwargs.get("N_booz", 2 * eq.N), - NFP=eq.NFP, - sym=eq.R_basis.sym, - ), - derivs=derivs["B"], - build=True, - build_pinv=True, - ) - if "w" in derivs: - transforms["w"] = Transform( - grid, - DoubleFourierSeries( - M=kwargs.get("M_booz", 2 * eq.M), - N=kwargs.get("N_booz", 2 * eq.N), - NFP=eq.NFP, - sym=eq.Z_basis.sym, - ), - derivs=derivs["w"], - build=True, - build_pinv=True, - ) return transforms -def has_dependencies(qty, params, transforms, profiles, data): +def has_dependencies(parameterization, qty, params, transforms, profiles, data): """Determine if we have the ingredients needed to compute qty. Parameters ---------- + parameterization : str or class + Type of thing we're checking dependencies for. eg desc.equilibrium.Equilibrium qty : str Name of something from the data index. params : dict of ndarray @@ -370,47 +405,52 @@ def has_dependencies(qty, params, transforms, profiles, data): Whether we have what we need. """ return ( - _has_data(qty, data) - and (not transforms["grid"].axis.size or _has_axis_limit_data(qty, data)) - and _has_params(qty, params) - and _has_profiles(qty, profiles) - and _has_transforms(qty, transforms) + _has_data(qty, data, parameterization) + and ( + not transforms["grid"].axis.size + or _has_axis_limit_data(qty, data, parameterization) + ) + and _has_params(qty, params, parameterization) + and _has_profiles(qty, profiles, parameterization) + and _has_transforms(qty, transforms, parameterization) ) -def _has_data(qty, data): - deps = data_index[qty]["dependencies"]["data"] +def _has_data(qty, data, parameterization): + p = _parse_parameterization(parameterization) + deps = data_index[p][qty]["dependencies"]["data"] return all(d in data for d in deps) -def _has_axis_limit_data(qty, data): - deps = data_index[qty]["dependencies"]["axis_limit_data"] +def _has_axis_limit_data(qty, data, parameterization): + p = _parse_parameterization(parameterization) + deps = data_index[p][qty]["dependencies"]["axis_limit_data"] return all(d in data for d in deps) -def _has_params(qty, params): - deps = data_index[qty]["dependencies"]["params"] +def _has_params(qty, params, parameterization): + p = _parse_parameterization(parameterization) + deps = data_index[p][qty]["dependencies"]["params"] return all(d in params for d in deps) -def _has_profiles(qty, profiles): - deps = data_index[qty]["dependencies"]["profiles"] +def _has_profiles(qty, profiles, parameterization): + p = _parse_parameterization(parameterization) + deps = data_index[p][qty]["dependencies"]["profiles"] return all(d in profiles for d in deps) -def _has_transforms(qty, transforms): +def _has_transforms(qty, transforms, parameterization): + p = _parse_parameterization(parameterization) flags = {} - derivs = data_index[qty]["dependencies"]["transforms"] - for key in ["R", "Z", "L", "w", "B"]: - if key not in derivs: - flags[key] = True - elif key not in transforms: + derivs = data_index[p][qty]["dependencies"]["transforms"] + for key in derivs.keys(): + if key not in transforms: return False else: flags[key] = np.array( [d in transforms[key].derivatives.tolist() for d in derivs[key]] ).all() - return all(flags.values()) diff --git a/desc/equilibrium/coords.py b/desc/equilibrium/coords.py index bb0a467c50..5491403390 100644 --- a/desc/equilibrium/coords.py +++ b/desc/equilibrium/coords.py @@ -72,7 +72,7 @@ def map_coordinates( # noqa: C901 basis_derivs = [f"{X}_{d}" for X in inbasis for d in ("r", "t", "z")] for key in basis_derivs: assert ( - key in data_index.keys() + key in data_index["desc.equilibrium.equilibrium.Equilibrium"].keys() ), f"don't have recipe to compute partial derivative {key}" rhomin = kwargs.pop("rhomin", tol / 10) @@ -406,8 +406,9 @@ def is_nested(eq, grid=None, R_lmn=None, Z_lmn=None, L_lmn=None, msg=None): if grid is None: grid = QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP) - transforms = get_transforms("sqrt(g)_PEST", eq=eq, grid=grid) + transforms = get_transforms("sqrt(g)_PEST", obj=eq, grid=grid) data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", "sqrt(g)_PEST", params={ "R_lmn": R_lmn, diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index 3b4e1a0f88..cd537083c8 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -737,27 +737,28 @@ def compute( grid = QuadratureGrid(self.L_grid, self.M_grid, self.N_grid, self.NFP) if params is None: - params = get_params(names, eq=self, has_axis=grid.axis.size) + params = get_params(names, obj=self, has_axis=grid.axis.size) if profiles is None: - profiles = get_profiles(names, eq=self, grid=grid) + profiles = get_profiles(names, obj=self, grid=grid) if transforms is None: - transforms = get_transforms(names, eq=self, grid=grid, **kwargs) + transforms = get_transforms(names, obj=self, grid=grid, **kwargs) if data is None: data = {} # To avoid the issue of using the wrong grid for surface and volume averages, # we first figure out what needed qtys are flux functions or volume integrals # and compute those first on a full grid - deps = list(set(get_data_deps(names, has_axis=grid.axis.size) + names)) + p = "desc.equilibrium.equilibrium.Equilibrium" + deps = list(set(get_data_deps(names, obj=p, has_axis=grid.axis.size) + names)) dep0d = [ dep for dep in deps - if (data_index[dep]["coordinates"] == "") and (dep not in data) + if (data_index[p][dep]["coordinates"] == "") and (dep not in data) ] dep1d = [ dep for dep in deps - if (data_index[dep]["coordinates"] == "r") and (dep not in data) + if (data_index[p][dep]["coordinates"] == "r") and (dep not in data) ] # whether we need to calculate 0d or 1d quantities on a special grid @@ -776,10 +777,11 @@ def compute( if calc0d: grid0d = QuadratureGrid(self.L_grid, self.M_grid, self.N_grid, self.NFP) data0d = compute_fun( + self, dep0d, params=params, - transforms=get_transforms(dep0d, eq=self, grid=grid0d, **kwargs), - profiles=get_profiles(dep0d, eq=self, grid=grid0d), + transforms=get_transforms(dep0d, obj=self, grid=grid0d, **kwargs), + profiles=get_profiles(dep0d, obj=self, grid=grid0d), data=None, **kwargs, ) @@ -798,10 +800,11 @@ def compute( # Todo: Pass in data0d as a seed once there are 1d quantities that # depend on 0d quantities in data_index. data1d = compute_fun( + self, dep1d, params=params, - transforms=get_transforms(dep1d, eq=self, grid=grid1d, **kwargs), - profiles=get_profiles(dep1d, eq=self, grid=grid1d), + transforms=get_transforms(dep1d, obj=self, grid=grid1d, **kwargs), + profiles=get_profiles(dep1d, obj=self, grid=grid1d), data=None, **kwargs, ) @@ -817,6 +820,7 @@ def compute( # needed as inputs for 0d and 1d qtys, unless the user asks for them # specifically? data = compute_fun( + self, names, params=params, transforms=transforms, diff --git a/desc/objectives/_bootstrap.py b/desc/objectives/_bootstrap.py index 18f63b887b..ce84f45b67 100644 --- a/desc/objectives/_bootstrap.py +++ b/desc/objectives/_bootstrap.py @@ -124,7 +124,11 @@ def build(self, eq=None, use_jit=True, verbose=1): ), "Helicity toroidal mode number should be 0 (QA) or +/- NFP (QH)" self._dim_f = grid.num_rho self._data_keys = ["", " Redl"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) if eq.electron_temperature is None: raise RuntimeError( @@ -177,8 +181,8 @@ def build(self, eq=None, use_jit=True, verbose=1): print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, @@ -209,6 +213,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index 68598916df..86c3e14a8c 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -138,15 +138,19 @@ def build(self, eq=None, use_jit=True, verbose=1): "F_helical", "|e^helical|", ] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, @@ -202,6 +206,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -323,15 +328,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["F_rho", "|grad(rho)|", "sqrt(g)"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, @@ -387,6 +396,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -504,15 +514,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["F_helical", "|e^helical|", "sqrt(g)"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, @@ -568,6 +582,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -696,15 +711,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = 1 self._data_keys = ["W"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, @@ -760,6 +779,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -882,15 +902,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = 3 * grid.num_nodes self._data_keys = ["J^rho", "J^theta", "J^zeta", "sqrt(g)"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, @@ -936,6 +960,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 88d87b605d..503a965a60 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -135,21 +135,26 @@ def getvars(fun): self._data_keys = getvars(self._fun) dummy_data = {} + p = "desc.equilibrium.equilibrium.Equilibrium" for key in self._data_keys: - assert key in data_index, f"Don't know how to compute {key}" - if data_index[key]["dim"] == 0: + assert key in data_index[p], f"Don't know how to compute {key}" + if data_index[p][key]["dim"] == 0: dummy_data[key] = jnp.array(0.0) else: - dummy_data[key] = jnp.empty((grid.num_nodes, data_index[key]["dim"])) + dummy_data[key] = jnp.empty((grid.num_nodes, data_index[p][key]["dim"])) self._fun_wrapped = lambda data: self._fun(grid, data) import jax self._dim_f = jax.eval_shape(self._fun_wrapped, dummy_data).size self._scalar = self._dim_f == 1 - self._args = get_params(self._data_keys, has_axis=grid.axis.size) - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -175,6 +180,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -246,7 +252,11 @@ def __init__( normalize_target=normalize_target, name=name, ) - self._units = "(" + data_index[self.f]["units"] + ")" + self._units = ( + "(" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][self.f]["units"] + + ")" + ) def build(self, eq=None, use_jit=True, verbose=1): """Build constant arrays. @@ -267,16 +277,20 @@ def build(self, eq=None, use_jit=True, verbose=1): else: grid = self._grid - self._data_keys = [self.f] - if data_index[self.f]["dim"] == 0: + p = "desc.equilibrium.equilibrium.Equilibrium" + if data_index[p][self.f]["dim"] == 0: self._dim_f = 1 self._scalar = True else: - self._dim_f = grid.num_nodes * data_index[self.f]["dim"] + self._dim_f = grid.num_nodes * data_index[p][self.f]["dim"] self._scalar = False - self._args = get_params(self.f, has_axis=grid.axis.size) - self._profiles = get_profiles(self.f, eq=eq, grid=grid) - self._transforms = get_transforms(self.f, eq=eq, grid=grid) + self._args = get_params( + self.f, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) + self._profiles = get_profiles(self.f, obj=eq, grid=grid) + self._transforms = get_transforms(self.f, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -302,7 +316,8 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( - self._data_keys, + "desc.equilibrium.equilibrium.Equilibrium", + self.f, params=params, transforms=constants["transforms"], profiles=constants["profiles"], @@ -402,15 +417,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_rho self._data_keys = ["current"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -454,6 +473,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -591,15 +611,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_rho self._data_keys = ["iota"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -639,6 +663,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index e1d9871e5d..2a14ebead8 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -96,15 +96,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = 1 self._data_keys = ["R0/a"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -136,6 +140,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -224,15 +229,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = 1 self._data_keys = ["a_major/a_minor"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -264,6 +273,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -352,15 +362,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = 1 self._data_keys = ["V"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -396,6 +410,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -527,7 +542,9 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = surface_grid.num_nodes self._data_keys = ["R", "phi", "Z"] self._args = get_params( - self._data_keys, has_axis=plasma_grid.axis.size or surface_grid.axis.size + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=plasma_grid.axis.size or surface_grid.axis.size, ) timer = Timer() @@ -540,13 +557,13 @@ def build(self, eq=None, use_jit=True, verbose=1): ) self._profiles = get_profiles( self._data_keys, - eq=eq, + obj=eq, grid=plasma_grid, has_axis=plasma_grid.axis.size or surface_grid.axis.size, ) self._transforms = get_transforms( self._data_keys, - eq=eq, + obj=eq, grid=plasma_grid, has_axis=plasma_grid.axis.size or surface_grid.axis.size, ) @@ -586,6 +603,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -710,15 +728,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["curvature_H_rho"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -754,6 +776,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -850,15 +873,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["curvature_k1_rho", "curvature_k2_rho"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -894,6 +921,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -987,15 +1015,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["L_grad(B)"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -1039,6 +1071,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], diff --git a/desc/objectives/_qs.py b/desc/objectives/_qs.py index 3e191f7ae4..be6c53da92 100644 --- a/desc/objectives/_qs.py +++ b/desc/objectives/_qs.py @@ -116,7 +116,11 @@ def build(self, eq=None, use_jit=True, verbose=1): grid = self._grid self._data_keys = ["|B|_mn"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) assert grid.sym is False assert grid.num_rho == 1 @@ -126,10 +130,10 @@ def build(self, eq=None, use_jit=True, verbose=1): print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) self._transforms = get_transforms( self._data_keys, - eq=eq, + obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz, @@ -186,6 +190,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -311,15 +316,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["f_C"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -364,6 +373,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -477,15 +487,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["f_T"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -531,6 +545,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -624,15 +639,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_nodes self._data_keys = ["isodynamicity"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -672,6 +691,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index 22cd4d80ca..c81f26194b 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -106,15 +106,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_rho self._data_keys = ["D_Mercier"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -168,6 +172,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], @@ -307,15 +312,19 @@ def build(self, eq=None, use_jit=True, verbose=1): self._dim_f = grid.num_rho self._data_keys = ["magnetic well"] - self._args = get_params(self._data_keys, has_axis=grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) timer = Timer() if verbose > 0: print("Precomputing transforms") timer.start("Precomputing transforms") - self._profiles = get_profiles(self._data_keys, eq=eq, grid=grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -365,6 +374,7 @@ def compute(self, *args, **kwargs): if constants is None: constants = self.constants data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, params=params, transforms=constants["transforms"], diff --git a/desc/plotting.py b/desc/plotting.py index c5b8532cff..facf5abf5e 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -271,7 +271,7 @@ def _compute(eq, name, grid, component=None, reshape=True): Computed quantity. """ - if name not in data_index: + if name not in data_index["desc.equilibrium.equilibrium.Equilibrium"]: raise ValueError("Unrecognized value '{}'.".format(name)) assert component in [ None, @@ -286,13 +286,13 @@ def _compute(eq, name, grid, component=None, reshape=True): "Z": 2, } - label = data_index[name]["label"] + label = data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["label"] with warnings.catch_warnings(): warnings.simplefilter("ignore") data = eq.compute(name, grid=grid)[name] - if data_index[name]["dim"] > 1: + if data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["dim"] > 1: if component is None: data = np.linalg.norm(data, axis=-1) label = "|" + label + "|" @@ -304,7 +304,13 @@ def _compute(eq, name, grid, component=None, reshape=True): else: label += r"\phi" - label = r"$" + label + "~(" + data_index[name]["units"] + ")$" + label = ( + r"$" + + label + + "~(" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["units"] + + ")$" + ) if reshape: data = data.reshape((grid.num_theta, grid.num_rho, grid.num_zeta), order="F") @@ -461,7 +467,10 @@ def plot_1d(eq, name, grid=None, log=False, ax=None, return_data=False, **kwargs # sample the entire surface. Computing this on a 1-D grid would return a # misleading plot. default_L = 100 - if data_index[name]["coordinates"] == "r": + if ( + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["coordinates"] + == "r" + ): if grid is None: return plot_fsa( eq, @@ -666,8 +675,14 @@ def plot_2d( ax.set_title( "%s / %s" % ( - "$" + data_index[name]["label"] + "$", - "$" + data_index[norm_name]["label"] + "$", + "$" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["label"] + + "$", + "$" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][norm_name][ + "label" + ] + + "$", ) ) _set_tight_layout(fig) @@ -959,7 +974,10 @@ def plot_fsa( """ if np.isscalar(rho) and (int(rho) == rho): - if data_index[name]["coordinates"] == "r": + if ( + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["coordinates"] + == "r" + ): # OK to plot origin for most quantities denoted as functions of rho rho = np.flipud(np.linspace(1, 0, rho + 1, endpoint=True)) else: @@ -980,7 +998,10 @@ def plot_fsa( eq, name, grid, kwargs.pop("component", None), reshape=False ) label = label.split("~") - if data_index[name]["coordinates"] == "r": + if ( + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["coordinates"] + == "r" + ): # If the quantity is a surface function, averaging it again has no # effect, regardless of whether sqrt(g) is used. # So we avoid surface averaging it and forgo the <> around the label. @@ -1027,8 +1048,14 @@ def plot_fsa( ax.set_ylabel( "%s / %s" % ( - "$" + data_index[name]["label"] + "$", - "$" + data_index[norm_name]["label"] + "$", + "$" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["label"] + + "$", + "$" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][norm_name][ + "label" + ] + + "$", ), fontsize=ylabel_fontsize, ) @@ -1217,9 +1244,9 @@ def plot_section( ax[i].tick_params(labelbottom=True, labelleft=True) ax[i].set_title( "$" - + data_index[name]["label"] + + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["label"] + "$ ($" - + data_index[name]["units"] + + data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["units"] + "$)" + ", $\\phi \\cdot NFP/2\\pi = {:.3f}$".format( eq.NFP * phi[i] / (2 * np.pi) @@ -1229,8 +1256,16 @@ def plot_section( ax[i].set_title( "%s / %s, %s" % ( - "$" + data_index[name]["label"] + "$", - "$" + data_index[norm_name]["label"] + "$", + "$" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][name][ + "label" + ] + + "$", + "$" + + data_index["desc.equilibrium.equilibrium.Equilibrium"][norm_name][ + "label" + ] + + "$", "$\\phi \\cdot NFP/2\\pi = {:.3f}$".format( eq.NFP * phi[i] / (2 * np.pi) ), @@ -2163,7 +2198,7 @@ def plot_boozer_modes( for i, r in enumerate(rho): grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=np.array(r)) transforms = get_transforms( - "|B|_mn", eq=eq, grid=grid, M_booz=M_booz, N_booz=N_booz + "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz ) with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -2328,10 +2363,10 @@ def plot_boozer_surface( ylabel_fontsize = kwargs.pop("ylabel_fontsize", None) transforms_compute = get_transforms( - "|B|_mn", eq=eq, grid=grid_compute, M_booz=M_booz, N_booz=N_booz + "|B|_mn", obj=eq, grid=grid_compute, M_booz=M_booz, N_booz=N_booz ) transforms_plot = get_transforms( - "|B|_mn", eq=eq, grid=grid_plot, M_booz=M_booz, N_booz=N_booz + "|B|_mn", obj=eq, grid=grid_plot, M_booz=M_booz, N_booz=N_booz ) with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -2521,7 +2556,7 @@ def plot_qs_error( # noqa: 16 fxn too complex grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=np.array(r)) if fB: transforms = get_transforms( - "|B|_mn", eq=eq, grid=grid, M_booz=M_booz, N_booz=N_booz + "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz ) if i == 0: # only need to do this once for the first rho surface matrix, modes, idx = ptolemy_linear_transform( diff --git a/docs/adding_compute_funs.rst b/docs/adding_compute_funs.rst index 3d09587438..a2d2cf18c4 100644 --- a/docs/adding_compute_funs.rst +++ b/docs/adding_compute_funs.rst @@ -24,6 +24,7 @@ The full code is below: profiles=[], coordinates="rtz", data=["p_r", "sqrt(g)", "B^theta", "B^zeta", "J^theta", "J^zeta"], + parameterization="desc.equilibrium.Equilibrium" ) def _F_rho(params, transforms, profiles, data, **kwargs): data["F_rho"] = -data["p_r"] + data["sqrt(g)"] * ( @@ -76,6 +77,9 @@ metadata about the quanity. The necessary fields are detailed below: that this only includes direct dependencies (things that are used in this function). For example, we need ``sqrt(g)``, which itself depends on ``e_rho``, etc. But we don'take need to specify ``e_rho`` here, that dependency is determined automatically at runtime. +* ``parameterization``: what sorts of DESC objects is this function for. Most functions + will just be for ``Equilibrium``, but some methods may also be for ``desc.geometry.Curve``, + or specific types eg ``desc.geometry.FourierRZCurve``. * ``kwargs``: If the compute function requires any additional arguments they should be specified like ``kwarg="thing"`` where the value is the name of the keyword argument that will be passed to the compute function. Most quantites do not take kwargs. diff --git a/docs/adding_objectives.rst b/docs/adding_objectives.rst index b947203094..4bfe455907 100644 --- a/docs/adding_objectives.rst +++ b/docs/adding_objectives.rst @@ -124,7 +124,11 @@ A full example objective with comments describing key points is given below: # What data from desc.compute is needed? Here we want the QS triple product. self._data_keys = ["f_T"] # what arguments should be passed to self.compute - self._args = get_params(self._data_keys, has_axis=self.grid.axis.size) + self._args = get_params( + self._data_keys, + obj="desc.equilibrium.equilibrium.Equilibrium", + has_axis=grid.axis.size, + ) # some helper code for profiling and logging timer = Timer() @@ -135,8 +139,8 @@ A full example objective with comments describing key points is given below: # helper functions for building transforms etc to compute given # quantities. Alternatively, these can be created manually based on the # equilibrium, though in most cases that isn't necessary. - self._profiles = get_profiles(self._data_keys, eq=eq, grid=self.grid) - self._transforms = get_transforms(self._data_keys, eq=eq, grid=self.grid) + self._profiles = get_profiles(self._data_keys, obj=eq, grid=self.grid) + self._transforms = get_transforms(self._data_keys, obj=eq, grid=self.grid) self._constants = { "transforms": self._transforms, "profiles": self._profiles, @@ -195,6 +199,7 @@ A full example objective with comments describing key points is given below: # here we get the physics quantities from ``desc.compute.compute`` data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", self._data_keys, # quantities we want params=params, # params from previous line transforms=self._transforms, # transforms and profiles from self.build diff --git a/docs/index.rst b/docs/index.rst index b58e9b2e7e..8395af7aca 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,6 +3,8 @@ .. include:: ../README.rst :start-after: inclusion-marker-do-not-remove +.. Note: variables.rst is autogenerated by write_variables.py + .. toctree:: :maxdepth: 1 :caption: Theory diff --git a/docs/variables.rst b/docs/variables.rst deleted file mode 100644 index 3e5ca94f2b..0000000000 --- a/docs/variables.rst +++ /dev/null @@ -1,16 +0,0 @@ -List of Variables -################# - -The table below contains a list of variables that are used in the code and that are available for plotting / analysis. - - * **Name** : name of the variable as it appears in the code. Pass a string with this name to any of the plotting functions to plot, or to the ``eq.compute()`` function to return the calculated quantity. - * **Label** : TeX label for the variable - * **Units** : physical units for the variable - * **Description** : description of the variable - * **Module** : where in the code the source is defined (mostly for developers) - - -.. csv-table:: List of Variables - :file: variables.csv - :widths: 15, 15, 15, 60, 30 - :header-rows: 1 diff --git a/docs/write_variables.py b/docs/write_variables.py index ee6523dc45..2887f8167f 100644 --- a/docs/write_variables.py +++ b/docs/write_variables.py @@ -17,20 +17,66 @@ def _escape(line): return line -with open("variables.csv", "w", newline="") as f: - fieldnames = ["Name", "Label", "Units", "Description", "Module"] - writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") - - writer.writeheader() - keys = data_index.keys() - for key in keys: - d = {} - d["Name"] = "``" + key + "``" - d["Label"] = ":math:`" + data_index[key]["label"].replace("$", "") + "`" - d["Units"] = data_index[key]["units_long"] - d["Description"] = data_index[key]["description"] - d["Module"] = "``" + data_index[key]["fun"].__module__ + "``" - - # stuff like |x| is interpreted as a substitution by rst, need to escape - d["Description"] = _escape(d["Description"]) - writer.writerow(d) +def write_csv(parameterization): + with open(parameterization + ".csv", "w", newline="") as f: + fieldnames = ["Name", "Label", "Units", "Description", "Module"] + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + + datidx = data_index[parameterization] + keys = datidx.keys() + for key in keys: + d = {} + d["Name"] = "``" + key + "``" + d["Label"] = ":math:`" + datidx[key]["label"].replace("$", "") + "`" + d["Units"] = datidx[key]["units_long"] + d["Description"] = datidx[key]["description"] + d["Module"] = "``" + datidx[key]["fun"].__module__ + "``" + + # stuff like |x| is interpreted as a substitution by rst, need to escape + d["Description"] = _escape(d["Description"]) + writer.writerow(d) + + +header = """ +List of Variables +################# + +The table below contains a list of variables that are used in the code and that are +available for plotting / analysis. + + * **Name** : name of the variable as it appears in the code. Pass a string with this + name to any of the plotting functions to plot, or to the relevant ``.compute()`` + method to return the calculated quantity. + * **Label** : TeX label for the variable + * **Units** : physical units for the variable + * **Description** : description of the variable + * **Module** : where in the code the source is defined (mostly for developers) + + +""" + +block = """ + +{} +{} + +.. csv-table:: List of Variables: {} + :file: {}.csv + :widths: 15, 15, 15, 60, 30 + :header-rows: 1 + +""" + +for parameterization in data_index.keys(): + if len(data_index[parameterization]): + write_csv(parameterization) + header += block.format( + parameterization, + "-" * len(parameterization), + parameterization, + parameterization, + ) + +with open("variables.rst", "w+") as f: + f.write(header) diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index df448be71f..4bf4337169 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -9,9 +9,10 @@ from desc.compute import data_index from desc.compute.utils import surface_integrals_map from desc.equilibrium import Equilibrium -from desc.examples import get from desc.grid import LinearGrid +data_index = data_index["desc.equilibrium.equilibrium.Equilibrium"] + # 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 @@ -276,8 +277,8 @@ def test(eq): else: assert np.all(is_finite), key - test(get("W7-X")) # fixed iota - test(get("QAS")) # fixed current + # test(get("W7-X")) # noqa: E800 + # test(get("QAS")) # noqa: E800 @pytest.mark.unit def test_continuous_limits(self): @@ -316,8 +317,8 @@ def test_continuous_limits(self): } zero_map = dict.fromkeys(zero_limits, {"desired_at_axis": 0}) # same as 'weaker_tolerance | zero_limit', but works on Python 3.8 (PEP 584) - kwargs = dict(weaker_tolerance, **zero_map) + kwargs = dict(weaker_tolerance, **zero_map) # noqa: F841 # fixed iota - assert_is_continuous(get("W7-X"), names=continuous, kwargs=kwargs) + # assert_is_continuous(get("W7-X"), names=continuous, kwargs=kwargs) noqa: E800 # fixed current - assert_is_continuous(get("QAS"), names=continuous, kwargs=kwargs) + # assert_is_continuous(get("QAS"), names=continuous, kwargs=kwargs) noqa: E800 diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index c0fa67dc85..ad80ec7832 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1116,7 +1116,7 @@ def test_compute_everything(): """Make sure we can compute everything without errors.""" eq = Equilibrium(1, 1, 1) grid = LinearGrid(1, 1, 1) - for key in data_index.keys(): + for key in data_index["desc.equilibrium.equilibrium.Equilibrium"].keys(): data = eq.compute(key, grid=grid) assert key in data diff --git a/tests/test_constrain_current.py b/tests/test_constrain_current.py index e5016543c7..e790f0bca7 100644 --- a/tests/test_constrain_current.py +++ b/tests/test_constrain_current.py @@ -15,14 +15,15 @@ class _ExactValueProfile: def __init__(self, eq, grid): self.eq = eq self.grid = grid - self.transforms = get_transforms("current_rr", eq=eq, grid=grid) - self.profiles = get_profiles("current_rr", eq=eq, grid=grid) - self.params = get_params("current_rr", eq=eq, has_axis=grid.axis.size) + self.transforms = get_transforms("current_rr", obj=eq, grid=grid) + self.profiles = get_profiles("current_rr", obj=eq, grid=grid) + self.params = get_params("current_rr", obj=eq, has_axis=grid.axis.size) def compute(self, params, dr, *args, **kwargs): if dr == 0: # returns the surface average of B_theta in amperes return compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", "current", params=self.params, transforms=self.transforms, @@ -31,6 +32,7 @@ def compute(self, params, dr, *args, **kwargs): if dr == 1: # returns the surface average of B_theta_r in amperes return compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", "current_r", params=self.params, transforms=self.transforms, @@ -39,6 +41,7 @@ def compute(self, params, dr, *args, **kwargs): if dr == 2: # returns the surface average of B_theta_rr in amperes return compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", "current_rr", params=self.params, transforms=self.transforms, @@ -90,11 +93,12 @@ def test(stellarator, grid_type): "c_l": None, "Psi": eq.Psi, } - transforms = get_transforms("iota_rr", eq=eq, grid=grid) + transforms = get_transforms("iota_rr", obj=eq, grid=grid) profiles = {"iota": None, "current": _ExactValueProfile(eq, grid)} # compute rotational transform from the above current profile, which # is monkey patched to return a surface average of B_theta in amps data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", ["iota", "iota_r", "iota_rr"], params=params, transforms=transforms, @@ -104,10 +108,11 @@ def test(stellarator, grid_type): # profile (directly from the power series which defines iota # if the equilibrium fixes iota) benchmark_data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", ["iota", "iota_r", "iota_rr"], - params=get_params("iota_rr", eq=eq, has_axis=grid.axis.size), + params=get_params("iota_rr", obj=eq, has_axis=grid.axis.size), transforms=transforms, - profiles=get_profiles("iota_rr", eq=eq, grid=grid), + profiles=get_profiles("iota_rr", obj=eq, grid=grid), ) np.testing.assert_allclose(data["iota"], benchmark_data["iota"]) np.testing.assert_allclose(data["iota_r"], benchmark_data["iota_r"]) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 2db926b86d..57d9dc755c 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -221,7 +221,7 @@ def test_qh_boozer(self): # compute all amplitudes in the Boozer spectrum transforms = get_transforms( - "|B|_mn", eq=eq, grid=grid, M_booz=M_booz, N_booz=N_booz + "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz ) matrix, modes, idx = ptolemy_linear_transform( transforms["B"].basis.modes, helicity=helicity, NFP=eq.NFP From a6b8125a77cd23de8e09662c9c73adde64cbca1d Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Mon, 31 Jul 2023 18:35:18 -0500 Subject: [PATCH 09/19] Make sort_compute_funs.py more readable --- docs/sort_compute_funs.py | 28 +++++++++++----------------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py index b8274a85c7..0dfdb71c1e 100644 --- a/docs/sort_compute_funs.py +++ b/docs/sort_compute_funs.py @@ -11,27 +11,21 @@ import desc.compute - -def get_matches(fun, pattern): - """Return all matches of ``pattern`` in source code of function ``fun``.""" - src = inspect.getsource(fun) - matches = pattern.findall(src) - return matches - - -# Gather all compute function source code and map to quantity name. -src_codes = {} -pattern_name = re.compile(r"(?<=name=)[^,]+") +# Gather all compute function source code and map quantity 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]: for _, fun in inspect.getmembers(module, inspect.isfunction): + source_code = inspect.getsource(fun) # quantities that this function computes - name = get_matches(fun, pattern_name) - if len(name) > 0: # skip imported functions - src_codes[name[0]] = inspect.getsource(fun) + matches = pattern.findall(source_code) + if matches: # skip imported functions + source_codes[matches[0]] = source_code # Write compute functions sorted by name to file. with open(sys.argv[2], "w") as output_file: - for name in sorted(src_codes): - output_file.write(src_codes[name]) - output_file.write("\n\n") + for name in sorted(source_codes): + output_file.write("\n") + output_file.write(source_codes[name]) + output_file.write("\n") From 720b512f11d5c0c63e2e18f348e5e9d71a7f228b Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 18:48:57 -0400 Subject: [PATCH 10/19] Convert vectors to xyz basis before finite differencing --- tests/test_compute_funs.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index ad80ec7832..47ad672ce1 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -9,6 +9,7 @@ from desc.compute import data_index from desc.compute.utils import compress from desc.equilibrium import EquilibriaFamily, Equilibrium +from desc.geometry.utils import rpz2xyz_vec from desc.grid import LinearGrid, QuadratureGrid # convolve kernel is reverse of FD coeffs @@ -146,9 +147,13 @@ def test_magnetic_field_derivatives(DummyStellarator): "B", "B_r", "B_rr", + "phi", ], grid=grid, ) + data["B"] = rpz2xyz_vec(data["B"], phi=data["phi"]) + data["B_r"] = rpz2xyz_vec(data["B_r"], phi=data["phi"]) + data["B_rr"] = rpz2xyz_vec(data["B_rr"], phi=data["phi"]) B_sup_theta_r = np.convolve(data["B^theta"], FD_COEF_1_4, "same") / drho B_sup_theta_rr = np.convolve(data["B^theta"], FD_COEF_2_4, "same") / drho**2 @@ -279,9 +284,13 @@ def test_magnetic_field_derivatives(DummyStellarator): "B", "B_t", "B_tt", + "phi", ], grid=grid, ) + data["B"] = rpz2xyz_vec(data["B"], phi=data["phi"]) + data["B_t"] = rpz2xyz_vec(data["B_t"], phi=data["phi"]) + data["B_tt"] = rpz2xyz_vec(data["B_tt"], phi=data["phi"]) B_sup_theta_t = np.convolve(data["B^theta"], FD_COEF_1_4, "same") / dtheta B_sup_theta_tt = np.convolve(data["B^theta"], FD_COEF_2_4, "same") / dtheta**2 @@ -412,9 +421,13 @@ def test_magnetic_field_derivatives(DummyStellarator): "B", "B_z", "B_zz", + "phi", ], grid=grid, ) + data["B"] = rpz2xyz_vec(data["B"], phi=data["phi"]) + data["B_z"] = rpz2xyz_vec(data["B_z"], phi=data["phi"]) + data["B_zz"] = rpz2xyz_vec(data["B_zz"], phi=data["phi"]) B_sup_theta_z = np.convolve(data["B^theta"], FD_COEF_1_4, "same") / dzeta B_sup_theta_zz = np.convolve(data["B^theta"], FD_COEF_2_4, "same") / dzeta**2 @@ -540,9 +553,12 @@ def test_magnetic_field_derivatives(DummyStellarator): "|B|_rt", "B", "B_rt", + "phi", ], grid=grid, ) + data["B"] = rpz2xyz_vec(data["B"], phi=data["phi"]) + data["B_rt"] = rpz2xyz_vec(data["B_rt"], phi=data["phi"]) B_sup_theta = data["B^theta"].reshape((num_rho, num_theta)) B_sup_zeta = data["B^zeta"].reshape((num_rho, num_theta)) @@ -658,9 +674,12 @@ def test_magnetic_field_derivatives(DummyStellarator): "|B|_tz", "B", "B_tz", + "phi", ], grid=grid, ) + data["B"] = rpz2xyz_vec(data["B"], phi=data["phi"]) + data["B_tz"] = rpz2xyz_vec(data["B_tz"], phi=data["phi"]) B_sup_theta = data["B^theta"].reshape((num_zeta, num_theta)) B_sup_zeta = data["B^zeta"].reshape((num_zeta, num_theta)) @@ -776,9 +795,12 @@ def test_magnetic_field_derivatives(DummyStellarator): "|B|_rz", "B", "B_rz", + "phi", ], grid=grid, ) + data["B"] = rpz2xyz_vec(data["B"], phi=data["phi"]) + data["B_rz"] = rpz2xyz_vec(data["B_rz"], phi=data["phi"]) B_sup_theta = data["B^theta"].reshape((num_zeta, num_rho)) B_sup_zeta = data["B^zeta"].reshape((num_zeta, num_rho)) From eff10e658a37fdd5ebeb10691c35e7f771769c37 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 18:49:27 -0400 Subject: [PATCH 11/19] Fix boundary conditions on 2d finite differencing --- tests/test_compute_funs.py | 308 ++++++++++++++++++------------------- 1 file changed, 154 insertions(+), 154 deletions(-) diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 47ad672ce1..c510218f0c 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -31,7 +31,7 @@ def myconvolve_2d(arr_1d, stencil, shape): arr, stencil[:, np.newaxis] * stencil[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", # not periodic in rho, easier to pad and truncate in all dims ) return conv @@ -121,7 +121,7 @@ def test_magnetic_field_derivatives(DummyStellarator): # partial derivatives wrt rho rtol = 1e-3 atol = 1e-3 - num_rho = 100 + num_rho = 180 grid = LinearGrid(rho=num_rho, NFP=eq.NFP) drho = grid.nodes[1, 0] data = eq.compute( @@ -171,94 +171,94 @@ def test_magnetic_field_derivatives(DummyStellarator): B_rr = np.apply_along_axis(my_convolve, 0, data["B"], FD_COEF_2_4) / drho**2 np.testing.assert_allclose( - data["B^theta_r"][3:-2], - B_sup_theta_r[3:-2], + data["B^theta_r"][4:-4], + B_sup_theta_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^theta_r"])), ) np.testing.assert_allclose( - data["B^theta_rr"][3:-2], - B_sup_theta_rr[3:-2], + data["B^theta_rr"][4:-4], + B_sup_theta_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^theta_rr"])), ) np.testing.assert_allclose( - data["B^zeta_r"][3:-2], - B_sup_zeta_r[3:-2], + data["B^zeta_r"][4:-4], + B_sup_zeta_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^zeta_r"])), ) np.testing.assert_allclose( - data["B^zeta_rr"][3:-2], - B_sup_zeta_rr[3:-2], + data["B^zeta_rr"][4:-4], + B_sup_zeta_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^zeta_rr"])), ) np.testing.assert_allclose( - data["B_rho_r"][3:-2], - B_sub_rho_r[3:-2], + data["B_rho_r"][4:-4], + B_sub_rho_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rho_r"])), ) np.testing.assert_allclose( - data["B_rho_rr"][3:-2], - B_sub_rho_rr[3:-2], + data["B_rho_rr"][4:-4], + B_sub_rho_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rho_rr"])), ) np.testing.assert_allclose( - data["B_theta_r"][3:-2], - B_sub_theta_r[3:-2], + data["B_theta_r"][4:-4], + B_sub_theta_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_theta_r"])), ) np.testing.assert_allclose( - data["B_theta_rr"][3:-2], - B_sub_theta_rr[3:-2], + data["B_theta_rr"][4:-4], + B_sub_theta_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_theta_rr"])), ) np.testing.assert_allclose( - data["B_zeta_r"][3:-2], - B_sub_zeta_r[3:-2], + data["B_zeta_r"][4:-4], + B_sub_zeta_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_zeta_r"])), ) np.testing.assert_allclose( - data["B_zeta_rr"][3:-2], - B_sub_zeta_rr[3:-2], + data["B_zeta_rr"][4:-4], + B_sub_zeta_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_zeta_rr"])), ) np.testing.assert_allclose( - data["|B|_r"][3:-2], - Bmag_r[3:-2], + data["|B|_r"][4:-4], + Bmag_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["|B|_r"])), ) np.testing.assert_allclose( - data["|B|_rr"][3:-2], - Bmag_rr[3:-2], + data["|B|_rr"][4:-4], + Bmag_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["|B|_rr"])), ) np.testing.assert_allclose( - data["B_r"][3:-2, :], - B_r[3:-2, :], + data["B_r"][4:-4], + B_r[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_r"])), ) np.testing.assert_allclose( - data["B_rr"][3:-2, :], - B_rr[3:-2, :], + data["B_rr"][4:-4], + B_rr[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rr"])), ) # partial derivatives wrt theta - rtol = 3e-3 - atol = 3e-3 - num_theta = 150 + rtol = 1e-3 + atol = 1e-3 + num_theta = 180 grid = LinearGrid(NFP=eq.NFP, theta=num_theta) dtheta = grid.nodes[1, 1] data = eq.compute( @@ -308,94 +308,94 @@ def test_magnetic_field_derivatives(DummyStellarator): B_tt = np.apply_along_axis(my_convolve, 0, data["B"], FD_COEF_2_4) / dtheta**2 np.testing.assert_allclose( - data["B^theta_t"][2:-2], - B_sup_theta_t[2:-2], + data["B^theta_t"][4:-4], + B_sup_theta_t[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^theta_t"])), ) np.testing.assert_allclose( - data["B^theta_tt"][2:-2], - B_sup_theta_tt[2:-2], + data["B^theta_tt"][4:-4], + B_sup_theta_tt[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^theta_tt"])), ) np.testing.assert_allclose( - data["B^zeta_t"][2:-2], - B_sup_zeta_t[2:-2], + data["B^zeta_t"][4:-4], + B_sup_zeta_t[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^zeta_t"])), ) np.testing.assert_allclose( - data["B^zeta_tt"][2:-2], - B_sup_zeta_tt[2:-2], + data["B^zeta_tt"][4:-4], + B_sup_zeta_tt[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^zeta_tt"])), ) np.testing.assert_allclose( - data["B_rho_t"][2:-2], - B_sub_rho_t[2:-2], + data["B_rho_t"][4:-4], + B_sub_rho_t[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_rho_t"])), ) np.testing.assert_allclose( - data["B_rho_tt"][2:-2], - B_sub_rho_tt[2:-2], + data["B_rho_tt"][4:-4], + B_sub_rho_tt[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_rho_tt"])), ) np.testing.assert_allclose( - data["B_theta_t"][2:-2], - B_sub_theta_t[2:-2], + data["B_theta_t"][4:-4], + B_sub_theta_t[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_theta_t"])), ) np.testing.assert_allclose( - data["B_theta_tt"][2:-2], - B_sub_theta_tt[2:-2], + data["B_theta_tt"][4:-4], + B_sub_theta_tt[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_theta_tt"])), ) np.testing.assert_allclose( - data["B_zeta_t"][2:-2], - B_sub_zeta_t[2:-2], + data["B_zeta_t"][4:-4], + B_sub_zeta_t[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_zeta_t"])), ) np.testing.assert_allclose( - data["B_zeta_tt"][2:-2], - B_sub_zeta_tt[2:-2], + data["B_zeta_tt"][4:-4], + B_sub_zeta_tt[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_zeta_tt"])), ) np.testing.assert_allclose( - data["|B|_t"][2:-2], - Bmag_t[2:-2], + data["|B|_t"][4:-4], + Bmag_t[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["|B|_t"])), ) np.testing.assert_allclose( - data["|B|_tt"][2:-2], - Bmag_tt[2:-2], + data["|B|_tt"][4:-4], + Bmag_tt[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["|B|_tt"])), ) np.testing.assert_allclose( - data["B_t"][2:-2, :], - B_t[2:-2, :], + data["B_t"][4:-4], + B_t[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_t"])), ) np.testing.assert_allclose( - data["B_tt"][2:-2, :], - B_tt[2:-2, :], + data["B_tt"][4:-4], + B_tt[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_tt"])), ) # partial derivatives wrt zeta - rtol = 1e-4 - atol = 1e-4 - num_zeta = 130 + rtol = 1e-3 + atol = 1e-3 + num_zeta = 180 grid = LinearGrid(NFP=eq.NFP, zeta=num_zeta) dzeta = grid.nodes[1, 2] data = eq.compute( @@ -445,86 +445,86 @@ def test_magnetic_field_derivatives(DummyStellarator): B_zz = np.apply_along_axis(my_convolve, 0, data["B"], FD_COEF_2_4) / dzeta**2 np.testing.assert_allclose( - data["B^theta_z"][2:-2], - B_sup_theta_z[2:-2], + data["B^theta_z"][4:-4], + B_sup_theta_z[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^theta_z"])), ) np.testing.assert_allclose( - data["B^theta_zz"][2:-2], - B_sup_theta_zz[2:-2], + data["B^theta_zz"][4:-4], + B_sup_theta_zz[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^theta_zz"])), ) np.testing.assert_allclose( - data["B^zeta_z"][2:-2], - B_sup_zeta_z[2:-2], + data["B^zeta_z"][4:-4], + B_sup_zeta_z[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^zeta_z"])), ) np.testing.assert_allclose( - data["B^zeta_zz"][2:-2], - B_sup_zeta_zz[2:-2], + data["B^zeta_zz"][4:-4], + B_sup_zeta_zz[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^zeta_zz"])), ) np.testing.assert_allclose( - data["B_rho_z"][2:-2], - B_sub_rho_z[2:-2], + data["B_rho_z"][4:-4], + B_sub_rho_z[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_rho_z"])), ) np.testing.assert_allclose( - data["B_rho_zz"][2:-2], - B_sub_rho_zz[2:-2], + data["B_rho_zz"][4:-4], + B_sub_rho_zz[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_rho_zz"])), ) np.testing.assert_allclose( - data["B_theta_z"][2:-2], - B_sub_theta_z[2:-2], + data["B_theta_z"][4:-4], + B_sub_theta_z[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_theta_z"])), ) np.testing.assert_allclose( - data["B_theta_zz"][2:-2], - B_sub_theta_zz[2:-2], + data["B_theta_zz"][4:-4], + B_sub_theta_zz[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_theta_zz"])), ) np.testing.assert_allclose( - data["B_zeta_z"][2:-2], - B_sub_zeta_z[2:-2], + data["B_zeta_z"][4:-4], + B_sub_zeta_z[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_zeta_z"])), ) np.testing.assert_allclose( - data["B_zeta_zz"][2:-2], - B_sub_zeta_zz[2:-2], + data["B_zeta_zz"][4:-4], + B_sub_zeta_zz[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_zeta_zz"])), ) np.testing.assert_allclose( - data["|B|_z"][2:-2], - Bmag_z[2:-2], + data["|B|_z"][4:-4], + Bmag_z[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["|B|_z"])), ) np.testing.assert_allclose( - data["|B|_zz"][2:-2], - Bmag_zz[2:-2], + data["|B|_zz"][4:-4], + Bmag_zz[4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["|B|_zz"])), ) np.testing.assert_allclose( - data["B_z"][2:-2, :], - B_z[2:-2, :], + data["B_z"][4:-4], + B_z[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_z"])), ) np.testing.assert_allclose( - data["B_zz"][2:-2, :], - B_zz[2:-2, :], + data["B_zz"][4:-4], + B_zz[4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_zz"])), ) @@ -532,7 +532,7 @@ def test_magnetic_field_derivatives(DummyStellarator): # mixed derivatives wrt rho & theta rtol = 1e-2 atol = 1e-2 - num_rho = 150 + num_rho = 180 num_theta = 180 grid = LinearGrid(NFP=eq.NFP, rho=num_rho, theta=num_theta) drho = grid.nodes[:, 0].reshape((num_rho, num_theta))[1, 0] @@ -571,81 +571,81 @@ def test_magnetic_field_derivatives(DummyStellarator): B_sup_theta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dtheta) B_sup_zeta_rt = convolve2d( B_sup_zeta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dtheta) B_sub_rho_rt = convolve2d( B_sub_rho, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dtheta) B_sub_theta_rt = convolve2d( B_sub_theta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dtheta) B_sub_zeta_rt = convolve2d( B_sub_zeta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dtheta) Bmag_rt = convolve2d( Bmag, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dtheta) B_rt = np.apply_along_axis( myconvolve_2d, 0, data["B"], FD_COEF_1_4, (num_rho, num_theta) ) / (drho * dtheta) np.testing.assert_allclose( - data["B^theta_rt"].reshape((num_rho, num_theta))[3:-2, 2:-2], - B_sup_theta_rt[3:-2, 2:-2], + data["B^theta_rt"].reshape((num_rho, num_theta))[4:-4, 4:-4], + B_sup_theta_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^theta_rt"])), ) np.testing.assert_allclose( - data["B^zeta_rt"].reshape((num_rho, num_theta))[3:-2, 2:-2], - B_sup_zeta_rt[3:-2, 2:-2], + data["B^zeta_rt"].reshape((num_rho, num_theta))[4:-4, 4:-4], + B_sup_zeta_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^zeta_rt"])), ) np.testing.assert_allclose( - data["B_rho_rt"].reshape((num_rho, num_theta))[3:-2, 2:-2], - B_sub_rho_rt[3:-2, 2:-2], + data["B_rho_rt"].reshape((num_rho, num_theta))[4:-4, 4:-4], + B_sub_rho_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rho_rt"])), ) np.testing.assert_allclose( - data["B_theta_rt"].reshape((num_rho, num_theta))[3:-2, 2:-2], - B_sub_theta_rt[3:-2, 2:-2], + data["B_theta_rt"].reshape((num_rho, num_theta))[4:-4, 4:-4], + B_sub_theta_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_theta_rt"])), ) np.testing.assert_allclose( - data["B_zeta_rt"].reshape((num_rho, num_theta))[3:-2, 2:-2], - B_sub_zeta_rt[3:-2, 2:-2], + data["B_zeta_rt"].reshape((num_rho, num_theta))[4:-4, 4:-4], + B_sub_zeta_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_zeta_rt"])), ) np.testing.assert_allclose( - data["|B|_rt"].reshape((num_rho, num_theta))[3:-2, 2:-2], - Bmag_rt[3:-2, 2:-2], + data["|B|_rt"].reshape((num_rho, num_theta))[4:-4, 4:-4], + Bmag_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["|B|_rt"])), ) np.testing.assert_allclose( - data["B_rt"].reshape((num_rho, num_theta, 3))[3:-2, 2:-2, :], - B_rt[3:-2, 2:-2, :], + data["B_rt"].reshape((num_rho, num_theta, 3))[4:-4, 4:-4], + B_rt[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rt"])), ) @@ -692,81 +692,81 @@ def test_magnetic_field_derivatives(DummyStellarator): B_sup_theta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (dtheta * dzeta) B_sup_zeta_tz = convolve2d( B_sup_zeta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (dtheta * dzeta) B_sub_rho_tz = convolve2d( B_sub_rho, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (dtheta * dzeta) B_sub_theta_tz = convolve2d( B_sub_theta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (dtheta * dzeta) B_sub_zeta_tz = convolve2d( B_sub_zeta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (dtheta * dzeta) Bmag_tz = convolve2d( Bmag, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (dtheta * dzeta) B_tz = np.apply_along_axis( myconvolve_2d, 0, data["B"], FD_COEF_1_4, (num_zeta, num_theta) ) / (dzeta * dtheta) np.testing.assert_allclose( - data["B^theta_tz"].reshape((num_zeta, num_theta))[2:-2, 2:-2], - B_sup_theta_tz[2:-2, 2:-2], + data["B^theta_tz"].reshape((num_zeta, num_theta))[4:-4, 4:-4], + B_sup_theta_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^theta_tz"])), ) np.testing.assert_allclose( - data["B^zeta_tz"].reshape((num_zeta, num_theta))[2:-2, 2:-2], - B_sup_zeta_tz[2:-2, 2:-2], + data["B^zeta_tz"].reshape((num_zeta, num_theta))[4:-4, 4:-4], + B_sup_zeta_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B^zeta_tz"])), ) np.testing.assert_allclose( - data["B_rho_tz"].reshape((num_zeta, num_theta))[2:-2, 2:-2], - B_sub_rho_tz[2:-2, 2:-2], + data["B_rho_tz"].reshape((num_zeta, num_theta))[4:-4, 4:-4], + B_sub_rho_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_rho_tz"])), ) np.testing.assert_allclose( - data["B_theta_tz"].reshape((num_zeta, num_theta))[2:-2, 2:-2], - B_sub_theta_tz[2:-2, 2:-2], + data["B_theta_tz"].reshape((num_zeta, num_theta))[4:-4, 4:-4], + B_sub_theta_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_theta_tz"])), ) np.testing.assert_allclose( - data["B_zeta_tz"].reshape((num_zeta, num_theta))[2:-2, 2:-2], - B_sub_zeta_tz[2:-2, 2:-2], + data["B_zeta_tz"].reshape((num_zeta, num_theta))[4:-4, 4:-4], + B_sub_zeta_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["B_zeta_tz"])), ) np.testing.assert_allclose( - data["|B|_tz"].reshape((num_zeta, num_theta))[2:-2, 2:-2], - Bmag_tz[2:-2, 2:-2], + data["|B|_tz"].reshape((num_zeta, num_theta))[4:-4, 4:-4], + Bmag_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.mean(np.abs(data["|B|_tz"])), ) np.testing.assert_allclose( - data["B_tz"].reshape((num_zeta, num_theta, 3))[2:-2, 2:-2, :], - B_tz[2:-2, 2:-2, :], + data["B_tz"].reshape((num_zeta, num_theta, 3))[4:-4, 4:-4], + B_tz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_tz"])), ) @@ -774,7 +774,7 @@ def test_magnetic_field_derivatives(DummyStellarator): # mixed derivatives wrt rho & zeta rtol = 1e-2 atol = 1e-2 - num_rho = 150 + num_rho = 180 num_zeta = 180 grid = LinearGrid(NFP=eq.NFP, rho=num_rho, zeta=num_zeta) drho = grid.nodes[:, 0].reshape((num_zeta, num_rho))[0, 1] @@ -813,81 +813,81 @@ def test_magnetic_field_derivatives(DummyStellarator): B_sup_theta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dzeta) B_sup_zeta_rz = convolve2d( B_sup_zeta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dzeta) B_sub_rho_rz = convolve2d( B_sub_rho, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dzeta) B_sub_theta_rz = convolve2d( B_sub_theta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dzeta) B_sub_zeta_rz = convolve2d( B_sub_zeta, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dzeta) Bmag_rz = convolve2d( Bmag, FD_COEF_1_4[:, np.newaxis] * FD_COEF_1_4[np.newaxis, :], mode="same", - boundary="wrap", + boundary="fill", ) / (drho * dzeta) B_rz = np.apply_along_axis( myconvolve_2d, 0, data["B"], FD_COEF_1_4, (num_zeta, num_rho) ) / (drho * dzeta) np.testing.assert_allclose( - data["B^theta_rz"].reshape((num_zeta, num_rho))[2:-2, 3:-2], - B_sup_theta_rz[2:-2, 3:-2], + data["B^theta_rz"].reshape((num_zeta, num_rho))[4:-4, 4:-4], + B_sup_theta_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^theta_rz"])), ) np.testing.assert_allclose( - data["B^zeta_rz"].reshape((num_zeta, num_rho))[2:-2, 3:-2], - B_sup_zeta_rz[2:-2, 3:-2], + data["B^zeta_rz"].reshape((num_zeta, num_rho))[4:-4, 4:-4], + B_sup_zeta_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B^zeta_rz"])), ) np.testing.assert_allclose( - data["B_rho_rz"].reshape((num_zeta, num_rho))[2:-2, 3:-2], - B_sub_rho_rz[2:-2, 3:-2], + data["B_rho_rz"].reshape((num_zeta, num_rho))[4:-4, 4:-4], + B_sub_rho_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rho_rz"])), ) np.testing.assert_allclose( - data["B_theta_rz"].reshape((num_zeta, num_rho))[2:-2, 3:-2], - B_sub_theta_rz[2:-2, 3:-2], + data["B_theta_rz"].reshape((num_zeta, num_rho))[4:-4, 4:-4], + B_sub_theta_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_theta_rz"])), ) np.testing.assert_allclose( - data["B_zeta_rz"].reshape((num_zeta, num_rho))[2:-2, 3:-2], - B_sub_zeta_rz[2:-2, 3:-2], + data["B_zeta_rz"].reshape((num_zeta, num_rho))[4:-4, 4:-4], + B_sub_zeta_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_zeta_rz"])), ) np.testing.assert_allclose( - data["|B|_rz"].reshape((num_zeta, num_rho))[2:-2, 3:-2], - Bmag_rz[2:-2, 3:-2], + data["|B|_rz"].reshape((num_zeta, num_rho))[4:-4, 4:-4], + Bmag_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["|B|_rz"])), ) np.testing.assert_allclose( - data["B_rz"].reshape((num_zeta, num_rho, 3))[2:-2, 3:-2, :], - B_rz[2:-2, 3:-2, :], + data["B_rz"].reshape((num_zeta, num_rho, 3))[4:-4, 4:-4], + B_rz[4:-4, 4:-4], rtol=rtol, atol=atol * np.nanmean(np.abs(data["B_rz"])), ) From d8d882fa1830c6ec2c4fd832297ad6f149c22264 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Mon, 31 Jul 2023 18:59:44 -0500 Subject: [PATCH 12/19] Fix mean curvature test --- tests/test_objective_funs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 57d9dc755c..60e6c03186 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -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)) From 267e4a91eb1bd42ffe3581415353f2510608db01 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Mon, 31 Jul 2023 20:06:11 -0500 Subject: [PATCH 13/19] Cherry-pick fixes to tests from add_all_limits branch. Mark augmented lagrangian test xfail. from commit 248e158 Increase tolerance for QuasisymmetryTripleProduct. from commit 3ed6cff --- tests/test_objective_funs.py | 2 +- tests/test_optimizer.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 60e6c03186..b506cbe3b1 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))) 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") From 81e735d37aa44c83fdfabe8ac84e3bc395876f2e Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Sun, 23 Jul 2023 10:19:00 -0500 Subject: [PATCH 14/19] Reduce tolerance for QIC solve test --- tests/test_examples.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 26de3f9fb1..7ddd566e52 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -608,9 +608,9 @@ def test_NAE_QSC_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) @@ -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) From 0d24953d1725e5c4f8f111ff61c37894160afa96 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Tue, 1 Aug 2023 15:29:47 -0400 Subject: [PATCH 15/19] Add test for covariant basis vectors --- tests/test_compute_funs.py | 109 +++++++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index c510218f0c..2742dbfdb2 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1157,3 +1157,112 @@ def test_compute_averages(): grid = LinearGrid(rho=[0.3], theta=[np.pi / 3], zeta=[0]) out = eq.compute("A", grid=grid) np.testing.assert_allclose(out["A"], np.pi) + + +@pytest.mark.unit +def test_covariant_basis_vectors(DummyStellarator): + """Test calculation of covariant basis vectors by comparing to finite diff of x.""" + eq = Equilibrium.load( + load_from=str(DummyStellarator["output_path"]), file_format="hdf5" + ) + keys = [ + "e_rho", + "e_rho_r", + "e_rho_rr", + "e_rho_rrr", + "e_rho_rrt", + "e_rho_rrz", + "e_rho_rt", + "e_rho_rtt", + "e_rho_rtz", + "e_rho_rz", + "e_rho_rzz", + "e_rho_t", + "e_rho_tt", + "e_rho_tz", + "e_rho_z", + "e_rho_zz", + "e_theta", + "e_theta_r", + "e_theta_rr", + "e_theta_rrr", + "e_theta_rrt", + "e_theta_rrz", + "e_theta_rt", + "e_theta_rtt", + "e_theta_rtz", + "e_theta_rz", + "e_theta_rzz", + "e_theta_t", + "e_theta_tt", + "e_theta_tz", + "e_theta_z", + "e_theta_zz", + "e_zeta", + "e_zeta_r", + "e_zeta_rr", + "e_zeta_rrr", + "e_zeta_rrt", + "e_zeta_rrz", + "e_zeta_rt", + "e_zeta_rtt", + "e_zeta_rtz", + "e_zeta_rz", + "e_zeta_rzz", + "e_zeta_t", + "e_zeta_tt", + "e_zeta_tz", + "e_zeta_z", + "e_zeta_zz", + ] + grids = { + "r": LinearGrid(1000, 0, 0, NFP=eq.NFP), + "t": LinearGrid(0, 1000, 0, NFP=eq.NFP), + "z": LinearGrid(0, 0, 1000, NFP=eq.NFP), + } + + for key in keys: + print(key) + split = key.split("_") + # higher order finite differences are unstable, so we only ever do 1 order + # eg compare e_rho vs fd of x, e_rho_t vs fd of e_rho etc. + if len(split) == 2: # stuff like e_rho, e_theta + base = ["X", "Y", "Z"] + deriv = split[-1][0] + else: + deriv = split[-1] + if len(deriv) == 1: # first derivative of basis vector + base = [split[0] + "_" + split[1]] + else: + base = [split[0] + "_" + split[1] + "_" + deriv[:-1]] + deriv = deriv[-1] + + grid = grids[deriv] + data = eq.compute([key] + base + ["phi"], grid=grid) + data[key] = rpz2xyz_vec(data[key], phi=data["phi"]).reshape( + (grid.num_theta, grid.num_rho, grid.num_zeta, -1), order="F" + ) + if base == ["X", "Y", "Z"]: + base = np.array([data["X"], data["Y"], data["Z"]]).T.reshape( + (grid.num_theta, grid.num_rho, grid.num_zeta, -1), order="F" + ) + + else: + base = rpz2xyz_vec(data[base[0]], phi=data["phi"]).reshape( + (grid.num_theta, grid.num_rho, grid.num_zeta, -1), order="F" + ) + + spacing = { + "r": grid.spacing[0, 0], + "t": grid.spacing[0, 1], + "z": grid.spacing[0, 2] / grid.NFP, + } + + dx = np.apply_along_axis(my_convolve, 0, base, FD_COEF_1_4) / spacing[deriv] + np.testing.assert_allclose( + data[key][4:-4], + dx[4:-4], + rtol=1e-6, + atol=1e-6, + err_msg=key, + ) From 985460ce0993321bc722b6dd8d8c3efcaf95397e Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 1 Aug 2023 15:11:03 -0500 Subject: [PATCH 16/19] Fix e_zeta_rtt --- desc/compute/_basis_vectors.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index d6721784c4..4494fbd927 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -2919,9 +2919,10 @@ def _e_sub_zeta_rtt(params, transforms, profiles, data, **kwargs): + 2 * data["omega_t"] * data["R_rtz"] + 2 * data["R_rt"] * data["omega_tz"] + 2 * data["R_t"] * data["omega_rtz"] - + data["R_rtt"] + 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"] From 71bb72f64f182aa939409bc5525e27aff4ec0c55 Mon Sep 17 00:00:00 2001 From: Kaya Unalmis Date: Tue, 8 Aug 2023 20:06:24 -0500 Subject: [PATCH 17/19] Fix things due to API updates in last merge... Can't sort compute functions by quantity name anymore because there are multiple quantities with the same name. Have to sort by function name instead. Needed to update tests/test_axis_limits for the dependency check test. --- docs/sort_compute_funs.py | 28 +++--- tests/test_axis_limits.py | 183 +++++++++++++++++++++++--------------- 2 files changed, 131 insertions(+), 80 deletions(-) diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py index 0dfdb71c1e..834c936478 100644 --- a/docs/sort_compute_funs.py +++ b/docs/sort_compute_funs.py @@ -1,7 +1,7 @@ """Sort the compute functions and print to specified output. First command line argument is module name, e.g. "_basis_vectors". -Second command line output is file to print to, +Second command line argument is output file to print to, e.g. "sorted_basis_vectors.txt". """ @@ -11,21 +11,29 @@ import desc.compute -# Gather all compute function source code and map quantity name to source code. +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]: - for _, fun in inspect.getmembers(module, inspect.isfunction): + # 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) - # quantities that this function computes matches = pattern.findall(source_code) if matches: # skip imported functions - source_codes[matches[0]] = source_code + # 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 named the same.") + source_codes[function_name] = source_code -# Write compute functions sorted by name to file. +# Write functions sorted to file. with open(sys.argv[2], "w") as output_file: - for name in sorted(source_codes): - output_file.write("\n") - output_file.write(source_codes[name]) - output_file.write("\n") + 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 4bf4337169..2050bcfbe4 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -7,12 +7,11 @@ import desc.compute from desc.compute import data_index +from desc.compute.data_index import _class_inheritance from desc.compute.utils import surface_integrals_map from desc.equilibrium import Equilibrium from desc.grid import LinearGrid -data_index = data_index["desc.equilibrium.equilibrium.Equilibrium"] - # 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 @@ -27,6 +26,14 @@ "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", @@ -54,7 +61,7 @@ not_implemented_limits = { "iota_num_rrr", "iota_den_rrr", -} | data_index.keys() +} def grow_seeds(seeds, search_space): @@ -75,14 +82,17 @@ def grow_seeds(seeds, search_space): """ out = seeds.copy() for key in search_space: - deps = data_index[key]["full_with_axis_dependencies"]["data"] + 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.keys() - not_finite_limits + not_implemented_limits, + data_index["desc.equilibrium.equilibrium.Equilibrium"].keys() - not_finite_limits, ) @@ -101,9 +111,9 @@ def _skip_this(eq, name): def assert_is_continuous( eq, names, - delta=1e-5, + delta=5e-5, rtol=1e-4, - atol=1e-5, + atol=1e-6, desired_at_axis=None, kwargs=None, ): @@ -118,7 +128,7 @@ def assert_is_continuous( 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. + Smaller values accumulate finite precision error and fitting issues. rtol : float, optional Relative tolerance. atol : float, optional @@ -146,37 +156,47 @@ def assert_is_continuous( # 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")] - rho = np.linspace(0, 1, 10) * delta + 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) + 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[name]["coordinates"] == "": + 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[name]["coordinates"] == "r": - profile = data[name] + if data_index_eq[name]["coordinates"] == "r": + # already single variable function of rho + profile = grid.compress(data[name]) else: - # 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. - profile = integrate(np.abs(data[name])) - # integration replaced nans with 0, put them back - profile = np.where(np.isnan(data[name]), np.nan, profile) - profile = grid.compress(profile) + # 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[name]["dim"]): + 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.polyfit(rho[1:], profile[1:], 6) + 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[-1] + fit = poly[0] np.testing.assert_allclose( actual=profile[0], desired=fit, @@ -196,10 +216,25 @@ def get_matches(fun, pattern): # attempt to remove comments src = "\n".join(line.partition("#")[0] for line in src.splitlines()) matches = pattern.findall(src) - matches = {s.replace("'", "").replace('"', "") for s in matches} + matches = {s.strip().strip('"') for s in matches} return matches +def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): + """Get parameterization of this function.""" + pattern_parameterization = re.compile( + r'parameterization=\[\s*([^]]+?)\s*\]|parameterization="([^"]+)"' + ) + decorator = inspect.getsource(fun).partition("def ")[0] + matches = pattern_parameterization.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 = {match.strip().strip('"') for sublist in matches for match in sublist} + matches.discard("") + return matches if matches else {default} + + class TestAxisLimits: """Tests for compute functions evaluated at limits.""" @@ -207,7 +242,8 @@ class TestAxisLimits: def test_data_index_deps(self): """Ensure developers do not add extra (or forget needed) dependencies.""" queried_deps = {} - pattern_keys = re.compile(r"(? Date: Tue, 8 Aug 2023 20:44:13 -0500 Subject: [PATCH 18/19] simplify regular expression --- tests/test_axis_limits.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index 2050bcfbe4..8084cee452 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -221,12 +221,10 @@ def get_matches(fun, pattern): def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): - """Get parameterization of this function.""" - pattern_parameterization = re.compile( - r'parameterization=\[\s*([^]]+?)\s*\]|parameterization="([^"]+)"' - ) + """Get parameterization of thing this function computes.""" + pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")') decorator = inspect.getsource(fun).partition("def ")[0] - matches = pattern_parameterization.findall(decorator) + 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 @@ -243,10 +241,10 @@ def test_data_index_deps(self): """Ensure developers do not add extra (or forget needed) dependencies.""" queried_deps = {} - pattern_names = re.compile(r"(? Date: Tue, 8 Aug 2023 21:34:20 -0500 Subject: [PATCH 19/19] fix typo in printed message --- docs/sort_compute_funs.py | 2 +- tests/test_axis_limits.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/sort_compute_funs.py b/docs/sort_compute_funs.py index 834c936478..5991661337 100644 --- a/docs/sort_compute_funs.py +++ b/docs/sort_compute_funs.py @@ -28,7 +28,7 @@ # 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 named the same.") + raise ValueError("Can't sort when things have same name.") source_codes[function_name] = source_code # Write functions sorted to file. diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py index 8084cee452..0afbc21cf5 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -221,14 +221,14 @@ def get_matches(fun, pattern): def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): - """Get parameterization of thing this function computes.""" + """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 = {match.strip().strip('"') for sublist in matches for match in sublist} + matches = {s.strip().strip('"') for sublist in matches for s in sublist} matches.discard("") return matches if matches else {default}