Adding Resolution to PickBins

src/cross_section_measurement/00_pick_bins.py

@@ -48,7 +48,10 @@
 from optparse import OptionParser
 from config.variable_binning import bin_edges_full, minimum_bin_width
 from tools.file_utilities import write_data_to_JSON
-from ROOT import TH1, TCanvas, TLine, gDirectory, TObjArray
+from ROOT import TH1, TCanvas, TLine, gDirectory, TObjArray, TColor, TLegend
+
+import tools.resolution as rs
+
 def main():
     '''
     Step 1: Get the 2D histogram for every sample (channel and/or centre of mass energy)
@@ -62,6 +65,8 @@ def main():
                       help = "Consider visible phase space or not" )
     parser.add_option( '-c', dest = "combined", action = "store_true",
                       help = "Combine channels" )
+    parser.add_option( '-r', dest = "redo_resolution", action = "store_true",
+                      help = "Recalculate the resolution plots" )
     ( options, _ ) = parser.parse_args()
 
     measurement_config = XSectionConfig(13)
@@ -86,6 +91,9 @@ def main():
             variableToUse = 'abs_%s' % variable
         histogram_information = get_histograms( variableToUse, options )
 
+        if options.redo_resolution:
+            rs.generate_resolution_plots(histogram_information, variable)
+
         if variable == 'HT':
             best_binning, histogram_information = get_best_binning( histogram_information , p_min, s_min, n_min, minimum_bin_width[variable], x_min=100. )
         elif variable == 'ST':
@@ -133,6 +141,7 @@ def main():
             outputInfo['p_i'] = info['p_i']
             outputInfo['s_i'] = info['s_i']
             outputInfo['N'] = info['N']
+            outputInfo['res'] = info['res']
             outputJsonFile = 'unfolding/13TeV/binningInfo_%s_%s_FullPS.txt' % ( variable, info['channel'] )
             if options.visiblePhaseSpace:
                 outputJsonFile = 'unfolding/13TeV/binningInfo_%s_%s_VisiblePS.txt' % ( variable, info['channel'] )
@@ -198,6 +207,8 @@ def get_histograms( variable, options ):
 
     return histogram_information
 
+
+
 def get_best_binning( histogram_information, p_min, s_min, n_min, min_width, x_min = None ):
     '''
     Step 1: Change the size of the first bin until it fulfils the minimal criteria
