Exemplo n.º 1
0
def compute_num_denom(type, bin, lep_cat, iLep, lepStr, base_cut, WP_tag,
                      WP_probe, FR_array, config, bin_NP):

    #print 'Inside compute_num_denom'

    this_base_cut = copy.deepcopy(base_cut)

    this_num_denom = 0.0

    if FR_array[bin][0] == 1 and lep_cat == 'ele_ele':
        return 0.0
    if FR_array[bin][0] == 0 and lep_cat == 'mu_mu':
        return 0.0

    ## Apply cuts for this bin, probe lepton
    #this_base_cut.append(lepStr+str(iLep+1)+'_lepMVA > '+str(WP_probe))
    this_base_cut.append(lepStr + str(iLep + 1) + '_lepCut >= ' +
                         str(WP_probe))
    this_base_cut.append(lepStr + str(iLep + 1) + '_isMuon == ' +
                         str(FR_array[bin][0]))
    this_base_cut.append(lepStr + str(iLep + 1) + '_pt > ' +
                         str(FR_array[bin][1]))
    this_base_cut.append(lepStr + str(iLep + 1) + '_pt < ' +
                         str(FR_array[bin][2]))
    this_base_cut.append('abs(' + lepStr + str(iLep + 1) + '_eta) > ' +
                         str(FR_array[bin][3]))
    this_base_cut.append('abs(' + lepStr + str(iLep + 1) + '_eta) < ' +
                         str(FR_array[bin][4]))
    this_base_cut.append(lepStr + str(iLep + 1) + '_jetBTagCSV > ' +
                         str(FR_array[bin][5]))
    this_base_cut.append(lepStr + str(iLep + 1) + '_jetBTagCSV < ' +
                         str(FR_array[bin][6]))

    #data_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #MC_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #QF_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #NP_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA < '+str(WP_tag)]
    data_cut = [lepStr + str(abs(iLep - 1) + 1) + '_lepCut >= ' + str(WP_tag)]
    MC_cut = [lepStr + str(abs(iLep - 1) + 1) + '_lepCut >= ' + str(WP_tag)]
    QF_cut = [lepStr + str(abs(iLep - 1) + 1) + '_lepCut >= ' + str(WP_tag)]
    NP_cut = [lepStr + str(abs(iLep - 1) + 1) + '_lepCut < ' + str(WP_tag)]

    ## Apply cuts for this bin, tag lepton
    if type == 'NP_tag':
        tag_lepton_bin_cut = [
            lepStr + str(abs(iLep - 1) + 1) + '_isMuon == ' +
            str(FR_array[bin_NP][0])
        ]
        tag_lepton_bin_cut[0] += (' && ' + lepStr + str(abs(iLep - 1) + 1) +
                                  '_pt > ' + str(FR_array[bin_NP][1]))
        tag_lepton_bin_cut[0] += (' && ' + lepStr + str(abs(iLep - 1) + 1) +
                                  '_pt < ' + str(FR_array[bin_NP][2]))
        tag_lepton_bin_cut[0] += (' && abs(' + lepStr +
                                  str(abs(iLep - 1) + 1) + '_eta) > ' +
                                  str(FR_array[bin_NP][3]))
        tag_lepton_bin_cut[0] += (' && abs(' + lepStr +
                                  str(abs(iLep - 1) + 1) + '_eta) < ' +
                                  str(FR_array[bin_NP][4]))
        tag_lepton_bin_cut[0] += (' && ' + lepStr + str(abs(iLep - 1) + 1) +
                                  '_jetBTagCSV > ' + str(FR_array[bin_NP][5]))
        tag_lepton_bin_cut[0] += (' && ' + lepStr + str(abs(iLep - 1) + 1) +
                                  '_jetBTagCSV < ' + str(FR_array[bin_NP][6]))

    for sample in config['samples']:

        sample_dict = config['samples'][sample] if config['samples'][
            sample] else {
                'systematics': ['common'],
                'weights': ['common']
            }

        ## Get basic sample information
        tree_sample = sample_dict.get('tree sample', sample)
        additional_cuts = sample_dict.get('additional cuts', [])
        cuts_to_remove = sample_dict.get('cuts to remove', [])
        sample_info = plot_helper.SampleInformation(tree_sample)

        ## Only compute relevant samples
        #print 'sample '+sample+', type '+type+', sample_type '+sample_info.sample_type
        if type == 'data' and ((not sample_info.sample_type == 'data') or
                               ('sideband' in sample)):
            continue
        if type == 'MC' and not sample_info.sample_type == 'MC':
            continue
        if type == 'QF' and not 'QF' in sample:
            continue
        if 'NP' in type and not 'NP' in sample:
            continue
        #print 'Proceeding with calculation'

        ## Get tree file and summaryTree
        if lep_cat == 'mu_mu' or lep_cat == 'mu_ele' or lep_cat == 'ele_ele':
            source_file_name = '%s/%s_%s_all.root' % (
                config['input_trees_directory'], tree_sample,
                config['SS label'])
        elif lep_cat == '3l':
            source_file_name = '%s/%s_%s_all.root' % (
                config['input_trees_directory'], tree_sample,
                config['3l label'])
        source_file = ROOT.TFile(source_file_name)
        tree = source_file.Get('summaryTree')

        draw_string_maker = plot_helper.DrawStringMaker()

        ## Apply proper selection criteria
        if sample in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(
                this_base_cut, data_cut)
        elif sample_info.sample_type == 'MC':
            draw_string_maker.append_selection_requirements(
                this_base_cut, MC_cut)
        elif 'QF_sideband' in sample and sample.replace(
                '_QF_sideband',
                '') in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(
                this_base_cut, QF_cut)
        elif 'NP_sideband' in sample and sample.replace(
                '_NP_sideband',
                '') in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(
                this_base_cut, NP_cut)
        else:
            #print sample+' not included in '+lep_cat+' num_denom'
            continue
        if type == 'NP_tag':
            draw_string_maker.append_selection_requirements(tag_lepton_bin_cut)

        draw_string_maker.remove_selection_requirements(cuts_to_remove)
        draw_string_maker.append_selection_requirements(additional_cuts)

        ## Are the QF weights implemented properly?
        weights = plot_helper.customize_list(
            config['weights'], sample_dict.get('weights', ['common']))

        ## Apply MC weights
        weights_cat = ['1.0']
        if (lep_cat == 'mu_mu' or lep_cat == 'mu_ele'
                or lep_cat == 'ele_ele') and sample_info.sample_type == 'MC':
            weights_cat = config['weights SS']
        elif lep_cat == '3l' and sample_info.sample_type == 'MC':
            weights_cat = config['weights 3l']

        if sample_info.sample_type == 'MC' and 'triggerSF' in weights:
            matched_SF = draw_string_maker.get_matched_SF(lep_cat)
            weights = [matched_SF if x == 'triggerSF' else x for x in weights]
        #draw_string_maker.multiply_by_factors(weights, [systematic_weight_string])
        draw_string_maker.multiply_by_factors(weights, weights_cat)

        ########################################
        ## ADD IN EFF and (1/EFF) WEIGHTS FOR MC
        ########################################

        if sample_info.sample_type not in [
                'MC', 'data'
        ] and 'sideband' not in sample_info.sample_type:
            sys.exit('Invalid sample_type must be data, sideband, or MC' %
                     (sample_info.sample_type))

        this_plot_integral = 0.0

        plot = plot_helper.Plot(sample, 0, tree, 'distribution',
                                config['distribution'],
                                draw_string_maker.draw_string)
        #if sample in config['lepton categories'][lep_cat]['data samples'] and plot.plot.Integral() == 0:
        #print draw_string_maker.draw_string
        #print draw_string_maker.draw_string

        if sample_info.sample_type == 'MC':
            plot.plot.Scale(sample_info.x_section * config['luminosity'] /
                            sample_info.num_generated)

        this_plot_integral = plot.plot.Integral()

        #this_plot_integral += plot.plot.Integral()*FR_hist_tight_mu.GetBinContent(ptBin+1, etaBin+1, csvBin+1)
        #this_plot_integral += plot.plot.Integral()*FR_hist_tight_ele.GetBinContent(ptBin+1, etaBin+1, csvBin+1)

        #print sample+' integral for '+lep_cat+', '+bin+', '+bin_NP+', lepton '+str(iLep+1)+', WP('+str(WP_probe)+','+str(WP_tag)+') = '+str(this_plot_integral)

        this_num_denom += this_plot_integral
        ## End loop over samples

    print 'num_denom for ' + type + ', ' + lep_cat + ', WP_probe ' + str(
        WP_probe) + ', ' + bin + ', ' + bin_NP + ', lepton ' + str(
            iLep + 1) + ' = ' + str(this_num_denom)
    return this_num_denom
