data-manipulation-utilities 0.2.6__py3-none-any.whl → 0.2.8.dev714__py3-none-any.whl

dmu/ml/cv_diagnostics.py

@@ -0,0 +1,224 @@
+'''
+Module containing CVDiagnostics class
+'''
+import os
+
+import numpy
+import matplotlib
+import matplotlib.pyplot as plt
+import pandas            as pnd
+
+from scipy.stats             import kendalltau
+from ROOT                    import RDataFrame, RDF
+from dmu.ml.cv_classifier    import CVClassifier
+from dmu.ml.cv_predict       import CVPredict
+from dmu.logging.log_store   import LogStore
+from dmu.plotting.plotter_1d import Plotter1D as Plotter
+
+NPA = numpy.ndarray
+Axis= matplotlib.axes._axes.Axes
+log = LogStore.add_logger('dmu:ml:cv_diagnostics')
+# -------------------------
+class CVDiagnostics:
+    '''
+    Class meant to rundiagnostics on classifier
+
+    Correlations
+    ------------------
+    Will calculate correlations between features + signal probability and some external target variable specified in the config
+    '''
+    # -------------------------
+    def __init__(self, models : list[CVClassifier], rdf : RDataFrame, cfg : dict):
+        self._l_model = models
+        self._cfg     = cfg
+        self._rdf     = rdf
+        self._target  = cfg['correlations']['target']['name']
+        self._l_feat  = self._get_features()
+        self._d_xlab  = self._get_xlabels()
+    # -------------------------
+    def _get_features(self) -> list[str]:
+        cfg   = self._l_model[0].cfg
+        l_var = cfg['training']['features']
+
+        return l_var
+    # -------------------------
+    def _get_xlabels(self) -> dict[str,str]:
+        cfg   = self._l_model[0].cfg
+        d_var = cfg['plotting']['features']['plots']
+
+        d_lab = { varname : d_field['labels'][0] for varname, d_field in d_var.items() }
+
+        target= self._cfg['correlations']['target']['name']
+        if 'overlay' not in self._cfg['correlations']['target']:
+            xlabel = target
+        else:
+            xlabel= self._cfg['correlations']['target']['overlay']['plots'][target]['labels'][0]
+
+        d_lab[target]  = xlabel
+        d_lab['score'] = 'score'
+
+        d_lab = { var_id : var_name.replace('MeV', '') for var_id, var_name in d_lab.items() }
+
+        return d_lab
+    # -------------------------
+    def _add_columns(self, rdf : RDataFrame) -> RDataFrame:
+        cfg    = self._l_model[0].cfg
+        d_def  = cfg['dataset']['define']
+        for var, expr in d_def.items():
+            rdf = rdf.Define(var, expr)
+
+        return rdf
+    # -------------------------
+    def _get_scores(self) -> NPA:
+        if 'score_from_rdf' not in self._cfg:
+            log.debug('Using score from model')
+            prd = CVPredict(models=self._l_model, rdf = self._rdf)
+
+            return prd.predict()
+
+        name = self._cfg['score_from_rdf']
+        log.debug(f'Picking up score from dataframe, column: {name}')
+        arr_score = self._rdf.AsNumpy([name])[name]
+
+        return arr_score
+    # -------------------------
+    def _get_arrays(self) -> dict[str, NPA]:
+        rdf   = self._add_columns(self._rdf)
+        l_col = [ name.c_str() for name in rdf.GetColumnNames() ]
+
+        missing= False
+        l_var  = self._l_feat + [self._target]
+        for var in l_var:
+            if var not in l_col:
+                log.error(f'{"Missing":<20}{var}')
+                missing=True
+
+        if missing:
+            raise ValueError('Columns missing')
+
+        d_var          = rdf.AsNumpy(l_var)
+        d_var['score'] = self._get_scores()
+
+        return d_var
+    # -------------------------
+    def _run_correlations(self, method : str, ax : Axis) -> Axis:
+        d_arr      = self._get_arrays()
+        arr_target = d_arr[self._target]
+
+        d_corr= {}
+        for name, arr_val in d_arr.items():
+            if name == self._target:
+                continue
+
+            d_corr[name] = self._calculate_correlations(var=arr_val, target=arr_target, method=method)
+
+        ax = self._plot_correlations(d_corr=d_corr, method=method, ax=ax)
+
+        return ax
+    # -------------------------
+    def _plot_correlations(self, d_corr : dict[str,float], method : str, ax : Axis) -> Axis:
+        df             = pnd.DataFrame.from_dict(d_corr, orient="index", columns=[method])
+        df['variable'] = df.index.map(self._d_xlab)
+
+        figsize = self._cfg['correlations']['figure']['size']
+        ax      = df.plot(x='variable', y=method,label=method, figsize=figsize, ax=ax)
+
+        # Needed to show all labels on x axis
+        plt.xticks(ticks=range(len(df)), labels=df.variable)
+        if 'xlabelsize' in self._cfg['correlations']['figure']:
+            xlabsize= self._cfg['correlations']['figure']['xlabelsize']
+        else:
+            xlabsize= 30
+
+        ax.tick_params(axis='x', labelsize=xlabsize)
+
+        return ax
+    # -------------------------
+    def _save_plot(self):
+        plot_dir = self._cfg['output']
+        os.makedirs(plot_dir, exist_ok=True)
+
+        plot_path = f'{plot_dir}/correlations.png'
+        log.info(f'Saving to: {plot_path}')
+
+        title = None
+        if 'title' in self._cfg['correlations']['figure']:
+            title = self._cfg['correlations']['figure']['title']
+
+        rotation=30
+        if 'rotate' in self._cfg['correlations']['figure']:
+            rotation = self._cfg['correlations']['figure']['rotate']
+
+        plt.ylim(-1, +1)
+        plt.title(title)
+        plt.xlabel('')
+        plt.ylabel('Correlation')
+        plt.grid()
+        plt.xticks(rotation=rotation)
+        plt.tight_layout()
+        plt.savefig(plot_path)
+        plt.close()
+    # -------------------------
+    def _remove_nans(self, var : NPA, tgt : NPA) -> tuple[NPA,NPA]:
+        arr_nan_var = numpy.isnan(var)
+        arr_nan_tgt = numpy.isnan(tgt)
+        arr_is_nan  = numpy.logical_or(arr_nan_var, arr_nan_tgt)
+        arr_not_nan = numpy.logical_not(arr_is_nan)
+
+        var         = var[arr_not_nan]
+        tgt         = tgt[arr_not_nan]
+
+        return var, tgt
+    # -------------------------
+    def _calculate_correlations(self, var : NPA, target : NPA, method : str) -> float:
+        var, target = self._remove_nans(var, target)
+
+        if method == 'Pearson':
+            mat         = numpy.corrcoef(var, target)
+
+            return mat[0,1]
+
+        if method == r'Kendall-$\tau$':
+            tau, _ = kendalltau(var, target)
+
+            return tau
+
+        raise NotImplementedError(f'Correlation coefficient {method} not implemented')
+    # -------------------------
+    def _plot_cutflow(self) -> None:
+        '''
+        Plot the 'mass' column for different values of working point
+        '''
+        if 'overlay' not in self._cfg['correlations']['target']:
+            log.debug('Not plotting cutflow of target distribution')
+            return
+
+        arr_score = self._get_scores()
+        arr_target= self._rdf.AsNumpy([self._target])[self._target]
+        arr_wp    = self._cfg['correlations']['target']['overlay']['wp']
+        rdf       = RDF.FromNumpy({'Score' : arr_score, self._target : arr_target})
+
+        d_rdf = {}
+        for wp in arr_wp:
+            name        = f'WP > {wp:.2}'
+            expr        = f'Score > {wp:.3}'
+            d_rdf[name] = rdf.Filter(expr)
+
+        cfg_target = self._cfg['correlations']['target']['overlay']
+
+        ptr=Plotter(d_rdf=d_rdf, cfg=cfg_target)
+        ptr.run()
+    # -------------------------
+    def run(self) -> None:
+        '''
+        Runs diagnostics
+        '''
+        if 'correlations' in self._cfg:
+            ax = None
+            for method in self._cfg['correlations']['methods']:
+                ax = self._run_correlations(method=method, ax=ax)
+
+            self._save_plot()
+
+            self._plot_cutflow()
+# -------------------------

