metadata.py

'''
Class to produce metadata for root to numpy conversion.

@author: hqu
'''

from __future__ import print_function

import os
import re
import json
import logging
import numpy as np
import pandas as pd

from helper import get_num_events

class Metadata(object):

    ''' Compile the metadata. '''

    def __init__(self,
                 inputdir,
                 input_filter=None,
                 treename='deepntuplizer/tree',
                 reweight_events=100000,
                 reweight_bins=[[200, 5000], [-999, 999]],
                 metadata_events=100000,
                 selection=None,
                 var_groups=None,  # {group_name:(regex, num)}
                 var_blacklist=None,
                 var_no_transform_branches=None,
                 label_list=None,
                 reweight_var=['fj_pt', 'fj_sdmass'],
                 reweight_classes=['fj_isTop', 'fj_isW', 'fj_isZ', 'fj_isH', 'fj_isQCD'],
                 reweight_method='flat',
                 scale_method='upper',
                 var_img='pfcand_ptrel',
                 var_pos=['pfcand_etarel', 'pfcand_phirel'],
                 n_pixels=64,
                 img_ranges=[[-0.8, 0.8], [-0.8, 0.8]],
                 ):
        self._inputdir = inputdir  # data members starting with '_' is not loaded from json

        self.input_filter = input_filter  # regex for filtering the input paths
        self.treename = treename
        self.reweight_var = reweight_var
        self.reweight_classes = reweight_classes
        self._reweight_events = reweight_events
        self._reweight_bins = reweight_bins
        self._metadata_events = metadata_events
        self.selection = selection
        self.var_groups = var_groups
        self.var_blacklist = var_blacklist
        self.var_no_transform_branches = var_no_transform_branches
        self.label_branches = label_list
        
        self.reweight_method = reweight_method.lower()
        if self.reweight_method not in ['none', 'flat', 'ref']:
            raise NotImplemented('reweight method %s not recognized' % reweight_method)

        self.scale_method = scale_method.lower()
        if self.scale_method not in ['upper', 'lower', 'average', 'max']:
            raise NotImplemented('scale method %s not recognized' % scale_method)

        self.var_img = var_img
        self.var_pos = var_pos
        self.n_pixels = n_pixels
        self.img_ranges = img_ranges

        self.inputfiles = None
        self.num_events = None

    def produceMetadata(self, filepath):
        logging.info('Start producing metadata...')
        # make file list
        self.updateFilelist()
        # make var list
        self._make_varlist()
        # make weights
        self._make_weights()
        # make transfromation info
        self._make_infos()
        # write metadata
        self.writeMetadata(filepath)

    def loadMetadata(self, filepath, override=True):
        with open(filepath) as metafile:
            md = json.load(metafile)
            for k in md:
                if k.startswith('_'): continue
                if override:
                    setattr(self, k, md[k])
                else:
                    if not hasattr(self, k) or getattr(self, k) is None:
                        setattr(self, k, md[k])
        logging.info('Metadata loaded from ' + filepath)

    def updateFilelist(self, test_sample=False):
        import re
        self.inputfiles = []
        self.num_events = []
        counter = 0
        for dp, dn, filenames in os.walk(self._inputdir, followlinks=True):
            if 'failed' in dp or 'ignore' in dp:
                continue
            if self.input_filter and re.search(self.input_filter, dp):
                logging.debug('Ignoring inputdir %s', dp)
                continue