Exemplo n.º 2
0
def make_histos(args, config, samples, lepton_categories, jet_tag_categories):
    for sample, sample_dict in samples.items():
        tree_sample = sample_dict.get('tree sample', sample)
        additional_cuts = sample_dict.get('additional cuts', [])
        cuts_to_remove = sample_dict.get('cuts to remove', [])
        sample_info = plot_helper.SampleInformation(tree_sample)

        for lepton_category in lepton_categories:
            lepton_category_cut_strings = config['lepton categories'][lepton_category].get('cuts', {}).values()
            if sample_info.sample_type == 'data' or 'sideband' in sample_info.sample_type:
                if any([x==sample for x in config['lepton categories'][lepton_category].get('excluded samples', [])]):
                    config['weights'].append('0') #So we get empty histograms for hadding to get the inclusive category
                    #continue
                if not plot_helper.is_matching_data_sample(config['lepton categories'][lepton_category]['data samples'], sample):
                    continue

            for jet_tag_category, jet_tag_category_cut_strings in jet_tag_categories.items():
                systematics_list = plot_helper.customize_systematics(config['systematics'], sample_dict.get('systematics', 'common'))
                output_file_name = '%s/%s/%s_%s_%s_%s.root' % (config['output directory'], lepton_category, lepton_category, jet_tag_category, sample, config['output label'])
                if args.limits:
                    output_file_name = '%s/%s/%s_%s_%s_%s.root' % (config['limits output directory'], lepton_category, lepton_category, jet_tag_category, sample, config['output label'])
                    if config['limits skip systematics']:
                        systematics_list = ['nominal']
                elif config['skip systematics']:
                    systematics_list = ['nominal']

                output_file = ROOT.TFile(output_file_name, 'RECREATE')

                for systematic in systematics_list:
                    print 'Beginning next loop iteration. Sample: %10s Jet tag category: %-10s  Lepton category: %-10s Systematic: %-10s' % (sample, jet_tag_category, lepton_category, systematic)

                    systematic_weight_string, systematic_label = plot_helper.get_systematic_info(systematic)
                    source_file_name = '%s/%s_%s_all.root' % (config['input_trees_directory'], tree_sample, config['label'])
                    if 'JES' in systematic or 'JER' in systematic:
                        source_file_name = '%s/%s_%s_%s_all.root' % (config['input_trees_directory'], tree_sample, config['label'], systematic)
                    if args.file:
                        source_file_name = args.file
                    source_file = ROOT.TFile(source_file_name)
                    tree = source_file.Get('summaryTree')

                    draw_string_maker = plot_helper.DrawStringMaker()
                    draw_string_maker.append_selection_requirements(config['common cuts'].values(),
                                                                    lepton_category_cut_strings,
                                                                    jet_tag_category_cut_strings,
                                                                    additional_cuts) #additional_cuts is empty by default

                    draw_string_maker.remove_selection_requirements(cuts_to_remove)

                    if not args.no_weights:
                        weights = plot_helper.customize_list(config['weights'], sample_dict.get('weights', ['common']))
                        if sample_info.sample_type == 'MC' and 'triggerSF' in weights:
                            matched_SF = draw_string_maker.get_matched_SF(lepton_category)
                            weights = [matched_SF if x=='triggerSF' else x for x in weights]
                        draw_string_maker.multiply_by_factors(weights, [systematic_weight_string])

                    if sample_info.sample_type not in ['MC', 'data'] and 'sideband' not in sample_info.sample_type:
                        sys.exit('Invalid sample_type must be data, sideband, or MC' % (sample_info.sample_type))

                    config = plot_helper.append_integral_histo(config)
                    distribution_items = config['distributions'].items()
                    if args.limits:
                        distribution_items = config['limits distributions'].items()
                    for distribution, parameters in distribution_items:
                        if sample not in parameters.get('samples', [sample]):
                            continue
                        draw_string_maker.remove_selection_requirements(parameters.get('cuts to remove', []))
                        draw_string_maker.append_selection_requirements(parameters.get('additional cuts', []))
                        plot_name = '%s%s' % (distribution, systematic_label)
                        plot = plot_helper.Plot(sample, output_file, tree, plot_name, parameters, draw_string_maker.draw_string)
                        if sample_info.sample_type == 'MC':
                            plot.plot.Scale(sample_info.x_section * config['luminosity'] / sample_info.num_generated)
                        output_file.Write()
                        if args.pdf:
                            plot.save_image('pdf')
                        if args.web:
                            plot.post_to_web(config, lepton_category)
                    source_file.Close() #end systematic
                config_file = ROOT.TObjString(args.config_file_name)
                output_file.cd()
                config_file.Write('config_file')
                output_file.Close() #end jet tag category
