data-manipulation-utilities 0.0.1__py3-none-any.whl

dmu/plotting/plotter_1d.py

@@ -0,0 +1,113 @@
+'''
+Module containing plotter class
+'''
+
+import numpy
+import matplotlib.pyplot as plt
+
+from dmu.logging.log_store import LogStore
+from dmu.plotting.plotter  import Plotter
+
+log = LogStore.add_logger('dmu:plotting:Plotter1D')
+# --------------------------------------------
+class Plotter1D(Plotter):
+    '''
+    Class used to plot columns in ROOT dataframes
+    '''
+    # --------------------------------------------
+    def __init__(self, d_rdf=None, cfg=None):
+        '''
+        Parameters:
+
+        d_rdf (dict): Dictionary mapping the kind of sample with the ROOT dataframe
+        cfg   (dict): Dictionary with configuration, e.g. binning, ranges, etc
+        '''
+
+        super().__init__(d_rdf=d_rdf, cfg=cfg)
+    #-------------------------------------
+    def _get_labels(self, var : str) -> tuple[str,str]:
+        if 'labels' not in self._d_cfg['plots'][var]:
+            return var, 'Entries'
+
+        xname, yname = self._d_cfg['plots'][var]['labels' ]
+
+        return xname, yname
+    #-------------------------------------
+    def _plot_var(self, var):
+        '''
+        Will plot a variable from a dictionary of dataframes
+        Parameters
+        --------------------
+        var   (str)  : name of column
+        '''
+        # pylint: disable=too-many-locals
+
+        d_cfg = self._d_cfg['plots'][var]
+
+        minx, maxx, bins = d_cfg['binning']
+        yscale           = d_cfg['yscale' ] if 'yscale' in d_cfg else 'linear'
+        xname, yname     = self._get_labels(var)
+
+        normalized=False
+        if 'normalized' in d_cfg:
+            normalized = d_cfg['normalized']
+
+        title = ''
+        if 'title'      in d_cfg:
+            title = d_cfg['title']
+
+        d_data = {}
+        for name, rdf in self._d_rdf.items():
+            d_data[name] = rdf.AsNumpy([var])[var]
+
+        if maxx <= minx + 1e-5:
+            log.info(f'Bounds not set for {var}, will calculated them')
+            minx, maxx = self._find_bounds(d_data = d_data, qnt=minx)
+            log.info(f'Using bounds [{minx:.3e}, {maxx:.3e}]')
+        else:
+            log.debug(f'Using bounds [{minx:.3e}, {maxx:.3e}]')
+
+        l_bc_all = []
+        d_wgt    = self._get_weights(var)
+        for name, arr_val in d_data.items():
+            arr_wgt    = d_wgt[name] if d_wgt is not None else None
+
+            self._print_weights(arr_wgt, var, name)
+            l_bc, _, _ = plt.hist(arr_val, weights=arr_wgt, bins=bins, range=(minx, maxx), density=normalized, histtype='step', label=name)
+            l_bc_all  += numpy.array(l_bc).tolist()
+
+            plt.yscale(yscale)
+            plt.xlabel(xname)
+            plt.ylabel(yname)
+
+        if yscale == 'linear':
+            plt.ylim(bottom=0)
+
+        max_y = max(l_bc_all)
+        plt.ylim(top=1.2 * max_y)
+        plt.title(title)
+    # --------------------------------------------
+    def _plot_lines(self, var : str):
+        '''
+        Will plot vertical lines for some variables
+
+        var (str) : name of variable
+        '''
+        if var in ['B_const_mass_M', 'B_M']:
+            plt.axvline(x=5280, color='r', label=r'$B^+$'   , linestyle=':')
+        elif var == 'Jpsi_M':
+            plt.axvline(x=3096, color='r', label=r'$J/\psi$', linestyle=':')
+    # --------------------------------------------
+    def run(self):
+        '''
+        Will run plotting
+        '''
+
+        fig_size = self._get_fig_size()
+        for var in self._d_cfg['plots']:
+            log.debug(f'Plotting: {var}')
+            plt.figure(var, figsize=fig_size)
+            self._plot_var(var)
+            self._plot_lines(var)
+            self._save_plot(var)
+# --------------------------------------------

