diff --git a/desc/compute/_core.py b/desc/compute/_core.py index 7bc6e8acb..c61b6974e 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -2764,6 +2764,25 @@ def _phi_rr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="phi_rrz", + label="\\partial_{\\rho \\rho \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, second derivative wrt radial coordinate " + "and first wrt DESC toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_rrz"], +) +def _phi_rrz(params, transforms, profiles, data, **kwargs): + data["phi_rrz"] = data["omega_rrz"] + return data + + @register_compute_fun( name="phi_rt", label="\\partial_{\\rho \\theta} \\phi", @@ -2783,6 +2802,25 @@ def _phi_rt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="phi_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, third derivative wrt radial, " + "poloidal, and toroidal coordinates", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_rtz"], +) +def _phi_rtz(params, transforms, profiles, data, **kwargs): + data["phi_rtz"] = data["omega_rtz"] + return data + + @register_compute_fun( name="phi_rz", label="\\partial_{\\rho \\zeta} \\phi", @@ -2802,6 +2840,25 @@ def _phi_rz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="phi_rzz", + label="\\partial_{\\rho \\zeta \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, first derivative wrt radial and " + "second derivative wrt DESC toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_rzz"], +) +def _phi_rzz(params, transforms, profiles, data, **kwargs): + data["phi_rzz"] = data["omega_rzz"] + return data + + @register_compute_fun( name="phi_t", label="\\partial_{\\theta} \\phi", @@ -2843,12 +2900,31 @@ def _phi_tt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="phi_ttz", + label="\\partial_{\\theta \\theta \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, second derivative wrt poloidal " + "coordinate and first derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_ttz"], +) +def _phi_ttz(params, transforms, profiles, data, **kwargs): + data["phi_ttz"] = data["omega_ttz"] + return data + + @register_compute_fun( name="phi_tz", label="\\partial_{\\theta \\zeta} \\phi", units="rad", units_long="radians", - description="Toroidal angle in lab frame, second derivative wrt poloidal and " + description="Toroidal angle in lab frame, derivative wrt poloidal and " "toroidal coordinate", dim=1, params=[], @@ -2862,6 +2938,25 @@ def _phi_tz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="phi_tzz", + label="\\partial_{\\theta \\zeta \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, derivative wrt poloidal coordinate and " + "second derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_tzz"], +) +def _phi_tzz(params, transforms, profiles, data, **kwargs): + data["phi_tzz"] = data["omega_tzz"] + return data + + @register_compute_fun( name="phi_z", label="\\partial_{\\zeta} \\phi", @@ -2903,6 +2998,25 @@ def _phi_zz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="phi_zzz", + label="\\partial_{\\zeta \\zeta \\zeta} \\phi", + units="rad", + units_long="radians", + description="Toroidal angle in lab frame, third derivative wrt toroidal " + "coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["omega_zzz"], +) +def _phi_zzz(params, transforms, profiles, data, **kwargs): + data["phi_zzz"] = data["omega_zzz"] + return data + + @register_compute_fun( name="rho", label="\\rho", @@ -2986,6 +3100,83 @@ def _theta_PEST_r(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="theta_PEST_rt", + label="\\partial_{\\rho \\theta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "radial and DESC poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_rt"], +) +def _theta_PEST_rt(params, transforms, profiles, data, **kwargs): + data["theta_PEST_rt"] = data["lambda_rt"] + return data + + +@register_compute_fun( + name="theta_PEST_rrt", + label="\\partial_{\\rho \\rho \\theta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, second " + "derivative wrt radial coordinate and first derivative wrt DESC poloidal " + "coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_rrt"], +) +def _theta_PEST_rrt(params, transforms, profiles, data, **kwargs): + data["theta_PEST_rrt"] = data["lambda_rrt"] + return data + + +@register_compute_fun( + name="theta_PEST_rtz", + label="\\partial_{\\rho \\theta \\zeta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "radial and DESC poloidal and toroidal coordinates", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_rtz"], +) +def _theta_PEST_rtz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_rtz"] = data["lambda_rtz"] + return data + + +@register_compute_fun( + name="theta_PEST_rtt", + label="\\partial_{\\rho \\theta \\theta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "radial coordinate once and DESC poloidal coordinate twice", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_rtt"], +) +def _theta_PEST_rtt(params, transforms, profiles, data, **kwargs): + data["theta_PEST_rtt"] = data["lambda_rtt"] + return data + + @register_compute_fun( name="theta_PEST_t", label="\\partial_{\\theta} \\vartheta", @@ -3024,6 +3215,25 @@ def _theta_PEST_tt(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="theta_PEST_ttt", + label="\\partial_{\\theta \\theta \\theta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, third " + "derivative wrt poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_ttt"], +) +def _theta_PEST_ttt(params, transforms, profiles, data, **kwargs): + data["theta_PEST_ttt"] = data["lambda_ttt"] + return data + + @register_compute_fun( name="theta_PEST_tz", label="\\partial_{\\theta \\zeta} \\vartheta", @@ -3043,6 +3253,25 @@ def _theta_PEST_tz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="theta_PEST_tzz", + label="\\partial_{\\theta \\zeta \\zeta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "poloidal coordinate once and toroidal coordinate twice", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_tzz"], +) +def _theta_PEST_tzz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_tzz"] = data["lambda_tzz"] + return data + + @register_compute_fun( name="theta_PEST_z", label="\\partial_{\\zeta} \\vartheta", @@ -3081,6 +3310,25 @@ def _theta_PEST_zz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="theta_PEST_ttz", + label="\\partial_{\\theta \\theta \\zeta} \\vartheta", + units="rad", + units_long="radians", + description="PEST straight field line poloidal angular coordinate, second " + "derivative wrt poloidal coordinate and derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["lambda_ttz"], +) +def _theta_PEST_ttz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_ttz"] = data["lambda_ttz"] + return data + + @register_compute_fun( name="zeta", label="\\zeta", diff --git a/desc/compute/_field.py b/desc/compute/_field.py index 31a9d58a1..1ca0adb1f 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -74,11 +74,12 @@ def _B_sup_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["psi_r/sqrt(g)", "iota", "lambda_z", "omega_z"], + data=["psi_r/sqrt(g)", "iota", "phi_z", "lambda_z"], ) def _B_sup_theta(params, transforms, profiles, data, **kwargs): + # Assumes θ = ϑ − λ. data["B^theta"] = data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] + data["iota"] * data["phi_z"] - data["lambda_z"] ) return data @@ -94,11 +95,12 @@ def _B_sup_theta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["psi_r/sqrt(g)", "iota", "lambda_t", "omega_t"], + data=["psi_r/sqrt(g)", "iota", "theta_PEST_t", "omega_t"], ) def _B_sup_zeta(params, transforms, profiles, data, **kwargs): + # Assumes ζ = ϕ − ω. data["B^zeta"] = data["psi_r/sqrt(g)"] * ( - -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + -data["iota"] * data["omega_t"] + data["theta_PEST_t"] ) return data @@ -224,21 +226,18 @@ def _psi_r_over_sqrtg_r(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_r", "iota", "iota_r", + "phi_rz", + "phi_z", "lambda_rz", "lambda_z", - "omega_rz", - "omega_z", ], ) def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): data["B^theta_r"] = data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_rz"] - + data["iota_r"] * data["omega_z"] - + data["iota_r"] + data["iota"] * data["phi_rz"] + + data["iota_r"] * data["phi_z"] - data["lambda_rz"] - ) + data["(psi_r/sqrt(g))_r"] * ( - data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] - ) + ) + data["(psi_r/sqrt(g))_r"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) return data @@ -261,8 +260,8 @@ def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_r", "iota", "iota_r", - "lambda_rt", - "lambda_t", + "theta_PEST_rt", + "theta_PEST_t", "omega_rt", "omega_t", ], @@ -271,9 +270,9 @@ def _B_sup_zeta_r(params, transforms, profiles, data, **kwargs): data["B^zeta_r"] = data["psi_r/sqrt(g)"] * ( -data["iota"] * data["omega_rt"] - data["iota_r"] * data["omega_t"] - + data["lambda_rt"] + + data["theta_PEST_rt"] ) + data["(psi_r/sqrt(g))_r"] * ( - -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + -data["iota"] * data["omega_t"] + data["theta_PEST_t"] ) return data @@ -352,16 +351,14 @@ def _psi_r_over_sqrtg_t(params, transforms, profiles, data, **kwargs): "iota", "lambda_tz", "lambda_z", - "omega_tz", - "omega_z", + "phi_tz", + "phi_z", ], ) def _B_sup_theta_t(params, transforms, profiles, data, **kwargs): data["B^theta_t"] = data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_tz"] - data["lambda_tz"] - ) + data["(psi_r/sqrt(g))_t"] * ( - data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] - ) + data["iota"] * data["phi_tz"] - data["lambda_tz"] + ) + data["(psi_r/sqrt(g))_t"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) return data @@ -383,17 +380,17 @@ def _B_sup_theta_t(params, transforms, profiles, data, **kwargs): "psi_r/sqrt(g)", "(psi_r/sqrt(g))_t", "iota", - "lambda_t", - "lambda_tt", + "theta_PEST_t", + "theta_PEST_tt", "omega_t", "omega_tt", ], ) def _B_sup_zeta_t(params, transforms, profiles, data, **kwargs): data["B^zeta_t"] = data["psi_r/sqrt(g)"] * ( - -data["iota"] * data["omega_tt"] + data["lambda_tt"] + -data["iota"] * data["omega_tt"] + data["theta_PEST_tt"] ) + data["(psi_r/sqrt(g))_t"] * ( - -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + -data["iota"] * data["omega_t"] + data["theta_PEST_t"] ) return data @@ -472,16 +469,14 @@ def _psi_r_over_sqrtg_z(params, transforms, profiles, data, **kwargs): "iota", "lambda_z", "lambda_zz", - "omega_z", - "omega_zz", + "phi_z", + "phi_zz", ], ) def _B_sup_theta_z(params, transforms, profiles, data, **kwargs): data["B^theta_z"] = data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_zz"] - data["lambda_zz"] - ) + data["(psi_r/sqrt(g))_z"] * ( - data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] - ) + data["iota"] * data["phi_zz"] - data["lambda_zz"] + ) + data["(psi_r/sqrt(g))_z"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) return data @@ -503,17 +498,17 @@ def _B_sup_theta_z(params, transforms, profiles, data, **kwargs): "psi_r/sqrt(g)", "(psi_r/sqrt(g))_z", "iota", - "lambda_t", - "lambda_tz", + "theta_PEST_t", + "theta_PEST_tz", "omega_t", "omega_tz", ], ) def _B_sup_zeta_z(params, transforms, profiles, data, **kwargs): data["B^zeta_z"] = data["psi_r/sqrt(g)"] * ( - -data["iota"] * data["omega_tz"] + data["lambda_tz"] + -data["iota"] * data["omega_tz"] + data["theta_PEST_tz"] ) + data["(psi_r/sqrt(g))_z"] * ( - -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + -data["iota"] * data["omega_t"] + data["theta_PEST_t"] ) return data @@ -650,31 +645,28 @@ def _psi_r_over_sqrtg_rr(params, transforms, profiles, data, **kwargs): "lambda_rrz", "lambda_rz", "lambda_z", - "omega_rrz", - "omega_rz", - "omega_z", + "phi_rrz", + "phi_rz", + "phi_z", ], ) def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): data["B^theta_rr"] = ( data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_rrz"] - + 2 * data["iota_r"] * data["omega_rz"] - + data["iota_rr"] * data["omega_z"] - + data["iota_rr"] + data["iota"] * data["phi_rrz"] + + 2 * data["iota_r"] * data["phi_rz"] + + data["iota_rr"] * data["phi_z"] - data["lambda_rrz"] ) + 2 * data["(psi_r/sqrt(g))_r"] * ( - data["iota"] * data["omega_rz"] - + data["iota_r"] * data["omega_z"] - + data["iota_r"] + data["iota"] * data["phi_rz"] + + data["iota_r"] * data["phi_z"] - data["lambda_rz"] ) - + data["(psi_r/sqrt(g))_rr"] - * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) + + data["(psi_r/sqrt(g))_rr"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) ) return data @@ -700,9 +692,9 @@ def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): "iota", "iota_r", "iota_rr", - "lambda_rrt", - "lambda_rt", - "lambda_t", + "theta_PEST_rrt", + "theta_PEST_rt", + "theta_PEST_t", "omega_rrt", "omega_rt", "omega_t", @@ -715,17 +707,17 @@ def _B_sup_zeta_rr(params, transforms, profiles, data, **kwargs): data["iota"] * data["omega_rrt"] + 2 * data["iota_r"] * data["omega_rt"] + data["iota_rr"] * data["omega_t"] - - data["lambda_rrt"] + - data["theta_PEST_rrt"] ) - 2 * data["(psi_r/sqrt(g))_r"] * ( data["iota"] * data["omega_rt"] + data["iota_r"] * data["omega_t"] - - data["lambda_rt"] + - data["theta_PEST_rt"] ) + data["(psi_r/sqrt(g))_rr"] - * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + * (-data["iota"] * data["omega_t"] + data["theta_PEST_t"]) ) return data @@ -820,19 +812,18 @@ def _psi_r_over_sqrtg_tt(params, transforms, profiles, data, **kwargs): "lambda_ttz", "lambda_tz", "lambda_z", - "omega_ttz", - "omega_tz", - "omega_z", + "phi_ttz", + "phi_tz", + "phi_z", ], ) def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): data["B^theta_tt"] = ( - data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) + data["psi_r/sqrt(g)"] * (data["iota"] * data["phi_ttz"] - data["lambda_ttz"]) + 2 * data["(psi_r/sqrt(g))_t"] - * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - + data["(psi_r/sqrt(g))_tt"] - * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) + * (data["iota"] * data["phi_tz"] - data["lambda_tz"]) + + data["(psi_r/sqrt(g))_tt"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) ) return data @@ -856,9 +847,9 @@ def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_t", "(psi_r/sqrt(g))_tt", "iota", - "lambda_t", - "lambda_tt", - "lambda_ttt", + "theta_PEST_t", + "theta_PEST_tt", + "theta_PEST_ttt", "omega_t", "omega_tt", "omega_ttt", @@ -866,12 +857,13 @@ def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): ) def _B_sup_zeta_tt(params, transforms, profiles, data, **kwargs): data["B^zeta_tt"] = ( - -data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_ttt"] - data["lambda_ttt"]) + -data["psi_r/sqrt(g)"] + * (data["iota"] * data["omega_ttt"] - data["theta_PEST_ttt"]) - 2 * data["(psi_r/sqrt(g))_t"] - * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + * (data["iota"] * data["omega_tt"] - data["theta_PEST_tt"]) + data["(psi_r/sqrt(g))_tt"] - * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + * (-data["iota"] * data["omega_t"] + data["theta_PEST_t"]) ) return data @@ -966,19 +958,18 @@ def _psi_r_over_sqrtg_zz(params, transforms, profiles, data, **kwargs): "lambda_z", "lambda_zz", "lambda_zzz", - "omega_z", - "omega_zz", - "omega_zzz", + "phi_z", + "phi_zz", + "phi_zzz", ], ) def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): data["B^theta_zz"] = ( - data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_zzz"] - data["lambda_zzz"]) + data["psi_r/sqrt(g)"] * (data["iota"] * data["phi_zzz"] - data["lambda_zzz"]) + 2 * data["(psi_r/sqrt(g))_z"] - * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) - + data["(psi_r/sqrt(g))_zz"] - * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) + * (data["iota"] * data["phi_zz"] - data["lambda_zz"]) + + data["(psi_r/sqrt(g))_zz"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) ) return data @@ -1002,9 +993,9 @@ def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_z", "(psi_r/sqrt(g))_zz", "iota", - "lambda_t", - "lambda_tz", - "lambda_tzz", + "theta_PEST_t", + "theta_PEST_tz", + "theta_PEST_tzz", "omega_t", "omega_tz", "omega_tzz", @@ -1012,12 +1003,13 @@ def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): ) def _B_sup_zeta_zz(params, transforms, profiles, data, **kwargs): data["B^zeta_zz"] = ( - -data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) + -data["psi_r/sqrt(g)"] + * (data["iota"] * data["omega_tzz"] - data["theta_PEST_tzz"]) - 2 * data["(psi_r/sqrt(g))_z"] - * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + * (data["iota"] * data["omega_tz"] - data["theta_PEST_tz"]) + data["(psi_r/sqrt(g))_zz"] - * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + * (-data["iota"] * data["omega_t"] + data["theta_PEST_t"]) ) return data @@ -1120,31 +1112,29 @@ def _psi_r_over_sqrtg_rt(params, transforms, profiles, data, **kwargs): "lambda_rz", "lambda_tz", "lambda_z", - "omega_rtz", - "omega_rz", - "omega_tz", - "omega_z", + "phi_rtz", + "phi_rz", + "phi_tz", + "phi_z", ], ) def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): data["B^theta_rt"] = ( data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_rtz"] - + data["iota_r"] * data["omega_tz"] + data["iota"] * data["phi_rtz"] + + data["iota_r"] * data["phi_tz"] - data["lambda_rtz"] ) + data["(psi_r/sqrt(g))_r"] - * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + * (data["iota"] * data["phi_tz"] - data["lambda_tz"]) + data["(psi_r/sqrt(g))_t"] * ( - data["iota"] * data["omega_rz"] - + data["iota_r"] * data["omega_z"] - + data["iota_r"] + data["iota"] * data["phi_rz"] + + data["iota_r"] * data["phi_z"] - data["lambda_rz"] ) - + data["(psi_r/sqrt(g))_rt"] - * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) + + data["(psi_r/sqrt(g))_rt"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) ) return data @@ -1170,10 +1160,10 @@ def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_t", "iota", "iota_r", - "lambda_rt", - "lambda_rtt", - "lambda_t", - "lambda_tt", + "theta_PEST_rt", + "theta_PEST_rtt", + "theta_PEST_t", + "theta_PEST_tt", "omega_rt", "omega_rtt", "omega_t", @@ -1186,18 +1176,18 @@ def _B_sup_zeta_rt(params, transforms, profiles, data, **kwargs): * ( data["iota"] * data["omega_rtt"] + data["iota_r"] * data["omega_tt"] - - data["lambda_rtt"] + - data["theta_PEST_rtt"] ) - data["(psi_r/sqrt(g))_r"] - * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + * (data["iota"] * data["omega_tt"] - data["theta_PEST_tt"]) - data["(psi_r/sqrt(g))_t"] * ( data["iota"] * data["omega_rt"] + data["iota_r"] * data["omega_t"] - - data["lambda_rt"] + - data["theta_PEST_rt"] ) + data["(psi_r/sqrt(g))_rt"] - * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + * (-data["iota"] * data["omega_t"] + data["theta_PEST_t"]) ) return data @@ -1307,21 +1297,20 @@ def _psi_r_over_sqrtg_tz(params, transforms, profiles, data, **kwargs): "lambda_tzz", "lambda_z", "lambda_zz", - "omega_tz", - "omega_tzz", - "omega_z", - "omega_zz", + "phi_tz", + "phi_tzz", + "phi_z", + "phi_zz", ], ) def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): data["B^theta_tz"] = ( - data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) + data["psi_r/sqrt(g)"] * (data["iota"] * data["phi_tzz"] - data["lambda_tzz"]) + data["(psi_r/sqrt(g))_t"] - * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) + * (data["iota"] * data["phi_zz"] - data["lambda_zz"]) + data["(psi_r/sqrt(g))_z"] - * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - + data["(psi_r/sqrt(g))_tz"] - * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) + * (data["iota"] * data["phi_tz"] - data["lambda_tz"]) + + data["(psi_r/sqrt(g))_tz"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) ) return data @@ -1346,10 +1335,10 @@ def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_tz", "(psi_r/sqrt(g))_z", "iota", - "lambda_t", - "lambda_tt", - "lambda_ttz", - "lambda_tz", + "theta_PEST_t", + "theta_PEST_tt", + "theta_PEST_ttz", + "theta_PEST_tz", "omega_t", "omega_tt", "omega_ttz", @@ -1358,13 +1347,14 @@ def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): ) def _B_sup_zeta_tz(params, transforms, profiles, data, **kwargs): data["B^zeta_tz"] = ( - -data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) + -data["psi_r/sqrt(g)"] + * (data["iota"] * data["omega_ttz"] - data["theta_PEST_ttz"]) - data["(psi_r/sqrt(g))_t"] - * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + * (data["iota"] * data["omega_tz"] - data["theta_PEST_tz"]) - data["(psi_r/sqrt(g))_z"] - * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + * (data["iota"] * data["omega_tt"] - data["theta_PEST_tt"]) + data["(psi_r/sqrt(g))_tz"] - * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + * (-data["iota"] * data["omega_t"] + data["theta_PEST_t"]) ) return data @@ -1473,31 +1463,29 @@ def _psi_r_over_sqrtg_rz(params, transforms, profiles, data, **kwargs): "lambda_rzz", "lambda_z", "lambda_zz", - "omega_rz", - "omega_rzz", - "omega_z", - "omega_zz", + "phi_rz", + "phi_rzz", + "phi_z", + "phi_zz", ], ) def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): data["B^theta_rz"] = ( data["psi_r/sqrt(g)"] * ( - data["iota"] * data["omega_rzz"] - + data["iota_r"] * data["omega_zz"] + data["iota"] * data["phi_rzz"] + + data["iota_r"] * data["phi_zz"] - data["lambda_rzz"] ) + data["(psi_r/sqrt(g))_r"] - * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) + * (data["iota"] * data["phi_zz"] - data["lambda_zz"]) + data["(psi_r/sqrt(g))_z"] * ( - data["iota"] * data["omega_rz"] - + data["iota_r"] * data["omega_z"] - + data["iota_r"] + data["iota"] * data["phi_rz"] + + data["iota_r"] * data["phi_z"] - data["lambda_rz"] ) - + data["(psi_r/sqrt(g))_rz"] - * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) + + data["(psi_r/sqrt(g))_rz"] * (data["iota"] * data["phi_z"] - data["lambda_z"]) ) return data @@ -1523,10 +1511,10 @@ def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): "(psi_r/sqrt(g))_z", "iota", "iota_r", - "lambda_rt", - "lambda_rtz", - "lambda_t", - "lambda_tz", + "theta_PEST_rt", + "theta_PEST_rtz", + "theta_PEST_t", + "theta_PEST_tz", "omega_rt", "omega_rtz", "omega_t", @@ -1539,18 +1527,18 @@ def _B_sup_zeta_rz(params, transforms, profiles, data, **kwargs): * ( data["iota"] * data["omega_rtz"] + data["iota_r"] * data["omega_tz"] - - data["lambda_rtz"] + - data["theta_PEST_rtz"] ) - data["(psi_r/sqrt(g))_r"] - * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + * (data["iota"] * data["omega_tz"] - data["theta_PEST_tz"]) - data["(psi_r/sqrt(g))_z"] * ( data["iota"] * data["omega_rt"] + data["iota_r"] * data["omega_t"] - - data["lambda_rt"] + - data["theta_PEST_rt"] ) + data["(psi_r/sqrt(g))_rz"] - * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + * (-data["iota"] * data["omega_t"] + data["theta_PEST_t"]) ) return data @@ -1654,6 +1642,25 @@ def _B_sub_zeta(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="B_phi|r,t", + label="B_{\\phi} = B \\dot \\mathbf{e}_{\\phi} |_{\\rho, \\theta}", + units="T \\cdot m", + units_long="Tesla * meters", + description="Covariant toroidal component of magnetic field in (ρ,θ,ϕ) " + "coordinates.", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B", "e_phi|r,t"], +) +def _B_sub_phi_rt(params, transforms, profiles, data, **kwargs): + data["B_phi|r,t"] = dot(data["B"], data["e_phi|r,t"]) + return data + + @register_compute_fun( name="B_rho_r", label="\\partial_{\\rho} B_{\\rho}", diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index fa4397689..cbc54561c 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -23,7 +23,7 @@ units="T \\cdot m}", units_long="Tesla * meters", description="Fourier coefficients for covariant poloidal component of " - + "magnetic field", + "magnetic field.", dim=1, params=[], transforms={"B": [[0, 0, 0]]}, @@ -39,25 +39,27 @@ def _B_theta_mn(params, transforms, profiles, data, **kwargs): return data +# TODO: do math to change definition of nu so that we can just use B_zeta_mn here @register_compute_fun( - name="B_zeta_mn", - label="B_{\\zeta, m, n}", + name="B_phi_mn", + label="B_{\\phi, m, n}", units="T \\cdot m}", units_long="Tesla * meters", description="Fourier coefficients for covariant toroidal component of " - + "magnetic field", + "magnetic field in (ρ,θ,ϕ) coordinates.", dim=1, params=[], transforms={"B": [[0, 0, 0]]}, profiles=[], coordinates="rtz", - data=["B_zeta"], + data=["B_phi|r,t"], M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", resolution_requirement="tz", + aliases="B_zeta_mn", # TODO: remove when phi != zeta ) -def _B_zeta_mn(params, transforms, profiles, data, **kwargs): - data["B_zeta_mn"] = transforms["B"].fit(data["B_zeta"]) +def _B_phi_mn(params, transforms, profiles, data, **kwargs): + data["B_phi_mn"] = transforms["B"].fit(data["B_phi|r,t"]) return data @@ -73,7 +75,7 @@ def _B_zeta_mn(params, transforms, profiles, data, **kwargs): transforms={"w": [[0, 0, 0]], "B": [[0, 0, 0]]}, profiles=[], coordinates="rtz", - data=["B_theta_mn", "B_zeta_mn"], + data=["B_theta_mn", "B_phi_mn"], M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) @@ -89,7 +91,7 @@ def _w_mn(params, transforms, profiles, data, **kwargs): num_t = (mask_t @ sign(wn)) * data["B_theta_mn"] den_t = mask_t @ jnp.abs(wm) - num_z = (mask_z @ sign(wm)) * data["B_zeta_mn"] + num_z = (mask_z @ sign(wm)) * data["B_phi_mn"] den_z = mask_z @ jnp.abs(NFP * wn) w_mn = jnp.where(mask_t.any(axis=0), mask_t.T @ safediv(num_t, den_t), w_mn) @@ -233,10 +235,10 @@ def _nu_z(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["theta", "lambda", "iota", "nu"], + data=["theta_PEST", "iota", "nu"], ) def _theta_B(params, transforms, profiles, data, **kwargs): - data["theta_B"] = data["theta"] + data["lambda"] + data["iota"] * data["nu"] + data["theta_B"] = data["theta_PEST"] + data["iota"] * data["nu"] return data @@ -251,10 +253,10 @@ def _theta_B(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["zeta", "nu"], + data=["phi", "nu"], ) def _zeta_B(params, transforms, profiles, data, **kwargs): - data["zeta_B"] = data["zeta"] + data["nu"] + data["zeta_B"] = data["phi"] + data["nu"] return data @@ -263,18 +265,20 @@ def _zeta_B(params, transforms, profiles, data, **kwargs): label="\\sqrt{g}_{B}", units="~", units_long="None", - description="Jacobian determinant of Boozer coordinates", + description="Jacobian determinant from Boozer to DESC coordinates", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["lambda_t", "lambda_z", "nu_t", "nu_z", "iota"], + data=["theta_PEST_t", "theta_PEST_z", "phi_t", "phi_z", "nu_t", "nu_z", "iota"], ) def _sqrtg_B(params, transforms, profiles, data, **kwargs): - data["sqrt(g)_B"] = (1 + data["lambda_t"]) * (1 + data["nu_z"]) + ( - data["iota"] - data["lambda_z"] - ) * data["nu_t"] + data["sqrt(g)_B"] = ( + data["theta_PEST_t"] * (data["phi_z"] + data["nu_z"]) + - data["theta_PEST_z"] * (data["phi_t"] + data["nu_t"]) + + data["iota"] * (data["nu_t"] * data["phi_z"] - data["nu_z"] * data["phi_t"]) + ) return data diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index 60a55125a..940a46395 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -764,6 +764,7 @@ def _iota(params, transforms, profiles, data, **kwargs): data["iota"] = profiles["iota"].compute(transforms["grid"], params["i_l"], dr=0) elif profiles["current"] is not None: # See the document attached to GitHub pull request #556 for the math. + # Assumes ζ = ϕ − ω and θ = ϑ − λ. data["iota"] = transforms["grid"].replace_at_axis( safediv(data["iota_num"], data["iota_den"]), lambda: safediv(data["iota_num_r"], data["iota_den_r"]), @@ -792,6 +793,7 @@ def _iota_r(params, transforms, profiles, data, **kwargs): ) elif profiles["current"] is not None: # See the document attached to GitHub pull request #556 for the math. + # Assumes ζ = ϕ − ω and θ = ϑ − λ. data["iota_r"] = transforms["grid"].replace_at_axis( safediv( data["iota_num_r"] * data["iota_den"] @@ -835,6 +837,7 @@ def _iota_rr(params, transforms, profiles, data, **kwargs): ) elif profiles["current"] is not None: # See the document attached to GitHub pull request #556 for the math. + # Assumes ζ = ϕ − ω and θ = ϑ − λ. data["iota_rr"] = transforms["grid"].replace_at_axis( safediv( data["iota_num_rr"] * data["iota_den"] ** 2 @@ -922,7 +925,7 @@ def _iota_num_current(params, transforms, profiles, data, **kwargs): iota = profiles["iota"].compute(transforms["grid"], params["i_l"], dr=0) data["iota_num current"] = iota * data["iota_den"] - data["iota_num vacuum"] elif profiles["current"] is not None: - # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current + # 4π² I = 4π² (μ₀ current / 2π) = 2π μ₀ current current = profiles["current"].compute(transforms["grid"], params["c_l"], dr=0) current_r = profiles["current"].compute(transforms["grid"], params["c_l"], dr=1) data["iota_num current"] = ( @@ -954,6 +957,7 @@ def _iota_num_current(params, transforms, profiles, data, **kwargs): ) def _iota_num_vacuum(params, transforms, profiles, data, **kwargs): """Vacuum contribution to the numerator of rotational transform formula.""" + # Assumes ζ = ϕ − ω and θ = ϑ − λ. iota_num_vacuum = transforms["grid"].replace_at_axis( safediv( data["lambda_z"] * data["g_tt"] - (1 + data["lambda_t"]) * data["g_tz"], @@ -1035,6 +1039,7 @@ def _iota_num_r_current(params, transforms, profiles, data, **kwargs): resolution_requirement="tz", ) def _iota_num_r_vacuum(params, transforms, profiles, data, **kwargs): + # Assumes ζ = ϕ − ω and θ = ϑ − λ. iota_num_vacuum = safediv( data["lambda_z"] * data["g_tt"] - (1 + data["lambda_t"]) * data["g_tz"], data["sqrt(g)"], @@ -1154,11 +1159,11 @@ def _iota_num_rr(params, transforms, profiles, data, **kwargs): Computes d2(𝛼+𝛽)/d𝜌2 as defined in the document attached to the description of GitHub pull request #556. 𝛼 supplements the rotational transform with an additional term to account for the enclosed net toroidal current. + Assumes ζ = ϕ − ω and θ = ϑ − λ. """ if profiles["iota"] is not None: data["iota_num_rr"] = jnp.nan * data["0"] elif profiles["current"] is not None: - # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current current = profiles["current"].compute(transforms["grid"], params["c_l"], dr=0) current_r = profiles["current"].compute(transforms["grid"], params["c_l"], dr=1) current_rr = profiles["current"].compute( @@ -1167,6 +1172,7 @@ def _iota_num_rr(params, transforms, profiles, data, **kwargs): current_rrr = profiles["current"].compute( transforms["grid"], params["c_l"], dr=3 ) + # 4π² I = 4π² (μ₀ current / 2π) = 2π μ₀ current alpha_rr = ( jnp.pi * mu_0 @@ -1283,6 +1289,7 @@ def _iota_num_rrr(params, transforms, profiles, data, **kwargs): Computes d3(𝛼+𝛽)/d𝜌3 as defined in the document attached to the description of GitHub pull request #556. 𝛼 supplements the rotational transform with an additional term to account for the enclosed net toroidal current. + Assumes ζ = ϕ − ω and θ = ϑ − λ. """ if profiles["iota"] is not None: data["iota_num_rrr"] = jnp.nan * data["0"] @@ -1298,7 +1305,7 @@ def _iota_num_rrr(params, transforms, profiles, data, **kwargs): current_rrrr = profiles["current"].compute( transforms["grid"], params["c_l"], dr=4 ) - # 4π^2 I = 4π^2 (mu_0 current / 2π) = 2π mu_0 current + # 4π² I = 4π² (μ₀ current / 2π) = 2π μ₀ current alpha_rrr = ( jnp.pi * mu_0 @@ -1402,14 +1409,14 @@ def _iota_den(params, transforms, profiles, data, **kwargs): """Denominator of rotational transform formula. Computes 𝛾 as defined in the document attached to the description - of GitHub pull request #556. + of GitHub pull request #556. Assumes ζ = ϕ − ω and θ = ϑ − λ. """ gamma = safediv( (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"], data["sqrt(g)"], ) # Assumes toroidal stream function behaves such that the magnetic axis limit - # of gamma is zero (as it would if omega = 0 identically). + # of γ is zero (as it would if ω = 0 identically). gamma = transforms["grid"].replace_at_axis( surface_integrals(transforms["grid"], gamma), 0 ) @@ -1447,7 +1454,7 @@ def _iota_den_r(params, transforms, profiles, data, **kwargs): """Denominator of rotational transform formula, first radial derivative. Computes d𝛾/d𝜌 as defined in the document attached to the description - of GitHub pull request #556. + of GitHub pull request #556. Assumes ζ = ϕ − ω and θ = ϑ − λ. """ gamma = safediv( (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"], @@ -1514,7 +1521,7 @@ def _iota_den_rr(params, transforms, profiles, data, **kwargs): """Denominator of rotational transform formula, second radial derivative. Computes d2𝛾/d𝜌2 as defined in the document attached to the description - of GitHub pull request #556. + of GitHub pull request #556. Assumes ζ = ϕ − ω and θ = ϑ − λ. """ gamma = safediv( (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"], @@ -1609,7 +1616,7 @@ def _iota_den_rrr(params, transforms, profiles, data, **kwargs): """Denominator of rotational transform formula, third radial derivative. Computes d3𝛾/d𝜌3 as defined in the document attached to the description - of GitHub pull request #556. + of GitHub pull request #556. Assumes ζ = ϕ − ω and θ = ϑ − λ. """ gamma = safediv( (1 + data["omega_z"]) * data["g_tt"] - data["omega_t"] * data["g_tz"], @@ -1675,9 +1682,12 @@ def _iota_den_rrr(params, transforms, profiles, data, **kwargs): axis_limit_data=["iota_rr", "psi_rr"], ) def _iota_psi(params, transforms, profiles, data, **kwargs): - # Existence of limit at magnetic axis requires ∂ᵨ iota = 0 at axis. - # Assume iota may be expanded as an even power series of ρ so that this - # condition is satisfied. + """∂ι/∂ψ. + + Existence of limit at magnetic axis requires ∂ι/∂ρ = 0 at axis. + Assume ι may be expanded as an even power series of ρ so that this + condition is satisfied. + """ data["iota_psi"] = transforms["grid"].replace_at_axis( safediv(data["iota_r"], data["psi_r"]), lambda: safediv(data["iota_rr"], data["psi_rr"]), diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 7591194c4..eef5bbf2f 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/test_axis_limits.py b/tests/test_axis_limits.py index de9ad815a..8c847ef3a 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -24,10 +24,12 @@ # functions tend toward zero as the magnetic axis is approached and that # d²ψ/(dρ)² and 𝜕√𝑔/𝜕𝜌 are both finite nonzero at the magnetic axis. # Also, dⁿψ/(dρ)ⁿ for n > 3 is assumed zero everywhere. -zero_limits = {"rho", "psi", "psi_r", "e_theta", "sqrt(g)", "B_t"} +zero_limits = {"rho", "psi", "psi_r", "psi_rrr", "e_theta", "sqrt(g)", "B_t"} + # These compute quantities require kinetic profiles, which are not defined for all # configurations (giving NaN values) not_continuous_limits = {"current Redl", "P_ISS04", "P_fusion", ""} + not_finite_limits = { "D_Mercier", "D_geodesic", diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py index aff0345e8..308138b26 100644 --- a/tests/test_compute_everything.py +++ b/tests/test_compute_everything.py @@ -80,8 +80,21 @@ def _compare_against_rpz(p, data, data_rpz, coordinate_conversion_func): def test_compute_everything(): """Test that the computations on this branch agree with those on master. - Also make sure we can compute everything without errors. Computed quantities - are both in "rpz" and "xyz" basis. + Also make sure we can compute everything without errors. + + Notes + ----- + This test will fail if the benchmark file has been updated on both + the local and upstream branches and git cannot resolve the merge + conflict. In that case, please regenerate the benchmark file. + Here are instructions for convenience. + + 1. Prepend true to the line near the end of this test. + ``if True or (not error_rpz and update_master_data_rpz):`` + 2. Run pytest -k test_compute_everything + 3. Revert 1. + 4. git add tests/inputs/master_compute_data_rpz.pkl + """ elliptic_cross_section_with_torsion = { "R_lmn": [10, 1, 0.2],