Exemplo n.º 3
0
def make_histos(args, config, samples, lepton_categories, jet_tag_categories):
    for sample, sample_dict in samples.items():
        tree_sample = sample_dict.get('tree sample', sample)
        additional_cuts = sample_dict.get('additional cuts', [])
        cuts_to_remove = sample_dict.get('cuts to remove', [])
        sample_info = plot_helper.SampleInformation(tree_sample)

        for lepton_category in lepton_categories:
            lepton_category_cut_strings = config['lepton categories'][
                lepton_category].get('cuts', {}).values()
            if sample_info.sample_type == 'data' or 'sideband' in sample_info.sample_type:
                if any([
                        x == sample for x in config['lepton categories']
                    [lepton_category].get('excluded samples', [])
                ]):
                    config['weights'].append(
                        '0'
                    )  #So we get empty histograms for hadding to get the inclusive category
                    #continue
                if not plot_helper.is_matching_data_sample(
                        config['lepton categories'][lepton_category]
                    ['data samples'], sample):
                    continue

            for jet_tag_category, jet_tag_category_cut_strings in jet_tag_categories.items(
            ):
                systematics_list = plot_helper.customize_systematics(
                    config['systematics'],
                    sample_dict.get('systematics', 'common'))
                output_file_name = '%s/%s/%s_%s_%s_%s.root' % (
                    config['output directory'], lepton_category,
                    lepton_category, jet_tag_category, sample,
                    config['output label'])
                if args.limits:
                    output_file_name = '%s/%s/%s_%s_%s_%s.root' % (
                        config['limits output directory'], lepton_category,
                        lepton_category, jet_tag_category, sample,
                        config['output label'])
                    if config['limits skip systematics']:
                        systematics_list = ['nominal']
                elif config['skip systematics']:
                    systematics_list = ['nominal']

                output_file = ROOT.TFile(output_file_name, 'RECREATE')

                for systematic in systematics_list:
                    print 'Beginning next loop iteration. Sample: %10s Jet tag category: %-10s  Lepton category: %-10s Systematic: %-10s' % (
                        sample, jet_tag_category, lepton_category, systematic)

                    systematic_weight_string, systematic_label = plot_helper.get_systematic_info(
                        systematic)
                    source_file_name = '%s/%s_%s_all.root' % (
                        config['input_trees_directory'], tree_sample,
                        config['label'])
                    if 'JES' in systematic or 'JER' in systematic:
                        source_file_name = '%s/%s_%s_%s_all.root' % (
                            config['input_trees_directory'], tree_sample,
                            config['label'], systematic)
                    if args.file:
                        source_file_name = args.file
                    source_file = ROOT.TFile(source_file_name)
                    tree = source_file.Get('summaryTree')

                    draw_string_maker = plot_helper.DrawStringMaker()
                    draw_string_maker.append_selection_requirements(
                        config['common cuts'].values(),
                        lepton_category_cut_strings,
                        jet_tag_category_cut_strings,
                        additional_cuts)  #additional_cuts is empty by default

                    draw_string_maker.remove_selection_requirements(
                        cuts_to_remove)

                    if not args.no_weights:
                        weights = plot_helper.customize_list(
                            config['weights'],
                            sample_dict.get('weights', ['common']))
                        if sample_info.sample_type == 'MC' and 'triggerSF' in weights:
                            matched_SF = draw_string_maker.get_matched_SF(
                                lepton_category)
                            weights = [
                                matched_SF if x == 'triggerSF' else x
                                for x in weights
                            ]
                        draw_string_maker.multiply_by_factors(
                            weights, [systematic_weight_string])

                    if sample_info.sample_type not in [
                            'MC', 'data'
                    ] and 'sideband' not in sample_info.sample_type:
                        sys.exit(
                            'Invalid sample_type must be data, sideband, or MC'
                            % (sample_info.sample_type))

                    config = plot_helper.append_integral_histo(config)
                    distribution_items = config['distributions'].items()
                    if args.limits:
                        distribution_items = config[
                            'limits distributions'].items()
                    for distribution, parameters in distribution_items:
                        if sample not in parameters.get('samples', [sample]):
                            continue
                        draw_string_maker.remove_selection_requirements(
                            parameters.get('cuts to remove', []))
                        draw_string_maker.append_selection_requirements(
                            parameters.get('additional cuts', []))
                        plot_name = '%s%s' % (distribution, systematic_label)
                        plot = plot_helper.Plot(sample, output_file, tree,
                                                plot_name, parameters,
                                                draw_string_maker.draw_string)
                        if sample_info.sample_type == 'MC':
                            plot.plot.Scale(sample_info.x_section *
                                            config['luminosity'] /
                                            sample_info.num_generated)
                        output_file.Write()
                        if args.pdf:
                            plot.save_image('pdf')
                        if args.web:
                            plot.post_to_web(config, lepton_category)
                    source_file.Close()  #end systematic
                config_file = ROOT.TObjString(args.config_file_name)
                output_file.cd()
                config_file.Write('config_file')
                output_file.Close()  #end jet tag category