#             if not test_sample and 'test_sample' in dp:
#                 # train/val samples
#                 continue
#             if test_sample and 'test_sample' not in dp:
#                 # test samples
#                 continue
            for f in filenames:
                if not f.endswith('.root'):
                    continue
                fullpath = os.path.realpath(os.path.join(dp, f))
                nevts = get_num_events(fullpath, self.treename)
                if nevts:
                    self.inputfiles.append(fullpath)
                    self.num_events.append(nevts)
                    counter += 1
                    if counter%10==0:
                        logging.debug('%d files processed...' % counter)
                else:
                    logging.warning('Ignore erroneous file %s' % fullpath)
        self._total_events = sum(self.num_events)
        logging.info('Created file list from directory %s\nFiles:%d, Events:%d' % (self._inputdir, len(self.inputfiles), self._total_events))
        return (self.inputfiles, self.num_events)

    def updateWeights(self, test_sample=False):
        if test_sample:
            return
        else:
            self._make_weights()

    def writeMetadata(self, filepath):
        with open(filepath, 'w') as metafile:
            json.dump(self.__dict__, metafile, indent=2, sort_keys=True)
        logging.info('Metadata written to ' + filepath)


    def _make_varlist(self):
        # get all branches and filter them using input variable list
        from root_numpy import root2array
        df = pd.DataFrame(root2array(self.inputfiles[0], treename=self.treename, stop=1))
        self._all_branches = df.columns.values.tolist()
        self.var_branches = []
        self.var_sizes = {}
        for k in self._all_branches:
            matched = False
            for v_group in self.var_groups:
                size = self.var_groups[v_group][1]
                for regex in self.var_groups[v_group][0]:
                    if re.match(regex, k):
                        self.var_branches.append(k)
                        self.var_sizes[k] = size
                        matched = True
                        break
                if matched: break
        for var in self.var_blacklist + self.label_branches + self.reweight_classes:
            try:
                self.var_branches.remove(var)
            except ValueError:
                pass
        logging.info('Training vars:\n' + '\n'.join(self.var_branches))
        # check no_transform vars
        _var_no_transform = []
        for v in self.var_no_transform_branches:
            if v in self._all_branches:
                _var_no_transform.append(v)
            else:
                logging.warning('No-transform var %s not found in the input. Will be ignored!' % v)
        self.var_no_transform_branches = _var_no_transform


    def _prepare_reweight_info(self, rec):
        ''' Produce metadata for reweighting. Goal:
            1) Produce flat pT spectrum.
            2) Balance the class weights on top of that
        '''
        class_events = {}
        result = {}
        for label in self.reweight_classes:
            pos = (rec[label] == 1)
            x = np.minimum(rec[self.reweight_var[0]][pos], max(self._reweight_bins[0]))
            y = np.minimum(rec[self.reweight_var[1]][pos], max(self._reweight_bins[1]))
#             class_events[label] = 0
            hist, x_edges, y_edges = np.histogram2d(x, y, bins=self._reweight_bins)
            hist = np.asfarray(hist, dtype=np.float32)
            result[label] = {'x_edges':x_edges.tolist(), 'y_edges':y_edges.tolist(), 'hist':hist, 'raw_hist':hist[:].tolist()}
            logging.debug('%s:\n%s' % (label, str(hist)))