dmu/plotting/plotter_2d.py

@@ -0,0 +1,87 @@
+'''
+Module containing Plotter2D class
+'''
+from typing import Union
+
+import hist
+import numpy
+import mplhep
+import matplotlib.pyplot as plt
+
+from hist                  import Hist
+from ROOT                  import RDataFrame
+from dmu.logging.log_store import LogStore
+from dmu.plotting.plotter  import Plotter
+
+log = LogStore.add_logger('dmu:plotting:Plotter2D')
+# --------------------------------------------
+class Plotter2D(Plotter):
+    '''
+    Class used to plot columns in ROOT dataframes
+    '''
+    # --------------------------------------------
+    def __init__(self, rdf=None, cfg=None):
+        '''
+        Parameters:
+
+        d_rdf (dict): Dictionary mapping the kind of sample with the ROOT dataframe
+        cfg   (dict): Dictionary with configuration, e.g. binning, ranges, etc
+        '''
+
+        if not isinstance(cfg, dict):
+            raise ValueError('Config dictionary not passed')
+
+        self._rdf   : RDataFrame = rdf
+        self._d_cfg : dict       = cfg
+
+        self._wgt : numpy.ndarray
+    # --------------------------------------------
+    def _get_axis(self, var : str):
+        [minx, maxx, nbins] = self._d_cfg['axes'][var]['binning']
+        label               = self._d_cfg['axes'][var][  'label']
+
+        axis = hist.axis.Regular(nbins, minx, maxx, name=label, label=label)
+
+        return axis
+    # --------------------------------------------
+    def _get_data(self, varx : str, vary : str) -> tuple[numpy.ndarray, numpy.ndarray]:
+        d_data = self._rdf.AsNumpy([varx, vary])
+        arr_x  = d_data[varx]
+        arr_y  = d_data[vary]
+
+        return arr_x, arr_y
+    # --------------------------------------------
+    def _get_dataset_weights(self, wgt_name : Union[str, None]) -> Union[numpy.ndarray, None]:
+        if wgt_name is None:
+            log.debug('Skipping weights')
+            return None
+
+        log.debug(f'Adding weights form column {wgt_name}')
+        arr_wgt  = self._rdf.AsNumpy([wgt_name])[wgt_name]
+
+        return arr_wgt
+    # --------------------------------------------
+    def _plot_vars(self, varx : str, vary : str, wgt_name : str) -> None:
+        log.info(f'Plotting {varx} vs {vary} with weights {wgt_name}')
+
+        ax_x         = self._get_axis(varx)
+        ax_y         = self._get_axis(vary)
+        arr_x, arr_y = self._get_data(varx, vary)
+
+        arr_w = self._get_dataset_weights(wgt_name)
+        hst   = Hist(ax_x, ax_y)
+        hst.fill(arr_x, arr_y, weight=arr_w)
+
+        mplhep.hist2dplot(hst)
+    # --------------------------------------------
+    def run(self):
+        '''
+        Will run plotting
+        '''
+
+        fig_size = self._get_fig_size()
+        for [varx, vary, wgt_name, plot_name] in self._d_cfg['plots_2d']:
+            plt.figure(plot_name, figsize=fig_size)
+            self._plot_vars(varx, vary, wgt_name)
+            self._save_plot(plot_name)
+# --------------------------------------------

dmu/rdataframe/atr_mgr.py

