How to use sub_key method in hypothesis
Best Python code snippet using hypothesis
ECRad_Results.py
Source:ECRad_Results.py 
1'''2Created on Dec 7, 20203@author: Severin Denk4Restructuring of the old ECRadresults class. Uses the OMFit style approach where the parent class is a dictionary5'''6from Distribution_Classes import Distribution7import numpy as np8np.set_printoptions(threshold=10)9import os10from Global_Settings import globalsettings11from scipy.io import savemat, loadmat12from scipy import constants as cnst13from ECRad_Config import ECRadConfig14from ECRad_Scenario import ECRadScenario15from Ndarray_Helper import ndarray_math_operation, ndarray_check_for_None16from netCDF4 import Dataset17import sys18class ECRadResults(dict):19 def __init__(self, lastused=False):20 self.result_keys = ["Trad", "resonance", "ray", "BPD", "weights", "dimensions"]21 self.shapes = {}22 self.units = {}23 self.scales = {}24 self.xaxis_link = {}25 self.legend_entries = {}26 self.labels = {}27 self.sub_keys = {}28 self.graph_style = {}29 self.failed_keys = {}30 # Mode order is mixed, X, O31 self.sub_keys["Trad"] = ["Trad", "tau", "T", \32 "Trad_second", "tau_second", "T_second"]33 self.shapes["Trad"] = ["N_time", "N_mode_mix", "N_ch"]34 self.units["Trad"] = {"Trad":"keV", "tau":"", "T":"", \35 "Trad_second":"keV", "tau_second":"", "T_second":""}36 self.scales["Trad"] = {"Trad":1.e-3, "tau":1.0, "T":1.0, \37 "Trad_second":1.e-3, "tau_second":1.0, "T_second":1.0}38 self.xaxis_link["Trad"] = ["resonance" ,"Trad"]39 self.legend_entries["Trad"] = {"Trad":r"$T_" + globalsettings.mathrm + r"{rad}$", "tau":r"$\tau$", "T":r"$T$", \40 "Trad_second":r"$T_" + globalsettings.mathrm + r"{rad,2nd\,model}$", \41 "tau_second":r"$\tau_" + globalsettings.mathrm + r"{2nd\,model}$", "T_second":r"$T_" + globalsettings.mathrm + r"{2nd\,model}$"}42 self.labels["Trad"] = {"Trad":r"$T_" + globalsettings.mathrm + r"{rad}$", "tau":r"$\tau$", "T":r"$T$", \43 "Trad_second":r"$T_" + globalsettings.mathrm + r"{rad}$", "tau_second":r"$\tau$", "T_second":r"$T$"}44 self.graph_style["Trad"] = "point"45 self.sub_keys["resonance"] = ["s_cold", "R_cold", "z_cold", \46 "rhop_cold", "rhot_cold", "s_warm", "R_warm", \47 "z_warm", "rhop_warm", "rhot_warm", \48 "s_warm_second", "R_warm_second", \49 "z_warm_second", "rhop_warm_second", \50 "rhot_warm_second"]51 # You want the channel as the inner most index since you want to plot as a function of channel52 self.shapes["resonance"] = ["N_time", "N_mode_mix", "N_ch"]53 self.units["resonance"] = {}54 self.scales["resonance"] = {}55 for sub_key in self.sub_keys["resonance"]:56 if(sub_key.startswith("rho")):57 self.units["resonance"][sub_key] = ""58 else:59 self.units["resonance"][sub_key] = "m"60 self.scales["resonance"][sub_key] = 1.061 self.xaxis_link["resonance"] = ["Trad", "resonance"]62 self.legend_entries["resonance"] = {"s_cold":r"$s_" + globalsettings.mathrm + r"{cold}$", \63 "R_cold":r"$R_" + globalsettings.mathrm + r"{cold}$", \64 "z_cold":r"$z_" + globalsettings.mathrm + r"{cold}$", \65 "rhop_cold":r"$\rho_" + globalsettings.mathrm + r"{pol,cold}$", \66 "rhot_cold":r"$\rho_" + globalsettings.mathrm + r"{tor,cold}$", \67 "s_warm":r"$s_" + globalsettings.mathrm + r"{warm}$", \68 "R_warm":r"$R_" + globalsettings.mathrm + r"{warm}$", \69 "z_warm":r"$z_" + globalsettings.mathrm + r"{warm}$", \70 "rhop_warm":r"$\rho_" + globalsettings.mathrm + r"{pol,warm}$", \71 "rhot_warmd":r"$\rho_" + globalsettings.mathrm + r"{tor,warm}$", \72 "s_warm_second":r"$s_" + globalsettings.mathrm + r"{warm,\,2nd\,model}$", \73 "R_warm_second":r"$R_" + globalsettings.mathrm + r"{warm,\,2nd\,model}$", \74 "z_warm_second":r"$z_" + globalsettings.mathrm + r"{warm,\,2nd\,model}$", \75 "rhop_warm_second":r"$\rho_" + globalsettings.mathrm + r"{pol,warm,\,2nd\,model}$", \76 "rhot_warm_second":r"$\rho_" + globalsettings.mathrm + r"{tor,warm,\,2nd\,model}$"}77 self.labels["resonance"] = {"s_cold":r"$s_" + globalsettings.mathrm + r"{cold}$", \78 "R_cold":r"$R_" + globalsettings.mathrm + r"{cold}$", \79 "z_cold":r"$z_" + globalsettings.mathrm + r"{cold}$", \80 "rhop_cold":r"$\rho_" + globalsettings.mathrm + r"{pol,cold}$", \81 "rhot_cold":r"$\rho_" + globalsettings.mathrm + r"{tor,cold}$", \82 "s_warm":r"$s_" + globalsettings.mathrm + r"{warm}$", \83 "R_warm":r"$R_" + globalsettings.mathrm + r"{warm}$", \84 "z_warm":r"$z_" + globalsettings.mathrm + r"{warm}$", \85 "rhop_warm":r"$\rho_" + globalsettings.mathrm + r"{pol,warm}$", \86 "rhot_warmd":r"$\rho_" + globalsettings.mathrm + r"{tor,warm}$", \87 "s_warm_second":r"$s_" + globalsettings.mathrm + r"{warm}$", \88 "R_warm_second":r"$R_" + globalsettings.mathrm + r"{warm}$", \89 "z_warm_second":r"$z_" + globalsettings.mathrm + r"{warm}$", \90 "rhop_warm_second":r"$\rho_" + globalsettings.mathrm + r"{pol,warm}$", \91 "rhot_warm_second":r"$\rho_" + globalsettings.mathrm + r"{tor,warm}$"}92 self.graph_style["resonance"] = "point"93 self.sub_keys["ray"] = ["s", "x", "y", "R", "z", \94 "Nx", "Ny", "Nz", \95 "Bx", "By", "Bz", \96 "H", "N", "Nc", \97 "X", "Y", "rhop", "rhot" ,\98 "Te", "ne", \99 "theta", "BPD", \100 "BPD_second", 101 "Trad", "Trad_second", "em", \102 "em_second", "ab", \103 "ab_second", \104 "T", "T_second", \105 "v_g_perp"]106 self.shapes["ray"] = ["N_time", "N_ch", "N_mode", "N_ray", "N_LOS"]107 self.legend_entries["ray"] = {"s":r"$s$", "x":r"$x$", "y":r"$y$", "R":r"$R$", "z":r"$z$", \108 "Nx":r"$N_" + globalsettings.mathrm + r"{x}$", "Ny":r"$N_" + globalsettings.mathrm + r"{y}$", \109 "Nz":r"$N_" + globalsettings.mathrm + r"{z}$", "Bx":r"$B_" + globalsettings.mathrm + r"{x}$", \110 "By":r"$B_" + globalsettings.mathrm + r"{y}$", "Bz":r"$B_" + globalsettings.mathrm + r"{z}$", \111 "H":r"Hamiltonian", "N":r"$N_" + globalsettings.mathrm + r"{ray}$", \112 "Nc":r"$N_" + globalsettings.mathrm + r"{disp}$", "X":r"Stix Parameter $X$", \113 "Y":r"Stix Parameter $Y$", "rhop":r"$\rho_" + globalsettings.mathrm + r"{pol}$", \114 "rhot":r"$\rho_" + globalsettings.mathrm + r"{tor}$", "Te":r"$T_" + globalsettings.mathrm + r"{e}$", \115 "ne":r"$n_" + globalsettings.mathrm + r"{e}$", "theta":r"$\theta$", \116 "BPD":"BPD", "BPD_second":"BPD$_" + globalsettings.mathrm + r"{2nd\,model}$", \117 "Trad":r"$T_" + globalsettings.mathrm + r"{rad}$", "Trad_second":r"$T_" + globalsettings.mathrm + r"{rad,\,2nd\,model}}$", \118 "em":r"$j$", \119 "em_second":r"$j_{" + globalsettings.mathrm + r"{2nd\,model}}$", "ab":r"$\alpha$", \120 "ab_second":r"$\alpha_{" + globalsettings.mathrm + r"{2nd\,model}}$", "T":r"$\mathcal{T}$", \121 "T_second":r"$\mathcal{T}_{" + globalsettings.mathrm + r"{2nd\,model}}$", "v_g_perp":r"$v_{" + globalsettings.mathrm + r"{g},\perp}$"}122 self.labels["ray"] = {"s":r"$s$", "x":r"$x$", "y":r"$y$", "R":r"$R$", "z":r"$z$", \123 "Nx":r"$N_" + globalsettings.mathrm + r"{x}$", "Ny":r"$N_" + globalsettings.mathrm + r"{y}$", \124 "Nz":r"$N_" + globalsettings.mathrm + r"{z}$", "Bx":r"$B_" + globalsettings.mathrm + r"{x}$", \125 "By":r"$B_" + globalsettings.mathrm + r"{y}$", "Bz":r"$B_" + globalsettings.mathrm + r"{z}$", \126 "H":r"Hamiltonian", "N":r"$N_" + globalsettings.mathrm + r"{ray}$", \127 "Nc":r"$N_" + globalsettings.mathrm + r"{disp}$", "X":r"Stix Parameter $X$", \128 "Y":r"Stix Parameter $Y$", "rhop":r"$\rho_" + globalsettings.mathrm + r"{pol}$", \129 "rhot":r"$\rho_" + globalsettings.mathrm + r"{tor}$", "Te":r"$T_" + globalsettings.mathrm + r"{e}$", \130 "ne":r"$n_" + globalsettings.mathrm + r"{e}$", "theta":r"$\theta$", \131 "BPD":"BPD", "BPD_second":"BPD", \132 "Trad":r"$T_" + globalsettings.mathrm + r"{rad}$", "Trad_second":r"$T_" + globalsettings.mathrm + r"{rad}$", \133 "em":r"$j_\omega$", \134 "em_second":r"$j_\omega$", "ab":r"$\alpha_\omega$", \135 "ab_second":r"$\alpha_\omega$", "T":r"$\mathcal{T}_\omega$", \136 "T_second":r"$\mathcal{T}_\omega$", "v_g_perp":r"$v_{" + globalsettings.mathrm + r"{g},\perp}$"}137 self.units["ray"] = {"s":"m", "x":"m", "y":"m", "R":"m", "z":"m", \138 "Nx":"", "Ny":"", "Nz":"", \139 "Bx":"T", "By":"T", "Bz":"T", \140 "H":"", "N":"", "Nc":"", \141 "X":"", "Y":"", \142 "rhop":"", "rhot":"",\143 "Te":r"keV", "ne":r"$10^{19}$m$^{-3}$", \144 "theta":r"$^\circ$", "BPD":r"m$^{-1}$", \145 "BPD_second":r"m$^{-1}$", \146 "Trad":r"keV", "Trad_second":r"keV", \147 "em":r"nW m$^{-3}$", \148 "em_second":r"nW m$^{-3}$", "ab":r"m$^{-1}$", \149 "ab_second":r"m$^{-1}$", \150 "T":"", "T_second":"", \151 "v_g_perp":""}152 self.scales["ray"] = {"s":1, "x":1, "y":1, "R":1, "z":1, \153 "Nx":1, "Ny":1, "Nz":1, \154 "Bx":1, "By":1, "Bz":1, \155 "H":1, "N":1, "Nc":1, \156 "X":1, "Y":1, \157 "rhop":1, "rhot":1, \158 "Te":1.e-3, "ne":1.e-19, \159 "theta":np.rad2deg(1), "BPD":1, \160 "BPD_second":1, \161 "Trad":1.e-3, "Trad_second":1.e-3, \162 "em":1.e9, \163 "em_second":1.e9, "ab":1, \164 "ab_second":1, \165 "T":1, "T_second":1, \166 "v_g_perp":1.0/cnst.speed_of_light}167 self.xaxis_link["ray"] = ["ray"]168 self.graph_style["ray"] = "line"169 self.sub_keys["BPD"] = ["rhop", "rhot", "BPD", "BPD_second"]170 self.scales["BPD"] = {"rhop":1.0, "rhot":1.0, "BPD":1.0, "BPD_second":1.0}171 self.labels["BPD"] = {"rhop":r"$\rho_" + globalsettings.mathrm + r"{pol}$", "rhot":r"$\rho_" + globalsettings.mathrm + r"{tor}$", \172 "BPD":"BPD","BPD_second":"BPD"}173 self.legend_entries["BPD"] = {"rhop":r"$\rho_" + globalsettings.mathrm + r"{pol}$", "rhot":r"$\rho_" + globalsettings.mathrm + r"{tor}$", \174 "BPD":"BPD","BPD_second":"BPD$_" + globalsettings.mathrm + r"{2nd\,model}$"}175 self.units["BPD"] = {"rhop":"", "rhot":"", "BPD":"m$^{-1}$", \176 "BPD_second":"m$^{-1}$"}177 self.xaxis_link["BPD"] = ["BPD"]178 self.shapes["BPD"] = ["N_time", "N_ch", "N_mode_mix", "N_BPD"]179 self.graph_style["BPD"] = "line"180 self.sub_keys["weights"] = ["mode_frac", "mode_frac_second", "ray_weights", "freq_weights"]181 self.units["weights"] = {}182 self.scales["weights"] = {}183 self.graph_style["weights"] = "point"184 for sub_key in self.sub_keys["weights"]:185 self.units["weights"][sub_key] = r'$\%$'186 self.scales["weights"][sub_key] = 1.e2187 self.labels["weights"] = {"mode_frac":r"Mode contribution", "mode_frac_second":"Mode contribution", \188 "ray_weights":"ray weight","freq_weights":"frequency weight"}189 self.legend_entries["weights"] = {"mode_frac":r"Mode contribution", "mode_frac_second":"Mode contribution", \190 "ray_weights":"ray weight","freq_weights":"frequency weight"}191 # The weights have a lot of different shapes need to store those by sub_key192 self.shapes["mode_frac"] = ["N_time", "N_mode", "N_ch"]193 self.shapes["mode_frac_second"] = ["N_time", "N_mode", "N_ch"]194 self.shapes["ray_weights"] = ["N_time", "N_ch", "N_ray"]195 self.shapes["freq_weights"] = ["N_time", "N_ch", "N_freq"]196 self.xaxis_link["weights"] = ["weights"]197 self.graph_style["BPD"] = "line"198 self["git"] = {"ECRad":"Unknoiwn", "GUI":"Unknoiwn", "Pylib":"Unknoiwn"}199 self["types"] = {}200 for key in self.result_keys:201 if(key == 'dimensions'):202 self["types"][key] = "int"203 else:204 self["types"][key] = "float"205 for key in self.units.keys():206 for sub_key in self.units[key].keys():207 if(len(self.units[key][sub_key]) > 0):208 self.units[key][sub_key] = "[" + self.units[key][sub_key] + "]"209 self.reset(not lastused)210 def reset(self, noLoad=True, light=False):211 # Does not reset Config or Scenario212 self.status = 0213 self.edition = 0214 self.init = False215 # Edition of this result216 # Remains zero until the result is saved217 # Time of the results, should be identical to self.Scenario.plasma_dict["time"]218 # This holds the same information as the scenario219 self.modes = None220 self.comment = ""221 self["dimensions"] = {}222 self["dimensions"]["N_LOS"] = []223 self["git"] = {}224 for key in self.result_keys:225 if(key in ["dimensions"]):226 continue227 self.failed_keys[key] = []228 self[key] = {}229 if(key in self.sub_keys):230 for sub_key in self.sub_keys[key]:231 self[key][sub_key] = []232 if(not light):233 self.Config = ECRadConfig(noLoad=noLoad)234 self.Scenario = ECRadScenario(noLoad=noLoad)235 self.data_origin = None236 237 def load(self, filename):238 if(filename is not None):239 ext = os.path.splitext(filename)[1]240 if(ext == ".mat"):241 self.from_mat(filename=filename)242 elif(ext == ".nc"):243 self.from_netcdf(filename=filename)244 else:245 print("Extension " + ext + " is unknown")246 raise(ValueError)247 def tidy_up(self, autosave=True):248 if(self.status != 0):249 return250 # Put everything into numpy arrays251 for key in self.result_keys:252 if(key in ["dimensions", "git", "types"]):253 continue254 for sub_key in self.sub_keys[key]:255 self[key][sub_key] = np.array(self[key][sub_key])256 # Autosave results257 self["git"]["ECRad"] = np.genfromtxt(os.path.join(globalsettings.ECRadRoot, "id"), dtype=str).item()258 self["git"]["GUI"] = np.genfromtxt(os.path.join(globalsettings.ECRadGUIRoot, "id"), dtype=str).item()259 self["git"]["Pylib"] = np.genfromtxt(os.path.join(globalsettings.ECRadPylibRoot, "id"), dtype=str).item()260 self.data_origin = "ECRad"261 if(autosave):262 self.autosave()263 264 def autosave(self,filename=None):265 self.to_netcdf(filename)266 def get_shape(self, key, start=None, stop=None,i_time=None, \267 i_ch=None, i_mode=None, i_ray=None):268 shape = ()269 for dim_ref in self.shapes[key][start:stop]:270 if(np.isscalar(self["dimensions"][dim_ref])):271 shape += (self["dimensions"][dim_ref],)272 else:273 # Cannot use numpy indexing here because the time dimension 274 # will be appended as we go275 shape += (self["dimensions"][dim_ref][i_time][i_ch, i_mode, i_ray],)276 return shape277 278 def get_index_reference(self, key, sub_key, ndim, index):279 # Retrieves the value this particular index refers to280 # used in plotting to label graphs.281 # The routine formats the quantity and returns a string282 # Note that the index should have the shape corresponding to the requested quantity283 if(key != "weights"):284 dim_ref = self.shapes[key][ndim]285 else:286 dim_ref = self.shapes[sub_key][ndim]287 if(dim_ref == "N_time"):288 return r"$t = $ " + "{0:1.3f}".format(self.Scenario["time"][index[ndim]]) + " s"289 elif(dim_ref == "N_ch"):290 return r"$f = $ " + "{0:3.1f}".format(self.Scenario["diagnostic"]["f"][index[0]][index[ndim]]/1.e9) + " GHz" 291 elif(dim_ref in ["N_mode"]):292 if(self.Config["Physics"]["considered_modes"] == 1):293 return "X-mode"294 elif(self.Config["Physics"]["considered_modes"] == 2):295 return "O-mode"296 else:297 if(index[ndim] == 0):298 return "X-mode"299 else:300 return "O-mode"301 elif(dim_ref == "N_mode_mix"):302 if(self.Config["Physics"]["considered_modes"] < 3):303 return ""304 else:305 if(index[ndim] == 1):306 return "X-mode"307 elif(index[ndim] == 2):308 return "O-mode"309 else:310 return ""311 elif(dim_ref == "N_ray"):312 if(index[ndim] == 0):313 return ""314 else:315 return r"ray \#" + str(index[ndim] + 1)316 else:317 raise ValueError("ECRadResults.get_index_reference could unexpected dim ref " + key + " " + sub_key + " " + str(ndim))318 319 def set_dimensions(self):320 # Sets the dimensions from Scenario and Configself["dimensions"]["N_time"] = len(self.Scenario.plasma_dict["time"])321 self["dimensions"]["N_time"] = self.Scenario["dimensions"]["N_time"]322 self["dimensions"]["N_ray"] = self.Config["Physics"]["N_ray"]323 self["dimensions"]["N_freq"] = self.Config["Physics"]["N_freq"]324 self["dimensions"]["N_BPD"] = self.Config["Numerics"]["N_BPD"]325 self["dimensions"]["N_ch"] = self.Scenario["dimensions"]["N_ch"]326 if(self.Config["Physics"]["considered_modes"] > 2):327 self["dimensions"]["N_mode"] = 2328 self["dimensions"]["N_mode_mix"] = 3329 modes_mix = ["", "X", "O"]330 modes = ["X", "O"]331 else:332 self["dimensions"]["N_mode"] = 1333 self["dimensions"]["N_mode_mix"] = 1334 modes = ["X"]335 if(self.Config["Physics"]["considered_modes"] == 2):336 modes = ["O"]337 modes_mix = [""]338 self["dimensions"]["N_ch"] = self.Scenario["dimensions"]["N_ch"]339 return modes, modes_mix340 def from_mat(self, filename):341 try:342 mdict = loadmat(filename, chars_as_strings=True, squeeze_me=True)343 except IOError as e:344 print(e)345 print("Error: " + filename + " does not exist")346 return347 self.Config.from_mat(mdict=mdict)348 self.Scenario.from_mat(mdict=mdict, load_plasma_dict=True)349 self.edition = mdict["edition"]350 # We need to do this song and dance because351 # 3D equlibria use rho tor instead of rho pol352 # This was not clearly indicated in the old353 # .mat files, but the new result files distinguish this.354 # Further cases are marked with #3D rhot355 if(self.Scenario["plasma"]["eq_dim"] == 3):356 rho = "rhot"357 else:358 rho = "rhop"359 if("comment" in mdict):360 self.comment = mdict["comment"]361 try:362 if(type(self.comment) == np.ndarray):363 if(len(self.comment) > 0):364 self.comment = self.comment[0]365 else:366 self.comment = ""367 except Exception as e:368 print("Failed to parse comment")369 print(e)370 self.comment = ""371 modes, modes_mix = self.set_dimensions()372 for key in ["Trad"]:373 for sub_key in self.sub_keys[key]:374 self[key][sub_key] = np.zeros(self.get_shape(key))375 for i_mode, mode in enumerate(modes_mix):376 if(sub_key == "T"):377 self[key][sub_key] = np.exp(-self["Trad"]["tau"])378 continue379 elif(sub_key == "T_second"):380 self[key][sub_key] = np.exp(-self["Trad"]["tau_second"])381 continue382 mdict_key = sub_key.replace("second","comp")383 self[key][sub_key][...,i_mode,:] = mdict[mode+mdict_key].reshape(self.get_shape(key,stop=1) + \384 self.get_shape(key,start=2)) *1.e3385 if(self["dimensions"]["N_mode"] > 1):386 mode_info_printed = False387 # mode resolved resonances are not available in .mat files388 for key in ["resonance"]:389 for sub_key in self.sub_keys[key]:390 self[key][sub_key] = np.zeros(self.get_shape(key))391 if(sub_key.endswith("second")):392 formatted_key = sub_key + "ary"393 else:394 formatted_key = sub_key395 #3D rhot396 if(rho == "rhot" and sub_key.startswith("rhot")):397 formatted_key = formatted_key.replace("rhot","rhop")398 elif(rho == "rhot" and formatted_key.startswith("rhop")):399 continue400 try:401 self[key][sub_key][...,0,:] = mdict[formatted_key].reshape(self.get_shape(key,stop=1) + \402 self.get_shape(key,start=2))403 if(not mode_info_printed):404 print("INFO:: No mode specific resonances in .mat files. Using mixed modes for all resonances.")405 mode_info_printed = True406 for imode in range(1,3):407 self[key][sub_key][...,imode,:] = self[key][sub_key][...,0,:]408 except KeyError:409 print("INFO: Couldn't load " + sub_key + " from result file")410 else:411 for key in ["resonance"]:412 for sub_key in self.sub_keys[key]:413 if(sub_key.endswith("second")):414 formatted_key = sub_key + "ary"415 else:416 formatted_key = sub_key417 #3D rhot418 if(rho == "rhot" and sub_key.startswith("rhot")):419 formatted_key = formatted_key.replace("rhot","rhop")420 elif(rho == "rhot" and formatted_key.startswith("rhop")):421 continue422 try:423 self[key][sub_key] = mdict[formatted_key].reshape(self.get_shape(key))424 except KeyError:425 print("INFO: Couldn't load " + sub_key + " from result file")426 self["dimensions"]["N_LOS"] = np.zeros(self.get_shape("ray", 0, -1),dtype=np.int) 427 # We need to fix the shape of the mdict ray info428 key = "ray"429 for sub_key in self.sub_keys[key]:430 for mode in modes:431 if(sub_key in ["em", "ab", "T", "BPD", "Trad", "Trad_second",\432 "em_second", "ab_second", "T_second", "BPD_second"]):433 mdict_key = "ray_" + sub_key434 #3D rhot435 elif(sub_key == "rhot" and rho == "rhot"):436 mdict_key = "rhop"437 elif(sub_key == "rhop" and rho == "rhot"):438 continue439 else:440 mdict_key = sub_key441 if(mdict_key+mode not in mdict.keys()):442 print("Cannot find " + key + "/" + sub_key + " in .mat")443 if(sub_key not in self.failed_keys[key]):444 self.failed_keys[key].append(sub_key)445 continue446 if(self["dimensions"]["N_time"] == 1):447 mdict[mdict_key+mode] = np.expand_dims(mdict[mdict_key+mode], 0)448 if(self["dimensions"]["N_ch"] == 1):449 mdict[mdict_key+mode] = np.expand_dims(mdict[mdict_key+mode],1)450 if(self["dimensions"]["N_ray"] == 1):451 mdict[mdict_key+mode] = np.expand_dims(mdict[mdict_key+mode],2)452 for key in ["ray", "BPD"]:453 for sub_key in self.sub_keys[key]:454 sub_key_error_printed = False455 if(key == "BPD"):456 mdict_key = None457 #3D rhot458 if(sub_key == "rhot" and self.Scenario["plasma"]["eq_dim"] == 2):459 continue460 elif(sub_key == "rhop" and self.Scenario["plasma"]["eq_dim"] == 3):461 continue462 elif(sub_key == "rhot" and self.Scenario["plasma"]["eq_dim"] == 3):463 # The BPD axis is mislabeled for 3D equilibria. It should be rhot464 mdict_key = "BPD" + "rhop"465 if(mdict_key is None):466 if(sub_key in ["rhop", "rhot"]):467 mdict_key = "BPD" + sub_key468 else:469 mdict_key = sub_key470 if(mdict_key+mode not in mdict.keys()):471 print("INFO: Cannot load " + key + "/" + sub_key)472 if(sub_key not in self.failed_keys[key]):473 self.failed_keys[key].append(sub_key)474 continue475 self[key][sub_key] = np.zeros(self.get_shape(key))476 for i_mode, mode in enumerate(modes):477 self[key][sub_key][...,i_mode,:] = mdict[mdict_key+mode].reshape(self.get_shape(key,stop=-2) + \478 (self["dimensions"]["N_BPD"],))479 else:480 if(sub_key in ["em", "ab", "T", "BPD", \481 "em_second", "ab_second", "T_second", "BPD_second"]):482 mdict_key = "ray_" + sub_key483 #3D rhot484 elif(sub_key == "rhot" and rho == "rhot"):485 mdict_key = "rhop"486 elif(sub_key == "rhop" and rho == "rhot"):487 continue488 else:489 mdict_key = sub_key490 if(mdict_key+mode not in mdict.keys()):491 if(sub_key !="R"):492 print("INFO: Cannot load " + key + "/" + sub_key)493 if(sub_key not in self.failed_keys[key]):494 self.failed_keys[key].append(sub_key)495 continue496 self["ray"][sub_key] = []497 for i_time in range(self["dimensions"]["N_time"]):498 self["ray"][sub_key].append([])499 for i_ch in range(self["dimensions"]["N_ch"]):500 self["ray"][sub_key][i_time].append([])501 for i_mode, mode in enumerate(modes):502 self["ray"][sub_key][i_time][i_ch].append([])503 for i_ray in range(self["dimensions"]["N_ray"]):504 try:505 self["ray"][sub_key][i_time][i_ch][i_mode].append( \506 mdict[mdict_key+mode][i_time][i_ch][i_ray])507 self["dimensions"]["N_LOS"][i_time][i_ch][i_mode][i_ray] = \508 len(self["ray"][sub_key][i_time][i_ch][i_mode][i_ray])509 except IndexError:510 if(not sub_key_error_printed):511 print("INFO: Failed to load {0:s} {1:s}".format(key,sub_key))512 print("INFO: For time index {0:d} channel {1:d} mode index {2:d} ray {3:d}".format(513 i_time, i_ch, i_mode, i_ray))514 sub_key_error_printed = True515 self["ray"][sub_key][i_time][i_ch][i_mode].append([])516 self["dimensions"]["N_LOS"][i_time][i_ch][i_mode][i_ray] = 0517 # Convert to ragged np array518 self["ray"][sub_key] = np.array(self["ray"][sub_key], dtype=np.object)519 self["weights"]["ray_weights"] = mdict["ray_weights"]520 self["weights"]["freq_weights"] = mdict["ray_weights"]521 if(self.Config["Physics"]["considered_modes"] > 2):522 self["weights"]["mode_frac"] = np.zeros(self.get_shape("mode_frac"))523 self["weights"]["mode_frac_second"] = np.zeros(self.get_shape("mode_frac_second"))524 self["weights"]["mode_frac"][...,0,:] = mdict["X_mode_frac"].reshape(self.get_shape("mode_frac",stop=1) + \525 self.get_shape("mode_frac",start=2))526 self["weights"]["mode_frac_second"][...,0,:] = mdict["X_mode_frac_comp"].reshape(self.get_shape("mode_frac_second",stop=1) + \527 self.get_shape("mode_frac_second",start=2))528 self["weights"]["mode_frac"][...,1,:] = 1 - mdict["X_mode_frac"].reshape(self.get_shape("mode_frac",stop=1) + \529 self.get_shape("mode_frac",start=2))530 self["weights"]["mode_frac_second"][...,1,:] = 1 - mdict["X_mode_frac_comp"].reshape(self.get_shape("mode_frac_second",stop=1) + \531 self.get_shape("mode_frac_second",start=2))532 else:533 self["weights"]["mode_frac"] = np.ones(self.get_shape("mode_frac"))534 self["weights"]["mode_frac_second"] = np.ones(self.get_shape("mode_frac_second"))535 self["ray"]["R"] = np.zeros(self.get_shape("ray", stop=-1), dtype=np.object)536 print("INFO: Fixing missing ray/R.")537 for itime in range(self["dimensions"]["N_time"]):538 for ich in range(self["dimensions"]["N_ch"]):539 for imode in range(self["dimensions"]["N_mode"]):540 for iray in range(self["dimensions"]["N_ray"]):541 if(self["dimensions"]["N_LOS"][i_time][i_ch][i_mode][i_ray] > 0):542 self["ray"]["R"][itime,ich,imode,iray] = np.sqrt(self["ray"]["x"][itime,ich,imode,iray]**2 + \543 self["ray"]["y"][itime,ich,imode,iray]**2)544 # We fix R later so we do not need to delete it545 self.failed_keys["ray"].remove("R")546 self["git"]["ECRad"] = mdict["ECRad_git_tag"]547 self["git"]["GUI"] = mdict["ECRadGUI_git_tag"]548 self["git"]["Pylib"] = mdict["ECRadPylib_git_tag"]549 self.data_origin = filename550 self.init = True551 return True552 553 def get_default_filename_and_edition(self, scratch=False, ed=None):554 if(scratch):555 dir = self.Config["Execution"]["scratch_dir"]556 else:557 dir = self.Config["Execution"]["working_dir"]558 diag_str = ""559 for key in self.Scenario["used_diags_dict"]:560 diag_str += key561 if(ed is None):562 ed = 1563 filename = os.path.join(dir, "ECRad_{0:5d}_{1:s}_ed{2:d}.nc".format(self.Scenario["shot"], diag_str, ed))564 while(os.path.exists(filename)):565 ed += 1566 filename = os.path.join(dir, "ECRad_{0:5d}_{1:s}_ed{2:d}.nc".format(self.Scenario["shot"], diag_str, ed))567 return filename, ed568 else:569 filename = os.path.join(dir, "ECRad_Results_{0:d}.nc".format(ed))570 return filename, 0571 def to_netcdf(self, filename=None, scratch=False, ed=None):572 if(filename is not None):573 rootgrp = Dataset(filename, "w", format="NETCDF4")574 else:575 filename, self.edition = self.get_default_filename_and_edition(scratch, ed=ed)576 rootgrp = Dataset(filename, "w", format="NETCDF4")577 rootgrp.createGroup("Results")578 self.Config.to_netcdf(rootgrp=rootgrp)579 self.Scenario.to_netcdf(rootgrp=rootgrp)580 for sub_key in self["dimensions"].keys():581 if(np.isscalar(self["dimensions"][sub_key])):582 rootgrp["Results"].createDimension(sub_key, self["dimensions"][sub_key])583 else:584 rootgrp["Results"].createDimension(sub_key, None) 585 for key in self.result_keys:586 if(key == "dimensions" or key == "ray"):587 continue588 dtype = "f8"589 if(self["types"][key] != "float"):590 dtype = "i8"591 try:592 for sub_key in self.sub_keys[key]:593 if(len(self[key][sub_key]) == 0):594 print("INFO: Not saving " + key + " " + sub_key + " because there is no data.")595 continue596 if(sub_key in self.failed_keys[key]):597 continue598 if(sub_key == "rhot" and self.Scenario["plasma"]["eq_dim"] == 2):599 continue600 elif(sub_key == "rhop" and self.Scenario["plasma"]["eq_dim"] == 3):601 continue602 if(key != "weights"):603 var = rootgrp["Results"].createVariable(key + "_" + sub_key,dtype, tuple(self.shapes[key]))604 else:605 var = rootgrp["Results"].createVariable(key + "_" + sub_key,dtype, tuple(self.shapes[sub_key]))606 var[:] = self[key][sub_key]607 except Exception as e:608 print(key, sub_key)609 raise e610 key = "ray"611 for sub_key in self.sub_keys["ray"]:612 if(sub_key in self.failed_keys[key]):613 continue614 elif(len(self[key][sub_key]) == 0):615 continue616 var = rootgrp["Results"].createVariable(key + "_" + sub_key,dtype, tuple(self.shapes[key]))617 for i_time in range(self["dimensions"]["N_time"]):618 for i_ch in range(self["dimensions"]["N_ch"]):619 for i_mode in range(self["dimensions"]["N_mode"]):620 for i_ray in range(self["dimensions"]["N_ray"]):621 var[i_time,i_ch,i_mode,i_ray,:] = self[key][sub_key][i_time,i_ch,i_mode,i_ray]622 # Get the shape information of the individual LOS length into the NETCDF file623 var = rootgrp["Results"].createVariable("dimensions" + "_" + "N_LOS", "i8", self.shapes["ray"][:-1])624 var[:] = self["dimensions"]["N_LOS"]625 rootgrp["Results"].comment = self.comment626 rootgrp["Results"].edition = self.edition627 rootgrp["Results"].ECRad_git_tag = self["git"]["ECRad"]628 rootgrp["Results"].ECRadGUI_git_tag = self["git"]["GUI"]629 rootgrp["Results"].ECRadPylib_git_tag = self["git"]["Pylib"]630 rootgrp.close()631 print("Created " + filename)632 633 def from_netcdf(self, filename):634 rootgrp = Dataset(filename, "r", format="NETCDF4")635 self.Config.from_netcdf(rootgrp=rootgrp)636 self.Scenario.from_netcdf(rootgrp=rootgrp)637 for sub_key in rootgrp["Results"].dimensions.keys():638 self["dimensions"][sub_key] = rootgrp["Results"].dimensions[sub_key].size639 self["dimensions"]["N_LOS"] = np.array(rootgrp["Results"]["dimensions" + "_" + "N_LOS"])640 for key in self.result_keys:641 if(key == "dimensions" or key == "ray"):642 continue643 for sub_key in self.sub_keys[key]:644 if(key + "_" + sub_key not in rootgrp["Results"].variables.keys()):645 print("INFO: Could not find " + key + " " + sub_key + " in the result file.")646 self.failed_keys[key].append(sub_key)647 continue648 if(key == "BPD"):649 if(sub_key == "rhot" and self.Scenario["plasma"]["eq_dim"] == 2):650 continue651 elif(sub_key == "rhop" and self.Scenario["plasma"]["eq_dim"] == 3):652 continue653 self[key][sub_key] = np.array(rootgrp["Results"][key + "_" + sub_key])654 655 key = "ray"656 for sub_key in self.sub_keys["ray"]:657 self["ray"][sub_key] = []658 if(key + "_" + sub_key not in rootgrp["Results"].variables.keys()):659 print("INFO: Cannot load " + key + "/" + sub_key)660 self.failed_keys[key].append(sub_key)661 continue662 for i_time in range(self["dimensions"]["N_time"]):663 self["ray"][sub_key].append([])664 for i_ch in range(self["dimensions"]["N_ch"]):665 self["ray"][sub_key][i_time].append([])666 for i_mode in range(self["dimensions"]["N_mode"]):667 self["ray"][sub_key][i_time][i_ch].append([])668 for i_ray in range(self["dimensions"]["N_ray"]):669 self["ray"][sub_key][i_time][i_ch][i_mode].append( \670 rootgrp["Results"][key+ "_" +sub_key][i_time,i_ch,i_mode,i_ray,\671 :self["dimensions"]["N_LOS"][i_time,i_ch,i_mode,i_ray]])672 self["ray"][sub_key] = np.array(self["ray"][sub_key], dtype=np.object)673 # Get the shape information of the individual LOS length into the NETCDF file674 self.comment = rootgrp["Results"].comment 675 self.edition = rootgrp["Results"].edition676 self["git"]["ECRad"] = rootgrp["Results"].ECRad_git_tag677 self["git"]["GUI"] = rootgrp["Results"].ECRadGUI_git_tag678 self["git"]["Pylib"] = rootgrp["Results"].ECRadPylib_git_tag679 self.data_origin = filename680 rootgrp.close()681 682if(__name__ == "__main__"):683 res = ECRadResults()684# res.from_mat("/mnt/c/Users/Severin/ECRad/ECRad_33585_EXT_ed1.mat")685# res.to_netcdf("/mnt/c/Users/Severin/ECRad/ECRad_33585_EXT_ed1.nc")686#"/mnt/c/Users/Severin/ECRad_regression/AUGX3/ECRad_32934_EXT_ed1.nc"687 # res.from_mat("/mnt/c/Users/Severin/ECRad/ECRad_35662_EXT_ed8.mat")688 # res.to_netcdf("/mnt/c/Users/Severin/ECRad/ECRad_35662_EXT_ed8.nc")689# res.reset()690 res.from_netcdf("/mnt/c/Users/Severin/ECRad/HFS_LHCD/ECRad_147634_EXT_ed1.nc")691 res.Scenario["plasma"]["dist_obj"].plot_Te_ne()692# res.reset()693# res.from_mat("/mnt/c/Users/Severin/ECRad_regression/W7X/ECRad_20180823016002_EXT_ed19.mat")694# res.to_netcdf("/mnt/c/Users/Severin/ECRad_regression/W7X/ECRad_20180823016002_EXT_Scenario.nc")695# res.reset()...
ECRad_Config.py
Source:ECRad_Config.py
1# -*- coding: utf-8 -*-2import os3from scipy.io import loadmat, savemat4from netCDF4 import Dataset5class ECRadConfig(dict):6 def __init__(self, noLoad = False):7 self.default_config_file = os.path.join(os.path.expanduser("~"), ".ECRad_GUI_Default.nc")8 if(noLoad):9 self.reset()10 else:11 try:12 self.reset()13 self.load(filename=self.default_config_file, default=True)14 except IOError:15 self.reset()16 def reset(self):17 self.main_keys = ["Physics", "Execution", "Numerics"]18 self.sub_keys = {}19 for key in self.main_keys:20 self[key] = {}21 self.sub_keys["Physics"] = ["dstf", "raytracing", "ripple", \22 "weak_rel", "considered_modes", \23 "N_freq", "N_ray", \24 "ratio_for_3rd_harm", "tau_ignore", "mode_conv" ,\25 "reflec_X", "reflec_O", \26 "R_shift", "z_shift", \27 "use_ext_rays"]28 self.sub_keys["Execution"] = ["working_dir", "scratch_dir", "extra_output", \29 "debug", "batch", "parallel" , \30 "parallel_cores", "wall_time", \31 "vmem"]32 self.sub_keys["Numerics"] = ["large_ds", "small_ds", "max_points_svec" , \33 "N_BPD"]34 self["Physics"]["dstf"] = "Th"35 self["Physics"]["raytracing"] = True36 self["Physics"]["ripple"] = True37 self["Physics"]["weak_rel"] = True38 self["Physics"]["N_freq"] = 139 self["Physics"]["N_ray"] = 140 self["Physics"]["ratio_for_3rd_harm"] = 0.441 self["Physics"]["tau_ignore"] = 1.e-842 self["Physics"]["considered_modes"] = 143 # 1 -> Only X44 # 2 -> Only O45 # 3 -> Both46 self["Physics"]["mode_conv"] = 0.047 self["Physics"]["reflec_X"] = 0.948 self["Physics"]["reflec_O"] = 0.949 self["Physics"]["R_shift"] = 0.050 self["Physics"]["z_shift"] = 0.051 self["Physics"]["use_ext_rays"] = False52 self["Execution"]["working_dir"] = ""53 self["Execution"]["scratch_dir"] = ""54 self["Execution"]["extra_output"] = True55 self["Execution"]["debug"] = False56 self["Execution"]["batch"] = True57 self["Execution"]["parallel"] = True58 self["Execution"]["parallel_cores"] = 3259 self["Execution"]["wall_time"] = 2.060 self["Execution"]["vmem"] = 3200061 self["Numerics"]["large_ds"] = 2.5e-362 self["Numerics"]["small_ds"] = 2.5e-463 self["Numerics"]["max_points_svec"] = 2000064 self["Numerics"]["N_BPD"] = 200065 self.types = {"working_dir":str, "scratch_dir":str, "dstf":str, "extra_output": "b", 66 "debug": "b", "batch": "b", "parallel": "b", 67 "parallel_cores": "i8", "wall_time": "f8", 68 "vmem": "i8", "raytracing": "b", "ripple": "b", 69 "weak_rel": "b", "N_freq" : "i8", "N_ray": "i8", 70 "ratio_for_3rd_harm": "f8", "tau_ignore" :"f8", "considered_modes" : "i8", 71 "mode_conv" : "f8", "reflec_X" : "f8","reflec_O" : "f8", 72 "R_shift" : "f8","z_shift" : "f8","large_ds" : "f8",73 "small_ds" : "f8","max_points_svec" : "i8","use_ext_rays" : "b", 74 "N_BPD" : "i8"}75 self.nice_labels = {"working_dir":"Working dir.", "scratch_dir":"Scratch dir.","dstf":"Distribution type", "extra_output": "Extra output", 76 "debug": "Debug", "batch": "Batch", "parallel": "Parallel", 77 "parallel_cores": "# cores", "wall_time": "wall time [h]", 78 "vmem": "virtual memory [MB]", "raytracing": "Raytracing", "ripple": "Magn. field Ripple", 79 "weak_rel": "Relativistic cor. for rt.", "N_freq" : "# frequencies", "N_ray": "# rays", 80 "ratio_for_3rd_harm": "omega_c/omega w. 3rd", "tau_ignore": "tau threshhold for computation of alpha/j",81 "considered_modes" : "Modes to consider", 82 "mode_conv" : "mode conv. ratio", "reflec_X" : "Wall refl. coeff. X-mode", 83 "reflec_O" : "Wall refl. coeff. O-mode", "use_ext_rays" : "Use ext rays", 84 "R_shift" : "R shift [m]","z_shift" : "z shift [m]", "large_ds" : "Large step size [m]",85 "small_ds" : "Small step size [m]","max_points_svec" : "Max points on LOS", 86 "N_BPD" : "Points for BPD"}87 88 def load(self, filename=None, mdict=None, rootgrp=None, default=False):89 if(filename is not None):90 ext = os.path.splitext(filename)[1]91 if(ext == ".mat"):92 self.from_mat(path_in=filename)93 elif(ext == ".nc"):94 self.from_netcdf(filename=filename,default=default)95 else:96 print("Extension " + ext + " is unknown")97 raise(ValueError)98 elif(mdict is not None):99 self.from_mat(mdict)100 elif(rootgrp is not None):101 self.from_netcdf(rootgrp=rootgrp)102 103 def from_mat(self, mdict=None, path_in=None):104 ext_mdict = False105 temp_config = None106 if(mdict is not None or path_in is not None):107 if(path_in is not None):108 mdict = loadmat(path_in, chars_as_strings=True, squeeze_me=True)109 else:110 ext_mdict = True111 else:112 raise ValueError("Either filename or mdict must be present")113 self.reset()114 key = "Execution"115 for sub_key in ["working_dir", "scratch_dir"]:116 if(os.path.isdir(mdict[sub_key])):117 self[key][sub_key] = mdict[sub_key]118 elif(not ext_mdict):119 self[key][sub_key] = mdict[sub_key]120 else:121 print("Warning working dir not imported, since it is not a valid directory")122 print("Falling back to last used working directory")123 temp_config = ECRadConfig()124 self[key][sub_key] = temp_config[key][sub_key]125 for key in self.main_keys:126 for sub_key in self.sub_keys[key]:127 if(sub_key in ["working_dir", "scratch_dir"]):128 continue129 else:130 if(sub_key in mdict.keys()):131 self[key][sub_key] = mdict[sub_key]132 else:133 print("Could not find " + sub_key + " in config file.")134 if(path_in is None and not ext_mdict):135 print("Successfully loaded last used configuration")136 return137 138 def to_netcdf(self, filename=None, rootgrp=None):139 if(filename is not None):140 rootgrp = Dataset(filename, "w", format="NETCDF4")141 rootgrp.createGroup("Config")142 rootgrp["Config"].createDimension('str_dim', 1)143 for key in self.main_keys:144 for sub_key in self.sub_keys[key]:145 if(self.types[sub_key] == str):146 var = rootgrp["Config"].createVariable(key + "_" + sub_key, self.types[sub_key], 'str_dim')147 var[0] = self[key][sub_key]148 else:149 var = rootgrp["Config"].createVariable(key + "_" + sub_key, self.types[sub_key])150 var[...] = self[key][sub_key]151 if(filename is not None):152 rootgrp.close()153 154 def from_netcdf(self, filename=None, rootgrp=None, default=False):155 self.reset()156 if(filename is not None):157 rootgrp = Dataset(filename, "r", format="NETCDF4")158 key = "Execution"159 for sub_key in ["working_dir", "scratch_dir"]:160 if(os.path.isdir(rootgrp["Config"][key + "_" + sub_key][0])):161 self[key][sub_key] = rootgrp["Config"][key + "_" + sub_key][0]162 elif(not default):163 print("Warning " + sub_key + " not imported, since it is not a valid directory")164 print("Falling back to last used " + sub_key)165 temp_config = ECRadConfig()166 self[key][sub_key] = temp_config["Execution"][sub_key]167 else:168 continue169 for key in self.main_keys:170 for sub_key in self.sub_keys[key]:171 try:172 if(sub_key in ["working_dir", "scratch_dir"]):173 continue174 if(self.types[sub_key] == "b"):175 self[key][sub_key] = bool(rootgrp["Config"][key + "_" + sub_key][...])176 elif(self.types[sub_key] == str):177 self[key][sub_key] = rootgrp["Config"][key + "_" + sub_key][0]178 else:179 self[key][sub_key] = rootgrp["Config"][key + "_" + sub_key][...].item()180 except IndexError:181 print("ERROR: Cannot find {0:s} in Config file.".format(key + "/" + sub_key))182 print("INFO: Using default value.")183 if(filename is not None):184 rootgrp.close()185 def autosave(self):186 config_file = os.path.join(os.path.expanduser("~"), ".ECRad_GUI_Default.nc")187 self.to_netcdf(filename=config_file)188if(__name__ == "__main__"):189 newConf = ECRadConfig(noLoad=True)190 newConf.load( filename="/mnt/c/Users/Severin/ECRad_regression/AUGX3/ECRad_32934_EXT_ed1.nc")...
test_winreg.py
Source:test_winreg.py
1# Test the windows specific win32reg module.2# Only win32reg functions not hit here: FlushKey, LoadKey and SaveKey3from _winreg import *4import os, sys5from test_support import verify, have_unicode6test_key_name = "SOFTWARE\\Python Registry Test Key - Delete Me"7test_data = [8 ("Int Value", 45, REG_DWORD),9 ("String Val", "A string value", REG_SZ,),10 ("StringExpand", "The path is %path%", REG_EXPAND_SZ),11 ("Multi-string", ["Lots", "of", "string", "values"], REG_MULTI_SZ),12 ("Raw Data", ("binary"+chr(0)+"data"), REG_BINARY),13]14if have_unicode:15 test_data+=[16 (unicode("Unicode Val"), unicode("A Unicode value"), REG_SZ,),17 ("UnicodeExpand", unicode("The path is %path%"), REG_EXPAND_SZ),18 ("Multi-unicode", [unicode("Lots"), unicode("of"), unicode("unicode"), unicode("values")], REG_MULTI_SZ),19 ("Multi-mixed", [unicode("Unicode"), unicode("and"), "string", "values"],REG_MULTI_SZ),20 ]21def WriteTestData(root_key):22 # Set the default value for this key.23 SetValue(root_key, test_key_name, REG_SZ, "Default value")24 key = CreateKey(root_key, test_key_name)25 # Create a sub-key26 sub_key = CreateKey(key, "sub_key")27 # Give the sub-key some named values28 for value_name, value_data, value_type in test_data:29 SetValueEx(sub_key, value_name, 0, value_type, value_data)30 # Check we wrote as many items as we thought.31 nkeys, nvalues, since_mod = QueryInfoKey(key)32 verify(nkeys==1, "Not the correct number of sub keys")33 verify(nvalues==1, "Not the correct number of values")34 nkeys, nvalues, since_mod = QueryInfoKey(sub_key)35 verify(nkeys==0, "Not the correct number of sub keys")36 verify(nvalues==len(test_data), "Not the correct number of values")37 # Close this key this way...38 # (but before we do, copy the key as an integer - this allows39 # us to test that the key really gets closed).40 int_sub_key = int(sub_key)41 CloseKey(sub_key)42 try:43 QueryInfoKey(int_sub_key)44 raise RuntimeError, "It appears the CloseKey() function does not close the actual key!"45 except EnvironmentError:46 pass47 # ... and close that key that way :-)48 int_key = int(key)49 key.Close()50 try:51 QueryInfoKey(int_key)52 raise RuntimeError, "It appears the key.Close() function does not close the actual key!"53 except EnvironmentError:54 pass55def ReadTestData(root_key):56 # Check we can get default value for this key.57 val = QueryValue(root_key, test_key_name)58 verify(val=="Default value", "Registry didn't give back the correct value")59 key = OpenKey(root_key, test_key_name)60 # Read the sub-keys61 sub_key = OpenKey(key, "sub_key")62 # Check I can enumerate over the values.63 index = 064 while 1:65 try:66 data = EnumValue(sub_key, index)67 except EnvironmentError:68 break69 verify(data in test_data, "Didn't read back the correct test data")70 index = index + 171 verify(index==len(test_data), "Didn't read the correct number of items")72 # Check I can directly access each item73 for value_name, value_data, value_type in test_data:74 read_val, read_typ = QueryValueEx(sub_key, value_name)75 verify(read_val==value_data and read_typ == value_type, \76 "Could not directly read the value" )77 sub_key.Close()78 # Enumerate our main key.79 read_val = EnumKey(key, 0)80 verify(read_val == "sub_key", "Read subkey value wrong")81 try:82 EnumKey(key, 1)83 verify(0, "Was able to get a second key when I only have one!")84 except EnvironmentError:85 pass86 key.Close()87def DeleteTestData(root_key):88 key = OpenKey(root_key, test_key_name, 0, KEY_ALL_ACCESS)89 sub_key = OpenKey(key, "sub_key", 0, KEY_ALL_ACCESS)90 # It is not necessary to delete the values before deleting91 # the key (although subkeys must not exist). We delete them92 # manually just to prove we can :-)93 for value_name, value_data, value_type in test_data:94 DeleteValue(sub_key, value_name)95 nkeys, nvalues, since_mod = QueryInfoKey(sub_key)96 verify(nkeys==0 and nvalues==0, "subkey not empty before delete")97 sub_key.Close()98 DeleteKey(key, "sub_key")99 try:100 # Shouldnt be able to delete it twice!101 DeleteKey(key, "sub_key")102 verify(0, "Deleting the key twice succeeded")103 except EnvironmentError:104 pass105 key.Close()106 DeleteKey(root_key, test_key_name)107 # Opening should now fail!108 try:109 key = OpenKey(root_key, test_key_name)110 verify(0, "Could open the non-existent key")111 except WindowsError: # Use this error name this time112 pass113def TestAll(root_key):114 WriteTestData(root_key)115 ReadTestData(root_key)116 DeleteTestData(root_key)117# Test on my local machine.118TestAll(HKEY_CURRENT_USER)119print "Local registry tests worked"120try:121 remote_name = sys.argv[sys.argv.index("--remote")+1]122except (IndexError, ValueError):123 remote_name = None124if remote_name is not None:125 try:126 remote_key = ConnectRegistry(remote_name, HKEY_CURRENT_USER)127 except EnvironmentError, exc:128 print "Could not connect to the remote machine -", exc.strerror129 remote_key = None130 if remote_key is not None:131 TestAll(remote_key)132 print "Remote registry tests worked"133else:134 print "Remote registry calls can be tested using",...
Automation Testing Tutorials
Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
LambdaTest Learning Hubs:
- JUnit Tutorial
- TestNG Tutorial
- Webdriver Tutorial
- WebDriverIO Tutorial
- Protractor Tutorial
- Selenium 4 Tutorial
- Jenkins Tutorial
- NUnit Tutorial
- Jest Tutorial
- Playwright Tutorial
- Cypress Tutorial
- PyTest Tutorial
YouTube
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Run hypothesis automation tests on LambdaTest cloud grid
Perform automation testing on 3000+ real desktop and mobile devices online.
Try LambdaTest Now !!
Get 100 minutes of automation test minutes FREE!!
Was this article helpful?
