diff --git a/desc/basis.py b/desc/basis.py index f3175c337e..c5d0cb4aed 100644 --- a/desc/basis.py +++ b/desc/basis.py @@ -40,6 +40,8 @@ def __init__(self): def _set_up(self): """Do things after loading or changing resolution.""" + # Also recreates any attributes not in _io_attrs on load from input file. + # See IOAble class docstring for more info. self._enforce_symmetry() self._sort_modes() self._create_idx() @@ -614,7 +616,7 @@ class ZernikePolynomial(Basis): For L>0, the indexing scheme defines order of the basis functions: ``'ansi'``: ANSI indexing fills in the pyramid with triangles of - decreasing size, ending in a triagle shape. For L == M, + decreasing size, ending in a triangle shape. For L == M, the traditional ANSI pyramid indexing is recovered. For L>M, adds rows to the bottom of the pyramid, increasing L while keeping M constant, giving a "house" shape. @@ -657,7 +659,7 @@ def _get_modes(self, L=-1, M=0, spectral_indexing="ansi"): For L>0, the indexing scheme defines order of the basis functions: ``'ansi'``: ANSI indexing fills in the pyramid with triangles of - decreasing size, ending in a triagle shape. For L == M, + decreasing size, ending in a triangle shape. For L == M, the traditional ANSI pyramid indexing is recovered. For L>M, adds rows to the bottom of the pyramid, increasing L while keeping M constant, giving a "house" shape. @@ -971,7 +973,7 @@ class FourierZernikeBasis(Basis): For L>0, the indexing scheme defines order of the basis functions: ``'ansi'``: ANSI indexing fills in the pyramid with triangles of - decreasing size, ending in a triagle shape. For L == M, + decreasing size, ending in a triangle shape. For L == M, the traditional ANSI pyramid indexing is recovered. For L>M, adds rows to the bottom of the pyramid, increasing L while keeping M constant, giving a "house" shape. @@ -1016,7 +1018,7 @@ def _get_modes(self, L=-1, M=0, N=0, spectral_indexing="ansi"): For L>0, the indexing scheme defines order of the basis functions: ``'ansi'``: ANSI indexing fills in the pyramid with triangles of - decreasing size, ending in a triagle shape. For L == M, + decreasing size, ending in a triangle shape. For L == M, the traditional ANSI pyramid indexing is recovered. For L>M, adds rows to the bottom of the pyramid, increasing L while keeping M constant, giving a "house" shape. @@ -1395,37 +1397,49 @@ def zernike_radial(r, l, m, dr=0): beta = 0 n = (l - m) // 2 s = (-1) ** n + jacobi_arg = 1 - 2 * r**2 if dr == 0: - out = r**m * _jacobi(n, alpha, beta, 1 - 2 * r**2, 0) + out = r**m * _jacobi(n, alpha, beta, jacobi_arg, 0) elif dr == 1: - f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 0) - df = _jacobi(n, alpha, beta, 1 - 2 * r**2, 1) + f = _jacobi(n, alpha, beta, jacobi_arg, 0) + df = _jacobi(n, alpha, beta, jacobi_arg, 1) out = m * r ** jnp.maximum(m - 1, 0) * f - 4 * r ** (m + 1) * df elif dr == 2: - f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 0) - df = _jacobi(n, alpha, beta, 1 - 2 * r**2, 1) - d2f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 2) + f = _jacobi(n, alpha, beta, jacobi_arg, 0) + df = _jacobi(n, alpha, beta, jacobi_arg, 1) + d2f = _jacobi(n, alpha, beta, jacobi_arg, 2) out = ( - m * (m - 1) * r ** jnp.maximum((m - 2), 0) * f - - 2 * 4 * m * r**m * df - + r**m * (16 * r**2 * d2f - 4 * df) + (m - 1) * m * r ** jnp.maximum(m - 2, 0) * f + - 4 * (2 * m + 1) * r**m * df + + 16 * r ** (m + 2) * d2f ) elif dr == 3: - f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 0) - df = _jacobi(n, alpha, beta, 1 - 2 * r**2, 1) - d2f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 2) - d3f = _jacobi(n, alpha, beta, 1 - 2 * r**2, 3) + f = _jacobi(n, alpha, beta, jacobi_arg, 0) + df = _jacobi(n, alpha, beta, jacobi_arg, 1) + d2f = _jacobi(n, alpha, beta, jacobi_arg, 2) + d3f = _jacobi(n, alpha, beta, jacobi_arg, 3) out = ( (m - 2) * (m - 1) * m * r ** jnp.maximum(m - 3, 0) * f - - 12 * (m - 1) * m * r ** jnp.maximum(m - 1, 0) * df - + 48 * r ** (m + 1) * d2f + - 12 * m**2 * r ** jnp.maximum(m - 1, 0) * df + + 48 * (m + 1) * r ** (m + 1) * d2f - 64 * r ** (m + 3) * d3f - + 48 * m * r ** (m + 1) * d2f - - 12 * m * r ** jnp.maximum(m - 1, 0) * df + ) + elif dr == 4: + f = _jacobi(n, alpha, beta, jacobi_arg, 0) + df = _jacobi(n, alpha, beta, jacobi_arg, 1) + d2f = _jacobi(n, alpha, beta, jacobi_arg, 2) + d3f = _jacobi(n, alpha, beta, jacobi_arg, 3) + d4f = _jacobi(n, alpha, beta, jacobi_arg, 4) + out = ( + (m - 3) * (m - 2) * (m - 1) * m * r ** jnp.maximum(m - 4, 0) * f + - 8 * m * (2 * m**2 - 3 * m + 1) * r ** jnp.maximum(m - 2, 0) * df + + 48 * (2 * m**2 + 2 * m + 1) * r**m * d2f + - 128 * (2 * m + 3) * r ** (m + 2) * d3f + + 256 * r ** (m + 4) * d4f ) else: raise NotImplementedError( - "Analytic radial derivatives of Zernike polynomials for order>3 " + "Analytic radial derivatives of Zernike polynomials for order>4 " + "have not been implemented." ) return s * jnp.where((l - m) % 2 == 0, out, 0) diff --git a/desc/compute/_core.py b/desc/compute/_core.py index 418efb34b9..a370fdb4e4 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -24,300 +24,1282 @@ 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=["R_lmn"], + transforms={"R": [[0, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R(params, transforms, profiles, data, **kwargs): + data["R"] = transforms["R"].transform(params["R_lmn"], 0, 0, 0) + return data + + +@register_compute_fun( + name="R_r", + label="\\partial_{\\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, first radial derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="R_rr", + label="\\partial_{\\rho \\rho} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second radial derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[2, 0, 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) + return data + + +@register_compute_fun( + 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=["R_lmn"], + transforms={"R": [[3, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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_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=["R_lmn"], + transforms={"R": [[2, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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_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_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=["R_lmn"], + transforms={"R": [[2, 0, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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_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=["R_lmn"], + transforms={"R": [[1, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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=["R_lmn"], + transforms={"R": [[1, 2, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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_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_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=["R_lmn"], + transforms={"R": [[1, 1, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_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="R_rz", + label="\\partial_{\\rho \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second derivative wrt radius " + "and toroidal angle", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 0, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_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", + dim=1, + params=["R_lmn"], + transforms={"R": [[1, 0, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_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_t", + label="\\partial_{\\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, first poloidal derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_t(params, transforms, profiles, data, **kwargs): + data["R_t"] = transforms["R"].transform(params["R_lmn"], 0, 1, 0) + return data + + +@register_compute_fun( + name="R_tt", + label="\\partial_{\\theta \\theta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second poloidal derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 2, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _R_tt(params, transforms, profiles, data, **kwargs): + data["R_tt"] = transforms["R"].transform(params["R_lmn"], 0, 2, 0) + return data + + +@register_compute_fun( + 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, third derivative wrt poloidal angle " + "twice and toroidal angle", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 2, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_tz", + label="\\partial_{\\theta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second derivative wrt poloidal " + "and toroidal angles", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 1, 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) + return data + + +@register_compute_fun( + name="R_tzz", + label="\\partial_{\\theta \\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third derivative wrt poloidal angle " + "and toroidal angle twice", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 1, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_z", + label="\\partial_{\\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, first toroidal derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 0, 1]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_zz", + label="\\partial_{\\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, second toroidal derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 0, 2]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} R", + units="m", + units_long="meters", + description="Major radius in lab frame, third toroidal derivative", + dim=1, + params=["R_lmn"], + transforms={"R": [[0, 0, 3]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="X", + label="X = R \\cos{\\phi}", + units="m", + units_long="meters", + description="Cartesian X coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "phi"], +) +def _X(params, transforms, profiles, data, **kwargs): + data["X"] = data["R"] * jnp.cos(data["phi"]) + 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", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "phi", "phi_r"], +) +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="X_t", + label="\\partial_{\\theta} X", + units="m", + units_long="meters", + description="Cartesian X coordinate, derivative wrt poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_t", "phi", "phi_t"], +) +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="X_z", + label="\\partial_{\\zeta} X", + units="m", + units_long="meters", + description="Cartesian X coordinate, derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_z", "phi", "phi_z"], +) +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="Y", + label="Y = R \\sin{\\phi}", + units="m", + units_long="meters", + description="Cartesian Y coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "phi"], +) +def _Y(params, transforms, profiles, data, **kwargs): + data["Y"] = data["R"] * jnp.sin(data["phi"]) + 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", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_r", "phi", "phi_r"], +) +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="Y_t", + label="\\partial_{\\theta} Y", + units="m", + units_long="meters", + description="Cartesian Y coordinate, derivative wrt poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_t", "phi", "phi_t"], +) +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="Y_z", + label="\\partial_{\\zeta} Y", + units="m", + units_long="meters", + description="Cartesian Y coordinate, derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["R", "R_z", "phi", "phi_z"], +) +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="Z", + label="Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[0, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z(params, transforms, profiles, data, **kwargs): + data["Z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 0) + return data + + +@register_compute_fun( + name="Z_r", + label="\\partial_{\\rho} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, first radial derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[1, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="Z_rr", + label="\\partial_{\\rho \\rho} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, second radial derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[2, 0, 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) + return data + + +@register_compute_fun( + name="Z_rrr", + label="\\partial_{\rho \\rho \\rho} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, third radial derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[3, 0, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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_rrt", + label="\\partial_{\\rho \\rho \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, third derivative, wrt radius " + "twice and poloidal angle", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[2, 1, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +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="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_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_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", dim=1, - params=[], - transforms={"grid": []}, + params=["Z_lmn"], + transforms={"Z": [[2, 0, 1]]}, profiles=[], - coordinates="r", + coordinates="rtz", data=[], ) -def _rho(params, transforms, profiles, data, **kwargs): - data["rho"] = transforms["grid"].nodes[:, 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="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", + 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=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[2, 0, 2]]}, profiles=[], - coordinates="r", - data=["0"], + coordinates="rtz", + data=[], ) -def _rho_r(params, transforms, profiles, data, **kwargs): - data["rho_r"] = jnp.ones_like(data["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="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", + name="Z_rt", + label="\\partial_{\\rho \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, second derivative wrt radius " + "and poloidal angle", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[1, 1, 0]]}, profiles=[], - coordinates="r", - data=["0"], + coordinates="rtz", + data=[], ) -def _rho_t(params, transforms, profiles, data, **kwargs): - data["rho_t"] = data["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="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", + 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", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[1, 2, 0]]}, profiles=[], - coordinates="r", - data=["0"], + coordinates="rtz", + data=[], ) -def _rho_z(params, transforms, profiles, data, **kwargs): - data["rho_z"] = data["0"] +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="theta", - label="\\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic)", + 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=[], - transforms={"grid": []}, + params=["Z_lmn"], + transforms={"Z": [[1, 3, 0]]}, profiles=[], - coordinates="t", + coordinates="rtz", data=[], ) -def _theta(params, transforms, profiles, data, **kwargs): - data["theta"] = transforms["grid"].nodes[:, 1] +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="theta_r", - label="\\partial_{\\rho} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - + "derivative wrt radial coordinate", + 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=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[1, 2, 1]]}, profiles=[], - coordinates="t", - data=["0"], + coordinates="rtz", + data=[], ) -def _theta_r(params, transforms, profiles, data, **kwargs): - data["theta_r"] = data["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="theta_t", - label="\\partial_{\\theta} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - + "derivative wrt poloidal coordinate", + 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", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[1, 1, 1]]}, 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 _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="theta_z", - label="\\partial_{\\zeta} \\theta", - units="rad", - units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - + "derivative wrt toroidal coordinate", + 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=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[1, 1, 2]]}, profiles=[], - coordinates="t", - data=["0"], + coordinates="rtz", + data=[], ) -def _theta_z(params, transforms, profiles, data, **kwargs): - data["theta_z"] = data["0"] +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="zeta", - label="\\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate", + name="Z_rz", + label="\\partial_{\\rho \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, second derivative wrt radius " + "and toroidal angle", dim=1, - params=[], - transforms={"grid": []}, + params=["Z_lmn"], + transforms={"Z": [[1, 0, 1]]}, profiles=[], - coordinates="z", + coordinates="rtz", data=[], ) -def _zeta(params, transforms, profiles, data, **kwargs): - data["zeta"] = transforms["grid"].nodes[:, 2] +def _Z_rz(params, transforms, profiles, data, **kwargs): + data["Z_rz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 1) 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="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", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[1, 0, 2]]}, profiles=[], - coordinates="z", - data=["0"], + coordinates="rtz", + data=[], ) -def _zeta_r(params, transforms, profiles, data, **kwargs): - data["zeta_r"] = data["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="zeta_t", - label="\\partial_{\\theta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt poloidal coordinate", + 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_t", + label="\\partial_{\\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, first poloidal derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[0, 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) + return data + + +@register_compute_fun( + name="Z_tt", + label="\\partial_{\\theta \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, second poloidal derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[0, 2, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_tt(params, transforms, profiles, data, **kwargs): + data["Z_tt"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 0) + return data + + +@register_compute_fun( + name="Z_ttt", + label="\\partial_{\\theta \\theta \\theta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, third poloidal derivative", + dim=1, + params=["Z_lmn"], + transforms={"Z": [[0, 3, 0]]}, + profiles=[], + coordinates="rtz", + data=[], +) +def _Z_ttt(params, transforms, profiles, data, **kwargs): + data["Z_ttt"] = transforms["Z"].transform(params["Z_lmn"], 0, 3, 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", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[0, 2, 1]]}, profiles=[], - coordinates="z", - data=["0"], + coordinates="rtz", + data=[], ) -def _zeta_t(params, transforms, profiles, data, **kwargs): - data["zeta_t"] = data["0"] +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="zeta_z", - label="\\partial_{\\zeta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt toroidal coordinate", + 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=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[0, 1, 1]]}, 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 _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="theta_PEST", - label="\\vartheta", - units="rad", - units_long="radians", - description="PEST straight field line poloidal angular coordinate", + 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", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[0, 1, 2]]}, 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 _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="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="Z_z", + label="\\partial_{\\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, first toroidal derivative", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[0, 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 _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="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="Z_zz", + label="\\partial_{\\zeta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, second toroidal derivative", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[0, 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 _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="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="Z_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} Z", + units="m", + units_long="meters", + description="Vertical coordinate in lab frame, third toroidal derivative", dim=1, - params=[], - transforms={}, + params=["Z_lmn"], + transforms={"Z": [[0, 0, 3]]}, profiles=[], coordinates="rtz", - data=["lambda_z"], + data=[], ) -def _theta_PEST_z(params, transforms, profiles, data, **kwargs): - data["theta_PEST_z"] = data["lambda_z"] +def _Z_zzz(params, transforms, profiles, data, **kwargs): + data["Z_zzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 3) return data @@ -350,10 +1332,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 +1354,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,1465 +1372,1458 @@ 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 @register_compute_fun( - name="R", - label="R", - units="m", - units_long="meters", - description="Major radius in lab frame", + name="lambda", + label="\\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 0, 0]]}, + params=["L_lmn"], + transforms={"L": [[0, 0, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R(params, transforms, profiles, data, **kwargs): - data["R"] = transforms["R"].transform(params["R_lmn"], 0, 0, 0) +def _lambda(params, transforms, profiles, data, **kwargs): + data["lambda"] = transforms["L"].transform(params["L_lmn"], 0, 0, 0) return data @register_compute_fun( - name="R_r", - label="\\partial_{\\rho} R", - units="m", - units_long="meters", - description="Major radius in lab frame, first radial derivative", + name="lambda_r", + label="\\partial_{\\rho} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, first radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 0, 0]]}, + params=["L_lmn"], + transforms={"L": [[1, 0, 0]]}, 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 _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="R_t", - label="\\partial_{\\theta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, first poloidal derivative", + name="lambda_rr", + label="\\partial_{\\rho \\rho} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 1, 0]]}, + params=["L_lmn"], + transforms={"L": [[2, 0, 0]]}, 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 _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="R_z", - label="\\partial_{\\zeta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, first toroidal derivative", + name="lambda_rrr", + label="\\partial_{\rho \\rho \\rho} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third radial derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 0, 1]]}, + params=["L_lmn"], + transforms={"L": [[3, 0, 0]]}, 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 _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="R_rr", - label="\\partial_{\\rho \\rho} R", - units="m", - units_long="meters", - description="Major radius in lab frame, second radial derivative", + 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=["R_lmn"], - transforms={"R": [[2, 0, 0]]}, + params=["L_lmn"], + transforms={"L": [[3, 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 _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="R_tt", - label="\\partial_{\\theta \\theta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, second poloidal derivative", + 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=["R_lmn"], - transforms={"R": [[0, 2, 0]]}, + params=["L_lmn"], + transforms={"L": [[3, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_tt(params, transforms, profiles, data, **kwargs): - data["R_tt"] = transforms["R"].transform(params["R_lmn"], 0, 2, 0) +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="R_zz", - label="\\partial_{\\zeta \\zeta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, second toroidal derivative", + name="lambda_rrt", + label="\\partial_{\\rho \\rho \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative, wrt radius twice " + "and poloidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 0, 2]]}, + params=["L_lmn"], + transforms={"L": [[2, 1, 0]]}, 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 _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="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", + name="lambda_rrz", + label="\\partial_{\\rho \\rho \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative, wrt radius twice " + "and toroidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 1, 0]]}, + params=["L_lmn"], + transforms={"L": [[2, 0, 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 _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="R_rz", - label="\\partial_{\\rho \\zeta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, second derivative wrt radius " - + "and toroidal angle", + name="lambda_rt", + label="\\partial_{\\rho \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second derivative wrt radius and " + "poloidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 0, 1]]}, + params=["L_lmn"], + transforms={"L": [[1, 1, 0]]}, 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 _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="R_tz", - label="\\partial_{\\theta \\zeta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, second derivative wrt poloidal " - + "and toroidal angles", + name="lambda_rtt", + label="\\partial_{\\rho \\theta \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative wrt radius and " + "poloidal angle twice", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 1, 1]]}, + params=["L_lmn"], + transforms={"L": [[1, 2, 0]]}, 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 _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="R_rrr", - label="\\partial_{\rho \\rho \\rho} R", - units="m", - units_long="meters", - description="Major radius in lab frame, third radial derivative", + name="lambda_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative wrt radius, poloidal " + " angle, and toroidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[3, 0, 0]]}, + params=["L_lmn"], + transforms={"L": [[1, 1, 1]]}, 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 _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="R_ttt", - label="\\partial_{\\theta \\theta \\theta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, third poloidal derivative", + name="lambda_rz", + label="\\partial_{\\rho \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second derivative wrt radius and " + "toroidal angle", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 3, 0]]}, + params=["L_lmn"], + transforms={"L": [[1, 0, 1]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_ttt(params, transforms, profiles, data, **kwargs): - data["R_ttt"] = transforms["R"].transform(params["R_lmn"], 0, 3, 0) +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="R_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, third toroidal derivative", + name="lambda_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third derivative wrt radius and " + "toroidal angle twice", dim=1, - params=["R_lmn"], - transforms={"R": [[0, 0, 3]]}, + params=["L_lmn"], + transforms={"L": [[1, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_zzz(params, transforms, profiles, data, **kwargs): - data["R_zzz"] = transforms["R"].transform(params["R_lmn"], 0, 0, 3) +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="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", + name="lambda_t", + label="\\partial_{\\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, first poloidal derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 1, 0]]}, + params=["L_lmn"], + transforms={"L": [[0, 1, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rrt(params, transforms, profiles, data, **kwargs): - data["R_rrt"] = transforms["R"].transform(params["R_lmn"], 2, 1, 0) +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="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", + name="lambda_tt", + label="\\partial_{\\theta \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second poloidal derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 2, 0]]}, + params=["L_lmn"], + transforms={"L": [[0, 2, 0]]}, profiles=[], coordinates="rtz", data=[], ) -def _R_rtt(params, transforms, profiles, data, **kwargs): - data["R_rtt"] = transforms["R"].transform(params["R_lmn"], 1, 2, 0) +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="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", + name="lambda_ttt", + label="\\partial_{\\theta \\theta \\theta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third poloidal derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[2, 0, 1]]}, + params=["L_lmn"], + transforms={"L": [[0, 3, 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 _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="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", + 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=["R_lmn"], - transforms={"R": [[1, 0, 2]]}, + params=["L_lmn"], + transforms={"L": [[0, 2, 1]]}, 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 _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="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", + 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=["R_lmn"], - transforms={"R": [[0, 2, 1]]}, + params=["L_lmn"], + transforms={"L": [[0, 1, 1]]}, 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 _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="R_tzz", - label="\\partial_{\\theta \\zeta \\zeta} R", - units="m", - units_long="meters", - description="Major radius in lab frame, third derivative wrt poloidal angle " - + "and toroidal angle twice", + 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=["R_lmn"], - transforms={"R": [[0, 1, 2]]}, + params=["L_lmn"], + transforms={"L": [[0, 1, 2]]}, 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 _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="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", + name="lambda_z", + label="\\partial_{\\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, first toroidal derivative", dim=1, - params=["R_lmn"], - transforms={"R": [[1, 1, 1]]}, + params=["L_lmn"], + transforms={"L": [[0, 0, 1]]}, 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 _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="Z", - label="Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame", + name="lambda_zz", + label="\\partial_{\\zeta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, second toroidal derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 0, 0]]}, + params=["L_lmn"], + transforms={"L": [[0, 0, 2]]}, profiles=[], coordinates="rtz", data=[], ) -def _Z(params, transforms, profiles, data, **kwargs): - data["Z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 0) +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="Z_r", - label="\\partial_{\\rho} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, first radial derivative", + name="lambda_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} \\lambda", + units="rad", + units_long="radians", + description="Poloidal stream function, third toroidal derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[1, 0, 0]]}, + params=["L_lmn"], + transforms={"L": [[0, 0, 3]]}, 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 _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="Z_t", - label="\\partial_{\\theta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, first poloidal derivative", + name="omega", + label="\\omega", + units="rad", + units_long="radians", + description="Toroidal stream function", dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 1, 0]]}, + params=[], # ["W_lmn"], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_t(params, transforms, profiles, data, **kwargs): - data["Z_t"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 0) +def _omega(params, transforms, profiles, data, **kwargs): + data["omega"] = data["0"] return data @register_compute_fun( - name="Z_z", - label="\\partial_{\\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, first toroidal derivative", + name="omega_r", + label="\\partial_{\\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, first radial derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_z(params, transforms, profiles, data, **kwargs): - data["Z_z"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 1) +def _omega_r(params, transforms, profiles, data, **kwargs): + data["omega_r"] = data["0"] return data @register_compute_fun( - name="Z_rr", - label="\\partial_{\\rho \\rho} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, second radial derivative", + name="omega_rr", + label="\\partial_{\\rho \\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second radial derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[2, 0, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rr(params, transforms, profiles, data, **kwargs): - data["Z_rr"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 0) +def _omega_rr(params, transforms, profiles, data, **kwargs): + data["omega_rr"] = data["0"] return data @register_compute_fun( - name="Z_tt", - label="\\partial_{\\theta \\theta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, second poloidal derivative", + name="omega_rrr", + label="\\partial_{\rho \\rho \\rho} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third radial derivative", dim=1, - params=["Z_lmn"], - transforms={"Z": [[0, 2, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_tt(params, transforms, profiles, data, **kwargs): - data["Z_tt"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 0) +def _omega_rrr(params, transforms, profiles, data, **kwargs): + data["omega_rrr"] = data["0"] return data @register_compute_fun( - name="Z_zz", - label="\\partial_{\\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, second toroidal derivative", + 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=["Z_lmn"], - transforms={"Z": [[0, 0, 2]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_zz(params, transforms, profiles, data, **kwargs): - data["Z_zz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 2) +def _omega_rrrr(params, transforms, profiles, data, **kwargs): + data["omega_rrrr"] = data["0"] return data @register_compute_fun( - name="Z_rt", - label="\\partial_{\\rho \\theta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, second derivative wrt radius " - + "and poloidal angle", + 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=["Z_lmn"], - transforms={"Z": [[1, 1, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rt(params, transforms, profiles, data, **kwargs): - data["Z_rt"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 0) +def _omega_rrrt(params, transforms, profiles, data, **kwargs): + data["omega_rrrt"] = data["0"] return data @register_compute_fun( - name="Z_rz", - label="\\partial_{\\rho \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, second derivative wrt radius " - + "and toroidal angle", + 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=["Z_lmn"], - transforms={"Z": [[1, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rz(params, transforms, profiles, data, **kwargs): - data["Z_rz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 1) +def _omega_rrrz(params, transforms, profiles, data, **kwargs): + data["omega_rrrz"] = data["0"] return data @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", + 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=["Z_lmn"], - transforms={"Z": [[0, 1, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_tz(params, transforms, profiles, data, **kwargs): - data["Z_tz"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 1) +def _omega_rrt(params, transforms, profiles, data, **kwargs): + data["omega_rrt"] = data["0"] return data @register_compute_fun( - name="Z_rrr", - label="\\partial_{\rho \\rho \\rho} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third radial derivative", + 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=["Z_lmn"], - transforms={"Z": [[3, 0, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rrr(params, transforms, profiles, data, **kwargs): - data["Z_rrr"] = transforms["Z"].transform(params["Z_lmn"], 3, 0, 0) +def _omega_rrtt(params, transforms, profiles, data, **kwargs): + data["omega_rrtt"] = data["0"] return data @register_compute_fun( - name="Z_ttt", - label="\\partial_{\\theta \\theta \\theta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third poloidal derivative", + 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=["Z_lmn"], - transforms={"Z": [[0, 3, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_ttt(params, transforms, profiles, data, **kwargs): - data["Z_ttt"] = transforms["Z"].transform(params["Z_lmn"], 0, 3, 0) +def _omega_rrtz(params, transforms, profiles, data, **kwargs): + data["omega_rrtz"] = data["0"] return data @register_compute_fun( - name="Z_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third toroidal derivative", + 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=["Z_lmn"], - transforms={"Z": [[0, 0, 3]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_zzz(params, transforms, profiles, data, **kwargs): - data["Z_zzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 0, 3) +def _omega_rrz(params, transforms, profiles, data, **kwargs): + data["omega_rrz"] = data["0"] return data @register_compute_fun( - name="Z_rrt", - label="\\partial_{\\rho \\rho \\theta} Z", - units="m", - units_long="meters", - description="Vertical coordinate in lab frame, third derivative, wrt radius " - + "twice and poloidal angle", + 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=["Z_lmn"], - transforms={"Z": [[2, 1, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rrt(params, transforms, profiles, data, **kwargs): - data["Z_rrt"] = transforms["Z"].transform(params["Z_lmn"], 2, 1, 0) +def _omega_rrzz(params, transforms, profiles, data, **kwargs): + data["omega_rrzz"] = data["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", + 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=["Z_lmn"], - transforms={"Z": [[1, 2, 0]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rtt(params, transforms, profiles, data, **kwargs): - data["Z_rtt"] = transforms["Z"].transform(params["Z_lmn"], 1, 2, 0) +def _omega_rt(params, transforms, profiles, data, **kwargs): + data["omega_rt"] = data["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", + 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=["Z_lmn"], - transforms={"Z": [[2, 0, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rrz(params, transforms, profiles, data, **kwargs): - data["Z_rrz"] = transforms["Z"].transform(params["Z_lmn"], 2, 0, 1) +def _omega_rtt(params, transforms, profiles, data, **kwargs): + data["omega_rtt"] = data["0"] 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="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=["Z_lmn"], - transforms={"Z": [[1, 0, 2]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rzz(params, transforms, profiles, data, **kwargs): - data["Z_rzz"] = transforms["Z"].transform(params["Z_lmn"], 1, 0, 2) +def _omega_rttt(params, transforms, profiles, data, **kwargs): + data["omega_rttt"] = data["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="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=["Z_lmn"], - transforms={"Z": [[0, 2, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_ttz(params, transforms, profiles, data, **kwargs): - data["Z_ttz"] = transforms["Z"].transform(params["Z_lmn"], 0, 2, 1) +def _omega_rttz(params, transforms, profiles, data, **kwargs): + data["omega_rttz"] = data["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="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=["Z_lmn"], - transforms={"Z": [[0, 1, 2]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_tzz(params, transforms, profiles, data, **kwargs): - data["Z_tzz"] = transforms["Z"].transform(params["Z_lmn"], 0, 1, 2) +def _omega_rtz(params, transforms, profiles, data, **kwargs): + data["omega_rtz"] = data["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="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=["Z_lmn"], - transforms={"Z": [[1, 1, 1]]}, + params=[], # ["W_lmn"] + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["0"], ) -def _Z_rtz(params, transforms, profiles, data, **kwargs): - data["Z_rtz"] = transforms["Z"].transform(params["Z_lmn"], 1, 1, 1) +def _omega_rtzz(params, transforms, profiles, data, **kwargs): + data["omega_rtzz"] = data["0"] return data @register_compute_fun( - name="phi", - label="\\phi", + name="omega_rz", + label="\\partial_{\\rho \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame", + description="Toroidal stream function, second derivative wrt radius and " + "toroidal angle", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", - data=["zeta"], + data=["0"], ) -def _phi(params, transforms, profiles, data, **kwargs): - data["phi"] = data["zeta"] +def _omega_rz(params, transforms, profiles, data, **kwargs): + data["omega_rz"] = data["0"] return data @register_compute_fun( - name="phi_r", - label="\\partial_{\\rho} \\phi", + name="omega_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, derivative wrt radial coordinate", + description="Toroidal stream function, third derivative wrt radius and " + "toroidal angle twice", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_r(params, transforms, profiles, data, **kwargs): - data["phi_r"] = data["0"] +def _omega_rzz(params, transforms, profiles, data, **kwargs): + data["omega_rzz"] = data["0"] return data @register_compute_fun( - name="phi_t", - label="\\partial_{\\theta} \\phi", + name="omega_rzzz", + label="\\partial_{\\rho \\zeta \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, derivative wrt poloidal coordinate", + description="Toroidal stream function, third derivative wrt radius and " + "toroidal angle thrice", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_t(params, transforms, profiles, data, **kwargs): - data["phi_t"] = data["0"] +def _omega_rzzz(params, transforms, profiles, data, **kwargs): + data["omega_rzzz"] = data["0"] return data @register_compute_fun( - name="phi_z", - label="\\partial_{\\zeta} \\phi", + name="omega_t", + label="\\partial_{\\theta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, derivative wrt toroidal coordinate", + description="Toroidal stream function, first poloidal derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_z(params, transforms, profiles, data, **kwargs): - data["phi_z"] = jnp.ones_like(data["0"]) +def _omega_t(params, transforms, profiles, data, **kwargs): + data["omega_t"] = data["0"] return data @register_compute_fun( - name="phi_rr", - label="\\partial_{\\rho \\rho} \\phi", + name="omega_tt", + label="\\partial_{\\theta \\theta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt radial coordinate", + description="Toroidal stream function, second poloidal derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_rr(params, transforms, profiles, data, **kwargs): - data["phi_rr"] = data["0"] +def _omega_tt(params, transforms, profiles, data, **kwargs): + data["omega_tt"] = data["0"] return data @register_compute_fun( - name="phi_rt", - label="\\partial_{\\rho \\theta} \\phi", + name="omega_ttt", + label="\\partial_{\\theta \\theta \\theta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt radial and " - + "poloidal coordinate", + description="Toroidal stream function, third poloidal derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_rt(params, transforms, profiles, data, **kwargs): - data["phi_rt"] = data["0"] +def _omega_ttt(params, transforms, profiles, data, **kwargs): + data["omega_ttt"] = data["0"] return data @register_compute_fun( - name="phi_rz", - label="\\partial_{\\rho \\zeta} \\phi", + name="omega_ttz", + label="\\partial_{\\theta \\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt radial and " - + "toroidal coordinate", + description="Toroidal stream function, third derivative wrt poloidal angle " + "twice and toroidal angle", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_rz(params, transforms, profiles, data, **kwargs): - data["phi_rz"] = data["0"] +def _omega_ttz(params, transforms, profiles, data, **kwargs): + data["omega_ttz"] = data["0"] return data @register_compute_fun( - name="phi_tt", - label="\\partial_{\\theta \\theta} \\phi", + name="omega_tz", + label="\\partial_{\\theta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt poloidal " - + "coordinate", + description="Toroidal stream function, second derivative wrt poloidal and " + "toroidal angles", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_tt(params, transforms, profiles, data, **kwargs): - data["phi_tt"] = data["0"] +def _omega_tz(params, transforms, profiles, data, **kwargs): + data["omega_tz"] = data["0"] return data @register_compute_fun( - name="phi_tz", - label="\\partial_{\\theta \\zeta} \\phi", + name="omega_tzz", + label="\\partial_{\\theta \\zeta \\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt poloidal and " - + "toroidal coordinate", + description="Toroidal stream function, third derivative wrt poloidal angle " + "and toroidal angle twice", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_tz(params, transforms, profiles, data, **kwargs): - data["phi_tz"] = data["0"] +def _omega_tzz(params, transforms, profiles, data, **kwargs): + data["omega_tzz"] = data["0"] return data @register_compute_fun( - name="phi_zz", - label="\\partial_{\\zeta \\zeta} \\phi", + name="omega_z", + label="\\partial_{\\zeta} \\omega", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt toroidal " - + "coordinate", + description="Toroidal stream function, first toroidal derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", data=["0"], ) -def _phi_zz(params, transforms, profiles, data, **kwargs): - data["phi_zz"] = data["0"] +def _omega_z(params, transforms, profiles, data, **kwargs): + data["omega_z"] = data["0"] return data @register_compute_fun( - name="X", - label="X = R \\cos{\\phi}", - units="m", - units_long="meters", - description="Cartesian X coordinate", + name="omega_zz", + label="\\partial_{\\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, second toroidal derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", - data=["R", "phi"], + data=["0"], ) -def _X(params, transforms, profiles, data, **kwargs): - data["X"] = data["R"] * jnp.cos(data["phi"]) +def _omega_zz(params, transforms, profiles, data, **kwargs): + data["omega_zz"] = data["0"] 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="omega_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} \\omega", + units="rad", + units_long="radians", + description="Toroidal stream function, third toroidal derivative", dim=1, - params=[], + params=[], # ["W_lmn"] transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_r", "phi", "phi_r"], + data=["0"], ) -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 _omega_zzz(params, transforms, profiles, data, **kwargs): + data["omega_zzz"] = data["0"] 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="phi", + label="\\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "phi", "phi_t"], + data=["zeta", "omega"], ) -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 _phi(params, transforms, profiles, data, **kwargs): + data["phi"] = data["zeta"] + data["omega"] 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="phi_r", + label="\\partial_{\\rho} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, derivative wrt radial coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "phi", "phi_z"], + data=["omega_r"], ) -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 _phi_r(params, transforms, profiles, data, **kwargs): + data["phi_r"] = data["omega_r"] return data @register_compute_fun( - name="Y", - label="Y = R \\sin{\\phi}", - units="m", - units_long="meters", - description="Cartesian Y coordinate", + name="phi_rr", + label="\\partial_{\\rho \\rho} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, second derivative wrt radial coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "phi"], + data=["omega_rr"], ) -def _Y(params, transforms, profiles, data, **kwargs): - data["Y"] = data["R"] * jnp.sin(data["phi"]) +def _phi_rr(params, transforms, profiles, data, **kwargs): + data["phi_rr"] = data["omega_rr"] 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="phi_rt", + label="\\partial_{\\rho \\theta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, second derivative wrt radial and " + "poloidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_r", "phi", "phi_r"], + data=["omega_rt"], ) -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 _phi_rt(params, transforms, profiles, data, **kwargs): + data["phi_rt"] = data["omega_rt"] 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="phi_rz", + label="\\partial_{\\rho \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, second derivative wrt radial and " + "toroidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_t", "phi", "phi_t"], + data=["omega_rz"], ) -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 _phi_rz(params, transforms, profiles, data, **kwargs): + data["phi_rz"] = data["omega_rz"] 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="phi_t", + label="\\partial_{\\theta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, derivative wrt poloidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["R", "R_z", "phi", "phi_z"], + data=["omega_t"], ) -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 _phi_t(params, transforms, profiles, data, **kwargs): + data["phi_t"] = data["omega_t"] return data @register_compute_fun( - name="lambda", - label="\\lambda", + name="phi_tt", + label="\\partial_{\\theta \\theta} \\phi", units="rad", units_long="radians", - description="Poloidal stream function", + description="Toroidal angle in lab frame, second derivative wrt poloidal " + "coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["omega_tt"], ) -def _lambda(params, transforms, profiles, data, **kwargs): - data["lambda"] = transforms["L"].transform(params["L_lmn"], 0, 0, 0) +def _phi_tt(params, transforms, profiles, data, **kwargs): + data["phi_tt"] = data["omega_tt"] return data @register_compute_fun( - name="lambda_r", - label="\\partial_{\\rho} \\lambda", + name="phi_tz", + label="\\partial_{\\theta \\zeta} \\phi", units="rad", units_long="radians", - description="Poloidal stream function, first radial derivative", + description="Toroidal angle in lab frame, second derivative wrt poloidal and " + "toroidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 0, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["omega_tz"], ) -def _lambda_r(params, transforms, profiles, data, **kwargs): - data["lambda_r"] = transforms["L"].transform(params["L_lmn"], 1, 0, 0) +def _phi_tz(params, transforms, profiles, data, **kwargs): + data["phi_tz"] = data["omega_tz"] return data @register_compute_fun( - name="lambda_t", - label="\\partial_{\\theta} \\lambda", + name="phi_z", + label="\\partial_{\\zeta} \\phi", units="rad", units_long="radians", - description="Poloidal stream function, first poloidal derivative", + description="Toroidal angle in lab frame, derivative wrt toroidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 1, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["omega_z"], ) -def _lambda_t(params, transforms, profiles, data, **kwargs): - data["lambda_t"] = transforms["L"].transform(params["L_lmn"], 0, 1, 0) +def _phi_z(params, transforms, profiles, data, **kwargs): + data["phi_z"] = 1 + data["omega_z"] return data @register_compute_fun( - name="lambda_z", - label="\\partial_{\\zeta} \\lambda", + name="phi_zz", + label="\\partial_{\\zeta \\zeta} \\phi", units="rad", units_long="radians", - description="Poloidal stream function, first toroidal derivative", + description="Toroidal angle in lab frame, second derivative wrt toroidal " + "coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 1]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["omega_zz"], ) -def _lambda_z(params, transforms, profiles, data, **kwargs): - data["lambda_z"] = transforms["L"].transform(params["L_lmn"], 0, 0, 1) +def _phi_zz(params, transforms, profiles, data, **kwargs): + data["phi_zz"] = data["omega_zz"] return data @register_compute_fun( - name="lambda_rr", - label="\\partial_{\\rho \\rho} \\lambda", - units="rad", - units_long="radians", - description="Poloidal stream function, second radial derivative", + name="rho", + label="\\rho", + units="~", + units_long="None", + description="Radial coordinate, proportional to the square root " + + "of the toroidal flux", dim=1, - params=["L_lmn"], - transforms={"L": [[2, 0, 0]]}, + params=[], + transforms={"grid": []}, profiles=[], - coordinates="rtz", + coordinates="r", data=[], ) -def _lambda_rr(params, transforms, profiles, data, **kwargs): - data["lambda_rr"] = transforms["L"].transform(params["L_lmn"], 2, 0, 0) +def _rho(params, transforms, profiles, data, **kwargs): + data["rho"] = transforms["grid"].nodes[:, 0] return data @register_compute_fun( - name="lambda_tt", - label="\\partial_{\\theta \\theta} \\lambda", - units="rad", - units_long="radians", - description="Poloidal stream function, second poloidal 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=["L_lmn"], - transforms={"L": [[0, 2, 0]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="r", + data=["0"], ) -def _lambda_tt(params, transforms, profiles, data, **kwargs): - data["lambda_tt"] = transforms["L"].transform(params["L_lmn"], 0, 2, 0) +def _rho_r(params, transforms, profiles, data, **kwargs): + data["rho_r"] = jnp.ones_like(data["0"]) return data @register_compute_fun( - name="lambda_zz", - label="\\partial_{\\zeta \\zeta} \\lambda", - units="rad", - units_long="radians", - description="Poloidal stream function, second toroidal derivative", + 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=["L_lmn"], - transforms={"L": [[0, 0, 2]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="r", + data=["0"], ) -def _lambda_zz(params, transforms, profiles, data, **kwargs): - data["lambda_zz"] = transforms["L"].transform(params["L_lmn"], 0, 0, 2) +def _rho_t(params, transforms, profiles, data, **kwargs): + data["rho_t"] = data["0"] return data @register_compute_fun( - name="lambda_rt", - label="\\partial_{\\rho \\theta} \\lambda", - units="rad", - units_long="radians", - description="Poloidal stream function, second derivative wrt radius and " - + "poloidal angle", + 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=["L_lmn"], - transforms={"L": [[1, 1, 0]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="r", + data=["0"], ) -def _lambda_rt(params, transforms, profiles, data, **kwargs): - data["lambda_rt"] = transforms["L"].transform(params["L_lmn"], 1, 1, 0) +def _rho_z(params, transforms, profiles, data, **kwargs): + data["rho_z"] = data["0"] return data @register_compute_fun( - name="lambda_rz", - label="\\partial_{\\rho \\zeta} \\lambda", + name="theta", + label="\\theta", units="rad", units_long="radians", - description="Poloidal stream function, second derivative wrt radius and " - + "toroidal angle", + description="Poloidal angular coordinate (geometric, not magnetic)", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 0, 1]]}, + params=[], + transforms={"grid": []}, profiles=[], - coordinates="rtz", + coordinates="t", data=[], ) -def _lambda_rz(params, transforms, profiles, data, **kwargs): - data["lambda_rz"] = transforms["L"].transform(params["L_lmn"], 1, 0, 1) +def _theta(params, transforms, profiles, data, **kwargs): + data["theta"] = transforms["grid"].nodes[:, 1] return data @register_compute_fun( - name="lambda_tz", - label="\\partial_{\\theta \\zeta} \\lambda", + name="theta_PEST", + label="\\vartheta", units="rad", units_long="radians", - description="Poloidal stream function, second derivative wrt poloidal and " - + "toroidal angles", + description="PEST straight field line poloidal angular coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 1, 1]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["theta", "lambda"], ) -def _lambda_tz(params, transforms, profiles, data, **kwargs): - data["lambda_tz"] = transforms["L"].transform(params["L_lmn"], 0, 1, 1) +def _theta_PEST(params, transforms, profiles, data, **kwargs): + data["theta_PEST"] = (data["theta"] + data["lambda"]) % (2 * jnp.pi) return data @register_compute_fun( - name="lambda_rrr", - label="\\partial_{\rho \\rho \\rho} \\lambda", + name="theta_PEST_r", + label="\\partial_{\\rho} \\vartheta", units="rad", units_long="radians", - description="Poloidal stream function, third radial derivative", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "radial coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[3, 0, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["lambda_r"], ) -def _lambda_rrr(params, transforms, profiles, data, **kwargs): - data["lambda_rrr"] = transforms["L"].transform(params["L_lmn"], 3, 0, 0) +def _theta_PEST_r(params, transforms, profiles, data, **kwargs): + data["theta_PEST_r"] = data["lambda_r"] return data @register_compute_fun( - name="lambda_ttt", - label="\\partial_{\\theta \\theta \\theta} \\lambda", + name="theta_PEST_t", + label="\\partial_{\\theta} \\vartheta", units="rad", units_long="radians", - description="Poloidal stream function, third poloidal derivative", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "poloidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 3, 0]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["lambda_t"], ) -def _lambda_ttt(params, transforms, profiles, data, **kwargs): - data["lambda_ttt"] = transforms["L"].transform(params["L_lmn"], 0, 3, 0) +def _theta_PEST_t(params, transforms, profiles, data, **kwargs): + data["theta_PEST_t"] = 1 + data["lambda_t"] return data @register_compute_fun( - name="lambda_zzz", - label="\\partial_{\\zeta \\zeta \\zeta} \\lambda", + name="theta_PEST_z", + label="\\partial_{\\zeta} \\vartheta", units="rad", units_long="radians", - description="Poloidal stream function, third toroidal derivative", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "toroidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 0, 3]]}, + params=[], + transforms={}, profiles=[], coordinates="rtz", - data=[], + data=["lambda_z"], ) -def _lambda_zzz(params, transforms, profiles, data, **kwargs): - data["lambda_zzz"] = transforms["L"].transform(params["L_lmn"], 0, 0, 3) +def _theta_PEST_z(params, transforms, profiles, data, **kwargs): + data["theta_PEST_z"] = data["lambda_z"] return data @register_compute_fun( - name="lambda_rrt", - label="\\partial_{\\rho \\rho \\theta} \\lambda", + name="theta_r", + label="\\partial_{\\rho} \\theta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative, wrt radius twice " - + "and poloidal angle", + description="Poloidal angular coordinate (geometric, not magnetic), " + "derivative wrt radial coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[2, 1, 0]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="t", + data=["0"], ) -def _lambda_rrt(params, transforms, profiles, data, **kwargs): - data["lambda_rrt"] = transforms["L"].transform(params["L_lmn"], 2, 1, 0) +def _theta_r(params, transforms, profiles, data, **kwargs): + data["theta_r"] = data["0"] return data @register_compute_fun( - name="lambda_rtt", - label="\\partial_{\\rho \\theta \\theta} \\lambda", + name="theta_t", + label="\\partial_{\\theta} \\theta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt radius and " - + "poloidal angle twice", + description="Poloidal angular coordinate (geometric, not magnetic), " + "derivative wrt poloidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 2, 0]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="t", + data=["0"], ) -def _lambda_rtt(params, transforms, profiles, data, **kwargs): - data["lambda_rtt"] = transforms["L"].transform(params["L_lmn"], 1, 2, 0) +def _theta_t(params, transforms, profiles, data, **kwargs): + data["theta_t"] = jnp.ones_like(data["0"]) return data @register_compute_fun( - name="lambda_rrz", - label="\\partial_{\\rho \\rho \\zeta} \\lambda", + name="theta_z", + label="\\partial_{\\zeta} \\theta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative, wrt radius twice " - + "and toroidal angle", + description="Poloidal angular coordinate (geometric, not magnetic), " + "derivative wrt toroidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[2, 0, 1]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="t", + data=["0"], ) -def _lambda_rrz(params, transforms, profiles, data, **kwargs): - data["lambda_rrz"] = transforms["L"].transform(params["L_lmn"], 2, 0, 1) +def _theta_z(params, transforms, profiles, data, **kwargs): + data["theta_z"] = data["0"] return data @register_compute_fun( - name="lambda_rzz", - label="\\partial_{\\rho \\zeta \\zeta} \\lambda", + name="zeta", + label="\\zeta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt radius and " - + "toroidal angle twice", + description="Toroidal angular coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 0, 2]]}, + params=[], + transforms={"grid": []}, profiles=[], - coordinates="rtz", + coordinates="z", data=[], ) -def _lambda_rzz(params, transforms, profiles, data, **kwargs): - data["lambda_rzz"] = transforms["L"].transform(params["L_lmn"], 1, 0, 2) +def _zeta(params, transforms, profiles, data, **kwargs): + data["zeta"] = transforms["grid"].nodes[:, 2] return data @register_compute_fun( - name="lambda_ttz", - label="\\partial_{\\theta \\theta \\zeta} \\lambda", + name="zeta_r", + label="\\partial_{\\rho} \\zeta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt poloidal angle " - + "twice and toroidal angle", + description="Toroidal angular coordinate derivative, wrt radial coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 2, 1]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="z", + data=["0"], ) -def _lambda_ttz(params, transforms, profiles, data, **kwargs): - data["lambda_ttz"] = transforms["L"].transform(params["L_lmn"], 0, 2, 1) +def _zeta_r(params, transforms, profiles, data, **kwargs): + data["zeta_r"] = data["0"] return data @register_compute_fun( - name="lambda_tzz", - label="\\partial_{\\theta \\zeta \\zeta} \\lambda", + name="zeta_t", + label="\\partial_{\\theta} \\zeta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt poloidal angle " - + "and toroidal angle twice", + description="Toroidal angular coordinate, derivative wrt poloidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[0, 1, 2]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="z", + data=["0"], ) -def _lambda_tzz(params, transforms, profiles, data, **kwargs): - data["lambda_tzz"] = transforms["L"].transform(params["L_lmn"], 0, 1, 2) +def _zeta_t(params, transforms, profiles, data, **kwargs): + data["zeta_t"] = data["0"] return data @register_compute_fun( - name="lambda_rtz", - label="\\partial_{\\rho \\theta \\zeta} \\lambda", + name="zeta_z", + label="\\partial_{\\zeta} \\zeta", units="rad", units_long="radians", - description="Poloidal stream function, third derivative wrt radius, poloidal " - + " angle, and toroidal angle", + description="Toroidal angular coordinate, derivative wrt toroidal coordinate", dim=1, - params=["L_lmn"], - transforms={"L": [[1, 1, 1]]}, + params=[], + transforms={}, profiles=[], - coordinates="rtz", - data=[], + coordinates="z", + data=["0"], ) -def _lambda_rtz(params, transforms, profiles, data, **kwargs): - data["lambda_rtz"] = transforms["L"].transform(params["L_lmn"], 1, 1, 1) +def _zeta_z(params, transforms, profiles, data, **kwargs): + data["zeta_z"] = jnp.ones_like(data["0"]) return data diff --git a/desc/io/equilibrium_io.py b/desc/io/equilibrium_io.py index 9ceb29cb96..858bb1620f 100644 --- a/desc/io/equilibrium_io.py +++ b/desc/io/equilibrium_io.py @@ -68,7 +68,7 @@ def load(load_from, file_format=None): ) else: raise ValueError("Unknown file format: {}".format(file_format)) - # to set other secondary stuff that wasnt saved possibly: + # to set other secondary stuff that wasn't saved possibly: if hasattr(obj, "_set_up"): obj._set_up() return obj diff --git a/desc/transform.py b/desc/transform.py index 63453d6e35..dd2ab14eb3 100644 --- a/desc/transform.py +++ b/desc/transform.py @@ -58,9 +58,9 @@ def __init__( self._rcond = rcond if rcond is not None else "auto" if ( - not np.all(self.grid.nodes[:, 2] == 0) + np.any(self.grid.nodes[:, 2] != 0) and self.basis.N != 0 - and not (self.grid.NFP == self.basis.NFP) + and self.grid.NFP != self.basis.NFP and grid.node_pattern != "custom" ): warnings.warn( @@ -72,29 +72,19 @@ def __init__( ) ) + self._built = False + self._built_pinv = False self._derivatives = self._get_derivatives(derivs) self._sort_derivatives() self._method = method - - self._built = False - self._built_pinv = False - self._set_up() + # assign according to logic in setter function + self.method = method + self._matrices = self._get_matrices() if build: self.build() if build_pinv: self.build_pinv() - def _set_up(self): - - self.method = self._method - self._matrices = { - "direct1": { - i: {j: {k: {} for k in range(4)} for j in range(4)} for i in range(4) - }, - "fft": {i: {j: {} for j in range(4)} for i in range(4)}, - "direct2": {i: {} for i in range(4)}, - } - def _get_derivatives(self, derivs): """Get array of derivatives needed for calculating objective function. @@ -139,6 +129,18 @@ def _sort_derivatives(self): ) self._derivatives = self.derivatives[sort_idx] + def _get_matrices(self): + """Get matrices to compute all derivatives.""" + n = np.amax(self.derivatives) + 1 + matrices = { + "direct1": { + i: {j: {k: {} for k in range(n)} for j in range(n)} for i in range(n) + }, + "fft": {i: {j: {} for j in range(n)} for i in range(n)}, + "direct2": {i: {} for i in range(n)}, + } + return matrices + def _check_inputs_fft(self, grid, basis): """Check that inputs are formatted correctly for fft method.""" if grid.num_nodes == 0 or basis.num_modes == 0: @@ -371,7 +373,7 @@ def build(self): if self.method == "direct1": for d in self.derivatives: - self._matrices["direct1"][d[0]][d[1]][d[2]] = self.basis.evaluate( + self.matrices["direct1"][d[0]][d[1]][d[2]] = self.basis.evaluate( self.grid.nodes, d, unique=True ) @@ -405,7 +407,7 @@ def build_pinv(self): rcond = None if self.rcond == "auto" else self.rcond if self.method == "direct1": A = self.basis.evaluate(self.grid.nodes, np.array([0, 0, 0])) - self._matrices["pinv"] = ( + self.matrices["pinv"] = ( scipy.linalg.pinv(A, rcond=rcond) if A.size else np.zeros_like(A.T) ) elif self.method == "direct2": @@ -472,41 +474,36 @@ def transform(self, c, dr=0, dt=0, dz=0): return np.zeros(self.grid.num_nodes) if self.method == "direct1": - A = self.matrices["direct1"][dr][dt][dz] + A = self.matrices["direct1"].get(dr, {}).get(dt, {}).get(dz, {}) if isinstance(A, dict): raise ValueError( colored("Derivative orders are out of initialized bounds", "red") ) - return jnp.matmul(A, c) + return A @ c elif self.method == "direct2": - A = self.matrices["fft"][dr][dt] - B = self.matrices["direct2"][dz] - + A = self.matrices["fft"].get(dr, {}).get(dt, {}) + B = self.matrices["direct2"].get(dz, {}) if isinstance(A, dict) or isinstance(B, dict): raise ValueError( colored("Derivative orders are out of initialized bounds", "red") ) c_mtrx = jnp.zeros((self.num_lm_modes * self.num_n_modes,)) c_mtrx = put(c_mtrx, self.fft_index, c).reshape((-1, self.num_n_modes)) - - cc = jnp.matmul(A, c_mtrx) - return jnp.matmul(cc, B.T).flatten(order="F") + cc = A @ c_mtrx + return (cc @ B.T).flatten(order="F") elif self.method == "fft": - A = self.matrices["fft"][dr][dt] + A = self.matrices["fft"].get(dr, {}).get(dt, {}) if isinstance(A, dict): raise ValueError( colored("Derivative orders are out of initialized bounds", "red") ) - # reshape coefficients c_mtrx = jnp.zeros((self.num_lm_modes * self.num_n_modes,)) c_mtrx = put(c_mtrx, self.fft_index, c).reshape((-1, self.num_n_modes)) - # differentiate c_diff = c_mtrx[:, :: (-1) ** dz] * self.dk**dz * (-1) ** (dz > 1) - # re-format in complex notation c_real = jnp.pad( (self.num_z_nodes / 2) @@ -521,10 +518,9 @@ def transform(self, c, dr=0, dt=0, dz=0): jnp.fliplr(jnp.conj(c_real)), ) ) - # transform coefficients c_fft = jnp.real(jnp.fft.ifft(c_cplx)) - return jnp.matmul(A, c_fft).flatten(order="F") + return (A @ c_fft).flatten(order="F") def fit(self, x): """Transform from physical domain to spectral using weighted least squares fit. @@ -731,27 +727,20 @@ def change_derivatives(self, derivs, build=True): self._sort_derivatives() if len(derivs_to_add): - # if we actually added derivatives and didn't build them, then its not built + # if we actually added derivatives and didn't build them, then it's not + # built self._built = False if build: - # we don't update self._built here because it is still built from before + # we don't update self._built here because it is still built from before, # but it still might have unbuilt matrices from new derivatives self.build() @property def matrices(self): """dict: transform matrices such that x=A*c.""" - return self.__dict__.setdefault( - "_matrices", - { - "direct1": { - i: {j: {k: {} for k in range(4)} for j in range(4)} - for i in range(4) - }, - "fft": {i: {j: {} for j in range(4)} for i in range(4)}, - "direct2": {i: {} for i in range(4)}, - }, - ) + if not hasattr(self, "_matrices"): + self._matrices = self._get_matrices() + return self._matrices @property def num_nodes(self): diff --git a/setup.cfg b/setup.cfg index 87f6c23f57..0870f8b3b8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -47,6 +47,8 @@ markers= filterwarnings= error ignore::pytest.PytestUnraisableExceptionWarning + ignore::RuntimeWarning:desc.compute + # Ignore division by zero warnings. [flake8] # Primarily ignoring whitespace, indentation, and commenting etiquette that black does not catch diff --git a/tests/test_basis.py b/tests/test_basis.py index 269a8f47bc..680f754a87 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -142,44 +142,72 @@ def test_chebyshev(self): chebyshev(r[:, np.newaxis], l, dr=1) @pytest.mark.unit - def test_zernike_radial(self): + def test_zernike_radial(self): # noqa: C901 """Test zernike_radial function, comparing to analytic formulas.""" + # https://en.wikipedia.org/wiki/Zernike_polynomials#Radial_polynomials + + def Z3_1(x, dx=0): + if dx == 0: + return 3 * x**3 - 2 * x + if dx == 1: + return 9 * x**2 - 2 + if dx == 2: + return 18 * x + if dx == 3: + return np.full_like(x, 18) + if dx >= 4: + return np.zeros_like(x) + + def Z4_2(x, dx=0): + if dx == 0: + return 4 * x**4 - 3 * x**2 + if dx == 1: + return 16 * x**3 - 6 * x + if dx == 2: + return 48 * x**2 - 6 + if dx == 3: + return 96 * x + if dx == 4: + return np.full_like(x, 96) + if dx >= 5: + return np.zeros_like(x) + + def Z6_2(x, dx=0): + if dx == 0: + return 15 * x**6 - 20 * x**4 + 6 * x**2 + if dx == 1: + return 90 * x**5 - 80 * x**3 + 12 * x + if dx == 2: + return 450 * x**4 - 240 * x**2 + 12 + if dx == 3: + return 1800 * x**3 - 480 * x + if dx == 4: + return 5400 * x**2 - 480 + if dx == 5: + return 10800 * x + if dx == 6: + return np.full_like(x, 10800) + if dx >= 7: + return np.zeros_like(x) + l = np.array([3, 4, 6]) m = np.array([1, 2, 2]) r = np.linspace(0, 1, 11) # rho coordinates - - # correct value functions - def Z3_1(x): - return 3 * x**3 - 2 * x - - def Z4_2(x): - return 4 * x**4 - 3 * x**2 - - def Z6_2(x): - return 15 * x**6 - 20 * x**4 + 6 * x**2 - - # correct derivative functions - def dZ3_1(x): - return 9 * x**2 - 2 - - def dZ4_2(x): - return 16 * x**3 - 6 * x - - def dZ6_2(x): - return 90 * x**5 - 80 * x**3 + 12 * x - - correct_vals = np.array([Z3_1(r), Z4_2(r), Z6_2(r)]).T - correct_ders = np.array([dZ3_1(r), dZ4_2(r), dZ6_2(r)]).T - - values1 = zernike_radial(r[:, np.newaxis], l, m, 0) - derivs1 = zernike_radial(r[:, np.newaxis], l, m, 1) - values2 = zernike_radial_poly(r[:, np.newaxis], l, m, 0) - derivs2 = zernike_radial_poly(r[:, np.newaxis], l, m, 1) - - np.testing.assert_allclose(values1, correct_vals, atol=1e-8) - np.testing.assert_allclose(derivs1, correct_ders, atol=1e-8) - np.testing.assert_allclose(values2, correct_vals, atol=1e-8) - np.testing.assert_allclose(derivs2, correct_ders, atol=1e-8) + max_dr = 4 + desired = { + dr: np.array([Z3_1(r, dr), Z4_2(r, dr), Z6_2(r, dr)]).T + for dr in range(max_dr + 1) + } + radial = { + dr: zernike_radial(r[:, np.newaxis], l, m, dr) for dr in range(max_dr + 1) + } + radial_poly = { + dr: zernike_radial_poly(r[:, np.newaxis], l, m, dr) + for dr in range(max_dr + 1) + } + for dr in range(max_dr + 1): + np.testing.assert_allclose(radial[dr], desired[dr], err_msg=dr) + np.testing.assert_allclose(radial_poly[dr], desired[dr], err_msg=dr) @pytest.mark.unit def test_fourier(self): diff --git a/tests/test_curves.py b/tests/test_curves.py index bfc8ebb5cd..ab8f767990 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -256,7 +256,7 @@ def test_misc(self): def test_asserts(self): """Test error checking when creating FourierXYZCurve.""" c = FourierXYZCurve() - with pytest.raises(KeyError): + with pytest.raises(ValueError): c.compute_coordinates(dt=4) with pytest.raises(TypeError): c.grid = [1, 2, 3]