dmu/ml/cv_performance.py

@@ -0,0 +1,58 @@
+'''
+This module contains the class CVPerformance
+'''
+# pylint: disable=too-many-positional-arguments, too-many-arguments
+
+from ROOT                  import RDataFrame
+from dmu.ml.cv_classifier  import CVClassifier
+from dmu.ml.cv_predict     import CVPredict
+from dmu.ml.train_mva      import TrainMva
+from dmu.logging.log_store import LogStore
+
+log=LogStore.add_logger('dmu:ml:cv_performance')
+# -----------------------------------------------------
+class CVPerformance:
+    '''
+    This class is meant to:
+
+    - Compare the classifier performance, through the ROC curve, of a model, for a given background and signal sample
+    '''
+    # ---------------------------
+    def plot_roc(
+            self,
+            name  : str,
+            color : str,
+            sig   : RDataFrame,
+            bkg   : RDataFrame,
+            model : list[CVClassifier] ) -> float:
+        '''
+        Method in charge of picking up model and data and plotting ROC curve
+
+        Parameters
+        --------------------------
+        name : Label of combination, used for plots
+        sig  : ROOT dataframe storing signal samples
+        bkg  : ROOT dataframe storing background samples
+        model: List of instances of the CVClassifier
+
+        Returns
+        --------------------------
+        Area under the ROC curve
+        '''
+        log.info(f'Loading {name}')
+
+        cvp_sig = CVPredict(models=model, rdf=sig)
+        arr_sig = cvp_sig.predict()
+
+        cvp_bkg = CVPredict(models=model, rdf=bkg)
+        arr_bkg = cvp_bkg.predict()
+
+        _, _, auc = TrainMva.plot_roc_from_prob(
+                arr_sig_prb=arr_sig,
+                arr_bkg_prb=arr_bkg,
+                kind       =   name,
+                color      =  color, # This should allow the function to pick kind
+                ifold      =    999) # for the label
+
+        return auc
+# -----------------------------------------------------

