def make_correlation_plot_from_file( channel, variable, fit_variables, CoM, title, x_title, y_title, x_limits, y_limits, rebin = 1, save_folder = 'plots/fitchecks/', save_as = ['pdf', 'png'] ):
# global b_tag_bin
    parameters = ["TTJet", "SingleTop", "V+Jets", "QCD"]
    parameters_latex = []
    for template in parameters:
        parameters_latex.append(samples_latex[template])
        
    input_file = open( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), "r" )
    # cycle through the lines in the file
    for line_number, line in enumerate( input_file ):
        # for now, only make plots for the fits for the central measurement
        if "central" in line:
            # matrix we want begins 11 lines below the line with the measurement ("central")
            line_number = line_number + 11
            break
    input_file.close()
    
    #Note: For some reason, the fit outputs the correlation matrix with the templates in the following order:
    #parameter1: QCD
    #parameter2: SingleTop
    #parameter3: TTJet
    #parameter4: V+Jets
        
    for variable_bin in variable_bins_ROOT[variable]:
        weights = {}
        if channel == 'electron':
            #formula to calculate the number of lines below "central" to access in each loop
            number_of_lines_down = (variable_bins_ROOT[variable].index( variable_bin ) * 12)

            #Get QCD correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down )
            weights["QCD_QCD"] = matrix_line.split()[2]
            weights["QCD_SingleTop"] = matrix_line.split()[3]
            weights["QCD_TTJet"] = matrix_line.split()[4]
            weights["QCD_V+Jets"] = matrix_line.split()[5]

            #Get SingleTop correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down + 1 )
            weights["SingleTop_QCD"] = matrix_line.split()[2]
            weights["SingleTop_SingleTop"] = matrix_line.split()[3]
            weights["SingleTop_TTJet"] = matrix_line.split()[4]
            weights["SingleTop_V+Jets"] = matrix_line.split()[5]

            #Get TTJet correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down + 2 )
            weights["TTJet_QCD"] = matrix_line.split()[2]
            weights["TTJet_SingleTop"] = matrix_line.split()[3]            
            weights["TTJet_TTJet"] = matrix_line.split()[4]
            weights["TTJet_V+Jets"] = matrix_line.split()[5]

            #Get V+Jets correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down + 3 )
            weights["V+Jets_QCD"] = matrix_line.split()[2]
            weights["V+Jets_SingleTop"] = matrix_line.split()[3]
            weights["V+Jets_TTJet"] = matrix_line.split()[4]
            weights["V+Jets_V+Jets"] = matrix_line.split()[5]

        if channel == 'muon':
            #formula to calculate the number of lines below "central" to access in each bin loop
            number_of_lines_down =  ( len( variable_bins_ROOT [variable] ) * 12 ) + ( variable_bins_ROOT[variable].index( variable_bin ) * 12 )
            
            #Get QCD correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down )
            weights["QCD_QCD"] = matrix_line.split()[2]
            weights["QCD_SingleTop"] = matrix_line.split()[3]
            weights["QCD_TTJet"] = matrix_line.split()[4]
            weights["QCD_V+Jets"] = matrix_line.split()[5]

            #Get SingleTop correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down + 1 )
            weights["SingleTop_QCD"] = matrix_line.split()[2]
            weights["SingleTop_SingleTop"] = matrix_line.split()[3]
            weights["SingleTop_TTJet"] = matrix_line.split()[4]
            weights["SingleTop_V+Jets"] = matrix_line.split()[5]

            #Get TTJet correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down + 2 )
            weights["TTJet_QCD"] = matrix_line.split()[2]
            weights["TTJet_SingleTop"] = matrix_line.split()[3]
            weights["TTJet_TTJet"] = matrix_line.split()[4]
            weights["TTJet_V+Jets"] = matrix_line.split()[5]
            
            #Get V+Jets correlations
            matrix_line = linecache.getline( "logs/01_%s_fit_%dTeV_%s.log" % ( variable, CoM, fit_variables ), line_number + number_of_lines_down + 3 )
            weights["V+Jets_QCD"] = matrix_line.split()[2]
            weights["V+Jets_SingleTop"] = matrix_line.split()[3]
            weights["V+Jets_TTJet"] = matrix_line.split()[4]
            weights["V+Jets_V+Jets"] = matrix_line.split()[5]

        #Create histogram
        histogram_properties = Histogram_properties()
        histogram_properties.title = title
        histogram_properties.name = 'Correlations_' + channel + '_' + variable + '_' + variable_bin
        histogram_properties.y_axis_title = y_title
        histogram_properties.x_axis_title = x_title
        histogram_properties.y_limits = y_limits
        histogram_properties.x_limits = x_limits
        histogram_properties.mc_error = 0.0
        histogram_properties.legend_location = 'upper right'

        #initialise 2D histogram
        a = Hist2D( 4, 0, 4, 4, 0, 4 )
        #fill histogram
        for i in range( len( parameters ) ):
            for j in range( len( parameters ) ):
                a.fill( float( i ), float( j ), float( weights["%s_%s" % ( parameters[i], parameters[j] )] ) )
        # create figure
        plt.figure( figsize = CMS.figsize, dpi = CMS.dpi, facecolor = CMS.facecolor )
        # make subplot(?) 
        fig, ax = plt.subplots( nrows = 1, ncols = 1 )
        rplt.hist2d( a )
        plt.subplots_adjust( right = 0.8 )

        #Set labels and formats for titles and axes
        plt.ylabel( histogram_properties.y_axis_title )
        plt.xlabel( histogram_properties.x_axis_title )
        plt.title( histogram_properties.title )
        x_limits = histogram_properties.x_limits
        y_limits = histogram_properties.y_limits
        ax.set_xticklabels( parameters_latex )
        ax.set_yticklabels( parameters_latex )
        ax.set_xticks( [0.5, 1.5, 2.5, 3.5] )
        ax.set_yticks( [0.5, 1.5, 2.5, 3.5] )
        plt.setp( ax.get_xticklabels(), visible = True )
        plt.setp( ax.get_yticklabels(), visible = True )

        #create and draw colour bar to the right of the main plot
        im = rplt.imshow( a, axes = ax, vmin = -1.0, vmax = 1.0 )
        #set location and dimensions (left, lower, width, height)
        cbar_ax = fig.add_axes( [0.85, 0.10, 0.05, 0.8] )
        fig.colorbar( im, cax = cbar_ax )

        for xpoint in range( len( parameters ) ):
            for ypoint in range( len( parameters ) ):
                correlation_value = weights["%s_%s" % ( parameters[xpoint], parameters[ypoint] )]
                ax.annotate( correlation_value, xy = ( xpoint + 0.5, ypoint + 0.5 ), ha = 'center', va = 'center', bbox = dict( fc = 'white', ec = 'none' ) )
        for save in save_as:
            plt.savefig( save_folder + histogram_properties.name + '.' + save )
        plt.close(fig)
    plt.close('all')
