root_helpers

ROOT required

This module requires PyROOT. See checking ROOT availability.

analysis_helpers.root_helpers

is_root_available()

Return whether PyROOT can be imported.

Source code in src/analysis_helpers/root_helpers.py

def is_root_available():
    """Return whether PyROOT can be imported."""
    return _HAS_ROOT

require_root()

Raise a clear exception if PyROOT is not available.

Source code in src/analysis_helpers/root_helpers.py

def require_root():
    """Raise a clear exception if PyROOT is not available."""
    if not _HAS_ROOT:
        raise ImportError(
            "PyROOT is not available. Install and configure ROOT to use this function."
        )

LoadCompiledLibraries(libraries=None)

Load C++ Compiled libraries

Parameters:

Name	Type	Description	Default
libraries	list	A list of C++ compiled library paths. Defaults to None.	None

Source code in src/analysis_helpers/root_helpers.py

def LoadCompiledLibraries(libraries=None):
    """Load C++ Compiled libraries

    Args:
        libraries (list, optional): A list of C++ compiled library paths. Defaults to None.
    """
    if libraries is None: 
        return
    require_root()
    for lib in libraries:
        r.gSystem.Load(lib)
    return

DefineTree(variables, name='tree', title='a tree for my variable')

Given a set of variables in the format of a dictionary, a ROOT TTree object is created. The variables dictionary should have the form

variables = { <str>vname: <str>vtype }

where * vname: the name of the variable, * vtype: the type of the variable, one of the ROOT2ArrayTypes.keys()

Parameters:

Name	Type	Description	Default
variables	dict	a dictionary with variable names and their respective type.	required
name	str	the name of the ROOT TTree object. Defaults to 'tree'.	'tree'
title	str	the title of the ROOT TTree objects. Defaults to 'a tree for my variable'.	'a tree for my variable'

Source code in src/analysis_helpers/root_helpers.py

def DefineTree(variables, name='tree', title= 'a tree for my variable'):
    """Given a set of variables in the format of a dictionary, a ROOT TTree object is created. 
    The `variables` dictionary should have the form
    ```
    variables = { <str>vname: <str>vtype }
    ```
    where 
    * `vname`: the name of the variable,
    * `vtype`: the type of the variable, one of the `ROOT2ArrayTypes.keys()`

    Args:
        variables (dict): a dictionary with variable names and their respective type.
        name (str, optional): the name of the ROOT TTree object. Defaults to 'tree'.
        title (str, optional): the title of the ROOT TTree objects. Defaults to 'a tree for my variable'.
    """    
    require_root()
    t = r.TTree(name,title)
    t.SetDirectory(0)
    Vars = {}
    for vname, vtype in variables.items():
        Vars.update( { vname : array( ROOT2ArrayTypes[vtype],[0]) } )
    for var,args in Vars.items():
        print(var,args,args[0],type(args[0]).__name__,args.typecode)
        t.Branch(var, args, '%s/%s' % (var, Array2ROOTTypes[args.typecode]) )
#        Vars.update( { vname : array('i' if vtype == 'I' else 'f',[0]) })
#    for var,args in Vars.items(): t.Branch(var, args, '%s/%s' % (var, 'I' if type(args[0]) == type(0) else 'F'))
    t.Print()
    return dict(tree=t, vars=Vars)

GetTreeVariables(t)

Returns the dictionary of variables and types from TTree branches

Parameters:

Name	Type	Description	Default
t	TTree	The ROOT TTree object	required

Source code in src/analysis_helpers/root_helpers.py

def GetTreeVariables(t):
    """Returns the dictionary of variables and types from TTree branches

    Args:
        t (ROOT.TTree): The ROOT TTree object
    """
    require_root()
    tvars = {}
    for l in t.GetListOfLeaves():
        tvars.update( {l.GetName(): Array2ROOTTypes[ROOT2ArrayTypes_Long[l.GetTypeName()]] } )
    return tvars

SaveTree(t, fname)

Save a TTree into a file

Parameters:

Name	Type	Description	Default
t	TTree	A ROOT TTree object	required
fname	str	The file name	required

Source code in src/analysis_helpers/root_helpers.py

def SaveTree(t, fname):
    """Save a TTree into a file

    Args:
        t (ROOT.TTree): A ROOT TTree object
        fname (str): The file name
    """    
    require_root()
    # Save a tree
    f = r.TFile(fname, 'recreate')
    f.cd()
    t.Write()
    t.SetDirectory(0)
    f.Close()
    print('tree '+t.GetName()+' is saved to file '+fname)
    return

SaveTrees(trees, fname)

Save many TTrees into a file

Parameters:

Name	Type	Description	Default
trees	list	a list of TTree objects	required
fname	str	the file name	required

Source code in src/analysis_helpers/root_helpers.py