def compute_num_denom(bin, lep_cat, iLep, lepStr, base_cut, WP_tag, WP_probe, FR_array, FR_hist_tight_mu, FR_hist_tight_ele, config):

    #print 'Inside compute_num_denom'

    this_base_cut = copy.deepcopy(base_cut)

    this_num_denom = 0.0

    if FR_array[bin][0] == 1 and lep_cat == 'ele_ele':
        return 0.0
    if FR_array[bin][0] == 0 and lep_cat == 'mu_mu':
        return 0.0

    #this_base_cut.append(lepStr+str(iLep+1)+'_lepMVA > '+str(WP_probe))
    this_base_cut.append(lepStr+str(iLep+1)+'_lepCut >= '+str(WP_probe))
    this_base_cut.append(lepStr+str(iLep+1)+'_isMuon == '+str(FR_array[bin][0]))
    this_base_cut.append(lepStr+str(iLep+1)+'_pt > '+str(FR_array[bin][1]))
    this_base_cut.append(lepStr+str(iLep+1)+'_pt < '+str(FR_array[bin][2]))
    this_base_cut.append('abs('+lepStr+str(iLep+1)+'_eta) > '+str(FR_array[bin][3]))
    this_base_cut.append('abs('+lepStr+str(iLep+1)+'_eta) < '+str(FR_array[bin][4]))
    this_base_cut.append(lepStr+str(iLep+1)+'_jetBTagCSV > '+str(FR_array[bin][5]))
    this_base_cut.append(lepStr+str(iLep+1)+'_jetBTagCSV < '+str(FR_array[bin][6]))

    #data_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #MC_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #QF_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #NP_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA < '+str(WP_tag)]
    data_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut >= '+str(WP_tag)]
    MC_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut >= '+str(WP_tag)]
    QF_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut >= '+str(WP_tag)]
    NP_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut < '+str(WP_tag)]

    these_integrals = {}
    pool = multiprocessing.Pool(processes=10)
    work = []
    
    for sample in config['samples']:

        sample_dict = config['samples'][sample] if config['samples'][sample] else {'systematics':['common'], 'weights':['common']}
        
        tree_sample = sample_dict.get('tree sample', sample)
        additional_cuts = sample_dict.get('additional cuts', [])
        cuts_to_remove = sample_dict.get('cuts to remove', [])
        sample_info = plot_helper.SampleInformation(tree_sample)
        
        if lep_cat == 'mu_mu' or lep_cat == 'mu_ele' or lep_cat == 'ele_ele':
            source_file_name = '%s/%s_%s_all.root' % (config['input_trees_directory'], tree_sample, config['SS label'])
        elif lep_cat == '3l':
            source_file_name = '%s/%s_%s_all.root' % (config['input_trees_directory'], tree_sample, config['3l label'])
        source_file = ROOT.TFile(source_file_name)
        tree = source_file.Get('summaryTree')

        draw_string_maker = plot_helper.DrawStringMaker()

        if sample in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(this_base_cut, data_cut)
        elif sample_info.sample_type == 'MC':
            draw_string_maker.append_selection_requirements(this_base_cut, MC_cut)
        elif 'QF_sideband' in sample and sample.replace('_QF_sideband','') in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(this_base_cut, QF_cut)
        elif 'NP_sideband' in sample and sample.replace('_NP_sideband','') in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(this_base_cut, NP_cut)
        else:
            #print sample+' not included in '+lep_cat+' num_denom'
            continue

        draw_string_maker.remove_selection_requirements(cuts_to_remove)
        draw_string_maker.append_selection_requirements(additional_cuts)

        ## Are the QF weights implemented properly?
        weights = plot_helper.customize_list(config['weights'], sample_dict.get('weights', ['common']))

        weights_cat = ['1.0']
        if (lep_cat == 'mu_mu' or lep_cat == 'mu_ele' or lep_cat == 'ele_ele') and sample_info.sample_type == 'MC':
            weights_cat = config['weights SS']
        elif lep_cat == '3l' and sample_info.sample_type == 'MC':
            weights_cat = config['weights 3l']

        if sample_info.sample_type == 'MC' and 'triggerSF' in weights:
            matched_SF = draw_string_maker.get_matched_SF(lep_cat)
            weights = [matched_SF if x=='triggerSF' else x for x in weights]
        #draw_string_maker.multiply_by_factors(weights, [systematic_weight_string])
        draw_string_maker.multiply_by_factors(weights, weights_cat)

        if sample_info.sample_type not in ['MC', 'data'] and 'sideband' not in sample_info.sample_type:
            sys.exit('Invalid sample_type must be data, sideband, or MC' % (sample_info.sample_type))

        #this_plot_integral = 0.0

        if 'NP_sideband' in sample:
            #print sample
            for isMuon in range(2):
                if isMuon == 1:
                    for ptBin in range(FR_hist_tight_mu.GetNbinsX()):
                        for etaBin in range(FR_hist_tight_mu.GetNbinsY()):
                            for csvBin in range(FR_hist_tight_mu.GetNbinsZ()):
                                tag_lepton_bin_cut = [lepStr+str(abs(iLep-1)+1)+'_isMuon == '+str(isMuon)]
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_pt > '+str(FR_hist_tight_mu.GetXaxis().GetBinLowEdge(ptBin+1)))
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_pt < '+str(FR_hist_tight_mu.GetXaxis().GetBinLowEdge(ptBin+1)+FR_hist_tight_mu.GetXaxis().GetBinWidth(ptBin+1)))
                                tag_lepton_bin_cut[0] += (' && abs('+lepStr+str(abs(iLep-1)+1)+'_eta) > '+str(FR_hist_tight_mu.GetYaxis().GetBinLowEdge(etaBin+1)))
                                tag_lepton_bin_cut[0] += (' && abs('+lepStr+str(abs(iLep-1)+1)+'_eta) < '+str(FR_hist_tight_mu.GetYaxis().GetBinLowEdge(etaBin+1)+FR_hist_tight_mu.GetYaxis().GetBinWidth(etaBin+1)))
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_jetBTagCSV > '+str(FR_hist_tight_mu.GetZaxis().GetBinLowEdge(csvBin+1)))
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_jetBTagCSV < '+str(FR_hist_tight_mu.GetZaxis().GetBinLowEdge(csvBin+1)+FR_hist_tight_mu.GetZaxis().GetBinWidth(csvBin+1)))
                                #print tag_lepton_bin_cut
                                #print FR_hist_tight_mu.GetBinContent(ptBin+1, etaBin+1)
                                draw_string_maker.append_selection_requirements(tag_lepton_bin_cut)
                                these_integrals['%s_%d_%d_%d_%d' % (sample, isMuon, ptBin, etaBin, csvBin)] = -99.0
                                