Ejemplo n.º 2
0
def make_correlation_plot_from_file(channel,
                                    variable,
                                    fit_variables,
                                    CoM,
                                    title,
                                    x_title,
                                    y_title,
                                    x_limits,
                                    y_limits,
                                    rebin=1,
                                    save_folder='plots/fitchecks/',
                                    save_as=['pdf', 'png']):
    # global b_tag_bin
    parameters = ["TTJet", "SingleTop", "V+Jets", "QCD"]
    parameters_latex = []
    for template in parameters:
        parameters_latex.append(samples_latex[template])

    input_file = open(
        "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables), "r")
    # cycle through the lines in the file
    for line_number, line in enumerate(input_file):
        # for now, only make plots for the fits for the central measurement
        if "central" in line:
            # matrix we want begins 11 lines below the line with the measurement ("central")
            line_number = line_number + 11
            break
    input_file.close()

    #Note: For some reason, the fit outputs the correlation matrix with the templates in the following order:
    #parameter1: QCD
    #parameter2: SingleTop
    #parameter3: TTJet
    #parameter4: V+Jets

    for variable_bin in variable_bins_ROOT[variable]:
        weights = {}
        if channel == 'electron':
            #formula to calculate the number of lines below "central" to access in each loop
            number_of_lines_down = (
                variable_bins_ROOT[variable].index(variable_bin) * 12)

            #Get QCD correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down)
            weights["QCD_QCD"] = matrix_line.split()[2]
            weights["QCD_SingleTop"] = matrix_line.split()[3]
            weights["QCD_TTJet"] = matrix_line.split()[4]
            weights["QCD_V+Jets"] = matrix_line.split()[5]

            #Get SingleTop correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down + 1)
            weights["SingleTop_QCD"] = matrix_line.split()[2]
            weights["SingleTop_SingleTop"] = matrix_line.split()[3]
            weights["SingleTop_TTJet"] = matrix_line.split()[4]
            weights["SingleTop_V+Jets"] = matrix_line.split()[5]

            #Get TTJet correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down + 2)
            weights["TTJet_QCD"] = matrix_line.split()[2]
            weights["TTJet_SingleTop"] = matrix_line.split()[3]
            weights["TTJet_TTJet"] = matrix_line.split()[4]
            weights["TTJet_V+Jets"] = matrix_line.split()[5]

            #Get V+Jets correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down + 3)
            weights["V+Jets_QCD"] = matrix_line.split()[2]
            weights["V+Jets_SingleTop"] = matrix_line.split()[3]
            weights["V+Jets_TTJet"] = matrix_line.split()[4]
            weights["V+Jets_V+Jets"] = matrix_line.split()[5]

        if channel == 'muon':
            #formula to calculate the number of lines below "central" to access in each bin loop
            number_of_lines_down = (len(variable_bins_ROOT[variable]) * 12) + (
                variable_bins_ROOT[variable].index(variable_bin) * 12)

            #Get QCD correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down)
            weights["QCD_QCD"] = matrix_line.split()[2]
            weights["QCD_SingleTop"] = matrix_line.split()[3]
            weights["QCD_TTJet"] = matrix_line.split()[4]
            weights["QCD_V+Jets"] = matrix_line.split()[5]

            #Get SingleTop correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down + 1)
            weights["SingleTop_QCD"] = matrix_line.split()[2]
            weights["SingleTop_SingleTop"] = matrix_line.split()[3]
            weights["SingleTop_TTJet"] = matrix_line.split()[4]
            weights["SingleTop_V+Jets"] = matrix_line.split()[5]

            #Get TTJet correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down + 2)
            weights["TTJet_QCD"] = matrix_line.split()[2]
            weights["TTJet_SingleTop"] = matrix_line.split()[3]
            weights["TTJet_TTJet"] = matrix_line.split()[4]
            weights["TTJet_V+Jets"] = matrix_line.split()[5]

            #Get V+Jets correlations
            matrix_line = linecache.getline(
                "logs/01_%s_fit_%dTeV_%s.log" % (variable, CoM, fit_variables),
                line_number + number_of_lines_down + 3)
            weights["V+Jets_QCD"] = matrix_line.split()[2]
            weights["V+Jets_SingleTop"] = matrix_line.split()[3]
            weights["V+Jets_TTJet"] = matrix_line.split()[4]
            weights["V+Jets_V+Jets"] = matrix_line.split()[5]

        #Create histogram
        histogram_properties = Histogram_properties()
        histogram_properties.title = title
        histogram_properties.name = 'Correlations_' + channel + '_' + variable + '_' + variable_bin
        histogram_properties.y_axis_title = y_title
        histogram_properties.x_axis_title = x_title
        histogram_properties.y_limits = y_limits
        histogram_properties.x_limits = x_limits
        histogram_properties.mc_error = 0.0
        histogram_properties.legend_location = 'upper right'

        #initialise 2D histogram
        a = Hist2D(4, 0, 4, 4, 0, 4)
        #fill histogram
        for i in range(len(parameters)):
            for j in range(len(parameters)):
                a.fill(
                    float(i), float(j),
                    float(weights["%s_%s" % (parameters[i], parameters[j])]))
        # create figure
        plt.figure(figsize=CMS.figsize, dpi=CMS.dpi, facecolor=CMS.facecolor)
        # make subplot(?)
        fig, ax = plt.subplots(nrows=1, ncols=1)
        rplt.hist2d(a)
        plt.subplots_adjust(right=0.8)

        #Set labels and formats for titles and axes
        plt.ylabel(histogram_properties.y_axis_title)
        plt.xlabel(histogram_properties.x_axis_title)
        plt.title(histogram_properties.title)
        x_limits = histogram_properties.x_limits
        y_limits = histogram_properties.y_limits
        ax.set_xticklabels(parameters_latex)
        ax.set_yticklabels(parameters_latex)
        ax.set_xticks([0.5, 1.5, 2.5, 3.5])
        ax.set_yticks([0.5, 1.5, 2.5, 3.5])
        plt.setp(ax.get_xticklabels(), visible=True)
        plt.setp(ax.get_yticklabels(), visible=True)

        #create and draw colour bar to the right of the main plot
        im = rplt.imshow(a, axes=ax, vmin=-1.0, vmax=1.0)
        #set location and dimensions (left, lower, width, height)
        cbar_ax = fig.add_axes([0.85, 0.10, 0.05, 0.8])
        fig.colorbar(im, cax=cbar_ax)

        for xpoint in range(len(parameters)):
            for ypoint in range(len(parameters)):
                correlation_value = weights["%s_%s" % (parameters[xpoint],
                                                       parameters[ypoint])]
                ax.annotate(correlation_value,
                            xy=(xpoint + 0.5, ypoint + 0.5),
                            ha='center',
                            va='center',
                            bbox=dict(fc='white', ec='none'))
        for save in save_as:
            plt.savefig(save_folder + histogram_properties.name + '.' + save)
        plt.close(fig)
    plt.close('all')
    b_tag_bin = '0btag'
    control_region = 'topReconstruction/backgroundShape/mttbar_3jets_conversions_withMETAndAsymJets_' + b_tag_bin
    histograms = get_histograms_from_files([control_region], histogram_files)
    prepare_histograms(histograms, rebin=50)

    histograms_to_draw = [
        histograms['data'][control_region], histograms['QCD'][control_region],
        histograms['ZJets'][control_region],
        histograms['WJets'][control_region],
        histograms['SingleTop'][control_region],
        histograms['TTJet'][control_region]
    ]
    histogram_lables = [
        'data', 'QCD', samples_latex['ZJets'], samples_latex['WJets'],
        'Single-Top', samples_latex['TTJet']
    ]
    histogram_colors = ['black', 'yellow', 'blue', 'green', 'magenta', 'red']

    histogram_properties = Histogram_properties()
    histogram_properties.name = 'Mttbar'
    histogram_properties.title = 'CMS Preliminary, $\mathcal{L}$ = 5.1 fb$^{-1}$ at $\sqrt{s}$ = 7 TeV \n e+jets, $\geq$4 jets, ' + b_tag_bins_latex[
        b_tag_bin]
    histogram_properties.x_axis_title = '$m_{\mathrm{t}\\bar{\mathrm{t}}}$ [GeV]'
    histogram_properties.y_axis_title = 'Events/(50 GeV)'
    histogram_properties.x_limits = [300, 1800]
    histogram_properties.mc_error = 0.15
    histogram_properties.mc_errors_label = '$\mathrm{t}\\bar{\mathrm{t}}$ uncertainty'

    make_data_mc_comparison_plot(histograms_to_draw, histogram_lables,
                                 histogram_colors, histogram_properties)
def do_shape_check(channel, control_region_1, control_region_2, variable, normalisation, title, x_title, y_title, x_limits, y_limits,
                   name_region_1='conversions' , name_region_2='non-isolated electrons', name_region_3='fit results', rebin=1):
    global b_tag_bin
    # QCD shape comparison
    if channel == 'electron':
        histograms = get_histograms_from_files([control_region_1, control_region_2], histogram_files)
        
        region_1 = histograms[channel][control_region_1].Clone() - histograms['TTJet'][control_region_1].Clone() - histograms['V+Jets'][control_region_1].Clone() - histograms['SingleTop'][control_region_1].Clone()
        region_2 = histograms[channel][control_region_2].Clone() - histograms['TTJet'][control_region_2].Clone() - histograms['V+Jets'][control_region_2].Clone() - histograms['SingleTop'][control_region_2].Clone()
        
        region_1.Rebin(rebin)
        region_2.Rebin(rebin)
        
        histogram_properties = Histogram_properties()
        histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + variable + '_' + b_tag_bin
        histogram_properties.title = title + ', ' + b_tag_bins_latex[b_tag_bin]
        histogram_properties.x_axis_title = x_title
        histogram_properties.y_axis_title = 'arbitrary units/(0.1)'
        histogram_properties.x_limits = x_limits
        histogram_properties.y_limits = y_limits[0]
        histogram_properties.mc_error = 0.0
        histogram_properties.legend_location = 'upper right'
        make_control_region_comparison(region_1, region_2,
                                       name_region_1=name_region_1, name_region_2=name_region_2,
                                       histogram_properties=histogram_properties, save_folder=output_folder)
        
        # QCD shape comparison to fit results
        histograms = get_histograms_from_files([control_region_1], histogram_files)
        
        region_1_tmp = histograms[channel][control_region_1].Clone() - histograms['TTJet'][control_region_1].Clone() - histograms['V+Jets'][control_region_1].Clone() - histograms['SingleTop'][control_region_1].Clone()
        region_1 = rebin_asymmetric(region_1_tmp, bin_edges_vis[variable])
        
        fit_results_QCD = normalisation[variable]['QCD']
        region_2 = value_error_tuplelist_to_hist(fit_results_QCD, bin_edges_vis[variable])
        
        histogram_properties = Histogram_properties()
        histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + variable + '_fits_with_conversions_' + b_tag_bin
        histogram_properties.title = title + ', ' + b_tag_bins_latex[b_tag_bin]
        histogram_properties.x_axis_title = x_title
        histogram_properties.y_axis_title = 'arbitrary units/(0.1)'
        histogram_properties.x_limits = x_limits
        histogram_properties.y_limits = y_limits[1]
        histogram_properties.mc_error = 0.0
        histogram_properties.legend_location = 'upper right'
        make_control_region_comparison(region_1, region_2,
                                       name_region_1=name_region_1, name_region_2=name_region_3,
                                       histogram_properties=histogram_properties, save_folder=output_folder)
    
    histograms = get_histograms_from_files([control_region_2], histogram_files)
    
    region_1_tmp = histograms[channel][control_region_2].Clone() - histograms['TTJet'][control_region_2].Clone() - histograms['V+Jets'][control_region_2].Clone() - histograms['SingleTop'][control_region_2].Clone()
    region_1 = rebin_asymmetric(region_1_tmp, bin_edges_vis[variable])    
    
    fit_results_QCD = normalisation[variable]['QCD']
    region_2 = value_error_tuplelist_to_hist(fit_results_QCD, bin_edges_vis[variable])
    
    histogram_properties = Histogram_properties()
    histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + variable + '_fits_with_noniso_' + b_tag_bin
    histogram_properties.title = title + ', ' + b_tag_bins_latex[b_tag_bin]
    histogram_properties.x_axis_title = x_title
    histogram_properties.y_axis_title = 'arbitrary units/(0.1)'
    histogram_properties.x_limits = x_limits
    histogram_properties.y_limits = y_limits[1]
    histogram_properties.mc_error = 0.0
    histogram_properties.legend_location = 'upper right'
    make_control_region_comparison(region_1, region_2,
                                   name_region_1=name_region_2, name_region_2=name_region_3,
                                   histogram_properties=histogram_properties, save_folder=output_folder)
 
 qcd_predicted_mc = histograms['QCD'][control_region]
 
 histograms_to_draw = [histograms['data'][control_region], qcd_predicted_mc,
                       histograms['V+Jets'][control_region],
                       histograms['SingleTop'][control_region], histograms['TTJet'][control_region]]
 histogram_lables = ['data', 'QCD', 'V+Jets', 'Single-Top', samples_latex['TTJet']]
 histogram_colors = ['black', 'yellow', 'green', 'magenta', 'red']
 
 histogram_properties = Histogram_properties()
 histogram_properties.name = 'EPlusJets_BJets_invmass_' + b_tag_bin
 histogram_properties.title = e_title + ', ' + b_tag_bins_latex[b_tag_bin]
 histogram_properties.x_axis_title = '$M_{\mathrm{b}\\bar{\mathrm{b}}}$'
 histogram_properties.y_axis_title = 'Normalised events/(20 GeV)'
 histogram_properties.x_limits = [0, 800]
 histogram_properties.mc_error = 0.15
 make_data_mc_comparison_plot(histograms_to_draw, histogram_lables, histogram_colors,
                              histogram_properties, save_folder = output_folder, show_ratio = False)
 histogram_properties.name += '_with_ratio'
 make_data_mc_comparison_plot(histograms_to_draw, histogram_lables, histogram_colors,
                              histogram_properties, save_folder = output_folder, show_ratio = True)
 
 #bjet invariant mass
 b_tag_bin = '3btags'
 control_region = 'TTbar_plus_X_analysis/EPlusJets/Ref selection/bjet_invariant_mass_' + b_tag_bin
 
 histograms = get_histograms_from_files([control_region], histogram_files)
 prepare_histograms(histograms, rebin=10, scale_factor = measurement_config.luminosity_scale)
 
 qcd_predicted_mc = histograms['QCD'][control_region]
 
def make_ttbarReco_plot(
    channel,
    x_axis_title,
    y_axis_title,
    signal_region_tree,
    control_region_tree,
    branchName,
    name_prefix,
    x_limits,
    nBins,
    use_qcd_data_region=False,
    y_limits=[],
    y_max_scale=1.2,
    rebin=1,
    legend_location=(0.98, 0.78),
    cms_logo_location='right',
    log_y=False,
    legend_color=False,
    ratio_y_limits=[0.3, 1.7],
    normalise=False,
):
    global output_folder, measurement_config, category, normalise_to_fit
    global preliminary, norm_variable, sum_bins, b_tag_bin, histogram_files

    # Input files, normalisations, tree/region names
    qcd_data_region = ''
    title = title_template % (measurement_config.new_luminosity / 1000.,
                              measurement_config.centre_of_mass_energy)
    normalisation = None
    if channel == 'electron':
        histogram_files['data'] = measurement_config.data_file_electron_trees
        histogram_files[
            'QCD'] = measurement_config.electron_QCD_MC_category_templates_trees[
                category]
        if normalise_to_fit:
            normalisation = normalisations_electron[norm_variable]
        if use_qcd_data_region:
            qcd_data_region = 'QCDConversions'
    if channel == 'muon':
        histogram_files['data'] = measurement_config.data_file_muon_trees
        histogram_files[
            'QCD'] = measurement_config.muon_QCD_MC_category_templates_trees[
                category]
        if normalise_to_fit:
            normalisation = normalisations_muon[norm_variable]
        if use_qcd_data_region:
            qcd_data_region = 'QCD non iso mu+jets ge3j'

    histograms = get_histograms_from_trees(
        trees=[signal_region_tree, control_region_tree],
        branch=branchName,
        weightBranch='1',
        files=histogram_files,
        nBins=nBins,
        xMin=x_limits[0],
        xMax=x_limits[-1])

    selection = 'SolutionCategory == 0'
    histogramsNoSolution = get_histograms_from_trees(
        trees=[signal_region_tree],
        branch=branchName,
        weightBranch='1',
        selection=selection,
        files=histogram_files,
        nBins=nBins,
        xMin=x_limits[0],
        xMax=x_limits[-1])

    selection = 'SolutionCategory == 1'
    histogramsCorrect = get_histograms_from_trees(trees=[signal_region_tree],
                                                  branch=branchName,
                                                  weightBranch='1',
                                                  selection=selection,
                                                  files=histogram_files,
                                                  nBins=nBins,
                                                  xMin=x_limits[0],
                                                  xMax=x_limits[-1])

    selection = 'SolutionCategory == 2'
    histogramsNotSL = get_histograms_from_trees(trees=[signal_region_tree],
                                                branch=branchName,
                                                weightBranch='1',
                                                selection=selection,
                                                files=histogram_files,
                                                nBins=nBins,
                                                xMin=x_limits[0],
                                                xMax=x_limits[-1])

    selection = 'SolutionCategory == 3'
    histogramsNotReco = get_histograms_from_trees(trees=[signal_region_tree],
                                                  branch=branchName,
                                                  weightBranch='1',
                                                  selection=selection,
                                                  files=histogram_files,
                                                  nBins=nBins,
                                                  xMin=x_limits[0],
                                                  xMax=x_limits[-1])

    selection = 'SolutionCategory > 3'
    histogramsWrong = get_histograms_from_trees(trees=[signal_region_tree],
                                                branch=branchName,
                                                weightBranch='1',
                                                selection=selection,
                                                files=histogram_files,
                                                nBins=nBins,
                                                xMin=x_limits[0],
                                                xMax=x_limits[-1])

    # Split histograms up into signal/control (?)
    signal_region_hists = {}
    inclusive_control_region_hists = {}
    for sample in histograms.keys():
        signal_region_hists[sample] = histograms[sample][signal_region_tree]
        if use_qcd_data_region:
            inclusive_control_region_hists[sample] = histograms[sample][
                control_region_tree]

    prepare_histograms(histograms,
                       rebin=1,
                       scale_factor=measurement_config.luminosity_scale)
    prepare_histograms(histogramsNoSolution,
                       rebin=1,
                       scale_factor=measurement_config.luminosity_scale)
    prepare_histograms(histogramsCorrect,
                       rebin=1,
                       scale_factor=measurement_config.luminosity_scale)
    prepare_histograms(histogramsNotSL,
                       rebin=1,
                       scale_factor=measurement_config.luminosity_scale)
    prepare_histograms(histogramsNotReco,
                       rebin=1,
                       scale_factor=measurement_config.luminosity_scale)
    prepare_histograms(histogramsWrong,
                       rebin=1,
                       scale_factor=measurement_config.luminosity_scale)

    qcd_from_data = signal_region_hists['QCD']

    # Which histograms to draw, and properties
    histograms_to_draw = [
        signal_region_hists['data'], qcd_from_data,
        signal_region_hists['V+Jets'], signal_region_hists['SingleTop'],
        histogramsNoSolution['TTJet'][signal_region_tree],
        histogramsNotSL['TTJet'][signal_region_tree],
        histogramsNotReco['TTJet'][signal_region_tree],
        histogramsWrong['TTJet'][signal_region_tree],
        histogramsCorrect['TTJet'][signal_region_tree]
    ]
    histogram_lables = [
        'data',
        'QCD',
        'V+Jets',
        'Single-Top',
        samples_latex['TTJet'] + ' - no solution',
        samples_latex['TTJet'] + ' - not SL',
        samples_latex['TTJet'] + ' - not reconstructible',
        samples_latex['TTJet'] + ' - wrong reco',
        samples_latex['TTJet'] + ' - correct',
    ]
    histogram_colors = [
        'black', 'yellow', 'green', 'magenta', 'black', 'burlywood',
        'chartreuse', 'blue', 'red'
    ]

    histogram_properties = Histogram_properties()
    histogram_properties.name = name_prefix + b_tag_bin
    if category != 'central':
        histogram_properties.name += '_' + category
    histogram_properties.title = title
    histogram_properties.x_axis_title = x_axis_title
    histogram_properties.y_axis_title = y_axis_title
    histogram_properties.x_limits = x_limits
    histogram_properties.y_limits = y_limits
    histogram_properties.y_max_scale = y_max_scale
    histogram_properties.xerr = None
    # workaround for rootpy issue #638
    histogram_properties.emptybins = True
    if b_tag_bin:
        histogram_properties.additional_text = channel_latex[
            channel] + ', ' + b_tag_bins_latex[b_tag_bin]
    else:
        histogram_properties.additional_text = channel_latex[channel]
    histogram_properties.legend_location = legend_location
    histogram_properties.cms_logo_location = cms_logo_location
    histogram_properties.preliminary = preliminary
    histogram_properties.set_log_y = log_y
    histogram_properties.legend_color = legend_color
    if ratio_y_limits:
        histogram_properties.ratio_y_limits = ratio_y_limits

    if normalise_to_fit:
        histogram_properties.mc_error = get_normalisation_error(normalisation)
        histogram_properties.mc_errors_label = 'fit uncertainty'
    else:
        histogram_properties.mc_error = mc_uncertainty
        histogram_properties.mc_errors_label = 'MC unc.'

    # Actually draw histograms
    make_data_mc_comparison_plot(
        histograms_to_draw,
        histogram_lables,
        histogram_colors,
        histogram_properties,
        save_folder=output_folder,
        show_ratio=False,
        normalise=normalise,
    )
    histogram_properties.name += '_with_ratio'
    loc = histogram_properties.legend_location
    # adjust legend location as it is relative to canvas!
    histogram_properties.legend_location = (loc[0], loc[1] + 0.05)
    make_data_mc_comparison_plot(
        histograms_to_draw,
        histogram_lables,
        histogram_colors,
        histogram_properties,
        save_folder=output_folder,
        show_ratio=True,
        normalise=normalise,
    )
def do_shape_check(channel,
                   control_region_1,
                   control_region_2,
                   variable,
                   normalisation,
                   title,
                   x_title,
                   y_title,
                   x_limits,
                   y_limits,
                   name_region_1='conversions',
                   name_region_2='non-isolated electrons',
                   name_region_3='fit results',
                   rebin=1):
    global b_tag_bin
    # QCD shape comparison
    if channel == 'electron':
        histograms = get_histograms_from_files(
            [control_region_1, control_region_2], histogram_files)

        region_1 = histograms[channel][control_region_1].Clone(
        ) - histograms['TTJet'][control_region_1].Clone(
        ) - histograms['V+Jets'][control_region_1].Clone(
        ) - histograms['SingleTop'][control_region_1].Clone()
        region_2 = histograms[channel][control_region_2].Clone(
        ) - histograms['TTJet'][control_region_2].Clone(
        ) - histograms['V+Jets'][control_region_2].Clone(
        ) - histograms['SingleTop'][control_region_2].Clone()

        region_1.Rebin(rebin)
        region_2.Rebin(rebin)

        histogram_properties = Histogram_properties()
        histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + variable + '_' + b_tag_bin
        histogram_properties.title = title + ', ' + b_tag_bins_latex[b_tag_bin]
        histogram_properties.x_axis_title = x_title
        histogram_properties.y_axis_title = 'arbitrary units/(0.1)'
        histogram_properties.x_limits = x_limits
        histogram_properties.y_limits = y_limits[0]
        histogram_properties.mc_error = 0.0
        histogram_properties.legend_location = 'upper right'
        make_control_region_comparison(
            region_1,
            region_2,
            name_region_1=name_region_1,
            name_region_2=name_region_2,
            histogram_properties=histogram_properties,
            save_folder=output_folder)

        # QCD shape comparison to fit results
        histograms = get_histograms_from_files([control_region_1],
                                               histogram_files)

        region_1_tmp = histograms[channel][control_region_1].Clone(
        ) - histograms['TTJet'][control_region_1].Clone(
        ) - histograms['V+Jets'][control_region_1].Clone(
        ) - histograms['SingleTop'][control_region_1].Clone()
        region_1 = rebin_asymmetric(region_1_tmp, bin_edges_vis[variable])

        fit_results_QCD = normalisation[variable]['QCD']
        region_2 = value_error_tuplelist_to_hist(fit_results_QCD,
                                                 bin_edges_vis[variable])

        histogram_properties = Histogram_properties()
        histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + variable + '_fits_with_conversions_' + b_tag_bin
        histogram_properties.title = title + ', ' + b_tag_bins_latex[b_tag_bin]
        histogram_properties.x_axis_title = x_title
        histogram_properties.y_axis_title = 'arbitrary units/(0.1)'
        histogram_properties.x_limits = x_limits
        histogram_properties.y_limits = y_limits[1]
        histogram_properties.mc_error = 0.0
        histogram_properties.legend_location = 'upper right'
        make_control_region_comparison(
            region_1,
            region_2,
            name_region_1=name_region_1,
            name_region_2=name_region_3,
            histogram_properties=histogram_properties,
            save_folder=output_folder)

    histograms = get_histograms_from_files([control_region_2], histogram_files)

    region_1_tmp = histograms[channel][control_region_2].Clone(
    ) - histograms['TTJet'][control_region_2].Clone(
    ) - histograms['V+Jets'][control_region_2].Clone(
    ) - histograms['SingleTop'][control_region_2].Clone()
    region_1 = rebin_asymmetric(region_1_tmp, bin_edges_vis[variable])

    fit_results_QCD = normalisation[variable]['QCD']
    region_2 = value_error_tuplelist_to_hist(fit_results_QCD,
                                             bin_edges_vis[variable])

    histogram_properties = Histogram_properties()
    histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + variable + '_fits_with_noniso_' + b_tag_bin
    histogram_properties.title = title + ', ' + b_tag_bins_latex[b_tag_bin]
    histogram_properties.x_axis_title = x_title
    histogram_properties.y_axis_title = 'arbitrary units/(0.1)'
    histogram_properties.x_limits = x_limits
    histogram_properties.y_limits = y_limits[1]
    histogram_properties.mc_error = 0.0
    histogram_properties.legend_location = 'upper right'
    make_control_region_comparison(region_1,
                                   region_2,
                                   name_region_1=name_region_2,
                                   name_region_2=name_region_3,
                                   histogram_properties=histogram_properties,
                                   save_folder=output_folder)
def make_plot( channel, x_axis_title, y_axis_title,
              signal_region_tree,
              control_region_tree,
              branchName,
              name_prefix, x_limits, nBins,
              use_qcd_data_region = False,
              compare_qcd_signal_with_data_control = False,
              y_limits = [],
              y_max_scale = 1.3,
              rebin = 1,
              legend_location = ( 0.98, 0.78 ), cms_logo_location = 'right',
              log_y = False,
              legend_color = False,
              ratio_y_limits = [0.3, 2.5],
              normalise = False,
              ):
    global output_folder, measurement_config, category, normalise_to_fit
    global preliminary, norm_variable, sum_bins, b_tag_bin, histogram_files

    controlToCompare = []
    if 'electron' in channel :
        controlToCompare =  ['QCDConversions', 'QCD non iso e+jets']
    elif 'muon' in channel :
        controlToCompare =  ['QCD iso > 0.3', 'QCD 0.12 < iso <= 0.3']

    histogramsToCompare = {}
    for qcd_data_region in controlToCompare:
        print 'Doing ',qcd_data_region
        # Input files, normalisations, tree/region names
        title = title_template % ( measurement_config.new_luminosity, measurement_config.centre_of_mass_energy )
        normalisation = None
        weightBranchSignalRegion = 'EventWeight'
        if 'electron' in channel:
            histogram_files['data'] = measurement_config.data_file_electron_trees
            histogram_files['QCD'] = measurement_config.electron_QCD_MC_category_templates_trees[category]
            if normalise_to_fit:
                normalisation = normalisations_electron[norm_variable]
            # if use_qcd_data_region:
            #     qcd_data_region = 'QCDConversions'
            #     # qcd_data_region = 'QCD non iso e+jets'
            if not 'QCD' in channel and not 'NPU' in branchName:
                weightBranchSignalRegion += ' * ElectronEfficiencyCorrection'
        if 'muon' in channel:
            histogram_files['data'] = measurement_config.data_file_muon_trees
            histogram_files['QCD'] = measurement_config.muon_QCD_MC_category_templates_trees[category]
            if normalise_to_fit:
                normalisation = normalisations_muon[norm_variable]
            # if use_qcd_data_region:
            #     qcd_data_region = 'QCD iso > 0.3'
            if not 'QCD' in channel and not 'NPU' in branchName:
                weightBranchSignalRegion += ' * MuonEfficiencyCorrection'

        if not "_NPUNoWeight" in name_prefix:
            weightBranchSignalRegion += ' * PUWeight'

        if not "_NBJetsNoWeight" in name_prefix:
            weightBranchSignalRegion += ' * BJetWeight'

        selection = '1'
        if branchName == 'abs(lepton_eta)' :
            selection = 'lepton_eta > -10'
        else:
            selection = '%s >= 0' % branchName
        # if 'QCDConversions' in signal_region_tree:
        #     selection += '&& isTightElectron'
        # print selection
        histograms = get_histograms_from_trees( trees = [signal_region_tree, control_region_tree], branch = branchName, weightBranch = weightBranchSignalRegion, files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1], selection = selection )
        histograms_QCDControlRegion = None
        if use_qcd_data_region:
            qcd_control_region = signal_region_tree.replace( 'Ref selection', qcd_data_region )
            histograms_QCDControlRegion = get_histograms_from_trees( trees = [qcd_control_region], branch = branchName, weightBranch = 'EventWeight', files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1], selection = selection )

        # Split histograms up into signal/control (?)
        signal_region_hists = {}
        control_region_hists = {}
        for sample in histograms.keys():
            signal_region_hists[sample] = histograms[sample][signal_region_tree]

            if compare_qcd_signal_with_data_control:
                if sample is 'data':
                    signal_region_hists[sample] = histograms[sample][control_region_tree]
                elif sample is 'QCD' :
                    signal_region_hists[sample] = histograms[sample][signal_region_tree]
                else:
                    del signal_region_hists[sample]

            if use_qcd_data_region:
                control_region_hists[sample] = histograms_QCDControlRegion[sample][qcd_control_region]

        # Prepare histograms
        if normalise_to_fit:
            # only scale signal region to fit (results are invalid for control region)
            prepare_histograms( signal_region_hists, rebin = rebin,
                                scale_factor = measurement_config.luminosity_scale,
                                normalisation = normalisation )
        elif normalise_to_data:
            totalMC = 0
            for sample in signal_region_hists:
                if sample is 'data' : continue
                totalMC += signal_region_hists[sample].Integral()
            newScale = signal_region_hists['data'].Integral() / totalMC

            prepare_histograms( signal_region_hists, rebin = rebin,
                                scale_factor = newScale,
                               )
        else:
            print measurement_config.luminosity_scale
            prepare_histograms( signal_region_hists, rebin = rebin,
                                scale_factor = measurement_config.luminosity_scale )
            prepare_histograms( control_region_hists, rebin = rebin,
                                scale_factor = measurement_config.luminosity_scale )

        # Use qcd from data control region or not
        qcd_from_data = None
        if use_qcd_data_region:
            qcd_from_data = clean_control_region( control_region_hists,

                              subtract = ['TTJet', 'V+Jets', 'SingleTop'] )
            # Normalise control region correctly
            nBins = signal_region_hists['QCD'].GetNbinsX()
            n, error = signal_region_hists['QCD'].integral(0,nBins+1,error=True)
            n_qcd_predicted_mc_signal = ufloat( n, error)

            n, error = control_region_hists['QCD'].integral(0,nBins+1,error=True)
            n_qcd_predicted_mc_control = ufloat( n, error)

            n, error = qcd_from_data.integral(0,nBins+1,error=True)
            n_qcd_control_region = ufloat( n, error)

            if not n_qcd_control_region == 0:
                dataDrivenQCDScale = n_qcd_predicted_mc_signal / n_qcd_predicted_mc_control
                print 'Overall scale : ',dataDrivenQCDScale
                qcd_from_data.Scale( dataDrivenQCDScale.nominal_value )
                signalToControlScale = n_qcd_predicted_mc_signal / n_qcd_control_region
                dataToMCscale = n_qcd_control_region / n_qcd_predicted_mc_control
                print "Signal to control :",signalToControlScale
                print "QCD scale : ",dataToMCscale
        else:
            qcd_from_data = signal_region_hists['QCD']

        # Which histograms to draw, and properties
        histograms_to_draw = []
        histogram_lables = []
        histogram_colors = []

        if compare_qcd_signal_with_data_control :
            histograms_to_draw = [signal_region_hists['data'], qcd_from_data ]
            histogram_lables = ['data', 'QCD']
            histogram_colors = ['black', 'yellow']
        else :
            histograms_to_draw = [signal_region_hists['data'], qcd_from_data,
                                  signal_region_hists['V+Jets'],
                                  signal_region_hists['SingleTop'],
                                  signal_region_hists['TTJet']]
            histogram_lables = ['data', 'QCD', 'V+Jets', 'Single-Top', samples_latex['TTJet']]
            histogram_colors = [colours['data'], colours['QCD'], colours['V+Jets'], colours['Single-Top'], colours['TTJet'] ]

        
        print list(qcd_from_data.y())
        histogramsToCompare[qcd_data_region] = qcd_from_data

    print histogramsToCompare
    histogram_properties = Histogram_properties()
    histogram_properties.name = 'QCD_control_region_comparison_' + channel + '_' + branchName
    histogram_properties.title = title
    histogram_properties.x_axis_title = x_axis_title
    histogram_properties.y_axis_title = y_axis_title
    histogram_properties.x_limits = x_limits
    histogram_properties.y_limits = y_limits
    histogram_properties.mc_error = 0.0
    histogram_properties.legend_location = ( 0.98, 0.78 )
    histogram_properties.ratio_y_limits = ratio_y_limits
    if 'electron' in channel:
        make_control_region_comparison(histogramsToCompare['QCDConversions'], histogramsToCompare['QCD non iso e+jets'],
                                       name_region_1='Conversions', name_region_2='Non Iso',
                                       histogram_properties=histogram_properties, save_folder=output_folder)
    elif 'muon' in channel:
        make_control_region_comparison(histogramsToCompare['QCD iso > 0.3'], histogramsToCompare['QCD 0.12 < iso <= 0.3'],
                                       name_region_1='QCD iso > 0.3', name_region_2='QCD 0.12 < iso <= 0.3',
                                       histogram_properties=histogram_properties, save_folder=output_folder)