def SaveTrees(trees, fname):
    """Save many TTrees into a file

    Args:
        trees (list): a list of TTree objects
        fname (str): the file name
    """    
    require_root()
    # Save trees
    f = r.TFile(fname, 'recreate')
    f.cd()
    for t in trees:
        t.Write()
        t.SetDirectory(0)
        print('tree '+t.GetName()+' is saved to file '+fname)
    f.Close()
    return

ReduceTree(t, fname, nentries=None, cut='', firstentry=0, options='')

Reduce a TTree to a new one and save it to a file

Parameters:

Name	Type	Description	Default
t	TTree	A ROOT TTree object	required
fname	str	The file name	required
nentries	int	The number of entries to copy. Defaults to r.kMaxEntries.	None
cut	str	A selection to be made on the ROOT TTree branches. Defaults to ''.	''
firstentry	int	The first entry to copy. Defaults to 0.	0
options	str	Options to run the ROOT.TTree.CopyTree() function. Defaults to ''.	''

Source code in src/analysis_helpers/root_helpers.py

def ReduceTree(t, fname, nentries=None, cut='', firstentry=0, options=''):
    """Reduce a TTree to a new one and save it to a file

    Args:
        t (ROOT.TTree): A ROOT TTree object
        fname (str): The file name
        nentries (int, optional): The number of entries to copy. Defaults to r.kMaxEntries.
        cut (str, optional): A selection to be made on the ROOT TTree branches. Defaults to ''.
        firstentry (int, optional): The first entry to copy. Defaults to 0.
        options (str, optional): Options to run the `ROOT.TTree.CopyTree()` function. Defaults to ''.
    """     
    require_root()
    if nentries is None:
        nentries = r.TTree.kMaxEntries

    # Create a new tree
    f = r.TFile(fname,'recreate')
    t_new = t.CopyTree(cut,options,nentries,firstentry)
    t_new.Write()
    t_new.SetDirectory(0)
    f.Close()
    print('tree '+t.GetName()+' is reduced to file '+fname)
    return

ConvertTChain2DataFrame(ch, columns=None)

Convert a TChain to pandas.DataFrame

Parameters:

Name	Type	Description	Default
ch	TChain	a ROOT TChain object	required
columns	list	list of variables to save in the data frame. Defaults to None.	None

Returns:

Type	Description
	pandas.DataFrame: the output data frame

Source code in src/analysis_helpers/root_helpers.py

def ConvertTChain2DataFrame(ch,columns=None):
    """Convert a TChain to pandas.DataFrame

    Args:
        ch (r.TChain): a ROOT TChain object
        columns (list, optional): list of variables to save in the data frame. Defaults to None.

    Returns:
        pandas.DataFrame: the output data frame
    """    
    require_root()
    if columns is None: columns = [i.GetName() for i in ch.GetListOfLeaves()]
    rdf = r.RDataFrame(ch)
    df  = pd.DataFrame(rdf.AsNumpy(columns=columns))
    return df

Canvas(name='c1', width=600, height=600, title='')

A dictionary with a ROOT TCanvas and extra objects

Parameters:

Name	Type	Description	Default
name	str	the name of the TCanvas. Defaults to 'c1'.	'c1'
width	int	the width of the TCanvas. Defaults to 600.	600
height	int	the height of the TCanvas. Defaults to 600.	600
title	str	the title of the TCanvas. Defaults to ''.	''

Returns:

Name	Type	Description
dict		{'canvas': ROOT.TCanvas, 'extra': }

Source code in src/analysis_helpers/root_helpers.py

def Canvas(name='c1', width=600, height=600, title=''):
    """A dictionary with a ROOT TCanvas and extra objects

    Args:
        name (str, optional): the name of the TCanvas. Defaults to 'c1'.
        width (int, optional): the width of the TCanvas. Defaults to 600.
        height (int, optional): the height of the TCanvas. Defaults to 600.
        title (str, optional): the title of the TCanvas. Defaults to ''.

    Returns:
        dict: {'canvas': ROOT.TCanvas, 'extra': <list>}
    """
    require_root()
    if title=='':title=name
    return {'canvas': r.TCanvas(name, title, width, height), 'extra':[]}

SaveCansAsPdf(pdfName, cans)

Save a list of TCanvas objects into a single PDF file

Source code in src/analysis_helpers/root_helpers.py

def SaveCansAsPdf(pdfName, cans):
    """Save a list of TCanvas objects into a single PDF file
    """
    require_root()
    # formerly known as SavePdf
    print('Saving',[can.GetName() for can in cans])
    idx,idxmax = 0, len(cans)-1
    for can in cans:
        if idx==0: can.Print(pdfName+".pdf(")
        elif idx==idxmax: can.Print(pdfName+".pdf)")
        else: can.Print(pdfName+'.pdf')
        idx+=1
    return

TTree2Array(tree, leaves=None)

Transform a ROOT TTree into a numpy array

Parameters:

Name	Type	Description	Default
tree	TTree	the ROOT TTree object	required
leaves	list	the list of leaves to be transformed. Defaults to all leaves.	None