#                                 output_file = ROOT.TFile('temp_'+str(FR_hist_tight_mu.GetNbinsX())+'_'+str(FR_hist_tight_mu.GetNbinsY())+'_'+str(FR_hist_tight_mu.GetNbinsZ())+'.root', 'RECREATE')
#                                 plot = plot_helper.Plot(sample, output_file, tree, 'distribution', config['distribution'], draw_string_maker.draw_string)
#                                 #print draw_string_maker.draw_string
#                                 #print plot.plot.Integral()
#                                 this_plot_integral += plot.plot.Integral()*FR_hist_tight_mu.GetBinContent(ptBin+1, etaBin+1, csvBin+1)

                                thread_index = FR_hist_tight_mu.GetNbinsX()*100 + FR_hist_tight_mu.GetNbinsY()*10 + FR_hist_tight_mu.GetNbinsZ()
                                scale = FR_hist_tight_mu.GetBinContent(ptBin+1, etaBin+1, csvBin+1)
                                #these_integrals['%s_%d_%d_%d_%d' % (sample, isMuon, ptBin, etaBin, csvBin)] = get_one_integral(thread_index, sample, tree, config['distribution'], draw_string_maker.draw_string, scale)
                                work.append( ('%s_%d_%d_%d_%d' % (sample, isMuon, ptBin, etaBin, csvBin), pool.apply_async(get_one_integral, [thread_index, sample, config['distribution'], draw_string_maker.draw_string, scale, source_file_name]) ) )
                                
                                draw_string_maker.remove_selection_requirements(tag_lepton_bin_cut)
                                