def plot_fit_variable(histograms,
                      fit_variable,
                      variable,
                      bin_range,
                      fit_variable_distribution,
                      qcd_fit_variable_distribution,
                      title,
                      save_path,
                      channel='electron'):
    global fit_variable_properties, b_tag_bin, save_as, b_tag_bin_ctl
    histograms_ = deepcopy(histograms)
    mc_uncertainty = 0.10
    prepare_histograms(histograms_,
                       rebin=fit_variable_properties[fit_variable]['rebin'],
                       scale_factor=measurement_config.luminosity_scale)

    ######################################
    # plot the control regions as they are
    ######################################
    histogram_properties = Histogram_properties()
    histogram_properties.x_axis_title = fit_variable_properties[fit_variable][
        'x-title']
    histogram_properties.y_axis_title = fit_variable_properties[fit_variable][
        'y-title']
    histogram_properties.x_limits = [
        fit_variable_properties[fit_variable]['min'],
        fit_variable_properties[fit_variable]['max']
    ]
    histogram_properties.y_max_scale = 2

    histogram_lables = [
        'data', 'QCD', 'V+Jets', 'Single-Top', samples_latex['TTJet']
    ]
    histogram_colors = ['black', 'yellow', 'green', 'magenta', 'red']
    #     qcd_from_data = histograms_['data'][qcd_fit_variable_distribution].Clone()
    # clean against other processes
    histograms_for_cleaning = {
        'data': histograms_['data'][qcd_fit_variable_distribution],
        'V+Jets': histograms_['V+Jets'][qcd_fit_variable_distribution],
        'SingleTop': histograms_['SingleTop'][qcd_fit_variable_distribution],
        'TTJet': histograms_['TTJet'][qcd_fit_variable_distribution]
    }
    qcd_from_data = clean_control_region(
        histograms_for_cleaning, subtract=['TTJet', 'V+Jets', 'SingleTop'])

    histograms_to_draw = [
        histograms_['data'][qcd_fit_variable_distribution],
        histograms_['QCD'][qcd_fit_variable_distribution],
        histograms_['V+Jets'][qcd_fit_variable_distribution],
        histograms_['SingleTop'][qcd_fit_variable_distribution],
        histograms_['TTJet'][qcd_fit_variable_distribution]
    ]

    histogram_properties.title = title
    histogram_properties.additional_text = channel_latex[
        channel] + ', ' + b_tag_bins_latex[b_tag_bin_ctl]
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_%s_QCDConversions' % b_tag_bin_ctl
    make_data_mc_comparison_plot(
        histograms_to_draw,
        histogram_lables,
        histogram_colors,
        histogram_properties,
        save_folder=save_path + '/qcd/',
        show_ratio=False,
        save_as=save_as,
    )
    ######################################
    # plot QCD against data control region with TTJet, SingleTop and V+Jets removed
    ######################################
    histograms_to_draw = [
        qcd_from_data,
        histograms_['QCD'][qcd_fit_variable_distribution],
    ]
    histogram_properties.y_max_scale = 1.5
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_%s_QCDConversions_subtracted' % b_tag_bin_ctl
    make_data_mc_comparison_plot(
        histograms_to_draw,
        histogram_lables=['data', 'QCD'],
        histogram_colors=['black', 'yellow'],
        histogram_properties=histogram_properties,
        save_folder=save_path + '/qcd/',
        show_ratio=False,
        save_as=save_as,
    )
    ######################################
    # plot signal region
    ######################################
    # scale QCD to predicted
    n_qcd_predicted_mc = histograms_['QCD'][
        fit_variable_distribution].Integral()
    n_qcd_fit_variable_distribution = qcd_from_data.Integral()
    if not n_qcd_fit_variable_distribution == 0:
        qcd_from_data.Scale(1.0 / n_qcd_fit_variable_distribution *
                            n_qcd_predicted_mc)

    histograms_to_draw = [
        histograms_['data'][fit_variable_distribution], qcd_from_data,
        histograms_['V+Jets'][fit_variable_distribution],
        histograms_['SingleTop'][fit_variable_distribution],
        histograms_['TTJet'][fit_variable_distribution]
    ]

    histogram_properties.additional_text = channel_latex[
        channel] + ', ' + b_tag_bins_latex[b_tag_bin]
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_' + b_tag_bin
    make_data_mc_comparison_plot(
        histograms_to_draw,
        histogram_lables,
        histogram_colors,
        histogram_properties,
        save_folder=save_path,
        show_ratio=False,
        save_as=save_as,
    )
    ######################################
    # plot templates
    ######################################
    histogram_properties.mc_error = mc_uncertainty
    histogram_properties.mc_errors_label = '$\mathrm{t}\\bar{\mathrm{t}}$ uncertainty'
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_' + b_tag_bin + '_templates'
    histogram_properties.y_max_scale = 2
    # change histogram order for better visibility
    histograms_to_draw = [
        histograms_['TTJet'][fit_variable_distribution] +
        histograms_['SingleTop'][fit_variable_distribution],
        histograms_['TTJet'][fit_variable_distribution],
        histograms_['SingleTop'][fit_variable_distribution],
        histograms_['V+Jets'][fit_variable_distribution], qcd_from_data
    ]
    histogram_lables = [
        'QCD', 'V+Jets', 'Single-Top', samples_latex['TTJet'],
        samples_latex['TTJet'] + ' + ' + 'Single-Top'
    ]
    histogram_lables.reverse()
    # change QCD color to orange for better visibility
    histogram_colors = ['orange', 'green', 'magenta', 'red', 'black']
    histogram_colors.reverse()
    # plot template
    make_shape_comparison_plot(
        shapes=histograms_to_draw,
        names=histogram_lables,
        colours=histogram_colors,
        histogram_properties=histogram_properties,
        fill_area=False,
        alpha=1,
        save_folder=save_path,
        save_as=save_as,
    )