dmu/ml/cv_predict.py

@@ -1,8 +1,6 @@
 '''
 Module holding CVPredict class
 '''
-from typing import Optional
-
 import pandas as pnd
 import numpy
 import tqdm
@@ -21,41 +19,107 @@ class CVPredict:
     Class used to get classification probabilities from ROOT
     dataframe and a set of models. The models were trained with CVClassifier
     '''
-    def __init__(self, models : Optional[list] = None, rdf : Optional[RDataFrame] = None):
+    def __init__(
+            self,
+            rdf    : RDataFrame,
+            models : list[CVClassifier]):
         '''
         Will take a list of CVClassifier models and a ROOT dataframe
-        '''
-
-        if models is None:
-            raise ValueError('No list of models passed')
-
-        if rdf is None:
-            raise ValueError('No ROOT dataframe passed')
 
+        rdf   : ROOT dataframe where features will be extracted
+        models: List of models, one per fold
+        '''
         self._l_model   = models
         self._rdf       = rdf
+        self._nrows     : int
+        self._l_column  : list[str]
         self._d_nan_rep : dict[str,str]
 
-        self._arr_patch : numpy.ndarray
+        # Value of score used when no score has been assigned
+        self._dummy_score = -1.0
+
+        # name of column in ROOT dataframe where 1s will prevent prediction
+        self._skip_index_column = 'skip_mva_prediction'
+
+        # name of attribute of features dataframe where array of indices to skip are stored
+        self._index_skip  = 'skip_mva_prediction'
     # --------------------------------------------
     def _initialize(self):
+        self._rdf       = self._remove_periods(self._rdf)
         self._rdf       = self._define_columns(self._rdf)
         self._d_nan_rep = self._get_nan_replacements()
