Adds additional plots and implementation of mini-iso selection

.gitignore

@@ -0,0 +1,3 @@
+
+.ipynb_checkpoints/
+__pycache__/

HistCollection.h

@@ -3,59 +3,161 @@
 
 #include <string>
 #include <sstream>
+#include <iostream>
 #include <map>
+#include <cmath>
 
+#include <TCanvas.h>
 #include <TH1.h>
 #include <TH2.h>
 
+#define PI 3.14159
+#define PDG_ELECTRON = 11
+#define PDG_MUON = 13
+
+struct HistWithPlottingInfo{
+    TH1* hist;
+    std::string draw_string;
+    
+    HistWithPlottingInfo()
+      :hist(0),draw_string(""){};
+    HistWithPlottingInfo(TH1* hist, std::string draw_string)
+      :hist(hist),draw_string(draw_string){};
+};
 
 
 class HistCollection{
   private:
-    map<std::string, TH1*> histograms;
-    std::string prefix;
+    map<std::string, HistWithPlottingInfo> histograms;
+    std::string sample_name;
+    
+    const char* hist_id(const std::string &id){
+        return (sample_name+id).c_str();
+    }
     
     TH1D* book_histogram_1d(const std::string &id, const std::string &title,
-                            int nbins, double min, double max){
-        const char* full_id = (prefix+id).c_str();
-        TH1D *hist = new TH1D(full_id, title.c_str(), nbins, min, max);
-        histograms[id] = hist;
+                            int nbins, double min, double max,
+                            const std::string draw_string = ""){
+        TH1D *hist = new TH1D(hist_id(id), title.c_str(), nbins, min, max);
+        histograms[id] = HistWithPlottingInfo(hist, draw_string);
         return hist;
     };
 
     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){
-        const char* full_id = (prefix+id).c_str();
-        TH2D *hist = new TH2D(full_id, title.c_str(), nbins_x, min_x, max_x, nbins_y, min_y, max_y);
-        histograms[id] = hist;
+        TH2D *hist = new TH2D(hist_id(id), 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
     */
     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;
     
-    HistCollection(const std::string &filename)
-      : prefix(filename.substr(0, filename.size()-5))
+    TH1D *trilepton_iso_1;
+    TH1D *trilepton_iso_2;
+    TH1D *trilepton_iso_3;
+    
+    HistCollection(const std::string &sample_name)
+      : sample_name(sample_name)
     {
-//        this->prefix = prefix.substr(0, prefix.size()-5);
-        lepton_count = book_histogram_1d("lepton_count", "Good Lepton Count", 10, 0, 10);
-        delta_pt = book_histogram_1d("delta_pt", "{\\\\Delta P_T}", 100, -50, 50);
-        lepton_count_vs_delta_pt = book_histogram_2d("lepton_count_vs_delta_pt", "Good Lepton Count Vs {\\\\delta P_T}",
-                                                     10, 0, 10,
-                                                     100, -50, 50);
-    };
+        // ---------------------------------------------
+        lepton_count = book_histogram_1d("lepton_count", "Lepton Multiplicity",
+                                         8, 0, 8);
+        // ---------------------------------------------
+        lepton_count_pass_miniiso = book_histogram_1d("lepton_count_pass_miniiso",
+                                                      "Number of Leptons passing mini-isolation",
+                                                      8, 0, 8);
+        // ---------------------------------------------
+        delta_pt = book_histogram_1d("delta_pt", "{\\Delta P_T}",
+                                     100, -50, 50);
+        // ---------------------------------------------
+        lepton_count_vs_delta_pt = book_histogram_2d("lepton_count_vs_delta_pt",
+                                                     "Good Lepton Count Vs {\\Delta P_T}",
+                                                     10, 0, 10, 100, -50, 50);
+        // ---------------------------------------------
+        top_quark_count = book_histogram_1d("top_quark_count", "Top Quark Multiplicity",
+                                            4, 0, 4);
+        // ---------------------------------------------
+        jet_count_trilepton = book_histogram_1d("jet_count_trilepton",
+                                                "Jet Multiplicity - Trilepton",
+                                                13, 0, 13);
+        jet_count_ss_dilepton = book_histogram_1d("jet_count_ss_dilepton",
+                                                  "Jet Multiplicity - Same-Sign Dilepton",
+                                                  13, 0, 13);
+        jet_count_os_dilepton = book_histogram_1d("jet_count_os_dilepton",
+                                                  "Jet Multiplicity - Opposite-Sign Dilepton",
+                                                  13, 0, 13);
+        // ---------------------------------------------
+        jet_count = book_histogram_1d("jet_count", "Jet Multiplicity",
+                                      13, 0, 13);
+        // ---------------------------------------------
+        b_jet_discriminator = book_histogram_1d("b_jet_discriminator", "B-Jet Discriminator",
+                                                40, -1, 1);
+        // ---------------------------------------------
+        b_jet_count = book_histogram_1d("b_jet_count", "B-Jet Multiplicity",
+                                        10, 0, 10);
+        // ---------------------------------------------
+        b_jet_pt_all = book_histogram_1d("b_jet_pt_all", "B-Jet {P_T} - All",
+                                         50, 0, 2000);
+        // ---------------------------------------------
+        b_jet_pt_3_or_more = book_histogram_1d("b_jet_pt_3_or_more", "B-Jet  {P_T} - 3 or more",
+                                        50, 0, 2000);
+        // ---------------------------------------------
+        trilepton_iso_1 = book_histogram_1d("trilepton_iso_1", "Trilepton Isolation - First",
+                                            50, 0, PI);
+        // ---------------------------------------------
+        trilepton_iso_2 = book_histogram_1d("trilepton_iso_2", "Trilepton Isolation - Second",
+                                            50, 0, PI);
+        // ---------------------------------------------
+        trilepton_iso_3 = book_histogram_1d("trilepton_iso_3", "Trilepton Isolation - Third",
+                                            50, 0, PI);
+        // ---------------------------------------------
+    }
     
     HistCollection(): HistCollection("") {};
     
     ~HistCollection(){
         for(auto& p: histograms)
-            delete p.second;
-    };
+            delete p.second.hist;
+    }
+    
+    void draw(TCanvas* canvas,
+              std::pair<int, int> shape = {0,0}){
+        std::pair<int, int> null_shape(0,0);
+        if (shape == null_shape){
+            int n = (int)ceil(sqrt(histograms.size()));
+            shape.first = n;
+            shape.second = n;
+        }
+        canvas->Clear();
+        canvas->Divide(shape.first, shape.second);
+        int i = 1;
+        for(auto& p: histograms){
+            canvas->cd(i++);
+            p.second.hist->SetStats(false);
+            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;
@@ -63,8 +165,9 @@ class HistCollection{
             vec.push_back(p.first);
         return vec;
     }
+    
     TH1* get_by_id(const std::string &id){
-        TH1* hist = histograms[id];
+        TH1* hist = histograms[id].hist;
         return hist;
     }
 };

MiniTree.cpp

@@ -1,20 +1,18 @@
 #ifndef minitree_cxx
 #define minitree_cxx
+
 #include <cstdio>
 #include <string>
 #include <map>
 
-#include <TH1D.h>
-#include <TH2D.h>
 #include <TStyle.h>
-#include <TCanvas.h>
 
 #include "MiniTree.h"
 #include "HistCollection.h"
 
 
-HistCollection* build_histograms(const char* filename){
-    HistCollection *hc = new HistCollection(filename);
+HistCollection* build_histograms(const char* sample_name, const char* filename){
+    HistCollection *hc = new HistCollection(sample_name);
     TFile *f = TFile::Open(filename);
     TTree *tree = (TTree*) f->Get("tree");
     MiniTree minitree(tree);
@@ -26,9 +24,45 @@ HistCollection* build_histograms(const char* filename){
     return hc;
 }
 
+void set_branch_statuses(MiniTree *minitree){
+    auto on = [minitree](const char* bname){
+        minitree->fChain->SetBranchStatus(bname, true);
+    };
+    auto off = [minitree](const char* bname){
+        minitree->fChain->SetBranchStatus(bname, false);
+    };
+    off("*");
+    on("nLepGood");
+    on("LepGood_pt");
+    on("LepGood_eta");
+    on("LepGood_phi");
+    on("LepGood_mcPt");
+    on("LepGood_charge");
+    on("nJet");
+    on("Jet_btagCMVA");
+    on("Jet_pt");
+    on("nGenTop");
+}
 
 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){
+        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;
+    };
+    
     if (fChain == 0) return;
+    set_branch_statuses(this);
 
     Long64_t nentries = fChain->GetEntriesFast();
 
@@ -38,12 +72,44 @@ void MiniTree::Loop(HistCollection* hc){
         if (ientry < 0) break;
         nb = fChain->GetEntry(jentry);   nbytes += nb;
         // if (Cut(ientry) < 0) continue;
+        
+        // BEGIN ANALYSIS CODE
         hc->lepton_count->Fill(nLepGood);
+        hc->top_quark_count->Fill(nGenTop);
         for(int i=0; i<nLepGood; i++){
             double delta_pt_val = LepGood_pt[i] - LepGood_mcPt[i];
             hc->delta_pt->Fill(delta_pt_val);
             hc->lepton_count_vs_delta_pt->Fill(nLepGood, delta_pt_val);
         }
+        hc->jet_count->Fill(nJet);
+        int b_jet_count = 0;
+        for(int i=0; i<nJet; i++){
+            hc->b_jet_discriminator->Fill(Jet_btagCMVA[i]);
+            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){
+                hc->b_jet_pt_3_or_more->Fill(Jet_pt[i]);
+            }
+        }
+        int lepton_count_pass_miniiso= 0;
+        for(int i=0; i<nLepGood; i++){
+            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);
+        }
+        if (nLepGood == 2 && (LepGood_charge[0] * LepGood_charge[1]) == 1){
+            hc->jet_count_ss_dilepton->Fill(nJet);
+        }
+        if (nLepGood == 2 && (LepGood_charge[0] * LepGood_charge[1]) == -1){
+            hc->jet_count_os_dilepton->Fill(nJet);
+        }
     }
 }
 #endif // tree_cxx

