diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 305dc087fc..414c014182 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -602,7 +602,6 @@ def _e_sup_theta_rr(params, transforms, profiles, data, **kwargs): / data["sqrt(g)"] ** 2 + 2 * temp.T * data["sqrt(g)_r"] * data["sqrt(g)_r"] / data["sqrt(g)"] ** 3 ).T - return data @@ -1140,7 +1139,6 @@ def _e_sup_zeta_rz(params, transforms, profiles, data, **kwargs): / data["sqrt(g)"] ** 2 + 2 * temp.T * data["sqrt(g)_r"] * data["sqrt(g)_z"] / data["sqrt(g)"] ** 3 ).T - return data @@ -1184,7 +1182,6 @@ def _e_sup_zeta_t(params, transforms, profiles, data, **kwargs): * safediv(data["sqrt(g)_rt"], data["sqrt(g)_r"] ** 2) ).T, ) - return data @@ -1235,7 +1232,6 @@ def _e_sup_zeta_tt(params, transforms, profiles, data, **kwargs): / data["sqrt(g)"] ** 2 + 2 * temp.T * data["sqrt(g)_t"] * data["sqrt(g)_t"] / data["sqrt(g)"] ** 3 ).T - return data @@ -1291,7 +1287,6 @@ def _e_sup_zeta_tz(params, transforms, profiles, data, **kwargs): / data["sqrt(g)"] ** 2 + 2 * temp.T * data["sqrt(g)_t"] * data["sqrt(g)_z"] / data["sqrt(g)"] ** 3 ).T - return data @@ -1335,7 +1330,6 @@ def _e_sup_zeta_z(params, transforms, profiles, data, **kwargs): * safediv(data["sqrt(g)_rz"], data["sqrt(g)_r"] ** 2) ).T, ) - return data @@ -1386,16 +1380,15 @@ def _e_sup_zeta_zz(params, transforms, profiles, data, **kwargs): / data["sqrt(g)"] ** 2 + 2 * temp.T * data["sqrt(g)_z"] * data["sqrt(g)_z"] / data["sqrt(g)"] ** 3 ).T - return data @register_compute_fun( - name="e_phi", - label="\\mathbf{e}_{\\phi}", + name="e_phi|r,t", + label="\\mathbf{e}_{\\phi} |_{\\rho, \\theta}", units="m", units_long="meters", - description="Covariant cylindrical toroidal basis vector", + description="Covariant toroidal basis vector in (ρ,θ,ϕ) coordinates", dim=3, params=[], transforms={}, @@ -1407,11 +1400,13 @@ def _e_sup_zeta_zz(params, transforms, profiles, data, **kwargs): "desc.geometry.surface.FourierRZToroidalSurface", "desc.geometry.core.Surface", ], + aliases=["e_phi"], + # Our usual notation implies e_phi = (∂X/∂ϕ)|R,Z = R ϕ̂, but we need to alias e_phi + # to e_phi|r,t = (∂X/∂ϕ)|ρ,θ for compatibility with older versions of the code. ) -def _e_sub_phi(params, transforms, profiles, data, **kwargs): - # dX/dphi at const r,t = dX/dz * dz/dphi = dX/dz / (dphi/dz) - data["e_phi"] = (data["e_zeta"].T / data["phi_z"]).T - +def _e_sub_phi_rt(params, transforms, profiles, data, **kwargs): + # (∂X/∂ϕ)|ρ,θ = (∂X/∂ζ)|ρ,θ / (∂ϕ/∂ζ)|ρ,θ + data["e_phi|r,t"] = (data["e_zeta"].T / data["phi_z"]).T return data @@ -2434,27 +2429,81 @@ def _e_sub_theta_over_sqrt_g(params, transforms, profiles, data, **kwargs): safediv(data["e_theta"].T, data["sqrt(g)"]).T, lambda: safediv(data["e_theta_r"].T, data["sqrt(g)_r"]).T, ) - return data @register_compute_fun( name="e_theta_PEST", - label="\\mathbf{e}_{\\theta_{PEST}}", + label="\\mathbf{e}_{\\vartheta} |_{\\rho, \\phi} = \\mathbf{e}_{\\theta_{PEST}}", units="m", units_long="meters", - description="Covariant straight field line (PEST) poloidal basis vector", + description="Covariant poloidal basis vector in (ρ,ϑ,ϕ) coordinates or" + " straight field line PEST coordinates. ϕ increases counterclockwise" + " when viewed from above (cylindrical R,ϕ plane with Z out of page).", dim=3, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_theta", "theta_PEST_t"], + data=["e_theta", "theta_PEST_t", "e_zeta", "theta_PEST_z", "phi_t", "phi_z"], + aliases=["e_vartheta"], ) -def _e_sub_theta_pest(params, transforms, profiles, data, **kwargs): - # dX/dv at const r,z = dX/dt * dt/dv / dX/dt / dv/dt - data["e_theta_PEST"] = (data["e_theta"].T / data["theta_PEST_t"]).T +def _e_sub_vartheta_rp(params, transforms, profiles, data, **kwargs): + # constant ρ and ϕ + e_vartheta = ( + data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"] + ) / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"]) + data["e_theta_PEST"] = e_vartheta.T + return data + +@register_compute_fun( + name="e_phi|r,v", + label="\\mathbf{e}_{\\phi} |_{\\rho, \\vartheta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector in (ρ,ϑ,ϕ) coordinates or" + " straight field line PEST coordinates. ϕ increases counterclockwise" + " when viewed from above (cylindrical R,ϕ plane with Z out of page).", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "theta_PEST_t", "e_zeta", "theta_PEST_z", "phi_t", "phi_z"], +) +def _e_sub_phi_rv(params, transforms, profiles, data, **kwargs): + # constant ρ and ϑ + e_phi = ( + data["e_zeta"].T * data["theta_PEST_t"] + - data["e_theta"].T * data["theta_PEST_z"] + ) / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"]) + data["e_phi|r,v"] = e_phi.T + return data + + +@register_compute_fun( + name="e_rho|v,p", + label="\\mathbf{e}_{\\rho} |_{\\vartheta, \\phi}", + units="m", + units_long="meters", + description="Covariant radial basis vector in (ρ,ϑ,ϕ) coordinates or" + " straight field line PEST coordinates. ϕ increases counterclockwise" + " when viewed from above (cylindrical R,ϕ plane with Z out of page).", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_vartheta", "e_phi|r,v", "theta_PEST_r", "phi_r"], +) +def _e_sub_rho_vp(params, transforms, profiles, data, **kwargs): + # constant ϑ and ϕ + data["e_rho|v,p"] = ( + data["e_rho"].T + - data["e_vartheta"].T * data["theta_PEST_r"] + - data["e_phi|r,v"].T * data["phi_r"] + ).T return data diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index 42dd5c8218..3a842e9378 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -40,17 +40,21 @@ def _sqrtg(params, transforms, profiles, data, **kwargs): label="\\sqrt{g}_{PEST}", units="m^{3}", units_long="cubic meters", - description="Jacobian determinant of PEST flux coordinate system", + description="Jacobian determinant of (ρ,ϑ,ϕ) coordinate system or" + " straight field line PEST coordinates. ϕ increases counterclockwise" + " when viewed from above (cylindrical R,ϕ plane with Z out of page).", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["e_rho", "e_theta_PEST", "e_phi"], + data=["sqrt(g)", "theta_PEST_t", "phi_z", "theta_PEST_z", "phi_t"], ) def _sqrtg_pest(params, transforms, profiles, data, **kwargs): - data["sqrt(g)_PEST"] = dot( - data["e_rho"], cross(data["e_theta_PEST"], data["e_phi"]) + # Same as dot(data["e_rho|v,p"], cross(data["e_vartheta"], data["e_phi|r,v"])), but + # more efficient as it avoids computing radial derivatives of the stream functions. + data["sqrt(g)_PEST"] = data["sqrt(g)"] / ( + data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"] ) return data diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 5e76934ba6..5fb84fa060 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -3,6 +3,8 @@ from .data_index import register_compute_fun from .geom_utils import rpz2xyz +# TODO: review when zeta no longer equals phi + @register_compute_fun( name="x", @@ -25,6 +27,7 @@ def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"]) Z = transforms["Z"].transform(params["Z_lmn"]) + # TODO: change when zeta no longer equals phi phi = transforms["grid"].nodes[:, 2] coords = jnp.stack([R, phi, Z], axis=1) # default basis for "x" is rpz, the conversion will be done diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index b6c10ec359..f230d7a816 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/inputs/master_compute_data_xyz.pkl b/tests/inputs/master_compute_data_xyz.pkl index d998fa7d9b..0de07ede29 100644 Binary files a/tests/inputs/master_compute_data_xyz.pkl and b/tests/inputs/master_compute_data_xyz.pkl differ