Returns:

Name	Type	Description
array		The numpy array

Source code in src/analysis_helpers/root_helpers.py

def TTree2Array(tree, leaves=None):
    """Transform a ROOT TTree into a numpy array

    Args:
        tree (TTree): the ROOT TTree object
        leaves (list, optional): the list of leaves to be transformed. Defaults to all leaves.

    Returns:
        array: The numpy array
    """
    require_root()
    # tree_branches = [l.GetName() for l in tree.GetListOfLeaves()] if leaves is None else leaves
    # arr = np.asarray([[ev.__getattr__(tb) for tb in tree_branches] for ev in tree ])
    r.EnableImplicitMT()  # opzionale: multithreading
    rdf = r.RDataFrame(tree)
    d = rdf.AsNumpy(columns=leaves)
    arr = np.column_stack([d[leaf] for leaf in leaves])
    return arr

ROOT2MPLLineStyle(lineStyle)

A function to convert ROOT line styles to Matplotlib line styles

Parameters:

Name	Type	Description	Default
lineStyle	int	ROOT line style index	required

Returns:

Name	Type	Description
str		matplotlib line style

Source code in src/analysis_helpers/root_helpers.py

def ROOT2MPLLineStyle(lineStyle):
    """A function to convert ROOT line styles to Matplotlib line styles

    Args:
        lineStyle (int): ROOT line style index

    Returns:
        str: matplotlib line style
    """
    if lineStyle not in [1,2,3,4,5,6,7,8,9,10]:
        return 'solid'
    mpl_line_styles = {
        1: 'solid',
        2: 'dashed',
        3: 'dotted',
        4: 'densely dashdotted',
        5: 'dashdotted',
        6: 'densely dashdotdotted',
        7: 'densely dashed',
        8: 'dashdotdotted',
        9: 'long dash with offset',
        10: 'dashdot'
    }
    return mpl_line_styles[lineStyle]

ROOT2MPLColor(color)

A function to convert ROOT colors to Matplotlib colors

Parameters:

Name	Type	Description	Default
color	TColor	ROOT TColor object	required

Returns:

Name	Type	Description
str		matplotlib color name

Source code in src/analysis_helpers/root_helpers.py

def ROOT2MPLColor(color):
    """A function to convert ROOT colors to Matplotlib colors

    Args:
        color (TColor): ROOT TColor object

    Returns:
        str: matplotlib color name
    """    
    require_root()
    mpl_color_name = color.GetName().lower()[1:]
    mpl_color_name = 'tab:{}'.format(mpl_color_name) if mpl_color_name in ['blue','orange','green','red','purple','brown','pink','gray','olive','cyan'] else 'tab:blue'
    return mpl_color_name

ROOT2MPLText(text)

A function to convert ROOT text to Matplotlib text

Parameters:

Name	Type	Description	Default
text	str	the text to convert	required

Returns:

Name	Type	Description
str		matplotlib text

Source code in src/analysis_helpers/root_helpers.py

def ROOT2MPLText(text):
    """A function to convert ROOT text to Matplotlib text

    Args:
        text (str): the text to convert

    Returns:
        str: matplotlib text
    """
    mpl_text = text
    if '#' in text or '^' in text or '_' in text:
        mpl_text = text.replace('#','\\')
        mpl_text = f'${mpl_text}$'
    return mpl_text

save_np_histograms_uproot(histograms_dict, filename)

Save multiple numpy histograms to a ROOT file with uproot

Parameters:

Name	Type	Description	Default
histograms_dict	dict	A dictionary of histograms to save.	required
filename	str	The name of the output ROOT file.	required

Source code in src/analysis_helpers/root_helpers.py

def save_np_histograms_uproot(histograms_dict, filename):
    """Save multiple numpy histograms to a ROOT file with uproot

    Args:
        histograms_dict (dict): A dictionary of histograms to save.
        filename (str): The name of the output ROOT file.
    """
    with uproot.recreate(filename) as f:
        for name, (hist_values, bin_edges) in histograms_dict.items():
            # Convert numpy histogram to ROOT-compatible format
            f[name] = (hist_values, bin_edges)
    return

load_np_histograms_uproot(filename)

Load numpy histograms from a ROOT file with uproot

Parameters:

Name	Type	Description	Default
filename	str	The name of the input ROOT file.	required

Returns:

Name	Type	Description
dict		A dictionary of histograms.

Source code in src/analysis_helpers/root_helpers.py

def load_np_histograms_uproot(filename):
    """Load numpy histograms from a ROOT file with uproot

    Args:
        filename (str): The name of the input ROOT file.

    Returns:
        dict: A dictionary of histograms.
    """
    require_root()
    histograms = {}
    with uproot.open(filename) as f:
        for key in f.keys(cycle=None):
            hist = f[key]
            histograms[key] = hist.to_numpy()  # Returns (values, edges)
    return histograms