#!/usr/bin/env python
# coding: utf-8

# # Analysis of the di-muon spectrum using data from the CMS detector
# 
# This analysis takes data from the CMS experiment recorded in 2012 during Run B and C and extracts the di-muon spectrum. The di-muon spectrum is computed from the data by calculating the invariant mass of muon pairs with opposite charge. In the resulting plot, you are able to rediscover particle resonances in a wide energy range from the [eta meson](https://en.wikipedia.org/wiki/Eta_meson) at about 548 MeV up to the [Z boson](https://en.wikipedia.org/wiki/W_and_Z_bosons) at about 91 GeV.
# 
# The analysis code opens an interactive plot, which allows to zoom and navigate in the spectrum. Note that the bump at 30 GeV is not a resonance but an effect of the data taking due to the used trigger. The technical description of the dataset can be found in the respective record linked below.
# 
# The result of this analysis can be compared with [an official result of the CMS collaboration using data taken in 2010](https://cds.cern.ch/record/1456510), see the plot below:
# 
# ![](http://cds.cern.ch/record/1456510/files/pictures_samples_dimuonSpectrum_40pb-1_mod-combined.png)

# # Dataset description
# 
# The dataset consists of the following columns.
# 
# | Column name | Data type | Description |
# |-------------|-----------|-------------|
# | `nMuon` | `unsigned int` | Number of muons in this event |
# | `Muon_pt` | `float[nMuon]` | Transverse momentum of the muons (stored as an array of size `nMuon`) |
# | `Muon_eta` | `float[nMuon]` | Pseudorapidity of the muons |
# | `Muon_phi` | `float[nMuon]` | Azimuth of the muons |
# | `Muon_mass` | `float[nMuon]` | Mass of the muons |
# | `Muon_charge` | `int[nMuon]` | Charge of the muons (either 1 or -1) |

# # Data analysis
# 
# The following part of the notebook performs the data analysis of the dataset.

# In[ ]:


import ROOT


# ## Enable multi-threading
# 
# The default here is set to a single thread. You can choose the number of threads based on your system.

# In[ ]:


ROOT.ROOT.EnableImplicitMT(1)


# ## Create dataframe from NanoAOD files

# In[ ]:


df = ROOT.RDataFrame("Events", "root://eospublic.cern.ch//eos/opendata/cms/derived-data/AOD2NanoAODOutreachTool/Run2012BC_DoubleMuParked_Muons.root")


# ## Select events with exactly two muons

# In[ ]:


df_2mu = df.Filter("nMuon == 2", "Events with exactly two muons")


# ## Select events with two muons of opposite charge

# In[ ]:


df_os = df_2mu.Filter("Muon_charge[0] != Muon_charge[1]", "Muons with opposite charge")


# ## Compute invariant mass of the dimuon system
# 
# The following code just-in-time compiles the C++ function to compute the invariant mass, so that the function can be called in the Define node of the ROOT dataframe.

# In[ ]:


ROOT.gInterpreter.Declare(
"""
using namespace ROOT::VecOps;
float computeInvariantMass(RVec<float>& pt, RVec<float>& eta, RVec<float>& phi, RVec<float>& mass) {
    ROOT::Math::PtEtaPhiMVector m1(pt[0], eta[0], phi[0], mass[0]);
    ROOT::Math::PtEtaPhiMVector m2(pt[1], eta[1], phi[1], mass[1]);
    return (m1 + m2).mass();
}
""")
df_mass = df_os.Define("Dimuon_mass", "computeInvariantMass(Muon_pt, Muon_eta, Muon_phi, Muon_mass)")


# ## Book histogram of dimuon mass spectrum

# In[ ]:


bins = 30000 # Number of bins in the histogram
low = 0.25 # Lower edge of the histogram
up = 300.0 # Upper edge of the histogram
hist = df_mass.Histo1D(ROOT.RDF.TH1DModel("", "", bins, low, up), "Dimuon_mass")


# ## Request cut-flow report

# In[ ]:


report = df_mass.Report()


# ## Create canvas for plotting

# In[ ]:


ROOT.gStyle.SetOptStat(0)
ROOT.gStyle.SetTextFont(42)
c = ROOT.TCanvas("c", "", 800, 700)
c.SetLogx()
c.SetLogy();


# ## Draw histogram

# In[ ]:


hist.GetXaxis().SetTitle("m_{#mu#mu} (GeV)")
hist.GetXaxis().SetTitleSize(0.04)
hist.GetYaxis().SetTitle("N_{Events}")
hist.GetYaxis().SetTitleSize(0.04)
hist.SetStats(False)
hist.Draw();


# ## Draw labels

# In[ ]:


label = ROOT.TLatex()
label.SetTextAlign(22)
label.DrawLatex(0.55, 3.0e4, "#eta")
label.DrawLatex(0.77, 7.0e4, "#rho,#omega")
label.DrawLatex(1.20, 4.0e4, "#phi")
label.DrawLatex(4.40, 1.0e5, "J/#psi")
label.DrawLatex(4.60, 1.0e4, "#psi'")
label.DrawLatex(12.0, 2.0e4, "Y(1,2,3S)")
label.DrawLatex(91.0, 1.5e4, "Z")
label.SetNDC(True)
label.SetTextAlign(11)
label.SetTextSize(0.04)
label.DrawLatex(0.10, 0.92, "#bf{CMS Open Data}")
label.SetTextAlign(31)
label.DrawLatex(0.90, 0.92, "#sqrt{s} = 8 TeV, L_{int} = 11.6 fb^{-1}");


# ## Draw interactive canvas

# In[ ]:


c.Draw()


# ## Print cut-flow report

# In[ ]:


report.Print()