implements plotting multiple stack histograms in a single canvas

.gitignore

@@ -1,3 +1,8 @@
 
 .ipynb_checkpoints/
 __pycache__/
+
+# C++ Build Files
+*.d
+*.so
+*.pcm

HistCollection.h

@@ -18,7 +18,7 @@
 struct HistWithPlottingInfo{
     TH1* hist;
     std::string draw_string;
-    
+
     HistWithPlottingInfo()
       :hist(0),draw_string(""){};
     HistWithPlottingInfo(TH1* hist, std::string draw_string)
@@ -30,15 +30,16 @@ class HistCollection{
   private:
     map<std::string, HistWithPlottingInfo> histograms;
     std::string sample_name;
-    
-    const char* hist_id(const std::string &id){
-        return (sample_name+id).c_str();
+
+    std::string hist_id(const std::string &id){
+        return sample_name+"_"+id;
     }
-    
+
     TH1D* book_histogram_1d(const std::string &id, const std::string &title,
                             int nbins, double min, double max,
                             const std::string draw_string = ""){
-        TH1D *hist = new TH1D(hist_id(id), title.c_str(), nbins, min, max);
+        std::string full_id = hist_id(id);
+        TH1D *hist = new TH1D(full_id.c_str(), title.c_str(), nbins, min, max);
         histograms[id] = HistWithPlottingInfo(hist, draw_string);
         return hist;
     };
@@ -46,33 +47,33 @@ class HistCollection{
     TH2D* book_histogram_2d(const std::string id, const std::string title,
                             int nbins_x, double min_x, double max_x,
                             int nbins_y, double min_y, double max_y){
-        TH2D *hist = new TH2D(hist_id(id), title.c_str(), nbins_x, min_x, max_x, nbins_y, min_y, max_y);
+        std::string full_id = hist_id(id);
+        TH2D *hist = new TH2D(full_id.c_str(), title.c_str(), nbins_x, min_x, max_x, nbins_y, min_y, max_y);
         histograms[id] = HistWithPlottingInfo(hist, "COLZ");
         return hist;
     };
   public:
-    /* List of included histogram objects
-    */
+    // List of included histogram objects
     TH1D *lepton_count;
     TH1D *lepton_count_pass_miniiso;
     TH1D *delta_pt;
     TH2D *lepton_count_vs_delta_pt;
     TH1D *top_quark_count;
-    
+
     TH1D *jet_count_trilepton;
     TH1D *jet_count_ss_dilepton;
     TH1D *jet_count_os_dilepton;
     TH1D *jet_count;
-    
+
     TH1D *b_jet_discriminator;
     TH1D *b_jet_count;
     TH1D *b_jet_pt_all;
     TH1D *b_jet_pt_3_or_more;
-    
+
     TH1D *trilepton_iso_1;
     TH1D *trilepton_iso_2;
     TH1D *trilepton_iso_3;
-    
+
     HistCollection(const std::string &sample_name)
       : sample_name(sample_name)
     {
@@ -120,23 +121,23 @@ class HistCollection{
                                         50, 0, 2000);
         // ---------------------------------------------
         trilepton_iso_1 = book_histogram_1d("trilepton_iso_1", "Trilepton Isolation - First",
-                                            50, 0, PI);
+                                            50, 0, 2*PI);
         // ---------------------------------------------
         trilepton_iso_2 = book_histogram_1d("trilepton_iso_2", "Trilepton Isolation - Second",
-                                            50, 0, PI);
+                                            50, 0, 2*PI);
         // ---------------------------------------------
         trilepton_iso_3 = book_histogram_1d("trilepton_iso_3", "Trilepton Isolation - Third",
-                                            50, 0, PI);
+                                            50, 0, 2*PI);
         // ---------------------------------------------
     }
-    
+
     HistCollection(): HistCollection("") {};
-    
+
     ~HistCollection(){
         for(auto& p: histograms)
             delete p.second.hist;
     }
-    
+
     void draw(TCanvas* canvas,
               std::pair<int, int> shape = {0,0}){
         std::pair<int, int> null_shape(0,0);
@@ -154,18 +155,18 @@ class HistCollection{
             p.second.hist->Draw(p.second.draw_string.c_str());
         }
     }
-    
+
     std::string get_sample_name(){
         return sample_name;
     }
-    
+
     vector<std::string> get_ids(){
         vector<std::string> vec;
         for(auto& p: histograms)
             vec.push_back(p.first);
         return vec;
     }
-    
+
     TH1* get_by_id(const std::string &id){
         TH1* hist = histograms[id].hist;
         return hist;
@@ -173,4 +174,4 @@ class HistCollection{
 };
 
 
-#endif // histcollection_h
+#endif // histcollection_h

MiniTree.cpp

@@ -1,14 +1,18 @@
 #ifndef minitree_cxx
 #define minitree_cxx
 
+#include <algorithm>
+#include <functional>
 #include <cstdio>
 #include <string>
+#include <vector>
 #include <map>
 
 #include <TStyle.h>
 
 #include "MiniTree.h"
 #include "HistCollection.h"
+using namespace std;
 
 
 HistCollection* build_histograms(const char* sample_name, const char* filename){
@@ -20,7 +24,7 @@ HistCollection* build_histograms(const char* sample_name, const char* filename){
     delete tree;
     f->Close();
     delete f;
-    std::printf("Finished filling histograms for file %s\n", filename);
+    printf("Finished filling histograms for file %s\n", filename);
     return hc;
 }
 
@@ -44,24 +48,53 @@ void set_branch_statuses(MiniTree *minitree){
     on("nGenTop");
 }
 
+double delta_r2(double eta1, double eta2, double phi1, double phi2){
+    return pow(eta1 - eta2, 2) + pow(phi1 - phi2, 2);
+}
+
+double isolation(int n, int i, function<double(int, int)> delta_r2_calculator){
+    double min_iso = -1;
+    for (int j=0; j<n; j++){
+        if (i == j) continue;
+        double curr_iso = delta_r2_calculator(i, j);
+        if (curr_iso < min_iso || min_iso == -1){
+            min_iso = curr_iso;
+        }
+    }
+    if (min_iso == -1)
+        return min_iso;
+    else
+        return sqrt(min_iso);
+}
+
 void MiniTree::Loop(HistCollection* hc){
-    
-    // Implement the mini-isolation cut which utilizes a dynamic isolation cone size
-    // based on lepton pt
-    auto mini_isolation_cut = [this](int idx){
+
+    /* Define first a few little subroutines that will be used later in the
+     * main loop.
+     */
+    auto lepton_delta_r2 = [this](int i, int j){
+        return delta_r2(this->LepGood_eta[i], this->LepGood_eta[j],
+                        this->LepGood_phi[i], this->LepGood_phi[j]);
+    };
+    auto lepton_isolation = [this, lepton_delta_r2](int i){
+        return isolation(this->nLepGood, i, function<double(int, int)>(lepton_delta_r2));
+    };
+    auto mini_isolation_cut = [this, lepton_isolation](int i){
+        /* Implement the mini-isolation cut which utilizes a dynamic isolation cone size
+         * based on lepton pt
+         */
         double r0 = 0.5;
         double pt0 = 7.5; // GeV
-        double r_cut2 = pow(min(r0, pt0/this->LepGood_pt[idx]), 2); 
-        for (int j = 0; j < this->nLepGood; j++){
-            if (j == idx) continue;
-            double delta_r2 = pow(this->LepGood_eta[idx] - this->LepGood_eta[j], 2) + 
-                              pow(this->LepGood_phi[idx] - this->LepGood_phi[j], 2);
-            if (delta_r2 < r_cut2) return false;
-        }
-        return true;
+
+        double iso_cut = min(r0, pt0/this->LepGood_pt[i]);
+        double isolation = lepton_isolation(i);
+        return isolation > iso_cut;
     };
-    
+
     if (fChain == 0) return;
+    /* Main analysis code begins here
+     *
+     */
     set_branch_statuses(this);
 
     Long64_t nentries = fChain->GetEntriesFast();
@@ -71,9 +104,10 @@ void MiniTree::Loop(HistCollection* hc){
         Long64_t ientry = LoadTree(jentry);
         if (ientry < 0) break;
         nb = fChain->GetEntry(jentry);   nbytes += nb;
-        // if (Cut(ientry) < 0) continue;
-        
-        // BEGIN ANALYSIS CODE
+
+        /* Main analysis code begins here
+         *
+         */
         hc->lepton_count->Fill(nLepGood);
         hc->top_quark_count->Fill(nGenTop);
         for(int i=0; i<nLepGood; i++){
@@ -88,7 +122,7 @@ void MiniTree::Loop(HistCollection* hc){
             b_jet_count += (Jet_btagCMVA[i] > 0);
         }
         hc->b_jet_count->Fill(b_jet_count);
-        
+
         for(int i=0; i<nJet; i++){
             hc->b_jet_pt_all->Fill(Jet_pt[i]);
             if (b_jet_count >= 3){
@@ -100,7 +134,7 @@ void MiniTree::Loop(HistCollection* hc){
             lepton_count_pass_miniiso += mini_isolation_cut(i);
         }
         hc->lepton_count_pass_miniiso->Fill(lepton_count_pass_miniiso);
-        
+
         if (nLepGood == 3){
             hc->jet_count_trilepton->Fill(nJet);
         }
@@ -110,6 +144,16 @@ void MiniTree::Loop(HistCollection* hc){
         if (nLepGood == 2 && (LepGood_charge[0] * LepGood_charge[1]) == -1){
             hc->jet_count_os_dilepton->Fill(nJet);
         }
+
+        if (nLepGood == 3){
+            vector<double> isos = {lepton_isolation(0),
+                                   lepton_isolation(1),
+                                   lepton_isolation(2)};
+            sort(isos.begin(), isos.end()); //Sorts ascending by default
+            hc->trilepton_iso_1->Fill(isos[2]);
+            hc->trilepton_iso_2->Fill(isos[1]);
+            hc->trilepton_iso_3->Fill(isos[0]);
+        }
     }
 }
-#endif // tree_cxx
+#endif // tree_cxx

canvas_wrapper.py

@@ -59,7 +59,7 @@ png_formatter = get_ipython().display_formatter.formatters['image/png'] # noqa
 
 # Register ROOT types in ipython
 #
-# In [1]: canvas = rootnotes.canvas()
+# In [1]: canvas = canvas_wrapper.canvas()
 # In [2]: canvas
 # Out [2]: [image will be here]
 png_formatter.for_type(ROOT.TCanvas, _display_canvas)

utils.py

@@ -1,8 +1,10 @@
 import io
 import sys
+from math import ceil, sqrt
 import itertools as it
 import ROOT
 
+
 class OutputCapture:
     def __init__(self):
         self.my_stdout = io.StringIO()
@@ -28,17 +30,19 @@ class OutputCapture:
         self.stdout = None
         self.stderr = None
 
+
 def bin_range(n, end=None):
     if end is None:
         return range(1, n+1)
     else:
         return range(n+1, end+1)
-        
+
+
 def normalize_columns(hist2d):
     normHist = ROOT.TH2D(hist2d)
     cols, rows = hist2d.GetNbinsX(), hist2d.GetNbinsY()
     for col in bin_range(cols):
-        sum_ = 0;
+        sum_ = 0
         for row in bin_range(rows):
             sum_ += hist2d.GetBinContent(col, row)
         if sum_ == 0:
@@ -48,25 +52,60 @@ def normalize_columns(hist2d):
             normHist.SetBinContent(col, row, norm)
     return normHist
 
-def stack_hist(hists, labels=None, id_="stack", title="", enable_fill=False, normalize_to=0):
+
+def stack_hist(hists,
+               labels=None, id_=None,
+               title="", enable_fill=False,
+               normalize_to=0, draw=False,
+               draw_option="",
+               _stacks={}):
     """hists should be a list of TH1D objects
     returns a new stacked histogram
     """
-    colors = it.cycle([ROOT.kRed,ROOT.kBlue, ROOT.kGreen])
+    colors = it.cycle([ROOT.kRed, ROOT.kBlue, ROOT.kGreen])
     stack = ROOT.THStack(id_, title)
     if labels is None:
         labels = [hist.GetName() for hist in hists]
-    for i, (hist, label, color) in enumerate(zip(hists, labels, colors)):
+    if type(normalize_to) in (int, float):
+        normalize_to = [normalize_to]*len(hists)
+    if id_ is None:
+        id_ = hists[0].GetName() + "_stack"
+    ens = enumerate(zip(hists, labels, colors, normalize_to))
+    for i, (hist, label, color, norm) in ens:
         hist_copy = hist
-        hist_copy = hist.Clone(str(i)+"_clone")
+        hist_copy = hist.Clone(hist.GetName()+"_clone")
         hist_copy.SetTitle(label)
-        hist_copy.SetStats(False)
         if enable_fill:
             hist_copy.SetFillColorAlpha(color, 0.75)
             hist_copy.SetLineColorAlpha(color, 0.75)
-        if normalize_to:
+        if norm:
             integral = hist_copy.Integral()
-            hist_copy.Scale(normalize_to/integral, "nosw2")
+            hist_copy.Scale(norm/integral, "nosw2")
             hist_copy.SetStats(False)
         stack.Add(hist_copy)
-    return stack
+    if draw:
+        stack.Draw(draw_option)
+    _stacks[id_] = stack  # prevent stack from getting garbage collected
+                          # needed for multipad plots :/
+    return stack
+
+
+def stack_hist_array(canvas, histcollections, fields, titles,
+                     shape=None, **kwargs):
+    def get_hist_set(attrname):
+        hists, labels = zip(*[(getattr(h, attrname), h.get_sample_name())
+                              for h in histcollections])
+        return hists, labels
+    n_fields = len(fields)
+    if shape is None:
+        if n_fields <= 4:
+            shape = (1, n_fields)
+        else:
+            shape = (ceil(sqrt(n_fields)),)*2
+    canvas.Clear()
+    canvas.Divide(*shape)
+    for i, field, title in zip(bin_range(n_fields), fields, titles):
+        canvas.cd(i)
+        hists, labels = get_hist_set(field)
+        stack_hist(hists, labels, id_=field, title=title, draw=True, **kwargs)
+    canvas.cd(1).BuildLegend(0.75, 0.75, 0.95, 0.95, "")