#                                 output_file.Close()

                else:
                    for ptBin in range(FR_hist_tight_ele.GetNbinsX()):
                        for etaBin in range(FR_hist_tight_ele.GetNbinsY()):
                            for csvBin in range(FR_hist_tight_ele.GetNbinsZ()):
                                tag_lepton_bin_cut = [lepStr+str(abs(iLep-1)+1)+'_isMuon == '+str(isMuon)]
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_pt > '+str(FR_hist_tight_ele.GetXaxis().GetBinLowEdge(ptBin+1)))
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_pt < '+str(FR_hist_tight_ele.GetXaxis().GetBinLowEdge(ptBin+1)+FR_hist_tight_ele.GetXaxis().GetBinWidth(ptBin+1)))
                                tag_lepton_bin_cut[0] += (' && abs('+lepStr+str(abs(iLep-1)+1)+'_eta) > '+str(FR_hist_tight_ele.GetYaxis().GetBinLowEdge(etaBin+1)))
                                tag_lepton_bin_cut[0] += (' && abs('+lepStr+str(abs(iLep-1)+1)+'_eta) < '+str(FR_hist_tight_ele.GetYaxis().GetBinLowEdge(etaBin+1)+FR_hist_tight_ele.GetYaxis().GetBinWidth(etaBin+1)))
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_jetBTagCSV > '+str(FR_hist_tight_ele.GetZaxis().GetBinLowEdge(csvBin+1)))
                                tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_jetBTagCSV < '+str(FR_hist_tight_ele.GetZaxis().GetBinLowEdge(csvBin+1)+FR_hist_tight_ele.GetZaxis().GetBinWidth(csvBin+1)))
                                #print tag_lepton_bin_cut
                                #print FR_hist_tight_ele.GetBinContent(ptBin+1, etaBin+1)
                                draw_string_maker.append_selection_requirements(tag_lepton_bin_cut)
                                these_integrals['%s_%d_%d_%d_%d' % (sample, isMuon, ptBin, etaBin, csvBin)] = -99.0

#                                 output_file = ROOT.TFile('temp_'+str(FR_hist_tight_mu.GetNbinsX())+'_'+str(FR_hist_tight_mu.GetNbinsY())+'_'+str(FR_hist_tight_mu.GetNbinsZ())+'.root', 'RECREATE')
#                                 plot = plot_helper.Plot(sample, output_file, tree, 'distribution', config['distribution'], draw_string_maker.draw_string)
#                                 #print draw_string_maker.draw_string
#                                 #print plot.plot.Integral()
#                                 this_plot_integral += plot.plot.Integral()*FR_hist_tight_ele.GetBinContent(ptBin+1, etaBin+1, csvBin+1)

                                thread_index = FR_hist_tight_ele.GetNbinsX()*100 + FR_hist_tight_ele.GetNbinsY()*10 + FR_hist_tight_ele.GetNbinsZ()
                                scale = FR_hist_tight_ele.GetBinContent(ptBin+1, etaBin+1, csvBin+1)
                                #these_integrals['%s_%d_%d_%d_%d' % (sample, isMuon, ptBin, etaBin, csvBin)] = get_one_integral(thread_index, sample, tree, config['distribution'], draw_string_maker.draw_string, scale)
                                work.append( ('%s_%d_%d_%d_%d' % (sample, isMuon, ptBin, etaBin, csvBin), pool.apply_async(get_one_integral, [thread_index, sample, config['distribution'], draw_string_maker.draw_string, scale, source_file_name]) ) )
                                
                                draw_string_maker.remove_selection_requirements(tag_lepton_bin_cut)
                                
#                                 output_file.Close()
                            
        else:

            these_integrals[sample] = -99.0
            
