df102_NanoAODDimuonAnalysis-PurePython.py

import ROOT

# Enable multi-threading
ROOT.ROOT.EnableImplicitMT()

# Create list of input files
files = ROOT.std.vector("string")()
files.push_back("root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012B_DoubleMuParked.root")
files.push_back("root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012C_DoubleMuParked.root")

# Create dataframe from list of NanoAOD files
df = ROOT.ROOT.RDataFrame("Events", files);

# For simplicity, select only events with exactly two muons and require opposite charge
df_2mu = df.Filter("nMuon == 2", "Events with exactly two muons");
df_os  = df_2mu.Filter("Muon_charge[0] != Muon_charge[1]", "Muons with opposite charge");

# C++ Code to compute invariant mass of the dimuon system
compute_mass_code = """
using namespace ROOT::VecOps;
auto compute_mass = [](RVec<float> &pt, RVec<float> &eta, RVec<float> &phi, RVec<float> &mass) {
      // Compose four-vectors of both muons
      TLorentzVector p1, p2;
      p1.SetPtEtaPhiM(pt[0], eta[0], phi[0], mass[0]);
      p2.SetPtEtaPhiM(pt[1], eta[1], phi[1], mass[1]);
      // Add four-vectors to build dimuon system and return the invariant mass
      return (p1 + p2).M();
   };
"""

# Decleare function in the ROOT Interpreter to be accessible with JITted code
ROOT.gInterpreter.Declare(compute_mass_code)

df_mass = df_os.Define("Dimuon_mass1", "compute_mass(Muon_pt, Muon_eta, Muon_phi, Muon_mass)")

# Compute invariant mass of the dimuon system
h = df_mass.Histo1D(("Dimuon_mass1", "Dimuon_mass1", 30000, 0.25, 300), "Dimuon_mass1")

# Request cut-flow report
report = df_mass.Report()

# Produce plot
ROOT.gStyle.SetOptStat(0)
ROOT.gStyle.SetTextFont(42)

c = ROOT.TCanvas("c", "", 800, 700);
c.SetLogx(); c.SetLogy();
h.SetTitle("");
h.GetXaxis().SetTitle("m_{#mu#mu} (GeV)"); h.GetXaxis().SetTitleSize(0.04);
h.GetYaxis().SetTitle("N_{Events}"); h.GetYaxis().SetTitleSize(0.04);
h.Draw();

label = ROOT.TLatex()
label.SetNDC(True);

label.DrawLatex(0.175, 0.740, "#eta");
label.DrawLatex(0.205, 0.775, "#rho,#omega");
label.DrawLatex(0.270, 0.740, "#phi");
label.DrawLatex(0.400, 0.800, "J/#psi");
label.DrawLatex(0.415, 0.670, "#psi'");
label.DrawLatex(0.485, 0.700, "Y(1,2,3S)");
label.DrawLatex(0.755, 0.680, "Z");
label.SetTextSize(0.040); label.DrawLatex(0.100, 0.920, "#bf{CMS Open Data}");
label.SetTextSize(0.030); label.DrawLatex(0.630, 0.920, "#sqrt{s} = 8 TeV, L_{int} = 11.6 fb^{-1}");

c.SaveAs("dimuon_spectrum.pdf");

# Print cut-flow report
report.Print();