+        self._l_column  = [ name.c_str() for name in self._rdf.GetColumnNames() ]
+        self._nrows     = self._rdf.Count().GetValue()
+    # ----------------------------------
+    def _remove_periods(self, rdf : RDataFrame) -> RDataFrame:
+        '''
+        This will redefine all columns associated to friend trees as:
+
+        friend_preffix.branch_name -> friend_preffix.branch_name
+        '''
+        l_col = [ col.c_str() for col in rdf.GetColumnNames() ]
+        l_col = [ col         for col in l_col if '.' in col  ]
+
+        if len(l_col) == 0:
+            return rdf
+
+        log.debug(60 * '-')
+        log.debug('Renaming dotted columns')
+        log.debug(60 * '-')
+        for col in l_col:
+            new = col.replace('.', '_')
+            log.debug(f'{col:<50}{"->":10}{new:<20}')
+            rdf = rdf.Define(new, col)
+
+        return rdf
+    # --------------------------------------------
+    def _get_definitions(self) -> dict[str,str]:
+        '''
+        This method will search in the configuration the definitions used
+        on the dataframe before the dataframe was used to train the model.
+        '''
+        cfg   = self._l_model[0].cfg
+        d_def = {}
+        if 'define' in cfg['dataset']:
+            d_def_gen = cfg['dataset']['define'] # get generic definitions
+            d_def.update(d_def_gen)
+
+        sig_name = 'sig'
+        try:
+            # Get sample specific definitions. This will be taken from the signal section
+            # because predicted scores should come from features defined as for the signal.
+            d_def_sam = cfg['dataset']['samples'][sig_name]['definitions']
+        except KeyError:
+            log.debug(f'No sample specific definitions were found in: {sig_name}')
+            return d_def
+
+        log.info('Adding sample dependent definitions')
+        d_def.update(d_def_sam)
+
+        return d_def
     # --------------------------------------------
     def _define_columns(self, rdf : RDataFrame) -> RDataFrame:
-        cfg = self._l_model[0].cfg
-
-        if 'define' not in cfg['dataset']:
-            log.debug('No define section found in config, will not define extra columns')
+        d_def = self._get_definitions()
+        if len(d_def) == 0:
+            log.info('No definitions found')
             return self._rdf
 
-        d_def = cfg['dataset']['define']
+        dexc = None
         log.debug(60 * '-')
         log.info('Defining columns in RDF before evaluating classifier')
         log.debug(60 * '-')
         for name, expr in d_def.items():
+            expr = expr.replace('.', '_')
+
             log.debug(f'{name:<20}{"<---":20}{expr:<100}')
-            rdf = rdf.Define(name, expr)
+            try:
+                rdf = rdf.Define(name, expr)
+            except TypeError as exc:
+                log.error(f'Cannot define {name}={expr}')
+                dexc = exc
+
+        if dexc is not None:
+            raise TypeError('Could not define at least one column') from dexc
 
         return rdf
     # --------------------------------------------
@@ -68,21 +132,25 @@ class CVPredict:
 
         return cfg['dataset']['nan']
     # --------------------------------------------
-    def _replace_nans(self, df : pnd.DataFrame) -> pnd.DataFrame:
+    def _replace_nans(self, df_ft : pnd.DataFrame) -> pnd.DataFrame:
+        '''
+        Funtion replaces nans in user specified columns with user specified values
+        These NaNs are expected
+        '''
         if len(self._d_nan_rep) == 0:
             log.debug('Not doing any NaN replacement')
-            return df
+            return df_ft
 
         log.info(60 * '-')
         log.info('Doing NaN replacements')
         log.info(60 * '-')
         for var, val in self._d_nan_rep.items():
             log.info(f'{var:<20}{"--->":20}{val:<20.3f}')
-            df[var] = df[var].fillna(val)
+            df_ft[var] = df_ft[var].fillna(val)
 
-        return df
+        return df_ft
     # --------------------------------------------
-    def _get_df(self):
+    def _get_df(self) -> pnd.DataFrame:
         '''
         Will make ROOT rdf into dataframe and return it
         '''
@@ -90,11 +158,11 @@ class CVPredict:
         l_ft  = model.features
         d_data= self._rdf.AsNumpy(l_ft)
         df_ft = pnd.DataFrame(d_data)
-        df_ft = self._replace_nans(df_ft)
-        df_ft = ut.patch_and_tag(df_ft)
-
-        if 'patched_indices' in df_ft.attrs:
-            self._arr_patch = df_ft.attrs['patched_indices']
+        df_ft = self._replace_nans(df_ft=df_ft)
+        df_ft = self._tag_skipped(df_ft=df_ft)
+        df_ft = ut.tag_nans(
+            df      = df_ft,
+            indexes = self._index_skip)
 
         nfeat = len(l_ft)
         log.info(f'Found {nfeat} features')
@@ -103,6 +171,24 @@ class CVPredict:
 
         return df_ft
     # --------------------------------------------