#             output_file = ROOT.TFile('temp_'+str(FR_hist_tight_mu.GetNbinsX())+'_'+str(FR_hist_tight_mu.GetNbinsY())+'_'+str(FR_hist_tight_mu.GetNbinsZ())+'.root', 'RECREATE')
#             plot = plot_helper.Plot(sample, output_file, tree, 'distribution', config['distribution'], draw_string_maker.draw_string)
#             #if sample in config['lepton categories'][lep_cat]['data samples'] and plot.plot.Integral() == 0:
#                 #print draw_string_maker.draw_string
#             #print draw_string_maker.draw_string
#             if sample_info.sample_type == 'MC':
#                 plot.plot.Scale(sample_info.x_section * config['luminosity'] / sample_info.num_generated)
#             this_plot_integral = plot.plot.Integral()

            thread_index = FR_hist_tight_mu.GetNbinsX()*100 + FR_hist_tight_mu.GetNbinsY()*10 + FR_hist_tight_mu.GetNbinsZ()
            scale = 1.0
            if sample_info.sample_type == 'MC':
                scale = sample_info.x_section * config['luminosity'] / sample_info.num_generated
            #these_integrals[sample] = get_one_integral(thread_index, sample, tree, config['distribution'], draw_string_maker.draw_string, scale)
            work.append( (sample, pool.apply_async(get_one_integral, [thread_index, sample, config['distribution'], draw_string_maker.draw_string, scale, source_file_name]) ) )
            #for (id, res) in work:
                #these_integrals[id] = res.get()
                #print 'id %s in work yields %f' % (id, these_integrals[id])

#             output_file.Close()

        #print sample+' integral for '+lep_cat+', '+bin+', lepton '+str(iLep+1)+' = '+str(this_plot_integral)

#         if sample in config['lepton categories'][lep_cat]['data samples']:
#             this_num_denom += this_plot_integral
#         else:
#             this_num_denom -= this_plot_integral

    for (id, res) in work:
        these_integrals[id] = res.get()

    #pool.join()
    #pool.close()
    
    for sample in these_integrals:
        #print '%s integral is %f' % (sample, these_integrals[sample])
        if sample in config['lepton categories'][lep_cat]['data samples']:
            this_num_denom += these_integrals[sample]
        else:
            this_num_denom -= these_integrals[sample]

    print 'num_denom for '+lep_cat+', WP_probe '+str(WP_probe)+', '+bin+', lepton '+str(iLep+1)+' = '+str(this_num_denom)
    return this_num_denom