canvas_wrapper.py

@@ -0,0 +1,66 @@
+"""
+Helper module for displaying ROOT canvases in ipython notebooks
+
+Usage example:
+# Save this file as rootnotes.py to your working directory.
+import rootnotes
+c1 = rootnotes.canvas()
+fun1 = TF1( 'fun1', 'abs(sin(x)/x)', 0, 10)
+c1.SetGridx()
+c1.SetGridy()
+fun1.Draw()
+c1
+
+More examples: http://mazurov.github.io/webfest2013/
+
+@author alexander.mazurov@cern.ch
+@author andrey.ustyuzhanin@cern.ch
+@date 2013-08-09
+"""
+
+import ROOT
+ROOT.gROOT.SetBatch()
+
+import tempfile
+from IPython.core import display
+
+
+def canvas(name="icanvas", size=(800, 600)):
+    """Helper method for creating canvas"""
+
+    assert len(size) == 2
+    if len(size) != 2:
+        raise ValueError("Size must be a 2 element tuple")
+    # Check if icanvas already exists
+    canvas = ROOT.gROOT.FindObject(name)
+    if canvas:
+        return canvas
+    else:
+        return ROOT.TCanvas(name, name, size[0], size[1])
+
+
+def _display_canvas(canvas):
+    file = tempfile.NamedTemporaryFile(suffix=".png")
+    canvas.SaveAs(file.name)
+    ip_img = display.Image(filename=file.name, format='png', embed=True)
+    return ip_img._repr_png_()
+
+
+def _display_any(obj):
+    file = tempfile.NamedTemporaryFile(suffix=".png")
+    obj.Draw()
+    ROOT.gPad.SaveAs(file.name)
+    ip_img = display.Image(filename=file.name, format='png', embed=True)
+    return ip_img._repr_png_()
+
+
+# register display function with PNG formatter:
+png_formatter = get_ipython().display_formatter.formatters['image/png'] # noqa
+
+# Register ROOT types in ipython
+#
+# In [1]: canvas = rootnotes.canvas()
+# In [2]: canvas
+# Out [2]: [image will be here]
+png_formatter.for_type(ROOT.TCanvas, _display_canvas)
+png_formatter.for_type(ROOT.TF1, _display_any)

plotter.py

@@ -1,7 +1 @@
-import ROOT
-
-
-canvas = ROOT.TCanvas("c1")
-def do_plots():
-    ROOT.lepton_count.Draw()
-    canvas.Draw()
+import ROOT

utils.py

@@ -1,5 +1,6 @@
 import io
 import sys
+import itertools as it
 import ROOT
 
 class OutputCapture:
@@ -46,3 +47,26 @@ def normalize_columns(hist2d):
             norm = hist2d.GetBinContent(col, row) / sum_
             normHist.SetBinContent(col, row, norm)
     return normHist
+
+def stack_hist(hists, labels=None, id_="stack", title="", enable_fill=False, normalize_to=0):
+    """hists should be a list of TH1D objects
+    returns a new stacked histogram
+    """
+    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)):
+        hist_copy = hist
+        hist_copy = hist.Clone(str(i)+"_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:
+            integral = hist_copy.Integral()
+            hist_copy.Scale(normalize_to/integral, "nosw2")
+            hist_copy.SetStats(False)
+        stack.Add(hist_copy)
+    return stack