@@ -0,0 +1,79 @@
+'''
+Module with AtrMgr class
+'''
+
+import os
+
+from ROOT import RDataFrame
+
+import dmu.generic.utilities as gut
+from dmu.logging.log_store import LogStore
+
+#TODO:Skip attributes that start with Take< in a betterway
+log = LogStore.add_logger('dmu:rdataframe:atr_mgr')
+#------------------------
+class AtrMgr:
+    '''
+    Class intended to store attributes of ROOT dataframes and attach them back after a Filtering, definition, redefinition, etc operation 
+    These operations create new dataframes and therefore drop attributes.
+    '''
+    #------------------------
+    def __init__(self, rdf : RDataFrame):
+        self.d_in_atr = {}
+
+        self._store_atr(rdf)
+    #------------------------
+    def _store_atr(self, rdf : RDataFrame):
+        self.d_in_atr = self._get_atr(rdf)
+    #------------------------
+    def _skip_attr(self, name : str) -> bool:
+        if name.startswith('__') and name.endswith('__'):
+            return True
+
+        return False
+    #------------------------
+    def _get_atr(self, rdf : RDataFrame):
+        l_atr = dir(rdf)
+        d_atr = {}
+        for atr in l_atr:
+            if self._skip_attr(atr):
+                continue
+
+            val = getattr(rdf, atr)
+            d_atr[atr] = val
+
+        return d_atr
+    #------------------------
+    def add_atr(self, rdf : RDataFrame) -> RDataFrame:
+        '''
+        Takes dataframe and adds back the attributes
+        '''
+        d_ou_atr = self._get_atr(rdf)
+
+        key_in_atr = set(self.d_in_atr.keys())
+        key_ou_atr = set(     d_ou_atr.keys())
+
+        key_to_add = key_in_atr.difference(key_ou_atr)
+
+        for key in key_to_add:
+            val = self.d_in_atr[key]
+            if key.startswith('Take<') and key.endswith('>'):
+                continue
+
+            log.info(f'Adding attribute {key}')
+            setattr(rdf, key, val)
+
+        return rdf
+    #------------------------
+    def to_json(self, json_path : str) -> None:
+        '''
+        Saves the attributes inside current instance to JSON. Takes JSON path as argument
+        '''
+        json_dir = os.path.dirname(json_path)
+        os.makedirs(json_dir, exist_ok=True)
+
+        t_type   = (list, str, int, float)
+        d_fl_atr = { key : val for key, val in self.d_in_atr.items() if isinstance(val, t_type) and isinstance(key, t_type) }
+
+        gut.dump_json(d_fl_atr, json_path)
+#------------------------

dmu/rdataframe/utilities.py

@@ -0,0 +1,72 @@
+'''
+Module containing utility functions to be used with ROOT dataframes
+'''
+
+import re
+from dataclasses import dataclass
+
+import awkward as ak
+import numpy
+
+from ROOT import RDataFrame
+
+from dmu.logging.log_store import LogStore
+
+log = LogStore.add_logger('dmu:rdataframe:utilities')
+
+# ---------------------------------------------------------------------
+@dataclass
+class Data:
+    '''
+    Class meant to store data that is shared
+    '''
+    l_good_type = [int, numpy.bool_, numpy.int32, numpy.uint32, numpy.int64, numpy.uint64, numpy.float32, numpy.float64]
+    d_cast_type = {'bool': numpy.int32}
+# ---------------------------------------------------------------------
+def add_column(rdf : RDataFrame, arr_val : numpy.ndarray | None, name : str, d_opt : dict | None = None):
+    '''
+    Will take a dataframe, an array of numbers and a string
+    Will add the array as a colunm to the dataframe
+
+    d_opt (dict) : Used to configure adding columns
+         exclude_re : Regex with patter of column names that we won't pick
+    '''
+
+    d_opt = {} if d_opt is None else d_opt
+    if arr_val is None:
+        raise ValueError('Array of values not introduced')
+
+    if 'exclude_re' not in d_opt:
+        d_opt['exclude_re'] = None
+
+    v_col_org = rdf.GetColumnNames()
+    l_col_org = [name.c_str() for name in v_col_org ]
+    l_col     = []
+
+    tmva_rgx  = r'tmva_\d+_\d+'
+
+    for col in l_col_org:
+        user_rgx = d_opt['exclude_re']
+        if user_rgx is not None and re.match(user_rgx, col):
+            log.debug(f'Dropping: {col}')
+            continue
+
+        if                          re.match(tmva_rgx, col):
+            log.debug(f'Dropping: {col}')
+            continue
+
+        log.debug(f'Picking: {col}')
+        l_col.append(col)
+
+    data  = ak.from_rdataframe(rdf, columns=l_col)
+    d_data= { col : data[col] for col in l_col }
+
+    if arr_val.dtype == 'object':
+        arr_val = arr_val.astype(float)
+
+    d_data[name] = arr_val
+
+    rdf = ak.to_rdataframe(d_data)
+
+    return rdf
+# ---------------------------------------------------------------------