def compute_num_denom(type, bin, lep_cat, iLep, lepStr, base_cut, WP_tag, WP_probe, FR_array, config, bin_NP):

    #print 'Inside compute_num_denom'

    this_base_cut = copy.deepcopy(base_cut)

    this_num_denom = 0.0

    if FR_array[bin][0] == 1 and lep_cat == 'ele_ele':
        return 0.0
    if FR_array[bin][0] == 0 and lep_cat == 'mu_mu':
        return 0.0

    ## Apply cuts for this bin, probe lepton
    #this_base_cut.append(lepStr+str(iLep+1)+'_lepMVA > '+str(WP_probe))
    this_base_cut.append(lepStr+str(iLep+1)+'_lepCut >= '+str(WP_probe))
    this_base_cut.append(lepStr+str(iLep+1)+'_isMuon == '+str(FR_array[bin][0]))
    this_base_cut.append(lepStr+str(iLep+1)+'_pt > '+str(FR_array[bin][1]))
    this_base_cut.append(lepStr+str(iLep+1)+'_pt < '+str(FR_array[bin][2]))
    this_base_cut.append('abs('+lepStr+str(iLep+1)+'_eta) > '+str(FR_array[bin][3]))
    this_base_cut.append('abs('+lepStr+str(iLep+1)+'_eta) < '+str(FR_array[bin][4]))
    this_base_cut.append(lepStr+str(iLep+1)+'_jetBTagCSV > '+str(FR_array[bin][5]))
    this_base_cut.append(lepStr+str(iLep+1)+'_jetBTagCSV < '+str(FR_array[bin][6]))

    #data_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #MC_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #QF_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA > '+str(WP_tag)]
    #NP_cut = [lepStr+str(abs(iLep-1)+1)+'_lepMVA < '+str(WP_tag)]
    data_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut >= '+str(WP_tag)]
    MC_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut >= '+str(WP_tag)]
    QF_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut >= '+str(WP_tag)]
    NP_cut = [lepStr+str(abs(iLep-1)+1)+'_lepCut < '+str(WP_tag)]

    ## Apply cuts for this bin, tag lepton
    if type == 'NP_tag':
        tag_lepton_bin_cut = [lepStr+str(abs(iLep-1)+1)+'_isMuon == '+str(FR_array[bin_NP][0])]
        tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_pt > '+str(FR_array[bin_NP][1]))
        tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_pt < '+str(FR_array[bin_NP][2]))
        tag_lepton_bin_cut[0] += (' && abs('+lepStr+str(abs(iLep-1)+1)+'_eta) > '+str(FR_array[bin_NP][3]))
        tag_lepton_bin_cut[0] += (' && abs('+lepStr+str(abs(iLep-1)+1)+'_eta) < '+str(FR_array[bin_NP][4]))
        tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_jetBTagCSV > '+str(FR_array[bin_NP][5]))
        tag_lepton_bin_cut[0] += (' && '+lepStr+str(abs(iLep-1)+1)+'_jetBTagCSV < '+str(FR_array[bin_NP][6]))

    for sample in config['samples']:

        sample_dict = config['samples'][sample] if config['samples'][sample] else {'systematics':['common'], 'weights':['common']}

        ## Get basic sample information
        tree_sample = sample_dict.get('tree sample', sample)
        additional_cuts = sample_dict.get('additional cuts', [])
        cuts_to_remove = sample_dict.get('cuts to remove', [])
        sample_info = plot_helper.SampleInformation(tree_sample)

        ## Only compute relevant samples
        #print 'sample '+sample+', type '+type+', sample_type '+sample_info.sample_type
        if type == 'data' and ((not sample_info.sample_type == 'data') or ('sideband' in sample)):
            continue
        if type == 'MC' and not sample_info.sample_type == 'MC':
            continue
        if type == 'QF' and not 'QF' in sample:
            continue
        if 'NP' in type and not 'NP' in sample:
            continue
        #print 'Proceeding with calculation'
        
        ## Get tree file and summaryTree
        if lep_cat == 'mu_mu' or lep_cat == 'mu_ele' or lep_cat == 'ele_ele':
            source_file_name = '%s/%s_%s_all.root' % (config['input_trees_directory'], tree_sample, config['SS label'])
        elif lep_cat == '3l':
            source_file_name = '%s/%s_%s_all.root' % (config['input_trees_directory'], tree_sample, config['3l label'])
        source_file = ROOT.TFile(source_file_name)
        tree = source_file.Get('summaryTree')

        draw_string_maker = plot_helper.DrawStringMaker()

        ## Apply proper selection criteria
        if sample in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(this_base_cut, data_cut)
        elif sample_info.sample_type == 'MC':
            draw_string_maker.append_selection_requirements(this_base_cut, MC_cut)
        elif 'QF_sideband' in sample and sample.replace('_QF_sideband','') in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(this_base_cut, QF_cut)
        elif 'NP_sideband' in sample and sample.replace('_NP_sideband','') in config['lepton categories'][lep_cat]['data samples']:
            draw_string_maker.append_selection_requirements(this_base_cut, NP_cut)
        else:
            #print sample+' not included in '+lep_cat+' num_denom'
            continue
        if type == 'NP_tag':
            draw_string_maker.append_selection_requirements(tag_lepton_bin_cut)

        draw_string_maker.remove_selection_requirements(cuts_to_remove)
        draw_string_maker.append_selection_requirements(additional_cuts)

        ## Are the QF weights implemented properly?
        weights = plot_helper.customize_list(config['weights'], sample_dict.get('weights', ['common']))

        ## Apply MC weights
        weights_cat = ['1.0']
        if (lep_cat == 'mu_mu' or lep_cat == 'mu_ele' or lep_cat == 'ele_ele') and sample_info.sample_type == 'MC':
            weights_cat = config['weights SS']
        elif lep_cat == '3l' and sample_info.sample_type == 'MC':
            weights_cat = config['weights 3l']

        if sample_info.sample_type == 'MC' and 'triggerSF' in weights:
            matched_SF = draw_string_maker.get_matched_SF(lep_cat)
            weights = [matched_SF if x=='triggerSF' else x for x in weights]
        #draw_string_maker.multiply_by_factors(weights, [systematic_weight_string])
        draw_string_maker.multiply_by_factors(weights, weights_cat)

        ########################################
        ## ADD IN EFF and (1/EFF) WEIGHTS FOR MC
        ########################################

        if sample_info.sample_type not in ['MC', 'data'] and 'sideband' not in sample_info.sample_type:
            sys.exit('Invalid sample_type must be data, sideband, or MC' % (sample_info.sample_type))

        this_plot_integral = 0.0

        plot = plot_helper.Plot(sample, 0, tree, 'distribution', config['distribution'], draw_string_maker.draw_string)
        #if sample in config['lepton categories'][lep_cat]['data samples'] and plot.plot.Integral() == 0:
            #print draw_string_maker.draw_string
        #print draw_string_maker.draw_string

        if sample_info.sample_type == 'MC':
            plot.plot.Scale(sample_info.x_section * config['luminosity'] / sample_info.num_generated)

        this_plot_integral = plot.plot.Integral()

        #this_plot_integral += plot.plot.Integral()*FR_hist_tight_mu.GetBinContent(ptBin+1, etaBin+1, csvBin+1)
        #this_plot_integral += plot.plot.Integral()*FR_hist_tight_ele.GetBinContent(ptBin+1, etaBin+1, csvBin+1)

        #print sample+' integral for '+lep_cat+', '+bin+', '+bin_NP+', lepton '+str(iLep+1)+', WP('+str(WP_probe)+','+str(WP_tag)+') = '+str(this_plot_integral)

        this_num_denom += this_plot_integral
        ## End loop over samples

    print 'num_denom for '+type+', '+lep_cat+', WP_probe '+str(WP_probe)+', '+bin+', '+bin_NP+', lepton '+str(iLep+1)+' = '+str(this_num_denom)
    return this_num_denom