def plot_fit_variable( histograms, fit_variable, variable, bin_range,
                      fit_variable_distribution, qcd_fit_variable_distribution,
                      title, save_path, channel = 'electron' ):
    global fit_variable_properties, b_tag_bin, save_as, b_tag_bin_ctl
    histograms_ = deepcopy( histograms )
    mc_uncertainty = 0.10
    prepare_histograms( histograms_, rebin = fit_variable_properties[fit_variable]['rebin'], scale_factor = measurement_config.luminosity_scale )
    
    ######################################
    # plot the control regions as they are
    ######################################
    histogram_properties = Histogram_properties()
    histogram_properties.x_axis_title = fit_variable_properties[fit_variable]['x-title']
    histogram_properties.y_axis_title = fit_variable_properties[fit_variable]['y-title']
    histogram_properties.x_limits = [fit_variable_properties[fit_variable]['min'], fit_variable_properties[fit_variable]['max']]
    histogram_properties.y_max_scale = 2

    histogram_lables = ['data', 'QCD', 'V+Jets', 'Single-Top', samples_latex['TTJet']]
    histogram_colors = ['black', 'yellow', 'green', 'magenta', 'red']
#     qcd_from_data = histograms_['data'][qcd_fit_variable_distribution].Clone()
    # clean against other processes
    histograms_for_cleaning = {'data':histograms_['data'][qcd_fit_variable_distribution],
                               'V+Jets':histograms_['V+Jets'][qcd_fit_variable_distribution],
                               'SingleTop':histograms_['SingleTop'][qcd_fit_variable_distribution],
                               'TTJet':histograms_['TTJet'][qcd_fit_variable_distribution]}
    qcd_from_data = clean_control_region( histograms_for_cleaning, subtract = ['TTJet', 'V+Jets', 'SingleTop'] )
    
    histograms_to_draw = [histograms_['data'][qcd_fit_variable_distribution],
                          histograms_['QCD'][qcd_fit_variable_distribution],
                          histograms_['V+Jets'][qcd_fit_variable_distribution],
                          histograms_['SingleTop'][qcd_fit_variable_distribution],
                          histograms_['TTJet'][qcd_fit_variable_distribution]]
    
    histogram_properties.title = title
    histogram_properties.additional_text = channel_latex[channel] + ', ' + b_tag_bins_latex[b_tag_bin_ctl]
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_%s_QCDConversions' % b_tag_bin_ctl
    make_data_mc_comparison_plot( histograms_to_draw, histogram_lables, histogram_colors,
                                 histogram_properties,
                                 save_folder = save_path + '/qcd/',
                                 show_ratio = False,
                                 save_as = save_as,
                                 )
    ######################################
    # plot QCD against data control region with TTJet, SingleTop and V+Jets removed
    ######################################
    histograms_to_draw = [qcd_from_data,
                          histograms_['QCD'][qcd_fit_variable_distribution],
                          ]
    histogram_properties.y_max_scale = 1.5
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_%s_QCDConversions_subtracted' % b_tag_bin_ctl
    make_data_mc_comparison_plot( histograms_to_draw,
                                  histogram_lables = ['data', 'QCD'],
                                  histogram_colors = ['black', 'yellow'],
                                  histogram_properties = histogram_properties,
                                  save_folder = save_path + '/qcd/',
                                  show_ratio = False,
                                  save_as = save_as,
                                  )
    ######################################
    # plot signal region
    ######################################
    # scale QCD to predicted
    n_qcd_predicted_mc = histograms_['QCD'][fit_variable_distribution].Integral()
    n_qcd_fit_variable_distribution = qcd_from_data.Integral()
    if not n_qcd_fit_variable_distribution == 0:
        qcd_from_data.Scale( 1.0 / n_qcd_fit_variable_distribution * n_qcd_predicted_mc )
    
    histograms_to_draw = [histograms_['data'][fit_variable_distribution], qcd_from_data,
                          histograms_['V+Jets'][fit_variable_distribution],
                          histograms_['SingleTop'][fit_variable_distribution],
                          histograms_['TTJet'][fit_variable_distribution]]
    
    histogram_properties.additional_text = channel_latex[channel] + ', ' + b_tag_bins_latex[b_tag_bin]
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_' + b_tag_bin
    make_data_mc_comparison_plot( histograms_to_draw,
                                  histogram_lables,
                                  histogram_colors,
                                  histogram_properties,
                                  save_folder = save_path,
                                  show_ratio = False,
                                  save_as = save_as,
                                 )
    ######################################
    # plot templates
    ######################################
    histogram_properties.mc_error = mc_uncertainty
    histogram_properties.mc_errors_label = '$\mathrm{t}\\bar{\mathrm{t}}$ uncertainty'
    histogram_properties.name = variable + '_' + bin_range + '_' + fit_variable + '_' + b_tag_bin + '_templates'
    histogram_properties.y_max_scale = 2
    # change histogram order for better visibility
    histograms_to_draw = [histograms_['TTJet'][fit_variable_distribution] + histograms_['SingleTop'][fit_variable_distribution],
                          histograms_['TTJet'][fit_variable_distribution],
                          histograms_['SingleTop'][fit_variable_distribution],
                          histograms_['V+Jets'][fit_variable_distribution],
                          qcd_from_data]
    histogram_lables = ['QCD', 'V+Jets', 'Single-Top', samples_latex['TTJet'], samples_latex['TTJet'] + ' + ' + 'Single-Top']
    histogram_lables.reverse()
    # change QCD color to orange for better visibility
    histogram_colors = ['orange', 'green', 'magenta', 'red', 'black']
    histogram_colors.reverse()
    # plot template
    make_shape_comparison_plot( shapes = histograms_to_draw,
                                names = histogram_lables,
                                colours = histogram_colors,
                                histogram_properties = histogram_properties,
                                fill_area = False,
                                alpha = 1,
                                save_folder = save_path,
                                save_as = save_as,
                                )
