import sys
from pydiana.shell.logging_theme import logging,coloredlogs
logging.config.dictConfig({'disable_existing_loggers': True,'version': 1})
from pydiana import ROOT,Diana
from pydiana.tools.math_tools import sortmultivec
import pandas as pd
from pydiana.tools.globals import qsampleinfoeventtype
logger = logging.getLogger(__name__)
coloredlogs.install(level='WARNING', logger=logger)
#XXX: Disable multiprocessing because it breaks I/O
if ROOT.IsImplicitMTEnabled():
ROOT.DisableImplicitMT()
special_characters={'@' :'__AT__',
'.' :'__PO__',
'::':'__4PO__',
'[' :'__SQOP__',
']' :'__SQCL__',
' ' :'__SPC__',
'?' :'__QUM__',
}
operation_characters={'/':'__DIV__',
'*':'__MUL__',
'-':'__MIN__',
'+':'__SUM__',
}
[docs]
def name2var(diana_name:str,conv_operations:bool=True,conv_spaces:bool=True)->str:
"""
Function that converts variables names from Diana ROOT files to acceptable variables for RDataFrames.
Since "@", "." and other characters present in Diana ROOT files cannot be used in RDataFrames column names, they are redefined.
Parameters
----------
diana_name (str): name of the variable in Diana ROOT files
Returns
-------
frame_name (str): name of the same variable (column) in the RDataFrame
"""
frame_name = diana_name
for c in list(special_characters.keys()):
if c== " " and not conv_spaces:
continue
char=special_characters[c]
frame_name = char.join(frame_name.split(c))
if conv_operations:
for c in list(operation_characters.keys()):
char=operation_characters[c]
frame_name = char.join(frame_name.split(c))
return frame_name
################################################################################################################
[docs]
def var2name(frame_name:str,conv_operations:bool=True)->str:
"""
Function that converts back the variables names used for RDataFrames in the real names from Diana ROOT files.
Since "@", "." and other characters present in Diana ROOT files cannot be used in RDataFrames column names, they are redefined.
Parameters
----------
frame_name (str): name of the variable (column) in the RDataFrame
Returns
-------
diana_name (str): real name of the same variable from Diana ROOT files
"""
diana_name=frame_name
for c in list(special_characters.keys()):
char=special_characters[c]
diana_name = c.join(diana_name.split(char))
if conv_operations:
for c in list(operation_characters.keys()):
char=operation_characters[c]
diana_name = c.join(diana_name.split(char))
return diana_name
################################################################################################################
[docs]
def clean_colnames(colnames):
"""
Function (internal) that returns only the useful part of a variable name
Parameters
----------
colnames (list): list of str, input colnames
Returns
-------
cnames2 (list): list of str, output colnames
"""
cnames2=[]
#take the qtree_Preprocess and similar from the cnames
for c in colnames:
c = str(c)
if 'qtree' in c:
c = ".".join(c.split(".")[1:])
append = True
if '::' in c:
append=False
if 'TObject' == c.split('.')[-1]:
append=False
if 'fBits' == c.split('.')[-1]:
append=False
if 'fUniqueID' == c.split('.')[-1]:
append=False
if 'fData' == c.split('.')[-1]:
append=False
if 'fSize' == c.split('.')[-1]:
append=False
if 'fStride' == c.split('.')[-1]:
append=False
if '.' == c[-1]:
append=False
if append:
cnames2.append(c)
return cnames2
################################################################################################################
[docs]
def make_expr(string:str,rdataframe,returnval='None',init_classes:bool=False,conv_spaces:bool=True):
"""
Function made to convert a string in a usable expression for the rdataframe
Parameters
----------
string : expression to evaluate
rdataframe: RDataFrame instance to be used to check column names
returnval: modes of the return val, it can either return a None value (if set to 'None') or just the converted string (if set to 'best')
init_classes: boolean for enabling the initialization of the classes used in the expression
conv_spaces (bool) : boolean for toggling on and of the conversion of the spaces in the expression
Returns
-------
string with the expression or None based on returnval
"""
#Search if there are special characters to be converted
found_special_char=False
for c in list(special_characters.keys()):
if c==" " and not conv_spaces:
continue
if c in string:
found_special_char = True
#If not return the string
if not found_special_char:
return string
# Else convert all special characters in the string
expr2eval = name2var(string,conv_operations=False,conv_spaces=conv_spaces)
#Find all column names present in expression
colnames=clean_colnames(rdataframe.GetColumnNames())
candidates = findVar(expr2eval,rdataframe)
splitted_expr=[expr2eval]
for c in operation_characters:
splitting= []
for ll in splitted_expr:
cc = ll.split(c)
splitting +=cc
splitted_expr=splitting
for ll in splitted_expr:
if ll!="":
cc= findVar(ll,rdataframe)
if cc is not None:
candidates +=cc
if not conv_spaces:
for ll in expr2eval.split(" "):
if ll!="":
cc= findVar(ll,rdataframe)
if cc is not None:
candidates +=cc
candidates = list(set(candidates))
#if there are characters to be converted but no matches were found in the column names
#then something went wrong
#XXX find a better solution
if candidates is None:
if returnval=='None':
return None
elif returnval=='best':
return expr2eval
#Initialize classes so that root knows what to do when applying the cut
for cc in candidates:
if init_classes:
coltype=rdataframe.GetColumnType(cc)
if "Q" == coltype[0]:
ROOT.std.vector(f"{coltype}")(1)
string = f"{cc}".join(string.split(var2name(cc)))
return string
################################################################################################################
[docs]
def eval_alias(alias:str,expr:str,rdataframe,redefine:bool=False,verbose:bool=False):
"""
Function for adding an alias to the dataframe.
Parameters
----------
alias : string with the name of the alias to use
expr : string with the expression used to define the alias
rdataframe : RDataFrame instance to be aliased
redefine: boolean for enabling redefining of previously existing aliases
verbose: boolean for enabling message printing
Returns
-------
returns RDataFrame with implemented aliases. If aliasing was not successful returns RDataFrame
"""
alias = name2var(alias)
alias_expr = make_expr(string=expr,
rdataframe=rdataframe,
returnval='None',
init_classes=True)
if alias_expr is None:
if verbose:
logger.warning(f"Could not evaluate {expr}, skipping")
return rdataframe
if not rdataframe.HasColumn(alias):
try:
rdataframe = rdataframe.Define(alias,alias_expr)
except:
logger.warning(f"Could not define {alias} = {expr}")
if not rdataframe.HasColumn(alias):
logger.error(f"Failed to define {alias} = {expr}")
elif redefine:
rdataframe = rdataframe.Redefine(alias,alias_expr)
elif verbose:
logger.warning(f"A column named {alias} is already present. Skipping.")
return rdataframe
################################################################################################################
[docs]
def eval_aliases(rdataframe,aliases:dict=None,redefine:bool=False,verbose:bool=False):
"""
Function for defining multiple aliases from dictionary.
Parameters
----------
rdataframe : RDataFrame instance to be aliased
aliases : dictionary of the aliases used in the notebook. The keys are the aliases and the values are the expression to alias (using the default c++ aliases is allowed)
redefine: boolean for enabling redefining of previously existing aliases
verbose: boolean for enabling message printing
Returns
-------
RDataFrame with implemented aliases. If no aliases are provided then returns same RDataFrame.
"""
if aliases is not None:
for var in aliases:
rdataframe = eval_alias(alias=var,
expr=aliases[var],
rdataframe=rdataframe,
redefine=redefine,
verbose=verbose)
return rdataframe
################################################################################################################
#########################################################
[docs]
def open_rdataframe(datachain,local_aliases:dict=None,redefine:bool=False,verbose:bool=False):
"""
Function that opens the RDataFrame starting from QChain
Parameters
----------
datachain: QChain with data
local_aliases: dictionary of the aliases used in the notebook. The keys are the aliases and the values are the expression to alias (using the default c++ aliases is allowed)
redefine: boolean for enabling redefining of previously existing aliases
verbose: boolean for enabling message printing
Returns
-------
aliasDF: Opened RDataFrame (with implemented aliases from Diana)
"""
dataframe = ROOT.RDataFrame(datachain)
colnames ={}
cnames = clean_colnames(dataframe.GetColumnNames())
for i in range(len(cnames)):
name =cnames[i]
toprint=True
if "::" in name:
toprint=False
if ".fBits" in name:
toprint=False
if ".TObject" in name:
toprint=False
if ".fUniqueID" in name:
toprint=False
if ".fData" in name:
toprint=False
if ".fSize" in name:
toprint=False
if '.' == name[-1]:
toprint=False
if toprint:
varname = name2var(name)
colnames[name] =varname
dataframe=dataframe.Alias(varname,name)
aliases = extract_aliases(datachain=datachain)
cnames2=clean_colnames(dataframe.GetColumnNames())
aliasDF=dataframe
alias2eval={}
for alias in aliases:
path = aliases[alias]
if path not in cnames2:
if "DAQ@Header." in path and ("QHeader" not in path):
#XXX quick fix for DAQ header. Not sure what caused this.
path = "DAQ@Header.QHeader.".join(path.split("DAQ@Header."))
alias2eval[alias]=path
else:
aliasDF=aliasDF.Alias(alias,path)
aliasDF = eval_aliases(rdataframe=aliasDF,
aliases = alias2eval,
redefine=redefine,
verbose=verbose,
)
aliasDF = eval_aliases(rdataframe=aliasDF,
aliases = local_aliases,
redefine=redefine,
verbose=verbose,
)
return aliasDF
################################################################################################################
[docs]
def findVar(expr:str,rdataframe):
"""
Function for finding variable in the rdataframe
Parameters
----------
var: stirng with the variable to find
rdataframe: RDataFrame in which to search
Returns
-------
list with the variables in the expression
"""
#Find all column names present in expression
colnames=clean_colnames(rdataframe.GetColumnNames())
candidates = {}
for c in colnames:
c= str(c)
if c in expr:
candidates[c]=len(c)
if len(candidates)==0:
return None
#Take the longest match with the variables
candidates = sortmultivec(list(candidates.keys()),list(candidates.values()),lead=1)
indexes2remove = []
for i,cc in enumerate(candidates[0]):
for j,cc2 in enumerate(candidates[0]):
if j!=i and cc in cc2:
indexes2remove.append(i)
candidates=[cc for i,cc in enumerate(candidates[0]) if i not in indexes2remove]
return candidates
################################################################################################################
[docs]
def rfilter(rdataframe,cuts:str=None,name='MyCut',verbose:bool=False):
"""
Function for evaluating a filtering expression on an rdataframe
Parameters
----------
cuts: string with the cuts expression in root style
rdataframe: RDataFrame to filter
name (str): name of the cut using when referrency to it inside ROOT
verbose (bool): print information while applying cuts
Returns
-------
filtered rdataframe if successul otherwise returns same dataframe
"""
if cuts is not None and cuts !='':
myexpr = make_expr(string =cuts, rdataframe=rdataframe)
if myexpr is not None:
result = rdataframe.Filter(myexpr,name)
if verbose:
result.Report().GetPtr().Print()
sys.stdout.flush()
return result
elif verbose:
logger.error(f"Could not build expression: {cuts} -> {myexpr}")
elif verbose:
logger.warning("Cut string is empty")
return rdataframe
################################################################################################################
[docs]
def convert2pandas(df,colnames:list,sep_field:str="",eventtype2string:bool=True):
"""
Function that checks if the inserted dataframe is a Pandas DataFrame and returns it.
If not, one implemented option is that df is a RDataFrame: in this case this function uses the provided array colnames
(considering only the not None elements in it) to build a Pandas DataFrame (needed in our analisys setup).
For graph_tools, can be also needed to save the column from which there is the separation of plots, so along with the
requested columns, also the sep_field column (if not exists return an error) will be saved in the output Pandas dataframe.
Upgradable with other types of dataframe in the future.
N.B.: this function automatically create a map frame_name --> diana_name;
the inserted colnames list can contain be provided in frame_name, diana_name or even both format mixed inside the list.
Parameters
----------
df: inserted dataframe to be converted (if needed) in a Pandas DataFrame
colnames (list): list of the column to be considered (only used if df is not a Pandas DataFrame)
sep_field (str): name of the column to be considered as sep_field from graph_tools functions
(only used if df is not a Pandas DataFrame)
eventtype2string (bool): if True it looks for the column EventType and traslates it to a string
using as blueprint the enumerator in QSampleInfo
Returns
-------
res_df: resulting Pandas DataFrame (equal to df if df is a Pandas DataFrame)
"""
if isinstance(df,pd.core.frame.DataFrame):
res_df = df
#elif type(df) == ROOT.RDataFrame:
else: #### XXX: a me non ha funzionato questo type, non so perché
colnames_dict = {name2var(c):var2name(c) for c in colnames if c is not None} ### form {frame_name : diana_name}
### define sep field column
if sep_field is not None and sep_field != "":
colnames_dict.update({name2var(sep_field):var2name(sep_field)})
res_df = df.AsNumpy(list(colnames_dict.keys()))
res_df = pd.DataFrame().from_dict(res_df)
res_df = res_df.rename(columns=colnames_dict)
### dovrebbe non servire ora
#if "Run" in list(colnames_dict.keys()): res_df["Run"] = res_df['Run'].astype(int).astype(str)
if "EventType" in list(res_df.columns) and eventtype2string:
res_df = res_df.assign(EventType = lambda x: qsampleinfoeventtype[x['EventType'][0]])
return res_df
################################################################################################################
[docs]
def read_varsRDataFrame(variables,datafile,cuts:str =None):
"""
Function for reading a tuple for a diana root file passing through an RDataFrame (retrocompatibility)
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 : QChain instance of the file/files to read
cuts : String containing the cuts to select the wanted variables
Returns
-------
Pandas dataframe with all the requested variables, or None if no events pass the cuts.
"""
#open rdataframe and insert all vars2read as alias
rdataframe = open_rdataframe(datachain=datafile,local_aliases=variables)
#filted rdataframe based on cuts
rdataframe = rfilter(rdataframe=rdataframe,cuts=cuts,name='OpeningCut')
#since variables at this point are all aliased inside rdataframe get them
#directly through the alias. The aliases need to be translated in acceptable
#names
vars2get = {v: name2var(v) for v in (variables.keys())}
df_total = rdataframe.AsNumpy(list(vars2get.values()))
#Save to a pandas dataframe and rename the columns with the aliases defined
#in the notebook
df_total = pd.DataFrame().from_dict(df_total)
df_total = df_total.rename(columns={vars2get[v]:v for v in vars2get})
del rdataframe
return df_total