dmu/rfile/rfprinter.py

@@ -0,0 +1,91 @@
+'''
+Module containing RFPrinter
+'''
+import os
+
+from ROOT import TFile
+
+from dmu.logging.log_store import LogStore
+
+log = LogStore.add_logger('dmu:rfprinter')
+#--------------------------------------------------
+class RFPrinter:
+    '''
+    Class meant to print summary of ROOT file
+    '''
+    #-----------------------------------------
+    def __init__(self, path : str):
+        '''
+        Takes path to root file
+        '''
+        if not os.path.isfile(path):
+            raise FileNotFoundError(f'Cannot find {path}')
+
+        self._root_path = path
+        self._text_path = path.replace('.root', '.txt')
+    #-----------------------------------------
+    def _get_trees(self, ifile):
+        '''
+        Takes TFile object, returns list of TTree objects
+        '''
+        l_key=ifile.GetListOfKeys()
+
+        l_tree=[]
+        for key in l_key:
+            obj=key.ReadObj()
+            if obj.InheritsFrom("TTree"):
+                fname=ifile.GetName()
+                tname=obj.GetName()
+
+                title=f'{fname}/{tname}'
+                obj.SetTitle(title)
+                l_tree.append(obj)
+            elif obj.InheritsFrom("TDirectory"):
+                l_tree+=self._get_trees(obj)
+
+        return l_tree
+    #---------------------------------
+    def _get_tree_info(self, tree):
+        '''
+        Takes ROOT tree, returns list of strings with information about tree
+        '''
+        l_branch= tree.GetListOfBranches()
+        l_line  = []
+        for branch in l_branch:
+            bname = branch.GetName()
+            leaf  = branch.GetLeaf(bname)
+            btype = leaf.GetTypeName()
+
+            l_line.append(f'{"":4}{bname:<100}{btype:<40}')
+
+        return l_line
+    #-----------------------------------------
+    def _save_info(self, l_info):
+        '''
+        Takes list of strings, saves it to text file
+        '''
+
+        with open(self._text_path, 'w', encoding='utf-8') as ofile:
+            for info in l_info:
+                ofile.write(f'{info}\n')
+
+        log.info(f'Saved to: {self._text_path}')
+    #-----------------------------------------
+    def save(self, to_screen=False):
+        '''
+        Will save a text file with the summary of the ROOT file contents
+
+        to_screen (bool) : If true, will print to screen, default=False
+        '''
+        l_info = []
+        log.info(f'Reading from : {self._root_path}')
+        with TFile.Open(self._root_path) as ifile:
+            l_tree = self._get_trees(ifile)
+            for tree in l_tree:
+                l_info+= self._get_tree_info(tree)
+
+        self._save_info(l_info)
+        if to_screen:
+            for info in l_info:
+                log.info(info)
+#-----------------------------------------

dmu/rfile/utilities.py

@@ -0,0 +1,34 @@
+'''
+Module with utilities needed to manipulate ROOT files
+'''
+from ROOT import TTree, TDirectoryFile
+
+def get_trees_from_file(ifile : TDirectoryFile) -> dict[str,TTree]:
+    '''
+    Picks up a TFile object
+    Returns a dictionary of trees, with the tree location as the key
+    Can search recursively within directories
+    '''
+    if not ifile.InheritsFrom('TDirectoryFile'):
+        str_type = str(type(ifile))
+        raise ValueError(f'Unrecognized object type {str_type}')
+
+    dir_name = ifile.GetName()
+
+    d_tree={}
+    l_key =ifile.GetListOfKeys()
+
+    for key in l_key:
+        obj=key.ReadObj()
+        if   obj.InheritsFrom('TDirectoryFile'):
+            d_tmp = get_trees_from_file(obj)
+            d_tree.update(d_tmp)
+        elif obj.InheritsFrom('TTree'):
+            obj_name = obj.GetName()
+            key      = f'{dir_name}/{obj_name}'
+
+            d_tree[key] = obj
+        else:
+            continue
+
+    return d_tree