def make_ttbarReco_plot( channel, x_axis_title, y_axis_title,
              signal_region_tree,
              control_region_tree,
              branchName,
              name_prefix, x_limits, nBins,
              use_qcd_data_region = False,
              y_limits = [],
              y_max_scale = 1.2,
              rebin = 1,
              legend_location = ( 0.98, 0.78 ), cms_logo_location = 'right',
              log_y = False,
              legend_color = False,
              ratio_y_limits = [0.3, 1.7],
              normalise = False,
              ):
    global output_folder, measurement_config, category, normalise_to_fit
    global preliminary, norm_variable, sum_bins, b_tag_bin, histogram_files

    # Input files, normalisations, tree/region names
    qcd_data_region = ''
    title = title_template % ( measurement_config.new_luminosity / 1000., measurement_config.centre_of_mass_energy )
    normalisation = None
    if channel == 'electron':
        histogram_files['data'] = measurement_config.data_file_electron_trees
        histogram_files['QCD'] = measurement_config.electron_QCD_MC_category_templates_trees[category]
        if normalise_to_fit:
            normalisation = normalisations_electron[norm_variable]
        if use_qcd_data_region:
            qcd_data_region = 'QCDConversions'
    if channel == 'muon':
        histogram_files['data'] = measurement_config.data_file_muon_trees
        histogram_files['QCD'] = measurement_config.muon_QCD_MC_category_templates_trees[category]
        if normalise_to_fit:
            normalisation = normalisations_muon[norm_variable]
        if use_qcd_data_region:
            qcd_data_region = 'QCD non iso mu+jets ge3j'

    histograms = get_histograms_from_trees( trees = [signal_region_tree, control_region_tree], branch = branchName, weightBranch = '1', files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1] )

    selection = 'SolutionCategory == 0'
    histogramsNoSolution = get_histograms_from_trees( trees = [signal_region_tree], branch = branchName, weightBranch = '1', selection = selection, files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1] )

    selection = 'SolutionCategory == 1'
    histogramsCorrect = get_histograms_from_trees( trees = [signal_region_tree], branch = branchName, weightBranch = '1', selection = selection, files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1] )

    selection = 'SolutionCategory == 2'
    histogramsNotSL = get_histograms_from_trees( trees = [signal_region_tree], branch = branchName, weightBranch = '1', selection = selection, files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1] )

    selection = 'SolutionCategory == 3'
    histogramsNotReco = get_histograms_from_trees( trees = [signal_region_tree], branch = branchName, weightBranch = '1', selection = selection, files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1] )

    selection = 'SolutionCategory > 3'
    histogramsWrong = get_histograms_from_trees( trees = [signal_region_tree], branch = branchName, weightBranch = '1', selection = selection, files = histogram_files, nBins = nBins, xMin = x_limits[0], xMax = x_limits[-1] )

    # Split histograms up into signal/control (?)
    signal_region_hists = {}
    inclusive_control_region_hists = {}
    for sample in histograms.keys():
        signal_region_hists[sample] = histograms[sample][signal_region_tree]
        if use_qcd_data_region:
            inclusive_control_region_hists[sample] = histograms[sample][control_region_tree]

    prepare_histograms( histograms, rebin = 1, scale_factor = measurement_config.luminosity_scale )
    prepare_histograms( histogramsNoSolution, rebin = 1, scale_factor = measurement_config.luminosity_scale )
    prepare_histograms( histogramsCorrect, rebin = 1, scale_factor = measurement_config.luminosity_scale )
    prepare_histograms( histogramsNotSL, rebin = 1, scale_factor = measurement_config.luminosity_scale )
    prepare_histograms( histogramsNotReco, rebin = 1, scale_factor = measurement_config.luminosity_scale )
    prepare_histograms( histogramsWrong, rebin = 1, scale_factor = measurement_config.luminosity_scale )

    qcd_from_data = signal_region_hists['QCD']

    # Which histograms to draw, and properties
    histograms_to_draw = [signal_region_hists['data'], qcd_from_data,
                          signal_region_hists['V+Jets'],
                          signal_region_hists['SingleTop'],
                          histogramsNoSolution['TTJet'][signal_region_tree],
                          histogramsNotSL['TTJet'][signal_region_tree],
                          histogramsNotReco['TTJet'][signal_region_tree],
                          histogramsWrong['TTJet'][signal_region_tree],
                          histogramsCorrect['TTJet'][signal_region_tree]
                          ]
    histogram_lables = ['data', 'QCD', 'V+Jets', 'Single-Top', 
                        samples_latex['TTJet'] + ' - no solution',
                        samples_latex['TTJet'] + ' - not SL',
                        samples_latex['TTJet'] + ' - not reconstructible',
                        samples_latex['TTJet'] + ' - wrong reco',
                        samples_latex['TTJet'] + ' - correct',
                        ]
    histogram_colors = ['black', 'yellow', 'green', 'magenta',
                        'black',
                        'burlywood',
                        'chartreuse',
                        'blue',
                        'red'
                        ]

    histogram_properties = Histogram_properties()
    histogram_properties.name = name_prefix + b_tag_bin
    if category != 'central':
        histogram_properties.name += '_' + category
    histogram_properties.title = title
    histogram_properties.x_axis_title = x_axis_title
    histogram_properties.y_axis_title = y_axis_title
    histogram_properties.x_limits = x_limits
    histogram_properties.y_limits = y_limits
    histogram_properties.y_max_scale = y_max_scale
    histogram_properties.xerr = None
    # workaround for rootpy issue #638
    histogram_properties.emptybins = True
    if b_tag_bin:
        histogram_properties.additional_text = channel_latex[channel] + ', ' + b_tag_bins_latex[b_tag_bin]
    else:
        histogram_properties.additional_text = channel_latex[channel]
    histogram_properties.legend_location = legend_location
    histogram_properties.cms_logo_location = cms_logo_location
    histogram_properties.preliminary = preliminary
    histogram_properties.set_log_y = log_y
    histogram_properties.legend_color = legend_color
    if ratio_y_limits:
        histogram_properties.ratio_y_limits = ratio_y_limits

    if normalise_to_fit:
        histogram_properties.mc_error = get_normalisation_error( normalisation )
        histogram_properties.mc_errors_label = 'fit uncertainty'
    else:
        histogram_properties.mc_error = mc_uncertainty
        histogram_properties.mc_errors_label = 'MC unc.'

    # Actually draw histograms
    make_data_mc_comparison_plot( histograms_to_draw, histogram_lables, histogram_colors,
                                 histogram_properties, save_folder = output_folder,
                                 show_ratio = False, normalise = normalise,
                                 )
    histogram_properties.name += '_with_ratio'
    loc = histogram_properties.legend_location
    # adjust legend location as it is relative to canvas!
    histogram_properties.legend_location = ( loc[0], loc[1] + 0.05 )
    make_data_mc_comparison_plot( histograms_to_draw, histogram_lables, histogram_colors,
                                 histogram_properties, save_folder = output_folder,
                                 show_ratio = True, normalise = normalise,
                                 )