@@ -206,6 +217,7 @@ def get_best_binning( histogram_information, p_min, s_min, n_min, min_width, x_m
     '''
     histograms = [info['hist'] for info in histogram_information]
     bin_edges = []
+    resolutions = []
     purities = {}
     stabilities = {}
 
@@ -219,7 +231,9 @@ def get_best_binning( histogram_information, p_min, s_min, n_min, min_width, x_m
         current_bin_end = current_bin_start
 
     while current_bin_end < n_bins:
-        current_bin_end, _, _, _ = get_next_end( histograms, current_bin_start, current_bin_end, p_min, s_min, n_min, min_width )
+        # bin_End, p, s, N_reco
+        current_bin_end, _, _, _, r = get_next_end( histograms, current_bin_start, current_bin_end, p_min, s_min, n_min, min_width )
+        resolutions.append(r)
         if not bin_edges:
             # if empty
             bin_edges.append( first_hist.GetXaxis().GetBinLowEdge( current_bin_start + 1 ) )
@@ -246,39 +260,9 @@ def get_best_binning( histogram_information, p_min, s_min, n_min, min_width, x_m
         info['p_i'] = purities
         info['s_i'] = stabilities
         info['N'] = n_events
-
-    return bin_edges, histogram_information
+        info['res'] = resolutions
 
-def get_bin_resolution( fine_binned_response , bin_middle ):
-    gen_resolution = 0
-    reco_resolution = 0
-
-    reco_array = TObjArray()
-    gen_array = TObjArray()
-
-    # produces 4 TH1D histograms
-    # 0 : constant
-    # 1 : mean
-    # 2 : sigma
-    # 3 : chi-squared of it
-    gen_slices = fine_binned_response.FitSlicesY(0, bin_middle, bin_middle, 1, "Q", gen_array)
-    # for hist in gen_array:
-    #     print (type(hist))
-    # print("Gen Sigma : ", gen_array[2].GetBinContent(bin_middle))
-    # print("Gen Mean : ", gen_array[1].GetBinContent(bin_middle))
-    # print("Gen Constant : ", gen_array[0].GetBinContent(bin_middle))
-    gen_resolution = gen_array[2].GetBinContent(bin_middle)
-
-    reco_slices = fine_binned_response.FitSlicesX(0, bin_middle, bin_middle, 1, "Q", reco_array) 
-    # for hist in reco_array:
-    #     print (type(hist))
-    # print("Reco Sigma : ", reco_array[2].GetBinContent(bin_middle))
-    # print("Reco Mean : ", reco_array[1].GetBinContent(bin_middle))
-    # print("Reco Constant : ", reco_array[0].GetBinContent(bin_middle))
-    reco_resolution = reco_array[2].GetBinContent(bin_middle)
-
-    resolution = max(gen_resolution, reco_resolution)
-    return resolution
+    return bin_edges, histogram_information
 
 def get_next_end( histograms, bin_start, bin_end, p_min, s_min, n_min, min_width ): 
     current_bin_start = bin_start
@@ -292,7 +276,11 @@ def get_next_end( histograms, bin_start, bin_end, p_min, s_min, n_min, min_width
 
         # keep the start bin the same but roll the end bin
         for bin_i in range ( current_bin_end, len( reco_i ) + 1 ):
-            binWidth = reco.GetXaxis().GetBinLowEdge(bin_i) - reco.GetXaxis().GetBinUpEdge(current_bin_start)
+            x_high = reco.GetXaxis().GetBinLowEdge(bin_i)
+            x_mid = reco.GetXaxis().GetBinCenter(int( (current_bin_start+current_bin_end)/2 ) )
+            x_low = reco.GetXaxis().GetBinUpEdge(current_bin_start)
+            binWidth = x_high - x_low
+
             if binWidth < min_width:
                 current_bin_end = bin_i
                 continue
@@ -325,22 +313,22 @@ def get_next_end( histograms, bin_start, bin_end, p_min, s_min, n_min, min_width
                 # The StdDev of Gaussian = Resolution.
                 # If Resolution < Bin width then we are all good
 
+                # NJets is not great at the moment for fitting guassians 
+                if (var=='NJets'):
+                    current_bin_end = bin_i
+                    break
+
                 # Choose the middle fine bin of the total bin width as the resolution
-                current_bin_mid = int((bin_start+current_bin_end)/2)
-
-                resolution = get_bin_resolution( gen_vs_reco_histogram, current_bin_mid)
-                # print ("Max Resolution : ", resolution)
-                # print ("Bin Width : ", binWidth)
-                if( abs(reco.GetXaxis().GetBinCenter(current_bin_mid) - reco.GetXaxis().GetBinLowEdge(bin_i)) > resolution and 
-                    abs(reco.GetXaxis().GetBinCenter(current_bin_mid) - reco.GetXaxis().GetBinUpEdge(current_bin_start)) > resolution 
-                    ):
+                res = rs.get_merged_bin_resolution('plots/resolutionStudies/resolution.root', var, x_low, x_high)
+                res = round( res, 3 )
+                if ( x_high - x_mid > res and x_mid - x_low > res ):
                     current_bin_end = bin_i
                     break
 
+
             # if it gets to the end, this is the best we can do
             current_bin_end = bin_i
-
-    return current_bin_end, p, s, n_reco
+    return current_bin_end, p, s, n_reco, res
 
 def print_console(info, old_purities, old_stabilities, print_old = False):
     print('CoM =', info['CoM'], 'channel =', info['channel'])

tools/resolution.py

@@ -0,0 +1,131 @@
+import numpy as np
+from ROOT import TH1, TCanvas, TLine, gDirectory, TObjArray, TColor, TLegend, TAxis, TFile, TObject
+from rootpy import asrootpy
+from rootpy.io import File
+from tools.file_utilities import make_folder_if_not_exists
+
+
+def get_bin_resolution( fine_binned_response , bin_number ):
+    '''
+    Return the resolution for the reco and gen distributions for each variable in each fine bin
+    '''
+    gen_resolution = 0
+    reco_resolution = 0
+
+    reco_array = TObjArray()
+    gen_array = TObjArray()
+
+    # produces 4 TH1D histograms
+    # 0 : constant
+    # 1 : mean
+    # 2 : sigma
+    # 3 : chi-squared of it
+    gen_slices = fine_binned_response.FitSlicesY(0, bin_number, bin_number, 1, "Q", gen_array)
+    # for hist in gen_array:
+    #     print (type(hist))
+    # print("Gen Sigma : ", gen_array[2].GetBinContent(bin_number))
+    # print("Gen Mean : ", gen_array[1].GetBinContent(bin_number))
+    # print("Gen Constant : ", gen_array[0].GetBinContent(bin_number))
+    gen_resolution = gen_array[2].GetBinContent(bin_number)
+
+    reco_slices = fine_binned_response.FitSlicesX(0, bin_number, bin_number, 1, "Q", reco_array) 
+    # for hist in reco_array:
+    #     print (type(hist))
+    # print("Reco Sigma : ", reco_array[2].GetBinContent(bin_number))
+    # print("Reco Mean : ", reco_array[1].GetBinContent(bin_number))
+    # print("Reco Constant : ", reco_array[0].GetBinContent(bin_number))
+    reco_resolution = reco_array[2].GetBinContent(bin_number)
+
+    # resolution = max(gen_resolution, reco_resolution)
+    # return resolution
+    return reco_resolution, gen_resolution
+
+
+def get_merged_bin_resolution(res_file, var, low_bin, high_bin):
+    '''
+    Return mean value of resolutions in fine bins of the merged bin.
+    '''
+
+    bin_contents = []
+
+    f = File( res_file )
+    res_hist = f.Get( 'res_r_'+var ).Clone()
+    # change scope from file to memory
+    res_hist.SetDirectory( 0 )
+    f.close()
+
+    low_bin_n = res_hist.GetXaxis().FindBin(low_bin)
+    high_bin_n = res_hist.GetXaxis().FindBin(high_bin)
+
+    for bin_i in range(low_bin_n, high_bin_n+1):
+        bin_content = res_hist.GetBinContent(bin_i)
+        # resolution couldnt be reconstructed (High GeV with low stats)
+        # remove these from list of resolutions
+        if bin_content == 0 : continue
+        bin_contents.append(bin_content)
+
+    # print(bin_contents)
+    res = np.mean(bin_contents)
+    return res
+
+
+def generate_resolution_plots(histogram_information, var):
+    '''
+    Save the reconstructed resolution to root file
+    Plot rexonstructed vs generated resolution
+    '''
+
+    histograms = [info['hist'] for info in histogram_information]
+    first_hist = histograms[0]
+    rebin_hist = first_hist.Rebin2D( 10, 10, "rebinned" ) # Make a coarser fine bin - mor stats for fit and saves compute time
+    res_r_hist = rebin_hist.ProjectionX().Clone('res_r_'+var) #Initialise resolution plots
+    res_g_hist = rebin_hist.ProjectionX().Clone('res_g_'+var)
+    res_r_g_bias_hist = rebin_hist.ProjectionX().Clone('res_r_g_bias')
+
+    n_bins = rebin_hist.GetNbinsX()
+    for n_bin in range (0, n_bins):
+        r_res, g_res = get_bin_resolution(rebin_hist, n_bin)
+        res_r_hist.SetBinContent(n_bin, r_res)
+        res_g_hist.SetBinContent(n_bin, g_res)
+        res_r_hist.SetBinError(n_bin, 0)
+        res_g_hist.SetBinError(n_bin, 0)
+        res_r_g_bias_hist.SetBinContent(n_bin, r_res-g_res)
+
+    make_folder_if_not_exists('plots/resolutionStudies/')
+
+    # c = TCanvas('c1', 'c1', 800, 600)
+    # c.SetLogy()
+    # c.SetGrid()
+    # # res_r_hist.GetXaxis().SetRangeUser(0, 1500)
+    # res_r_hist.SetTitle("Reco(Gen) resolution;"+var+"; Resolution (GeV)")
+    # res_r_hist.SetFillColor(3)
+    # res_r_hist.SetFillStyle(3003);
+    # res_r_hist.Draw("hist")
+    # res_g_hist.SetFillColor(2)
+    # res_g_hist.SetFillStyle(3003);
+    # res_g_hist.Draw("hist same")
+    # leg = TLegend(0.1,0.8,0.38,0.9);
+    # leg.AddEntry(res_r_hist,"Res of Reco","f")
+    # leg.AddEntry(res_g_hist,"Res of Gen","f")
+    # leg.Draw()
+    # c.Update()
+    # c.SaveAs("plots/resolutionStudies/Resolution_"+var+".png")
+
+    # c2 = TCanvas('c2', 'c2', 800, 600)
+    # c.SetLogy()
+    # c2.SetGrid()
+    # res_r_g_bias_hist.SetTitle("Resolution Bias;"+var+"; Reco-Gen Resolution(GeV)")
+    # # res_r_g_bias_hist.GetXaxis().SetRangeUser(0, 1500)
+    # res_r_g_bias_hist.Draw("hist")
+    # leg2 = TLegend(0.75,0.1,0.9,0.2);
+    # leg2.AddEntry(res_r_g_bias_hist,"Resolution Bias","f")
+    # leg2.Draw()
+    # c2.Update()
+    # c2.SaveAs("plots/resolutionStudies/ResolutionBias_"+var+".png")
+
+
+    f = TFile("plots/resolutionStudies/resolution.root", "update")
+    res_r_hist.Write(res_r_hist.GetName(),TObject.kOverwrite)
+    f.Close()
+
+    return