Source code for idstools.cocos

import logging
import traceback

try:
    import imaspy as imas
except ImportError:
    import imas
import numpy as np

# set cocos in the DD version from the environment
IDS_COCOS = int(imas.dd_zip.dd_etree().find("cocos").text)


logger = logging.getLogger(f"module.{__name__}")


[docs]class COCOS: """ COCOS module in Python. This module provides functionality related to coordinate conventions in tokamak physics. References: O. Sauter and S. Yu. Medvedev, "Tokamak Coordinate Conventions: COCOS",Comput. Physics Commun 84 (2013), 293. `cocos_module.f90 (CHEASE)`. """ def __init__(self, index=None, values=None): """ Initialize COCOS index using values, or values using COCOS index. Parameters ---------- index : dict, optional COCOS index dictionary with signs of Ip and B0. Example: ``{"COCOS": 11, "ipsign": 1, "b0sign": 1}`` values : dict, optional COCOS values dictionary with coordinate convention parameters. """ if (index is None) and (values is None): raise ValueError("Initialize COCOS with either index or values: both not given") return elif (index is not None) and (values is not None): raise ValueError("Initialize COCOS with either index or values: both given") return # in case of init. by index elif index is not None: COCOS = index["COCOS"] ipsign = index["ipsign"] b0sign = index["b0sign"] # # Parameters from Table I # if COCOS in [1, 11]: # ITER, Boozer are cocos=11 val = (+1, +1, +1, +1, -1) elif COCOS in [2, 12]: # CHEASE, ONETWO, Hinton-Hazeltine, LION is cocos=2 val = (+1, -1, +1, +1, -1) elif COCOS in [3, 13]: # Freidberg, CAXE, KINX are cocos=3 # EU-ITM up to end of 2011 is COCOS=13 val = (-1, +1, -1, -1, +1) elif COCOS in [4, 14]: # val = (-1, -1, -1, -1, +1) elif COCOS in [5, 15]: # val = (+1, +1, -1, -1, -1) elif COCOS in [6, 16]: # val = (+1, -1, -1, -1, -1) elif COCOS in [7, 17]: # TCV psitbx is cocos=7 val = (-1, +1, +1, +1, +1) elif COCOS in [8, 18]: # val = (-1, -1, +1, +1, +1) else: # Should not be here since all cases defined raise ValueError(f"error: COCOS = {COCOS} does not exist") return ( sigma_bp, sigma_rphi_z, sigma_rhothetaphi, sign_q_pos, sign_pprime_pos, ) = val # cocos=i or 10+i have similar coordinate conventions except psi/2pi for # cocos=i and psi for cocos=10+i if COCOS >= 11: exp_bp = 1 else: exp_bp = 0 theta_sign_clockwise = sigma_rphi_z * sigma_rhothetaphi self.COCOS = COCOS self.sigma_ip = ipsign self.sigma_b0 = b0sign self.exp_bp = exp_bp self.sigma_bp = sigma_bp self.sigma_rphi_z = sigma_rphi_z self.sigma_rhothetaphi = sigma_rhothetaphi self.sign_q_pos = sign_q_pos self.sign_pprime_pos = sign_pprime_pos self.theta_sign_clockwise = theta_sign_clockwise # in case of init. by values else: sigma_ip = values["ipsign"] sigma_b0 = values["b0sign"] exp_bp = values["exp_Bp"] sigma_bp = values["sigma_Bp"] sigma_rphi_z = values["sigma_RphiZ"] sigma_rhothetaphi = values["sigma_rhothetaphi"] sign_q_pos = values["sign_q_pos"] sign_pprime_pos = values["sign_pprime_pos"] val = ( sigma_bp, sigma_rphi_z, sigma_rhothetaphi, sign_q_pos, sign_pprime_pos, ) # # Parameters from Table I # if val == (+1, +1, +1, +1, -1): # ITER, Boozer are cocos=11 COCOS = [1, 11] elif val == (+1, -1, +1, +1, -1): # CHEASE, ONETWO, Hinton-Hazeltine, LION is cocos=2 COCOS = [2, 12] elif val == (-1, +1, -1, -1, +1): # Freidberg, CAXE, KINX are cocos=3 # EU-ITM up to end of 2011 is COCOS=13 COCOS = [3, 13] elif val == (-1, -1, -1, -1, +1): # COCOS = [4, 14] elif val == (+1, +1, -1, -1, -1): # COCOS = [5, 15] elif val == (+1, -1, -1, -1, -1): # COCOS = [6, 16] elif val == (-1, +1, +1, +1, +1): # COCOS = [7, 17] elif val == (-1, -1, +1, +1, +1): # COCOS = [8, 18] else: # Should not be here since all cases defined raise ValueError(f"error: COCOS Values not match {val}") return theta_sign_clockwise = sigma_rphi_z * sigma_rhothetaphi self.COCOS = COCOS[exp_bp] self.sigma_ip = sigma_ip self.sigma_b0 = sigma_b0 self.exp_bp = exp_bp self.sigma_bp = sigma_bp self.sigma_rphi_z = sigma_rphi_z self.sigma_rhothetaphi = sigma_rhothetaphi self.sign_q_pos = sign_q_pos self.sign_pprime_pos = sign_pprime_pos self.theta_sign_clockwise = theta_sign_clockwise
[docs] def get(self): """ Return COCOS index and values Returns ------- dict COCOS index and values in type dict """ return { "COCOS": self.COCOS, "sigma_Ip": self.sigma_ip, "sigma_B0": self.sigma_b0, "exp_Bp": self.exp_bp, "sigma_Bp": self.sigma_bp, "sigma_RphiZ": self.sigma_rphi_z, "sigma_rhothetaphi": self.sigma_rhothetaphi, "sign_q_pos": self.sign_q_pos, "sign_pprime_pos": self.sign_pprime_pos, "theta_sign_clockwise": self.theta_sign_clockwise, }
[docs] @classmethod def values_coefficients(self, COCOS_in, COCOS_out, ip_in, b0_in, ipsign_out, b0sign_out): """ Provide transformation values for a set of quantities for a given pair of input/output COCOS numbers. Parameters ---------- COCOS_in : int COCOS input COCOS_out : int COCOS output ip_in : float Plasma current (toroidal component) [A] b0_in : float Vacuum toroidal field [T] ipsign_out : int Desired sign of Ip in output b0sign_out : int Desired sign of B0 in output Returns ------- dict COCOS transformation values in type dict """ # Default outputs sigma_ip_eff = 1.0 sigma_b0_eff = 1.0 sigma_bp_eff = 1.0 sigma_rhothetaphi_eff = 1.0 sigma_rphi_z_eff = 1.0 exp_bp_eff = 1.0 fact_psi = 1.0 fact_q = 1.0 fact_dpsi = 1.0 fact_dtheta = 1.0 # Check inputs sigma_ip_in = np.sign(ip_in) sigma_b0_in = np.sign(b0_in) # Get COCOS related parameters c_v_i = COCOS(index={"COCOS": COCOS_in, "ipsign": sigma_ip_in, "b0sign": sigma_b0_in}).get() c_v_o = COCOS(index={"COCOS": COCOS_out, "ipsign": ipsign_out, "b0sign": b0sign_out}).get() # Define effective variables: sigma_Ip_eff, si1gma_B0_eff, sigma_Bp_eff, # exp_Bp_eff as in Appendix C sigma_rphi_z_eff = float(c_v_o["sigma_RphiZ"] * c_v_i["sigma_RphiZ"]) # sign(Ip) in output if ipsign_out == 0: sigma_ip_eff = sigma_rphi_z_eff # sign folllowing transformation else: sigma_ip_eff = sigma_ip_in * float(ipsign_out) # sigma_Ip_out = sigma_Ip_in * sigma_Ip_eff # sign(B0) in output if b0sign_out == 0: sigma_b0_eff = sigma_rphi_z_eff # sign folllowing transformation else: sigma_b0_eff = sigma_b0_in * float(b0sign_out) # sigma_B0_out = sigma_B0_in * sigma_B0_eff sigma_bp_eff = float(c_v_o["sigma_Bp"] * c_v_i["sigma_Bp"]) exp_bp_eff = float(c_v_o["exp_Bp"] - c_v_i["exp_Bp"]) sigma_rhothetaphi_eff = float(c_v_o["sigma_rhothetaphi"] * c_v_i["sigma_rhothetaphi"]) # # Note that sign(sigma_RphiZ*sigma_rhothetaphi) gives theta in clockwise or counter-clockwise respectively # Thus sigma_RphiZ_eff*sigma_rhothetaphi_eff negative if the direction of # theta has changed from cocos_in to _out # fact_psi = sigma_ip_eff * sigma_bp_eff * (2.0 * np.pi) ** exp_bp_eff fact_dpsi = sigma_ip_eff * sigma_bp_eff / (2.0 * np.pi) ** exp_bp_eff fact_q = sigma_ip_eff * sigma_b0_eff * sigma_rhothetaphi_eff fact_dtheta = sigma_rphi_z_eff * sigma_rhothetaphi_eff self.sigma_ip_eff = sigma_ip_eff self.sigma_b0_eff = sigma_b0_eff self.sigma_bp_eff = sigma_bp_eff self.sigma_rhothetaphi_eff = sigma_rhothetaphi_eff self.sigma_rphi_z_eff = sigma_rphi_z_eff self.exp_bp_eff = exp_bp_eff self.fact_psi = fact_psi self.fact_q = fact_q self.fact_dpsi = fact_dpsi self.fact_dtheta = fact_dtheta return { "sigma_Ip_eff": self.sigma_ip_eff, "sigma_B0_eff": self.sigma_b0_eff, "sigma_Bp_eff": self.sigma_bp_eff, "sigma_rhothetaphi_eff": self.sigma_rhothetaphi_eff, "sigma_RphiZ_eff": self.sigma_rphi_z_eff, "exp_Bp_eff": self.exp_bp_eff, "fact_psi": self.fact_psi, "fact_q": self.fact_q, "fact_dpsi": self.fact_dpsi, "fact_dtheta": self.fact_dtheta, }
[docs]def compute_COCOS(ids, itime=None, i1=0): """ Compute COCOS values using experimental data in IDS/equilibrium. Parameters ---------- ids : object IDS/equilibrium for COCOS estimation itime : int, optional Index of struct_array time_slice in IDS/equilibrium i1 : int, optional Index of struct_array profiles_2d in IDS/equilibrium. Default is 0. Returns ------- dict Dictionary with COCOS values """ # COCOS Values in the middle of time sequence if itime is None: itime = int(np.floor(float(len(ids.time_slice)) / 2.0)) # Sign(Ip) and Sign(B0) from input ipsign = np.sign(ids.time_slice[itime].global_quantities.ip) b0sign = np.sign(ids.vacuum_toroidal_field.b0[itime]) # 1 Eq.(22) dpsi = ids.time_slice[itime].profiles_1d.psi[-1] - ids.time_slice[itime].profiles_1d.psi[0] sigma_bp = np.sign(dpsi) * ipsign # 2 Eq.(22) q = ids.time_slice[itime].profiles_1d.q sign_q = np.sign(np.sum(np.sign(q))) sign_q_pos = sign_q * ipsign * b0sign # 3 Eq.(22) sigma_rhothetaphi = sign_q_pos # 4 Eq.(22) dpressure_dpsi = ids.time_slice[itime].profiles_1d.dpressure_dpsi sign_pprime_pos = np.sign(np.sum(np.sign(dpressure_dpsi))) * ipsign # 5 sigma_RphiZ from Eq.(19) bz = ids.time_slice[itime].profiles_2d[i1].b_field_z psi2d = ids.time_slice[itime].profiles_2d[i1].psi r2d = ids.time_slice[itime].profiles_2d[i1].r dpsi2d = np.gradient(psi2d) dr2d = np.gradient(r2d) dpsi2drdr = dpsi2d[0] / dr2d[0] / r2d # todo - reduce num. of data for COCOS discrimination # - compute rtol(s) instead of fixed ones. dim2 = ids.time_slice[itime].profiles_2d[i1].grid.dim2 z_axis = ids.time_slice[itime].global_quantities.magnetic_axis.z psi_axis = ids.time_slice[itime].profiles_1d.psi[0] # grid of magnetic axis in Z iz = np.argmin(np.abs(dim2 - z_axis)) # psi ref. inside LCFS psi_ref = psi_axis + dpsi * 0.5 # grids close to psi ref. rows, cols = np.where((np.isclose(psi2d, psi_ref, rtol=0.25)) & (bz != 0)) if (not rows.any()) or (not cols.any()): raise ValueError("COCOS discrimination failed, len(grids) and/or Bz is zero") # discard grids in private flux region) w = np.where(np.isclose(cols, iz, rtol=0.1)) twopi_exp_bp_sigma_rphi_z = np.zeros(bz.shape) twopi_exp_bp_sigma_rphi_z = -sigma_bp * dpsi2drdr[rows[w], cols[w]] / bz[rows[w], cols[w]] sigma_rphi_z = np.sign(np.sum(np.sign(twopi_exp_bp_sigma_rphi_z))) # 6 exp_Bp from Eq.(19) x = np.average(twopi_exp_bp_sigma_rphi_z * sigma_rphi_z) exp_bp = np.where(np.isclose(x, 2.0 * np.pi, rtol=0.5), 1, 0) # values = { "ipsign": int(ipsign), "b0sign": int(b0sign), "exp_Bp": int(exp_bp), "sigma_Bp": int(sigma_bp), "sigma_RphiZ": int(sigma_rphi_z), "sigma_rhothetaphi": int(sigma_rhothetaphi), "sign_q_pos": int(sign_q_pos), "sign_pprime_pos": int(sign_pprime_pos), } cocos = COCOS(values=values).get() return cocos
[docs]def ids_compute_cocos(ids, itime=None, i1=0): """ Function Interface for computing COCOS. Parameters ---------- ids : object IDS for cocos estimation itime : int, optional Index of struct_array time_slice in IDS/equilibrium i1 : int, optional Index of struct_array profiles_2d in IDS/equilibrium. Default is 0. Returns ------- int COCOS number computed """ key = "COCOS" if ids.metadata.name == "equilibrium": try: cocos = compute_COCOS(ids, itime, i1) except Exception as e: logger.debug(f"{e}") logger.error(traceback.format_exc()) else: exit(f"equilibrium instead of {ids.metadata.name}") return cocos[key]