#             if min(hist[-2:]) < 10:
#                 logging.warning('Not enough events for cateogry %s:\n%s' % (label, str(hist)))
#                 raise RuntimeError('Not enough events for cateogry %s:\n%s' % (label, str(hist)))
        if self.reweight_method == 'flat':
            for label in self.reweight_classes:
                hist = result[label]['hist']
                hist_non_zero = hist[hist > 0]
                min_val = np.min(hist_non_zero)
                med_val = np.median(hist)
                ref_val = np.percentile(hist_non_zero, 10)
                logging.debug('label:%s, median=%f, min=%f, ref=%f, ref/min=%f' % (label, med_val, min_val, ref_val, ref_val / min_val))
                class_events[label] = ref_val
                wgt = ref_val / hist  # will produce inf if hist[ix,iy]=0
                wgt[np.isinf(wgt)] = 0  # get rid of inf
                wgt = np.clip(wgt, 0, 5)
                result[label]['hist'] = wgt.tolist()
            min_nevt = min(class_events.values())
            for label in self.reweight_classes:
                class_wgt = float(min_nevt) / class_events[label]
                result[label]['class_wgt'] = class_wgt
        elif self.reweight_method == 'ref':
            # use class 0 as the reference
            # will get both shape wgt and class wgt at the same time
            hist_ref = result[self.reweight_classes[0]]['hist']
            upper_wgt = 1.0
            for label in self.reweight_classes:
                wgt = hist_ref / result[label]['hist']
                wgt[np.isinf(wgt)] = 0  # get rid of inf
                upper = np.percentile(wgt, 90)
                if upper > upper_wgt:
                    upper_wgt = upper
                result[label]['hist'] = wgt
            # rescale the weights to make them less than 1
            for label in self.reweight_classes:
                result[label]['hist'] = (result[label]['hist'] / upper_wgt).tolist()
                result[label]['class_wgt'] = 1
        return result

    def _make_weights(self):
        if self.reweight_method == 'none':
            logging.info('-- Reweighting is disabled --')
            return
        logging.info('Start making weights...\n Var: %s\n Classes: %s\n Selection: %s' % (str(self.reweight_var), str(self.reweight_classes), self.selection))
        # fraction of events to take from each file
        from root_numpy import root2array
        frac = 1.0
        if self._reweight_events > 0:
            frac = float(self._reweight_events) / self._total_events
        if frac < 1:
            pieces = []
            for fn, n in zip(self.inputfiles, self.num_events):
                a = root2array(fn, treename=self.treename, selection=self.selection, stop=int(frac * n),
                               branches=self.reweight_classes + self.reweight_var)
                pieces.append(a)
            rec = np.concatenate(pieces)
        else:

            rec = root2array(self.inputfiles, treename=self.treename, selection=self.selection,
                               branches=self.reweight_classes + self.reweight_var)
        logging.info('Use %d events to produce reweight info, selection:\n%s' % (rec.shape[0], self.selection))
        # get distribution for reweighting
        self.reweight_info = self._prepare_reweight_info(rec)
        logging.debug('Reweight info:\n' + str(self.reweight_info))

    def _make_infos(self):
        # make variables transformation infos
        from root_numpy import root2array
        frac = 1.0
        _inputfiles = self.inputfiles
        _num_events = self.num_events
        if self._metadata_events > 0:
            nfiles = int(5 * float(self._metadata_events) / self._total_events * len(self.inputfiles))
            file_inds = np.arange(len(self.inputfiles))
            np.random.shuffle(file_inds)
            file_inds = file_inds[:nfiles]
            _inputfiles = [self.inputfiles[i] for i in file_inds]
            _num_events = [self.num_events[i] for i in file_inds]
            frac = float(self._metadata_events) / sum(_num_events)

        first = True

        self.branches_info = {}
        for var in self.var_branches:
            var_size = self.var_sizes[var]
            pieces = []
            for fn, n in zip(_inputfiles, _num_events):
                v = root2array(fn, treename=self.treename, selection=self.selection,
                               branches=var, stop=int(frac * n))
                pieces.append(v)
            a = np.concatenate(pieces)
            if first:
                first = False
                logging.debug('Use %d events from %d files for var transform info' % (a.shape[0], len(_inputfiles)))
            size = None
            if a.dtype == np.object:
                if var_size:
                    size = var_size  # use given size if provided
                else:
                    lengths = [len(row) for row in a]
                    size = int(round(np.percentile(lengths, 95)))  # else get 95% percentile of the length
                a = np.nan_to_num(np.concatenate(a))  # then flatten vector vars for calculations
            else:
                a = np.nan_to_num(a)
            self.branches_info[var] = {
                'size'  : size,
                'median': float(np.percentile(a, 50)),  # need float otherwise cannot serialize to json
                'lower' : float(np.percentile(a, 16)),
                'upper' : float(np.percentile(a, 84)),
                'min'   : float(np.min(a)),
                'max'   : float(np.max(a)),
                'mean'  : float(np.mean(a)),
                'std'   : float(np.std(a)),
                }
            logging.debug(var + ': ' + str(self.branches_info[var]))