overlap-yields

#!/usr/bin/env python

# -- Create by Quentin Buat quentin(dot)buat(at)cern(dot)ch
# create an ntuple containing trees for data events in the overlap between cutbased and mva
# with classifier scores
# and event weights
from array import array
from mva.cmd import get_parser

# args = get_parser(actions=False).parse_args()
args = get_parser(actions=False).parse_args()

from mva.analysis import get_analysis
from mva.plotting import draw_channel_array
from mva.variables import VARIABLES,WEIGHTS
from mva import log

analysis = get_analysis(args)
output_suffix = analysis.get_suffix()


from rootpy.io import root_open
from rootpy.tree import Cut
from rootpy.plotting import Hist,Canvas
from rootpy.interactive import wait

import ROOT
ROOT.gROOT.SetBatch(0)

import numpy as np
from numpy.lib import recfunctions
from root_numpy import array2tree


fields = [ 'RunNumber','EventNumber','mmc1_mass','weight']

target_region = args.target_region
data    = analysis.data
signals = analysis.signals
backgrounds = analysis.backgrounds

yields_info_master = {}


# ------------------
class sample_info:
    def __init__(self,name,events,color=None):
        self.name = name
        self.events = events 
        self.color  = color
# loop over the overlap categories (defined in mva/categories/__init__.py)
for category in analysis.iter_categories(
        args.categories, args.controls, names=args.category_names):

    if category.analysis_control:
        continue
    
    yields_info_cat = []
    

    # get the data record array
    rec = data.merged_records(
        category, target_region, fields=fields, include_weight=False)

    yields_info_cat += [ sample_info('data', len(rec)) ]
    yields_signal = 0.

    # Iterate over all the signals+background samples
    for sample in signals+backgrounds:
        # get the sample record array
        rec_sample = sample.merged_records( category,
                                            target_region,
                                            fields=fields,
                                            include_weight=True )
        yields_info_cat += [ sample_info(sample.name,rec_sample['weight'].sum()) ]
        if 'Signal' in sample.name:
            yields_signal += rec_sample['weight'].sum()
    yields_info_cat += [ sample_info('Signal_all', yields_signal)]
    yields_info_master[category.name] = yields_info_cat




for cat in yields_info_master:
    log.info( '=======================================')
    log.info( cat + ' category' )
    log.info( '=======================================')
    for s_i in yields_info_master[cat]:    
        log.info( s_i.name+'\t --> \t'+str(s_i.events) )