diff --git a/desc/compute/_equil.py b/desc/compute/_equil.py index 7901810665..7cd01491ed 100644 --- a/desc/compute/_equil.py +++ b/desc/compute/_equil.py @@ -9,6 +9,7 @@ expensive computations. """ +from interpax import interp1d from scipy.constants import mu_0 from desc.backend import jnp @@ -803,3 +804,42 @@ def _beta_volpol(params, transforms, profiles, data, **kwargs): def _beta_voltor(params, transforms, profiles, data, **kwargs): data["_vol"] = jnp.abs(data["W_p"] / data["W_Btor"]) return data + + +@register_compute_fun( + name="P_ISS04", + label="P_{ISS04}", + units="W", + units_long="Watts", + description="Heating power required by the ISS04 energy confinement time scaling", + dim=0, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="", + data=["a", "iota", "rho", "R0", "W_p", "_vol", "<|B|>_axis"], + method="str: Interpolation method. Default 'cubic'.", + H_ISS04="float: ISS04 confinement enhancement factor. Default 1.", +) +def _P_ISS04(params, transforms, profiles, data, **kwargs): + rho = transforms["grid"].compress(data["rho"], surface_label="rho") + iota = transforms["grid"].compress(data["iota"], surface_label="rho") + fx = {} + if "iota_r" in data: + fx["fx"] = transforms["grid"].compress( + data["iota_r"] + ) # noqa: unused dependency + iota_23 = interp1d(2 / 3, rho, iota, method=kwargs.get("method", "cubic"), **fx) + data["P_ISS04"] = 1e6 * ( # MW -> W + jnp.abs(data["W_p"] / 1e6) # J -> MJ + / ( + 0.134 + * data["a"] ** 2.28 # m + * data["R0"] ** 0.64 # m + * (data["_vol"] / 1e19) ** 0.54 # 1/m^3 -> 1e19/m^3 + * data["<|B|>_axis"] ** 0.84 # T + * iota_23**0.41 + * kwargs.get("H_ISS04", 1) + ) + ) ** (1 / 0.39) + return data diff --git a/desc/compute/_field.py b/desc/compute/_field.py index f2e22d4236..31a9d58a18 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -24,7 +24,7 @@ @register_compute_fun( - name="B0", + name="psi_r/sqrt(g)", label="\\psi' / \\sqrt{g}", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -37,8 +37,8 @@ data=["psi_r", "sqrt(g)"], axis_limit_data=["psi_rr", "sqrt(g)_r"], ) -def _B0(params, transforms, profiles, data, **kwargs): - data["B0"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg(params, transforms, profiles, data, **kwargs): + data["psi_r/sqrt(g)"] = transforms["grid"].replace_at_axis( safediv(data["psi_r"], data["sqrt(g)"]), lambda: safediv(data["psi_rr"], data["sqrt(g)_r"]), ) @@ -74,10 +74,10 @@ def _B_sup_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "iota", "lambda_z", "omega_z"], + data=["psi_r/sqrt(g)", "iota", "lambda_z", "omega_z"], ) def _B_sup_theta(params, transforms, profiles, data, **kwargs): - data["B^theta"] = data["B0"] * ( + data["B^theta"] = data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -94,10 +94,10 @@ def _B_sup_theta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "iota", "lambda_t", "omega_t"], + data=["psi_r/sqrt(g)", "iota", "lambda_t", "omega_t"], ) def _B_sup_zeta(params, transforms, profiles, data, **kwargs): - data["B^zeta"] = data["B0"] * ( + data["B^zeta"] = data["psi_r/sqrt(g)"] * ( -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 ) return data @@ -178,7 +178,7 @@ def _B_Z(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_r", + name="(psi_r/sqrt(g))_r", label="\\partial_{\\rho} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -191,8 +191,8 @@ def _B_Z(params, transforms, profiles, data, **kwargs): data=["psi_r", "psi_rr", "sqrt(g)", "sqrt(g)_r"], axis_limit_data=["psi_rrr", "sqrt(g)_rr"], ) -def _B0_r(params, transforms, profiles, data, **kwargs): - data["B0_r"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_r(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_r"] = transforms["grid"].replace_at_axis( safediv( data["psi_rr"] * data["sqrt(g)"] - data["psi_r"] * data["sqrt(g)_r"], data["sqrt(g)"] ** 2, @@ -220,8 +220,8 @@ def _B0_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", "iota", "iota_r", "lambda_rz", @@ -231,12 +231,12 @@ def _B0_r(params, transforms, profiles, data, **kwargs): ], ) def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): - data["B^theta_r"] = data["B0"] * ( + data["B^theta_r"] = data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rz"] + data["iota_r"] * data["omega_z"] + data["iota_r"] - data["lambda_rz"] - ) + data["B0_r"] * ( + ) + data["(psi_r/sqrt(g))_r"] * ( data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -257,8 +257,8 @@ def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", "iota", "iota_r", "lambda_rt", @@ -268,11 +268,13 @@ def _B_sup_theta_r(params, transforms, profiles, data, **kwargs): ], ) def _B_sup_zeta_r(params, transforms, profiles, data, **kwargs): - data["B^zeta_r"] = data["B0"] * ( + data["B^zeta_r"] = data["psi_r/sqrt(g)"] * ( -data["iota"] * data["omega_rt"] - data["iota_r"] * data["omega_t"] + data["lambda_rt"] - ) + data["B0_r"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + ) + data["(psi_r/sqrt(g))_r"] * ( + -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + ) return data @@ -309,7 +311,7 @@ def _B_r(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_t", + name="(psi_r/sqrt(g))_t", label="\\partial_{\\theta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -322,8 +324,8 @@ def _B_r(params, transforms, profiles, data, **kwargs): data=["psi_r", "sqrt(g)", "sqrt(g)_t"], axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rt"], ) -def _B0_t(params, transforms, profiles, data, **kwargs): - data["B0_t"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_t(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_t"] = transforms["grid"].replace_at_axis( safediv(-data["psi_r"] * data["sqrt(g)_t"], data["sqrt(g)"] ** 2), lambda: safediv(-data["psi_rr"] * data["sqrt(g)_rt"], data["sqrt(g)_r"] ** 2), ) @@ -344,12 +346,20 @@ def _B0_t(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_t", "iota", "lambda_tz", "lambda_z", "omega_tz", "omega_z"], + data=[ + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_t", + "iota", + "lambda_tz", + "lambda_z", + "omega_tz", + "omega_z", + ], ) def _B_sup_theta_t(params, transforms, profiles, data, **kwargs): - data["B^theta_t"] = data["B0"] * ( + data["B^theta_t"] = data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_tz"] - data["lambda_tz"] - ) + data["B0_t"] * ( + ) + data["(psi_r/sqrt(g))_t"] * ( data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -369,12 +379,22 @@ def _B_sup_theta_t(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_t", "iota", "lambda_t", "lambda_tt", "omega_t", "omega_tt"], + data=[ + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_t", + "iota", + "lambda_t", + "lambda_tt", + "omega_t", + "omega_tt", + ], ) def _B_sup_zeta_t(params, transforms, profiles, data, **kwargs): - data["B^zeta_t"] = data["B0"] * ( + data["B^zeta_t"] = data["psi_r/sqrt(g)"] * ( -data["iota"] * data["omega_tt"] + data["lambda_tt"] - ) + data["B0_t"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + ) + data["(psi_r/sqrt(g))_t"] * ( + -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + ) return data @@ -411,7 +431,7 @@ def _B_t(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_z", + name="(psi_r/sqrt(g))_z", label="\\partial_{\\zeta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -424,8 +444,8 @@ def _B_t(params, transforms, profiles, data, **kwargs): data=["psi_r", "sqrt(g)", "sqrt(g)_z"], axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rz"], ) -def _B0_z(params, transforms, profiles, data, **kwargs): - data["B0_z"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_z(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_z"] = transforms["grid"].replace_at_axis( safediv(-data["psi_r"] * data["sqrt(g)_z"], data["sqrt(g)"] ** 2), lambda: safediv(-data["psi_rr"] * data["sqrt(g)_rz"], data["sqrt(g)_r"] ** 2), ) @@ -446,12 +466,20 @@ def _B0_z(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_z", "iota", "lambda_z", "lambda_zz", "omega_z", "omega_zz"], + data=[ + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_z", + "iota", + "lambda_z", + "lambda_zz", + "omega_z", + "omega_zz", + ], ) def _B_sup_theta_z(params, transforms, profiles, data, **kwargs): - data["B^theta_z"] = data["B0"] * ( + data["B^theta_z"] = data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_zz"] - data["lambda_zz"] - ) + data["B0_z"] * ( + ) + data["(psi_r/sqrt(g))_z"] * ( data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"] ) return data @@ -471,12 +499,22 @@ def _B_sup_theta_z(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "B0_z", "iota", "lambda_t", "lambda_tz", "omega_t", "omega_tz"], + data=[ + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_z", + "iota", + "lambda_t", + "lambda_tz", + "omega_t", + "omega_tz", + ], ) def _B_sup_zeta_z(params, transforms, profiles, data, **kwargs): - data["B^zeta_z"] = data["B0"] * ( + data["B^zeta_z"] = data["psi_r/sqrt(g)"] * ( -data["iota"] * data["omega_tz"] + data["lambda_tz"] - ) + data["B0_z"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + ) + data["(psi_r/sqrt(g))_z"] * ( + -data["iota"] * data["omega_t"] + data["lambda_t"] + 1 + ) return data @@ -556,7 +594,7 @@ def _B_z(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_rr", + name="(psi_r/sqrt(g))_rr", label="\\partial_{\\rho \\rho} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meters", @@ -569,8 +607,8 @@ def _B_z(params, transforms, profiles, data, **kwargs): data=["psi_r", "psi_rr", "psi_rrr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_rr"], axis_limit_data=["sqrt(g)_rrr"], ) -def _B0_rr(params, transforms, profiles, data, **kwargs): - data["B0_rr"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_rr(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_rr"] = transforms["grid"].replace_at_axis( safediv( data["psi_rrr"] * data["sqrt(g)"] ** 2 - 2 * data["psi_rr"] * data["sqrt(g)_r"] * data["sqrt(g)"] @@ -603,9 +641,9 @@ def _B0_rr(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", - "B0_rr", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", + "(psi_r/sqrt(g))_rr", "iota", "iota_r", "iota_rr", @@ -619,7 +657,7 @@ def _B0_rr(params, transforms, profiles, data, **kwargs): ) def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): data["B^theta_rr"] = ( - data["B0"] + data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rrz"] + 2 * data["iota_r"] * data["omega_rz"] @@ -628,14 +666,14 @@ def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): - data["lambda_rrz"] ) + 2 - * data["B0_r"] + * data["(psi_r/sqrt(g))_r"] * ( data["iota"] * data["omega_rz"] + data["iota_r"] * data["omega_z"] + data["iota_r"] - data["lambda_rz"] ) - + data["B0_rr"] + + data["(psi_r/sqrt(g))_rr"] * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -656,9 +694,9 @@ def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", - "B0_rr", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", + "(psi_r/sqrt(g))_rr", "iota", "iota_r", "iota_rr", @@ -672,7 +710,7 @@ def _B_sup_theta_rr(params, transforms, profiles, data, **kwargs): ) def _B_sup_zeta_rr(params, transforms, profiles, data, **kwargs): data["B^zeta_rr"] = ( - -data["B0"] + -data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rrt"] + 2 * data["iota_r"] * data["omega_rt"] @@ -680,13 +718,14 @@ def _B_sup_zeta_rr(params, transforms, profiles, data, **kwargs): - data["lambda_rrt"] ) - 2 - * data["B0_r"] + * data["(psi_r/sqrt(g))_r"] * ( data["iota"] * data["omega_rt"] + data["iota_r"] * data["omega_t"] - data["lambda_rt"] ) - + data["B0_rr"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + + data["(psi_r/sqrt(g))_rr"] + * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -730,7 +769,7 @@ def _B_rr(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_tt", + name="(psi_r/sqrt(g))_tt", label="\\partial_{\\theta \\theta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -743,8 +782,8 @@ def _B_rr(params, transforms, profiles, data, **kwargs): data=["psi_r", "sqrt(g)", "sqrt(g)_t", "sqrt(g)_tt"], axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rt", "sqrt(g)_rtt"], ) -def _B0_tt(params, transforms, profiles, data, **kwargs): - data["B0_tt"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_tt(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_tt"] = transforms["grid"].replace_at_axis( safediv( data["psi_r"] * (2 * data["sqrt(g)_t"] ** 2 - data["sqrt(g)"] * data["sqrt(g)_tt"]), @@ -774,9 +813,9 @@ def _B0_tt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_t", - "B0_tt", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_t", + "(psi_r/sqrt(g))_tt", "iota", "lambda_ttz", "lambda_tz", @@ -788,9 +827,11 @@ def _B0_tt(params, transforms, profiles, data, **kwargs): ) def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): data["B^theta_tt"] = ( - data["B0"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) - + 2 * data["B0_t"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - + data["B0_tt"] + data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_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"]) ) return data @@ -811,9 +852,9 @@ def _B_sup_theta_tt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_t", - "B0_tt", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_t", + "(psi_r/sqrt(g))_tt", "iota", "lambda_t", "lambda_tt", @@ -825,9 +866,12 @@ 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["B0"] * (data["iota"] * data["omega_ttt"] - data["lambda_ttt"]) - - 2 * data["B0_t"] * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) - + data["B0_tt"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + -data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_ttt"] - data["lambda_ttt"]) + - 2 + * data["(psi_r/sqrt(g))_t"] + * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + + data["(psi_r/sqrt(g))_tt"] + * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -871,7 +915,7 @@ def _B_tt(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_zz", + name="(psi_r/sqrt(g))_zz", label="\\partial_{\\zeta \\zeta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -884,8 +928,8 @@ def _B_tt(params, transforms, profiles, data, **kwargs): data=["psi_r", "sqrt(g)", "sqrt(g)_z", "sqrt(g)_zz"], axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rz", "sqrt(g)_rzz"], ) -def _B0_zz(params, transforms, profiles, data, **kwargs): - data["B0_zz"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_zz(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_zz"] = transforms["grid"].replace_at_axis( safediv( data["psi_r"] * (2 * data["sqrt(g)_z"] ** 2 - data["sqrt(g)"] * data["sqrt(g)_zz"]), @@ -915,9 +959,9 @@ def _B0_zz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_z", - "B0_zz", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_z", + "(psi_r/sqrt(g))_zz", "iota", "lambda_z", "lambda_zz", @@ -929,9 +973,11 @@ def _B0_zz(params, transforms, profiles, data, **kwargs): ) def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): data["B^theta_zz"] = ( - data["B0"] * (data["iota"] * data["omega_zzz"] - data["lambda_zzz"]) - + 2 * data["B0_z"] * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) - + data["B0_zz"] + data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_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"]) ) return data @@ -952,9 +998,9 @@ def _B_sup_theta_zz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_z", - "B0_zz", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_z", + "(psi_r/sqrt(g))_zz", "iota", "lambda_t", "lambda_tz", @@ -966,9 +1012,12 @@ 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["B0"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) - - 2 * data["B0_z"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - + data["B0_zz"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + -data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) + - 2 + * data["(psi_r/sqrt(g))_z"] + * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + + data["(psi_r/sqrt(g))_zz"] + * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -1012,7 +1061,7 @@ def _B_zz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_rt", + name="(psi_r/sqrt(g))_rt", label="\\partial_{\\rho\\theta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meters", @@ -1025,8 +1074,8 @@ def _B_zz(params, transforms, profiles, data, **kwargs): data=["psi_r", "psi_rr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_t", "sqrt(g)_rt"], axis_limit_data=["psi_rrr", "sqrt(g)_rr", "sqrt(g)_rrt"], ) -def _B0_rt(params, transforms, profiles, data, **kwargs): - data["B0_rt"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_rt(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_rt"] = transforms["grid"].replace_at_axis( safediv( -data["sqrt(g)"] * (data["psi_rr"] * data["sqrt(g)_t"] + data["psi_r"] * data["sqrt(g)_rt"]) @@ -1061,10 +1110,10 @@ def _B0_rt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", - "B0_rt", - "B0_t", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", + "(psi_r/sqrt(g))_rt", + "(psi_r/sqrt(g))_t", "iota", "iota_r", "lambda_rtz", @@ -1079,21 +1128,22 @@ def _B0_rt(params, transforms, profiles, data, **kwargs): ) def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): data["B^theta_rt"] = ( - data["B0"] + data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rtz"] + data["iota_r"] * data["omega_tz"] - data["lambda_rtz"] ) - + data["B0_r"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - + data["B0_t"] + + data["(psi_r/sqrt(g))_r"] + * (data["iota"] * data["omega_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["lambda_rz"] ) - + data["B0_rt"] + + data["(psi_r/sqrt(g))_rt"] * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -1114,10 +1164,10 @@ def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", - "B0_rt", - "B0_t", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", + "(psi_r/sqrt(g))_rt", + "(psi_r/sqrt(g))_t", "iota", "iota_r", "lambda_rt", @@ -1132,20 +1182,22 @@ def _B_sup_theta_rt(params, transforms, profiles, data, **kwargs): ) def _B_sup_zeta_rt(params, transforms, profiles, data, **kwargs): data["B^zeta_rt"] = ( - -data["B0"] + -data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rtt"] + data["iota_r"] * data["omega_tt"] - data["lambda_rtt"] ) - - data["B0_r"] * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) - - data["B0_t"] + - data["(psi_r/sqrt(g))_r"] + * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + - data["(psi_r/sqrt(g))_t"] * ( data["iota"] * data["omega_rt"] + data["iota_r"] * data["omega_t"] - data["lambda_rt"] ) - + data["B0_rt"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + + data["(psi_r/sqrt(g))_rt"] + * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -1196,7 +1248,7 @@ def _B_rt(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_tz", + name="(psi_r/sqrt(g))_tz", label="\\partial_{\\theta\\zeta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meter", @@ -1209,8 +1261,8 @@ def _B_rt(params, transforms, profiles, data, **kwargs): data=["psi_r", "sqrt(g)", "sqrt(g)_t", "sqrt(g)_z", "sqrt(g)_tz"], axis_limit_data=["psi_rr", "sqrt(g)_r", "sqrt(g)_rt", "sqrt(g)_rz", "sqrt(g)_rtz"], ) -def _B0_tz(params, transforms, profiles, data, **kwargs): - data["B0_tz"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_tz(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_tz"] = transforms["grid"].replace_at_axis( safediv( data["psi_r"] * ( @@ -1246,10 +1298,10 @@ def _B0_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_t", - "B0_tz", - "B0_z", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_t", + "(psi_r/sqrt(g))_tz", + "(psi_r/sqrt(g))_z", "iota", "lambda_tz", "lambda_tzz", @@ -1263,10 +1315,12 @@ def _B0_tz(params, transforms, profiles, data, **kwargs): ) def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): data["B^theta_tz"] = ( - data["B0"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) - + data["B0_t"] * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) - + data["B0_z"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - + data["B0_tz"] + data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_tzz"] - data["lambda_tzz"]) + + data["(psi_r/sqrt(g))_t"] + * (data["iota"] * data["omega_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"]) ) return data @@ -1287,10 +1341,10 @@ def _B_sup_theta_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_t", - "B0_tz", - "B0_z", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_t", + "(psi_r/sqrt(g))_tz", + "(psi_r/sqrt(g))_z", "iota", "lambda_t", "lambda_tt", @@ -1304,10 +1358,13 @@ 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["B0"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) - - data["B0_t"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - - data["B0_z"] * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) - + data["B0_tz"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + -data["psi_r/sqrt(g)"] * (data["iota"] * data["omega_ttz"] - data["lambda_ttz"]) + - data["(psi_r/sqrt(g))_t"] + * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + - data["(psi_r/sqrt(g))_z"] + * (data["iota"] * data["omega_tt"] - data["lambda_tt"]) + + data["(psi_r/sqrt(g))_tz"] + * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -1357,7 +1414,7 @@ def _B_tz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="B0_rz", + name="(psi_r/sqrt(g))_rz", label="\\partial_{\\rho\\zeta} (\\psi' / \\sqrt{g})", units="T \\cdot m^{-1}", units_long="Tesla / meters", @@ -1370,8 +1427,8 @@ def _B_tz(params, transforms, profiles, data, **kwargs): data=["psi_r", "psi_rr", "sqrt(g)", "sqrt(g)_r", "sqrt(g)_z", "sqrt(g)_rz"], axis_limit_data=["psi_rrr", "sqrt(g)_rr", "sqrt(g)_rrz"], ) -def _B0_rz(params, transforms, profiles, data, **kwargs): - data["B0_rz"] = transforms["grid"].replace_at_axis( +def _psi_r_over_sqrtg_rz(params, transforms, profiles, data, **kwargs): + data["(psi_r/sqrt(g))_rz"] = transforms["grid"].replace_at_axis( safediv( -data["sqrt(g)"] * (data["psi_rr"] * data["sqrt(g)_z"] + data["psi_r"] * data["sqrt(g)_rz"]) @@ -1406,10 +1463,10 @@ def _B0_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", - "B0_rz", - "B0_z", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", + "(psi_r/sqrt(g))_rz", + "(psi_r/sqrt(g))_z", "iota", "iota_r", "lambda_rz", @@ -1424,21 +1481,22 @@ def _B0_rz(params, transforms, profiles, data, **kwargs): ) def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): data["B^theta_rz"] = ( - data["B0"] + data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rzz"] + data["iota_r"] * data["omega_zz"] - data["lambda_rzz"] ) - + data["B0_r"] * (data["iota"] * data["omega_zz"] - data["lambda_zz"]) - + data["B0_z"] + + data["(psi_r/sqrt(g))_r"] + * (data["iota"] * data["omega_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["lambda_rz"] ) - + data["B0_rz"] + + data["(psi_r/sqrt(g))_rz"] * (data["iota"] * data["omega_z"] + data["iota"] - data["lambda_z"]) ) return data @@ -1459,10 +1517,10 @@ def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=[ - "B0", - "B0_r", - "B0_rz", - "B0_z", + "psi_r/sqrt(g)", + "(psi_r/sqrt(g))_r", + "(psi_r/sqrt(g))_rz", + "(psi_r/sqrt(g))_z", "iota", "iota_r", "lambda_rt", @@ -1477,20 +1535,22 @@ def _B_sup_theta_rz(params, transforms, profiles, data, **kwargs): ) def _B_sup_zeta_rz(params, transforms, profiles, data, **kwargs): data["B^zeta_rz"] = ( - -data["B0"] + -data["psi_r/sqrt(g)"] * ( data["iota"] * data["omega_rtz"] + data["iota_r"] * data["omega_tz"] - data["lambda_rtz"] ) - - data["B0_r"] * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) - - data["B0_z"] + - data["(psi_r/sqrt(g))_r"] + * (data["iota"] * data["omega_tz"] - data["lambda_tz"]) + - data["(psi_r/sqrt(g))_z"] * ( data["iota"] * data["omega_rt"] + data["iota_r"] * data["omega_t"] - data["lambda_rt"] ) - + data["B0_rz"] * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) + + data["(psi_r/sqrt(g))_rz"] + * (-data["iota"] * data["omega_t"] + data["lambda_t"] + 1) ) return data @@ -2200,6 +2260,33 @@ def _B_sub_zeta_rz(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="<|B|>_axis", + label="\\lange |\\mathbf{B}| \\rangle_{axis}", + units="T", + units_long="Tesla", + description="Average magnitude of magnetic field on the magnetic axis", + dim=0, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="", + data=["|B|", "sqrt(g)"], + axis_limit_data=["sqrt(g)_r"], + resolution_requirement="tz", +) +def _B_mag_axis(params, transforms, profiles, data, **kwargs): + # mask of indices of the innermost flux surface + mask = transforms["grid"].inverse_rho_idx == 0 + sqrt_g = transforms["grid"].replace_at_axis( + data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True + ) + data["<|B|>_axis"] = jnp.sum( + mask * data["|B|"] * sqrt_g * transforms["grid"].weights + ) / jnp.sum(mask * sqrt_g * transforms["grid"].weights) + return data + + @register_compute_fun( name="|B|^2", label="|\\mathbf{B}|^{2}", diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index 892800d48e..fa43976898 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -336,12 +336,22 @@ def _B_modes(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["iota", "B0", "B_theta", "B_zeta", "|B|_t", "|B|_z", "G", "I", "B*grad(|B|)"], + data=[ + "iota", + "psi_r/sqrt(g)", + "B_theta", + "B_zeta", + "|B|_t", + "|B|_z", + "G", + "I", + "B*grad(|B|)", + ], helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)", ) def _f_C(params, transforms, profiles, data, **kwargs): M, N = kwargs.get("helicity", (1, 0)) - data["f_C"] = (M * data["iota"] - N) * data["B0"] * ( + data["f_C"] = (M * data["iota"] - N) * data["psi_r/sqrt(g)"] * ( data["B_zeta"] * data["|B|_t"] - data["B_theta"] * data["|B|_z"] ) - (M * data["G"] + N * data["I"]) * data["B*grad(|B|)"] return data @@ -359,10 +369,10 @@ def _f_C(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["B0", "|B|_t", "|B|_z", "(B*grad(|B|))_t", "(B*grad(|B|))_z"], + data=["psi_r/sqrt(g)", "|B|_t", "|B|_z", "(B*grad(|B|))_t", "(B*grad(|B|))_z"], ) def _f_T(params, transforms, profiles, data, **kwargs): - data["f_T"] = data["B0"] * ( + data["f_T"] = data["psi_r/sqrt(g)"] * ( data["|B|_t"] * data["(B*grad(|B|))_z"] - data["|B|_z"] * data["(B*grad(|B|))_t"] ) diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index 5069d06ee2..84de48e576 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -217,6 +217,27 @@ def _Te_rr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="_vol", + label="\\langle n_e \\rangle_{vol}", + units="m^{-3}", + units_long="1 / cubic meters", + description="Volume average electron density", + dim=0, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="", + data=["ne", "sqrt(g)", "V"], + resolution_requirement="rtz", +) +def _ne_vol(params, transforms, profiles, data, **kwargs): + data["_vol"] = ( + jnp.sum(data["ne"] * data["sqrt(g)"] * transforms["grid"].weights) / data["V"] + ) + return data + + @register_compute_fun( name="ne", label="n_e", @@ -355,6 +376,44 @@ def _Ti_rr(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="ni", + label="n_i", + units="m^{-3}", + units_long="1 / cubic meters", + description="Ion density", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["ne", "Zeff"], +) +def _ni(params, transforms, profiles, data, **kwargs): + data["ni"] = data["ne"] / data["Zeff"] + return data + + +@register_compute_fun( + name="ni_r", + label="\\partial_{\\rho} n_i", + units="m^{-3}", + units_long="1 / cubic meters", + description="Ion density, first radial derivative", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["ne", "ne_r", "Zeff", "Zeff_r"], +) +def _ni_r(params, transforms, profiles, data, **kwargs): + data["ni_r"] = (data["ne_r"] * data["Zeff"] - data["ne"] * data["Zeff_r"]) / data[ + "Zeff" + ] ** 2 + return data + + @register_compute_fun( name="Zeff", label="Z_{eff}", @@ -412,7 +471,7 @@ def _Zeff_r(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=["pressure"], coordinates="r", - data=["Te", "ne", "Ti", "Zeff"], + data=["ne", "ni", "Te", "Ti"], ) def _p(params, transforms, profiles, data, **kwargs): if profiles["pressure"] is not None: @@ -421,7 +480,7 @@ def _p(params, transforms, profiles, data, **kwargs): ) else: data["p"] = elementary_charge * ( - data["ne"] * data["Te"] + data["Ti"] * data["ne"] / data["Zeff"] + data["ne"] * data["Te"] + data["ni"] * data["Ti"] ) return data @@ -437,7 +496,7 @@ def _p(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=["pressure"], coordinates="r", - data=["Te", "Te_r", "ne", "ne_r", "Ti", "Ti_r", "Zeff", "Zeff_r"], + data=["ne", "ne_r", "ni", "ni_r", "Te", "Te_r", "Ti", "Ti_r"], ) def _p_r(params, transforms, profiles, data, **kwargs): if profiles["pressure"] is not None: @@ -448,9 +507,8 @@ def _p_r(params, transforms, profiles, data, **kwargs): data["p_r"] = elementary_charge * ( data["ne_r"] * data["Te"] + data["ne"] * data["Te_r"] - + data["Ti_r"] * data["ne"] / data["Zeff"] - + data["Ti"] * data["ne_r"] / data["Zeff"] - - data["Ti"] * data["ne"] * data["Zeff_r"] / data["Zeff"] ** 2 + + data["ni_r"] * data["Ti"] + + data["ni"] * data["Ti_r"] ) return data diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 3ae2e59af0..0443a8a504 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -11,6 +11,7 @@ QuadraticFlux, ToroidalFlux, ) +from ._confinement import HeatingPowerISS04 from ._equilibrium import ( CurrentDensity, Energy, diff --git a/desc/objectives/_confinement.py b/desc/objectives/_confinement.py new file mode 100644 index 0000000000..b768fa5e35 --- /dev/null +++ b/desc/objectives/_confinement.py @@ -0,0 +1,197 @@ +"""Objectives for targeting energy confinement.""" + +from desc.compute.utils import _compute as compute_fun +from desc.compute.utils import get_profiles, get_transforms +from desc.grid import QuadratureGrid +from desc.utils import Timer, errorif + +from .normalization import compute_scaling_factors +from .objective_funs import _Objective + + +class HeatingPowerISS04(_Objective): + """Heating power required by the ISS04 energy confinement time scaling. + + τ_E = Wₚ / P = 0.134 H_ISS04 a²ᐧ²⁸ R⁰ᐧ⁶⁴ P⁻⁰ᐧ⁶¹ nₑ⁰ᐧ⁵⁴ B⁰ᐧ⁸⁴ 𝜄⁰ᐧ⁴¹ (s) + + References + ---------- + https://doi.org/10.1088/0029-5515/45/12/024. + Characterization of energy confinement in net-current free plasmas using the + extended International Stellarator Database. + H. Yamada et al. Nucl. Fusion 29 November 2005; 45 (12): 1684–1693. + + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. + target : {float, ndarray}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. + bounds : tuple of {float, ndarray}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to to Objective.dim_f. + Defaults to ``target=0``. + weight : {float, ndarray}, optional + Weighting to apply to the Objective, relative to other Objectives. + Must be broadcastable to to Objective.dim_f + normalize : bool, optional + Whether to compute the error in physical units or non-dimensionalize. + normalize_target : bool, optional + Whether target and bounds should be normalized before comparing to computed + values. If `normalize` is `True` and the target is in physical units, + this should also be set to True. + loss_function : {None, 'mean', 'min', 'max'}, optional + Loss function to apply to the objective values once computed. This loss function + is called on the raw compute value, before any shifting, scaling, or + normalization. Note: Has no effect for this objective. + deriv_mode : {"auto", "fwd", "rev"} + Specify how to compute jacobian matrix, either forward mode or reverse mode AD. + "auto" selects forward or reverse mode based on the size of the input and output + of the objective. Has no effect on self.grad or self.hess which always use + reverse mode and forward over reverse mode respectively. + H_ISS04 : float, optional + ISS04 confinement enhancement factor. Default = 1. + gamma : float, optional + Adiabatic (compressional) index. Default = 0. + grid : Grid, optional + Collocation grid used to compute the intermediate quantities. + Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP)``. + name : str, optional + Name of the objective function. + + """ + + _scalar = True + _units = "(W)" + _print_value_fmt = "Heating power: {:10.3e} " + + def __init__( + self, + eq, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + loss_function=None, + deriv_mode="auto", + H_ISS04=1, + gamma=0, + grid=None, + name="heating power", + ): + if target is None and bounds is None: + target = 0 + self._H_ISS04 = H_ISS04 + self._gamma = gamma + self._grid = grid + super().__init__( + things=eq, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + loss_function=loss_function, + deriv_mode=deriv_mode, + name=name, + ) + + def build(self, use_jit=True, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + eq = self.things[0] + errorif( + eq.electron_density is None, + ValueError, + "Equilibrium must have an electron density profile.", + ) + if self._grid is None: + self._grid = QuadratureGrid( + L=eq.L_grid, + M=eq.M_grid, + N=eq.N_grid, + NFP=eq.NFP, + ) + self._dim_f = 1 + self._data_keys = ["P_ISS04"] + + timer = Timer() + if verbose > 0: + print("Precomputing transforms") + timer.start("Precomputing transforms") + + self._constants = { + "profiles": get_profiles(self._data_keys, obj=eq, grid=self._grid), + "transforms": get_transforms(self._data_keys, obj=eq, grid=self._grid), + "H_ISS04": self._H_ISS04, + "gamma": self._gamma, + } + + timer.stop("Precomputing transforms") + if verbose > 1: + timer.disp("Precomputing transforms") + + if self._normalize: + scales = compute_scaling_factors(eq) + self._normalization = scales["W_p"] + + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params, constants=None): + """Compute heating power. + + Parameters + ---------- + params : dict + Dictionary of equilibrium or surface degrees of freedom, eg + Equilibrium.params_dict + constants : dict + Dictionary of constant data, eg transforms, profiles etc. Defaults to + self.constants + + Returns + ------- + P : float + Heating power required by the ISS04 energy confinement time scaling (W). + + """ + if constants is None: + constants = self.constants + data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._data_keys, + params=params, + transforms=constants["transforms"], + profiles=constants["profiles"], + gamma=constants["gamma"], + H_ISS04=constants["H_ISS04"], + ) + return data["P_ISS04"] + + @property + def H_ISS04(self): + """float: Confinement enhancement factor.""" + return self._H_ISS04 + + @H_ISS04.setter + def H_ISS04(self, new): + self._H_ISS04 = new + + @property + def gamma(self): + """float: Adiabatic (compressional) index.""" + return self._gamma + + @gamma.setter + def gamma(self, new): + self._gamma = new diff --git a/desc/plotting.py b/desc/plotting.py index fc199b38ac..f43f408af4 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -1200,10 +1200,12 @@ def plot_fsa( # noqa: C901 # Attempt to compute the magnetic axis limit. # Compute derivative depending on various naming schemes. # e.g. B -> B_r, V(r) -> V_r(r), S_r(r) -> S_rr(r) + # psi_r/sqrt(g) -> (psi_r/sqrt(g))_r schemes = ( name + "_r", name[:-3] + "_r" + name[-3:], name[:-3] + "r" + name[-3:], + "(" + name + ")_r", ) values_r = next( ( diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 8a9086d237..5e54166e05 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 5a1200d537..21f733d859 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -25,9 +25,9 @@ # 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"} -# "current Redl" needs special treatment because it is generally not defined for all -# configurations (giving NaN values), except it is always 0 at the magnetic axis -not_continuous_limits = {"current Redl"} +# "current Redl" and "P_ISS04" need special treatment because they are not defined for +# all configurations (giving NaN values) +not_continuous_limits = {"current Redl", "P_ISS04"} not_finite_limits = { "D_Mercier", "D_geodesic", @@ -127,6 +127,7 @@ def _skip_this(eq, name): or (eq.electron_temperature is None and "Te" in name) or (eq.electron_density is None and "ne" in name) or (eq.ion_temperature is None and "Ti" in name) + or (eq.electron_density is None and "ni" in name) or (eq.anisotropy is None and "beta_a" in name) or (eq.pressure is not None and " Redl" in name) or (eq.current is None and "iota_num" in name) @@ -259,7 +260,7 @@ class TestAxisLimits: @pytest.mark.unit def test_axis_limit_api(self): """Test that axis limit dependencies are computed only when necessary.""" - name = "B0" + name = "psi_r/sqrt(g)" deps = {"psi_r", "sqrt(g)"} axis_limit_deps = {"psi_rr", "sqrt(g)_r"} eq = Equilibrium() diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 8b39843553..6de443dad0 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -49,6 +49,7 @@ ForceBalance, ForceBalanceAnisotropic, GenericObjective, + HeatingPowerISS04, Isodynamicity, LinearObjectiveFromUser, MagneticWell, @@ -2118,6 +2119,7 @@ class TestComputeScalarResolution: CoilSetMinDistance, CoilTorsion, GenericObjective, + HeatingPowerISS04, Omnigenity, PlasmaCoilSetMinDistance, PlasmaVesselDistance, @@ -2177,6 +2179,28 @@ def test_compute_scalar_resolution_bootstrap(self): f[i] = obj.compute_scalar(obj.x()) np.testing.assert_allclose(f, f[-1], rtol=5e-2) + @pytest.mark.regression + def test_compute_scalar_resolution_heating_power(self): + """HeatingPowerISS04.""" + eq = self.eq.copy() + eq.electron_density = PowerSeriesProfile([1e19, 0, -1e19]) + eq.electron_temperature = PowerSeriesProfile([1e3, 0, -1e3]) + eq.ion_temperature = PowerSeriesProfile([1e3, 0, -1e3]) + eq.atomic_number = 1.0 + + f = np.zeros_like(self.res_array, dtype=float) + for i, res in enumerate(self.res_array): + grid = QuadratureGrid( + L=int(self.eq.L * res), + M=int(self.eq.M * res), + N=int(self.eq.N * res), + NFP=self.eq.NFP, + ) + obj = ObjectiveFunction(HeatingPowerISS04(eq=eq, grid=grid)) + obj.build(verbose=0) + f[i] = obj.compute_scalar(obj.x()) + np.testing.assert_allclose(f, f[-1], rtol=5e-2) + @pytest.mark.regression def test_compute_scalar_resolution_boundary_error(self): """BoundaryError.""" @@ -2446,6 +2470,8 @@ class TestObjectiveNaNGrad: CoilSetMinDistance, CoilTorsion, ForceBalanceAnisotropic, + HeatingPowerISS04, + Omnigenity, PlasmaCoilSetMinDistance, PlasmaVesselDistance, QuadraticFlux, @@ -2455,8 +2481,6 @@ class TestObjectiveNaNGrad: GenericObjective, LinearObjectiveFromUser, ObjectiveFromUser, - # TODO: add Omnigenity objective (see GH issue #943) - Omnigenity, ] other_objectives = list(set(objectives) - set(specials)) @@ -2495,6 +2519,22 @@ def test_objective_no_nangrad_bootstrap(self): g = obj.grad(obj.x(eq)) assert not np.any(np.isnan(g)), "redl bootstrap" + @pytest.mark.unit + def test_objective_no_nangrad_heating_power(self): + """HeatingPowerISS04.""" + eq = Equilibrium( + L=2, + M=2, + N=2, + electron_density=PowerSeriesProfile([1e19, 0, -1e19]), + electron_temperature=PowerSeriesProfile([1e3, 0, -1e3]), + current=PowerSeriesProfile([1, 0, -1]), + ) + obj = ObjectiveFunction(HeatingPowerISS04(eq)) + obj.build() + g = obj.grad(obj.x(eq)) + assert not np.any(np.isnan(g)), "heating power" + @pytest.mark.unit def test_objective_no_nangrad_boundary_error(self): """BoundaryError.""" diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 944ec56de1..c1bd782ac5 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -335,7 +335,7 @@ def test_plot_fsa_axis_limit(self): plot_data["<" + name + ">"], desired, equal_nan=False ) - name = "B0" + name = "psi_r/sqrt(g)" assert ( "<" + name + ">" not in data_index["desc.equilibrium.equilibrium.Equilibrium"] @@ -371,7 +371,7 @@ def test_plot_fsa_axis_limit(self): rho=rho, M=eq.M_grid, N=eq.N_grid, - with_sqrt_g=False, # Test that does not compute data_index["<|B|>"] + with_sqrt_g=False, # test that does not compute data_index["<|B|>"] return_data=True, ) data = eq.compute(names=name, grid=grid) diff --git a/tests/test_profiles.py b/tests/test_profiles.py index 201996c1a7..bf75489a6b 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -410,7 +410,6 @@ def test_solve_with_combined(self): Te = PowerSeriesProfile(2.0e3 * np.array([1, -1]), modes=[0, 2]) Ti = Te pressure = elementary_charge * (ne * Te + ne * Ti) - print("pressure params:", pressure.params) LM_resolution = 6 eq1 = Equilibrium( @@ -463,7 +462,6 @@ def test_kinetic_pressure(self): Te = PowerSeriesProfile(2.0e3 * np.array([1, -1]), modes=[0, 2]) Ti = Te pressure = elementary_charge * (ne * Te + ne * Ti) - print("pressure params:", pressure.params) LM_resolution = 6 eq1 = Equilibrium(