+    def _tag_skipped(self, df_ft : pnd.DataFrame) -> pnd.DataFrame:
+        '''
+        Will drop rows with features where column with name _skip_name (currently "_skip_mva_prediction") has values of 1
+        '''
+        if self._skip_index_column not in self._l_column:
+            log.debug(f'Not dropping any rows through: {self._skip_index_column}')
+            return df_ft
+
+        log.info(f'Dropping rows through: {self._skip_index_column}')
+        arr_drop                = self._rdf.AsNumpy([self._skip_index_column])[self._skip_index_column]
+
+        if self._index_skip in df_ft.attrs:
+            raise ValueError(f'Feature dataframe already contains attribute: {self._index_skip}')
+
+        df_ft.attrs[self._index_skip] = numpy.where(arr_drop == 1)[0]
+
+        return df_ft
+    # --------------------------------------------
     def _non_overlapping_hashes(self, model, df_ft):
         '''
         Will return True if hashes of model and data do not overlap
@@ -147,8 +233,8 @@ class CVPredict:
         '''
         Evaluate the dataset for one of the folds, by taking the model and the full dataset
         '''
-        s_dat_hash = set(df_ft.index)
-        s_mod_hash = model.hashes
+        s_dat_hash : set[str] = set(df_ft.index)
+        s_mod_hash : set[str] = model.hashes
 
         s_dif_hash = s_dat_hash - s_mod_hash
 
@@ -164,19 +250,29 @@ class CVPredict:
         d_prob = dict(zip(l_hash, l_prob))
         nfeat  = len(df_ft_group)
         nprob  = len(l_prob)
-        log.debug(f'{nfeat:<10}{"->":10}{nprob:<10}')
+
+        if nfeat != nprob:
+            raise ValueError(f'Number of features and probabilities do not agree: {nfeat} != {nprob}')
 
         return d_prob
     # --------------------------------------------
-    def _patch_probabilities(self, arr_prb : numpy.ndarray) -> numpy.ndarray:
-        if not hasattr(self, '_arr_patch'):
-            return arr_prb
+    def _predict_signal_probabilities(
+            self,
+            model : CVClassifier,
+            df_ft : pnd.DataFrame) -> numpy.ndarray:
+        '''
+        Takes model and features dataframe, returns array of signal probabilities
+        '''
+        if self._non_overlapping_hashes(model, df_ft):
+            log.debug('No intersecting hashes found between model and data')
+            arr_prb = model.predict_proba(df_ft)
+        else:
+            log.info('Intersecting hashes found between model and data')
+            arr_prb = self._predict_with_overlap(df_ft)
 
-        nentries = len(self._arr_patch)
-        log.warning(f'Patching {nentries} probabilities with -1')
-        arr_prb[self._arr_patch] = -1
+        arr_sig_prb = arr_prb.T[1]
 
-        return arr_prb
+        return arr_sig_prb
     # --------------------------------------------
     def predict(self) -> numpy.ndarray:
         '''
@@ -187,15 +283,22 @@ class CVPredict:
         df_ft = self._get_df()
         model = self._l_model[0]
 
-        if self._non_overlapping_hashes(model, df_ft):
-            log.debug('No intersecting hashes found between model and data')
-            arr_prb = model.predict_proba(df_ft)
-        else:
-            log.info('Intersecting hashes found between model and data')
-            arr_prb = self._predict_with_overlap(df_ft)
+        arr_keep = None
+        arr_skip = None
+        if self._index_skip in df_ft.attrs:
+            arr_skip = df_ft.attrs[self._index_skip]
+            df_ft    = df_ft.drop(arr_skip)
+            arr_keep = df_ft.index.to_numpy()
+
+        arr_sig_prb  = self._predict_signal_probabilities(
+                model = model,
+                df_ft = df_ft)
+
+        if arr_skip is None:
+            return arr_sig_prb
 
-        arr_prb = self._patch_probabilities(arr_prb)
-        arr_prb = arr_prb.T[1]
+        arr_all_sig_prb           = numpy.full(self._nrows, self._dummy_score)
+        arr_all_sig_prb[arr_keep] = arr_sig_prb
 
-        return arr_prb
+        return arr_all_sig_prb
 # ---------------------------------------