Source code for pydiana.tools.pydianaIO

from pydiana import ROOT,Diana,diana_install,diana_ext_install
import subprocess
import numpy as np
from tqdm import tqdm,trange
import copy
import pandas as pd
import warnings
import array
from pydiana.tools import pyRDataFrame as pyRD
from pydiana.tools.globals import qsampleinfoeventtype

ROOT.gInterpreter.ProcessLine(f"#include \"{diana_install}/pkg/base/QFileList.hh\"")
ROOT.gInterpreter.ProcessLine(f"#include \"{diana_install}/pkg/dianaevent/QAliases.hh\"")

#################
#               #
#   Data Files  #
#               #
#################

[docs] def parse_listfile(inputfile:str): """ Function for parsing a .list file (not a nested reading). Parameters ---------- inputfile (str): path to input .list file Returns ------- dict containing the following structure: {'datapath':<datapath written in file>,'listfiles':<list of .list files>,'rootfiles':<list of .root files>} """ datafile = open(inputfile, "r") filecontent={'datapath':"",'listfiles':[],'rootfiles':[]} lines = datafile.readlines() readfiles=False for ll in lines: ll = ll.split("\n")[0]#remove newline ll = ll.split("#")[0] #remove commented parts of lines if "DATAPATH" in ll: filecontent['datapath']=ll.split("DATAPATH")[-1] if filecontent['datapath'].endswith("/"): filecontent['datapath']=filecontent['datapath'][:-1] filecontent['datapath']="".join(filecontent['datapath'].split(" ")) elif "START"==ll: readfiles=True elif "END" ==ll: readfiles=False elif ll.endswith('.root') and readfiles: filecontent['rootfiles'].append("".join(ll.split(" "))) elif ll.endswith('.list') and readfiles: filecontent['listfiles'].append("".join(ll.split(" "))) return filecontent
[docs] def find_all_files(datafile:str): """ Function for finding all files to include in the analysis. Parameters ---------- datafile : String containing the path to the file Returns ------- list with all the files to use """ filelist=[] if datafile.endswith(".list"): listfile = ROOT.QFileList() listfile.Read(f"{datafile}") #filelist = np.asarray(listfile.GetListVec()) filelist = listfile.GetListVec() filelist = [''.join(ff) for ff in filelist] elif datafile.endswith(".root"): filelist = [datafile] return filelist
[docs] def extract_list(a): return [a[i] for i in range(a.size())]
[docs] def openfile(datafile:str): """ Function for opening Diana data files. If a .root file is provided that is the only file provided, if a .list file all the .root files linked (directly or indirectly) from the .list files will be included. Parameters ---------- datafile : String containing the path to the file Returns ------- datafiles : QChain instance used to read the file (analogous to a TChain or TTree from root, can be thought of as the pointer to the file) """ filelist = find_all_files(datafile=datafile) datafiles = ROOT.QChain("qtree") _=[datafiles.Add(str(i)) for i in filelist] return datafiles
[docs] def get_subtree(datafiles:ROOT.QChain,cuts : str): """ Method to get elements of tree that respect cuts Parameters ---------- datafile : QChain instance with cuts to be applied to cuts : string with cuts to be imposed Returns ------- QChain instance with elements that passed cut selection """ return datafiles.CopyTree(cuts)
[docs] def alias_reader(alias_files : list = None): """ Function for importing the aliases. If no file is specified in automatically reads the `cfg/aliases.txt` files present in DIANA_INSTALL and DIANA_EXT_INSTALL. Parameters ---------- alias_files : array-like with all the paths to the alias files to be used. If only one file to be used it can be inserted as a simple string. Returns ------- dictionary with keys as the aliases and the elements as the path to the variable """ if alias_files is None: alias_files = [diana_install+"/cfg/aliases.txt"] if diana_ext_install is not None: alias_files.append(diana_ext_install+"/cfg/aliases.txt") elif type(alias_files) == type(''): alias_files = [alias_files] alias_reader = Diana.QAliases() for af in alias_files: alias_reader.FillFromFile(af) alias_map = alias_reader.Get() ali_tab = np.asarray(alias_map) aliases = {al[0]:al[1].fLabel.GetStringLabel()+'.'+al[1].fPath for al in ali_tab} return aliases
default_aliases = alias_reader()
[docs] def alias_compare(variable:str,list1,dict_alias=default_aliases): """ Function for converting a variable alias to a complete variable path Parameters ---------- variable : string with the variable to convert list1 : array-like of lists with the datafile aliases dict_alias : dictionary of the aliases (keys = aliases, values = variable path) Returns ------- Complete path to variable or '' if not found. """ filevar_found=False aliasvar_found = False if variable in list1: # found variable as alias of datafile filevar_found = True if variable in dict_alias.keys(): # found variable as alias in aliases.txt file aliasvar_found = True path2var = '' if filevar_found and aliasvar_found: path2var = dict_alias[variable] return path2var
[docs] def var_autocomplete(variable:str,dict_content:dict,dict_alias=default_aliases,verbose=False): """ Function for searching the best possible match of inserted variable. Parameters ---------- variable : string with the variable to search datafile : QChain instance of the opened datafile dict_alias : dictionary of the aliases (keys = aliases, values = variable path) Returns ------- String with best possible match of variable """ variables = dict_content myvar='' if '@' in variable: tmpvar=variable if variable[-1]=='.': tmpvar+='fValue' elif variable[-7:]!=".f": tmpvar+=".fValue" condition=True ii=-1 while(condition): ii+=1 if ii>=(len(list(variables.keys()))): condition=False break else: i=list(variables.keys())[ii] for j in variables[i]: if variable in j: myvar = j if (variable ==j) or (j==tmpvar): condition=False break else: myvar=alias_compare(variable,variables['Aliases'],dict_alias = dict_alias) if myvar=='':#in case reduced files are used if variable in variables['Leaves']: myvar = variable if myvar=='': warnings.warn(f"Cannot find {variable}, check spelling and that variable is present in file. Use `get_all_content` function to inspect the available variables.") elif myvar!='' and myvar[-1]!='.' and verbose: print(f'Variable found as {variable} -> {myvar}') if myvar[-1] =='.': if myvar+'fValue' not in variables['Leaves']: suggested = [vv for vv in variables['Leaves'] if var in vv and vv[-1]!= '.'] warnings.warn(f"Best guess of variable '{myvar}.fValue' not found. Suggested variables are: {suggested}") myvar+="fValue" if verbose: print(f'Variable found as {myvar}') return myvar
[docs] def apply_cuts(cuts,datafile): """ Function for getting cuts on datafile. Readapted from ROOT's TTreePlayer->CopyTree. Parameters ---------- cuts : string with cuts datafile : QChain instance on which to evaluate the cuts Returns ------- Array with indexes of the events that pass the cuts """ selected = [] select = ROOT.TTreeFormula("Selection",cuts,datafile) tnumber=-1 for entry in range(datafile.GetEntries()): entrynumber = datafile.GetEntryNumber(entry) if(entrynumber<0): break localentry = datafile.LoadTree(entrynumber) if(localentry<0): break if(tnumber != datafile.GetTreeNumber()): tnumber = datafile.GetTreeNumber() select.UpdateFormulaLeaves() ndata = select.GetNdata() keep = False current=0 while(current<ndata and keep==False): keep = keep or (select.EvalInstance(current)!=0) current+=1 if keep == False: continue else: selected.append(entry) del select return np.array(selected)
[docs] def check_legacy(variable:str,legacy:bool = False): """ Function for checking if the variable needs to be read with the legacy method or not Parameters ---------- variable : string with the name of the variable legacy : bool that overides check Returns ------- True if it's a legacy variable, False if not """ passed = False if legacy: passed = True return passed check_characters = ['(',')','*','+','-',"/",'|','&','='] for cc in check_characters: if cc in variable: passed = True if passed: break return passed
[docs] def read_var(variable:str,datafile,cuts:str ='',dtypes=(float,),dict_alias=default_aliases,legacy=False,verbose=False): """ Function for reading a tuple for a diana root file. Parameters ---------- variable : string with the variable to read datafile : Either QChain instance of the file/files to read cuts : String containing the cuts to select the wanted variables dtypes : array-like of types used to convert the read values (the variables are read as floats by default and then are converted from the first to the last type indicated) dict_alias : dictionary of the aliases (keys = aliases, values = variable path) legacy : boolean indicating if one must use old_read_var function Returns ------- array with the requested variable, or None if variable not found or cuts not passed """ dict_alias = pyRD.extract_aliases(datachain=datafile) dict_content = get_all_content(datafile,True) vv = var_autocomplete(variable,dict_content,dict_alias,verbose) if check_legacy(vv,legacy): if "@" in variable: res = old_read_var(vv,datafile,cuts) else: res = old_read_var(variable,datafile,cuts) else: res = datafile.AsMatrix([vv]).transpose()[0] if cuts != '': selected = apply_cuts(cuts,datafile) if len(selected) ==0: return None res = res[selected] for i in dtypes: res = np.array(res,dtype=i) return res
[docs] def old_read_var(variable:str,datafile:ROOT.QChain,cuts:str ='',dtypes=(float,)): """ Function for reading a variable from the root file Parameters ---------- variable : string containing the name of the variable datafile : QChain instance (given from the `openfile` function) cuts : String containing the cuts to select the wanted variables dtypes : array-like of types used to convert the read values (the variables are read as floats by default and then are converted from the first to the last type indicated) Returns ------- tuple of the values of that variable """ Nentries = datafile.Draw(variable,cuts,"goff") if Nentries==0: warnings.warn("Empty buffer") return None res= np.frombuffer(datafile.GetV1(),count=Nentries) return copy.deepcopy(res)
[docs] def read_vars(variables,datafile,cuts:str ='',dict_alias=default_aliases,legacy=False,verbose=True,useTQDM=False): """ Function for reading a tuple for a diana root file. Parameters ---------- variables : array-like of strings containing the name of the variables to fetch. If dictionary, the values must be the names of the variables to search and the keys the name to be shown in the resulting dataframe datafile : Either QChain instance of the file/files to read cuts : String containing the cuts to select the wanted variables dict_alias : dictionary of the aliases (keys = aliases, values = variable path) legacy : boolean indicating if one must use old_read_var function Returns ------- Pandas dataframe with all the requested variables, or None if no events pass the cuts. """ dict_alias = pyRD.extract_aliases(datachain=datafile) dict_content = get_all_content(datafile,True) keys=[] orig_vars=[] if isinstance(variables,dict): keys = list(variables.keys())#aliases to use in notebook variables = list(variables.values())#path to variable else: keys = variables keys = np.array(keys) orig_vars = np.array(variables) variables = np.array([var_autocomplete(vv,dict_content,dict_alias,verbose) for vv in variables])#autocomplete name of the variables #separate variables that can be read in python style(new_vars) #wrt the ones to be read with c++(legacy_vars) legacy_vars = [i for i,vv in enumerate(variables) if check_legacy(vv,legacy)] new_vars =[i for i,vv in enumerate(variables) if not check_legacy(vv,legacy)] leg_vars={} if float(ROOT.__version__.split("/")[0])>=6.24: #Call only RDataFrame if root version is >=6.24 legacy_vars = [] new_vars =[i for i,vv in enumerate(variables)] for lv in tqdm(legacy_vars,disable=(not useTQDM),desc="Reading Legacy Vars"): if "@" in orig_vars[lv]: res = old_read_var(variables[lv],datafile,cuts) else: res = old_read_var(orig_vars[lv],datafile,cuts) if res is not None: leg_vars[keys[lv]] = res df = [] if len(new_vars)>0: if float(ROOT.__version__.split("/")[0])>=6.24: #Call RDataFrame if root version >=6.24 df.append(pyRD.read_varsRDataFrame(variables={keys[i]: variables[i] for i in new_vars}, datafile=datafile, cuts=cuts, )) else: selected =apply_cuts(cuts,datafile) if len(selected) ==0: return None nvv = datafile.AsMatrix(list(variables[new_vars]),return_labels=True) nvv = nvv[0][selected,:] df.append(pd.DataFrame(data=nvv,columns=list(keys[new_vars]))) if len(list(leg_vars.keys()))>0: df.append(pd.DataFrame(data=np.array(list(leg_vars.values())).transpose(),columns=list(leg_vars.keys()))) if len(df)>0: df = pd.concat(df,axis=1) else: df = df[0] return df
##################################### # # # Reading Average Pulse Files # # # #####################################
[docs] def get_average_pulse(ap_input : str,dataset:int,chan:int,owner:str='AveragePulses',extralabel : str = None,verbose:bool=False): """ Function for importing the average pulse. Parameters ---------- ap_input : String with the path to the average pulse chan : Integer indicating the channel number extralabel : string indicating the extralabel used for the average pulse Returns ------- vector with the average pulse in ADC units """ dm = Diana.QGlobalDataManager() aph=ROOT.QAveragePulseHandle(chan) aph.SetDataset(dataset) apowner = owner if extralabel is not None: apowner+='_'+extralabel err = dm.Get(apowner,aph,ap_input) if not aph.IsValid(): raise Exception(f"Average Pulse Handle is not valid {aph.GetError().ToString()}") return None if verbose: print("Average Pulse Number of Events: ",aph.Get().fNumEvents) print("Average Pulse Number of Runs: ",aph.Get().fSourceRuns.size()) ap = Diana.QVector(aph.Get()) ap=np.array(array.array("d",ap.GetArrayVector())) return ap
######################################### # # # Reading Average Noise PSD Files # # # #########################################
[docs] def get_average_noise(an_input : str,dataset:int,chan:int,owner:str='NoiseAvgPowerSpectrum',extralabel : str = None,verbose:bool=False): """ Function for importing the average noise power spectrum. Parameters ---------- an_input : String with the path to the average noise power spectrum chan : Integer indicating the channel number extralabel : string indicating the extralabel used for the average noise power spectrum Returns ------- array with the values of the noise power spectrum in ADC units """ dm = Diana.QGlobalDataManager() anh=ROOT.QAverageNoiseHandle(chan) anh.SetDataset(dataset) anowner = owner if extralabel is not None: anowner+='_'+extralabel dm.Get(anowner,anh,an_input) if not anh.IsValid(): raise Exception(f"Average Noise Handle is not valid {anh.GetError().ToString()}") return None if verbose: print("Average Noise Number of Events: ",anh.Get().fNumEvents) print("Average Noise Number of Runs: ",anh.Get().fSourceRuns.size()) an = Diana.QVector(anh.Get()) an=np.array(array.array("d",an.GetArrayVector())) return an
######################################### # # # Reading Noise Covariance # # # #########################################
[docs] def get_noise_covariance(covinput : str,dataset:int,owner:str='NoiseCrossPowerSpectrum',extralabel : str = None,verbose:bool=False): """ Function for importing the noise covariance. Parameters ---------- covinput : String with the path to the file containing the noise covariance matrix (usually same as average noise power spectrum) dataset : int with the dataset owner: owner of the noise covariance matrix extralabel : string indicating the extralabel used for the covariance matrix Returns ------- Instance of QChannelCovariance """ dm = Diana.QGlobalDataManager() hcov=Diana.GlobalHandle(ROOT.QChannelCovariance)('Covariance') hcov.SetDataset(dataset) if extralabel is not None: owner+='_'+extralabel dm.Get(owner,hcov,covinput) NoiseCovariance=None if not hcov.IsValid(): raise Exception(f"Noise Covariance Handle is not valid {hcov.GetError().ToString()}") else: NoiseCovariance = ROOT.QChannelCovariance(hcov.Get()) if verbose: print(f"Getting noise covariance status: {hcov.GetError().ToString()}") print(f"Mean Correlation Matrix:") dumpstr = ROOT.std.ostringstream() NoiseCovariance.DumpMeanCorrelationMatrix(dumpstr) cov=dumpstr.str() del dumpstr cov = cov.split("\n") cov = [a.split("\t") for a in cov] cov = [[a.split(" ") for a in aa if len(a)>0] for aa in cov][:-1] cov = [[[b for b in a if len(b)>0] for a in aa if len(a)>0] for aa in cov] channels = np.concatenate([[int(ch[0])] for ch in cov[0]],dtype=int) covdict = {int(a[0][0]): np.array(a[1],dtype=float) for a in cov[1:]} cov = pd.DataFrame() cov['channels'] = list(covdict.keys()) for ch in channels: cov[ch]=covdict[ch] cov.set_index('channels',drop=True,inplace=True) display(cov) del hcov del dm return NoiseCovariance
############################################### # # # Getting Channel Run Data and Run Data # # # ###############################################
[docs] def get_run_data(filename : str , run : int, rundataowner:str='DAQ'): """ Method for getting rundata. Parameters ---------- filename : String with path to file to be read run: Int with run number. rundataowner: string with run data owner Returns ------- Run data """ if '.list'==filename[-5:]: files =find_all_files(filename) files = [ff for ff in files if 'p001.root' in ff] filename = files[0] if '.root' not in filename: raise Exception("File must be a .root file produced from diana") return None dm = Diana.QGlobalDataManager() rnhandle=ROOT.QRunDataHandle(run) dm.Get(rundataowner,rnhandle,filename) if not rnhandle.IsValid(): raise Exception(f"Run Data Handle is not valid {rnhandle.GetError().ToString()}") return None rundata=ROOT.QRunData(rnhandle.Get()) return rundata
[docs] def get_durations(listfile:str,verbose:bool=False,printres:bool=False): """ Method for getting the duration of a datafile (taken from rundata). The result is returned in hours. """ files =find_all_files(listfile) files = [ff for ff in files if 'p001.root' in ff] run_number ={ff:int(ff.split('/')[-1].split('_')[1]) for ff in files} rundata = {run_number[ff]:get_run_data(ff,run_number[ff]) for ff in run_number} durations = {rn: rundata[rn].fDuration/3600 for rn in rundata } total_dur = np.sum(list(durations.values())) if printres: if verbose: _=[print(f'\t{rn}: {durations[rn]/24:.2f} d ({durations[rn]:.2f} h)') for rn in durations] print(f'Total Duration: {total_dur/24:.2f} d ({total_dur:.2f} h)') res = total_dur if verbose: res = [total_dur,durations] return res
[docs] def get_channel_run_data(filename : str , ch : int, run : int,rundataowner:str='DAQ'): """ Method for getting the channel rundata Parameters ---------- filename : String with path to file to be read ch: Int with channel number run: Int with run number. Returns ------- Channel run data """ rundata =get_run_data(filename=filename,run=run,rundataowner=rundataowner) chdata =ROOT.QChannelRunData(rundata.GetChannelRunData(ch)) return chdata
######################## # # # Getting Pulses # # # ########################
[docs] def get_custom_class_from_event(datafile:ROOT.QChain, owner:str,label:str,cut=None): """ Utility Function for getting classes from datafile with selection. Parameters ---------- datafile : QChain instance with the datafile owner : string with the owner of the pulse label : string with the label of the pulse cut: either array with indexes of pulses to keep or string with cut Returns ------- 2D array with the results. Returns none if no pulses are found """ if cut is not None: if isinstance(cut,str): cut=apply_cuts(cut,datafile) else: cut = range(datafile.GetEntries()) if len(cut) ==0: raise Exception('No pulses selected') return None if len(cut)>datafile.GetEntries() or max(cut)>datafile.GetEntries() or min(cut)<0: raise Exception('Incompatible indexes given') return None leafname = f"{owner}@{label}" if leafname not in get_all_leaves(datafile): raise Exception(f"Could not fine {leafname} in datafile") return None leaf = datafile.GetLeaf(leafname) leaftype=leaf.GetTypeName() outclass=getattr(ROOT,leaftype)() if datafile.SetBranchAddress(leafname,outclass) != 0 : raise Exception("Could not set branch address") return None result=[] for pulse_idx in tqdm(cut,leave=False,desc=f'Reading {leafname}'): datafile.GetEntry(pulse_idx) result.append(getattr(ROOT,leaftype)(outclass)) datafile.ResetBranchAddresses() return result
def __get_pulses__(datafile:ROOT.QChain,cut=None,owner:str="DAQ",label:str="Pulse"): """ Utility Function for getting pulses from datafile given an array of indexes. Parameters ---------- datafile : QChain instance with the datafile cut: either array with indexes of pulses to keep or string with cut owner : string with the owner of the pulse label : string with the label of the pulse Returns ------- 2D array with the pulses. Returns none if no pulses are found """ mypulses=get_custom_class_from_event(datafile=datafile,owner=owner,label=label,cut=cut) if mypulses is not None: if isinstance(mypulses[0],ROOT.QPulse): mypulses = [np.array(pulse.GetSamples().GetArrayVector()) for pulse in mypulses] mypulses=np.array(mypulses) elif isinstance(mypulses[0],ROOT.QVector): mypulses = [np.array(pulse.GetArrayVector()) for pulse in mypulses] mypulses=np.array(mypulses) return mypulses
[docs] def get_pulses(datafile:ROOT.QChain,cut:str=None,owner:str="DAQ",label:str="Pulse."): """ Function for getting pulses from Diana .root file Parameters ---------- datafile : QChain instance with the datafile cuts : string implementing the cuts owner : string with the owner of the pulse label : string with the label of the pulse Returns ------- 2D array with the pulses. Returns none if no pulses are found """ return __get_pulses__(datafile=datafile, cut=cut, owner=owner, label=label, )
[docs] def get_pulses_with_index(datafile:ROOT.QChain,indexes=None,owner:str="DAQ",label:str="Pulse."): """ Function for getting pulses with cuts and indexes. Parameters ---------- datafile : QChain instance with the datafile indexes : list of pulses to get withing the ones selected by the cuts owner : string with the owner of the pulse label : string with the label of the pulse Returns ------- 2D array with the pulses. Returns none if no pulses are found """ return __get_pulses__(datafile=datafile, cut=indexes, owner=owner, label=label, )
[docs] def get_random_pulses(datafile:ROOT.QChain,num:int, cuts:str="",owner:str="DAQ",label:str="Pulse."): """ Function for getting random pulses from dataset. Parameters ---------- datafile : QChain instance with the datafile cuts : string implementing the cuts num : number of random pulses to be fetched within the ones that pass the cuts owner : string with the owner of the pulse label : string with the label of the pulse Returns ------- 2D array with the pulses. Returns none if no pulses are found """ selected =np.array(apply_cuts(cuts,datafile)) if num>len(selected): print("Invalid number of pulses selected, insert a value smaller than", len(selected)) return None indexes = np.random.choice(selected,num,replace=False) return __get_pulses__(datafile=datafile, cut=indexes, owner=owner, label=label, )
[docs] def get_pulse_info(datafile:ROOT.QChain,cut=None,owner:str='DAQ'): """ Utility Function for getting classes from datafile with selection. Parameters ---------- datafile : QChain instance with the datafile owner : string with the owner of the pulse cut: either array with indexes of pulses to keep or string with cut Returns ------- 2D array with the pulse infos. Returns none if no pulses are found """ return get_custom_class_from_event(datafile=datafile, cut=cut, owner=owner, label="PulseInfo.")
######################### # # # Inspect ROOT File # # # #########################
[docs] def get_all_aliases(datafile:ROOT.QChain): """ Function for getting all the aliases in the .root files Parameters ---------- datafile : QChain instance of the opened file Returns ------- List with all the aliases """ return [i.GetName() for i in datafile.GetListOfAllAliases()]
[docs] def get_all_leaves(datafile:ROOT.QChain): """ Function for getting all the leaves in the .root files Parameters ---------- datafile : QChain instance of the opened file Returns ------- List with all the leaves """ return [i.GetName() for i in datafile.GetListOfAllLeaves()]
[docs] def get_all_branches(datafile:ROOT.QChain): """ Function for getting all the branches in the .root files Parameters ---------- datafile : QChain instance of the opened file Returns ------- List with all the branches """ return [i.GetName() for i in datafile.GetListOfAllBranches()]
[docs] def get_all_content(datafile:ROOT.QChain,verbose:bool=False): """ Function for getting all the info in the .root files Parameters ---------- datafile : QChain instance of the opened file Returns ------- List with all the info """ res={} res['Aliases']=get_all_aliases(datafile) res['Leaves']=get_all_leaves(datafile) if verbose: res['Branches']=get_all_branches(datafile) return res