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

# <img src="http://files.oproject.org/img/HeaderOpenData.png">
# # CMS Open Data Example #6: Fitting the J/$\psi$ fit
# 
# ## Import required modules

# In[4]:


from ROOT import TFile, TCanvas, TH1F, TLorentzVector
get_ipython().run_line_magic('jsroot', 'on')


# ## Read in Data from Input File

# In[5]:


file = TFile("../data/Dimuons.root","READ");


# ## Declare Histograms

# In[6]:


Jpsi = TH1F("Jpsi","#mu#mu mass;#mu#mu mass [GeV];Events", 1000, 2, 3.5)


# # Compute Invariant Mass of J/$\psi$

# In[7]:


for dimu in file.Dimuons:

    if dimu.Muon1_Global and dimu.Muon2_Global:     
    
        muon1 = TLorentzVector(dimu.Muon1_Px, dimu.Muon1_Py, dimu.Muon1_Pz, dimu.Muon1_Energy)
    
        muon2 = TLorentzVector(dimu.Muon2_Px, dimu.Muon2_Py, dimu.Muon2_Pz, dimu.Muon2_Energy)
    
        InvariantMass = (muon1 + muon2).M() 
        
        if InvariantMass > 2.0 and InvariantMass < 4.0:

            Jpsi.Fill(InvariantMass)


# ## Fitting the J/$\psi$ Peak

# In[8]:


from ROOT import TF1, kGreen, kRed

Canvas = TCanvas()


# Possible fitting options: <BR>
# "gaus" gaussian <BR>
# "pol1" 1st order polynomial <BR>

# In[9]:


Gaussian = TF1("Gaussian","gaus", 3.03, 3.16)
Background  = TF1("Background","pol1", 2, 3.5)

Gaussian.SetLineColor(kRed)
Background.SetLineColor(kGreen)

Jpsi.Fit(Gaussian,"R")
Jpsi.Fit(Background,"R+")

Canvas.Draw()


# In[3]:


from ROOT import TFile, TCanvas, TH1F, TLorentzVector
get_ipython().run_line_magic('jsroot', 'on')

file = TFile("../data/Dimuons.root","READ");

z_boson = TH1F("zboson","#mu#mu mass;#mu#mu mass [GeV];Events", 1000, 70, 120)

for dimu in file.Dimuons:

    if dimu.Muon1_Global and dimu.Muon2_Global:     
    
        muon1 = TLorentzVector(dimu.Muon1_Px, dimu.Muon1_Py, dimu.Muon1_Pz, dimu.Muon1_Energy)
    
        muon2 = TLorentzVector(dimu.Muon2_Px, dimu.Muon2_Py, dimu.Muon2_Pz, dimu.Muon2_Energy)
    
        InvariantMass = (muon1 + muon2).M() 
        
        if InvariantMass > 70.0 and InvariantMass < 120.0:

            z_boson.Fill(InvariantMass)
            
from ROOT import TF1, kGreen, kRed

Canvas = TCanvas()

Gaussian = TF1("Gaussian","gaus", 87, 95)
Background  = TF1("Background","pol1", 70, 120)

Gaussian.SetLineColor(kRed)
Background.SetLineColor(kGreen)

z_boson.Fit(Gaussian,"R")
z_boson.Fit(Background,"R+")

Canvas.Draw()


# In[1]:


from ROOT import TFile, TCanvas, TH1F, TLorentzVector
get_ipython().run_line_magic('jsroot', 'on')

file = TFile("../data/Dimuons.root","READ");

upsilon = TH1F("upsilon","#mu#mu mass;#mu#mu mass [GeV];Events", 1000, 8.0, 12.0)

for dimu in file.Dimuons:

    if dimu.Muon1_Global and dimu.Muon2_Global:     
    
        muon1 = TLorentzVector(dimu.Muon1_Px, dimu.Muon1_Py, dimu.Muon1_Pz, dimu.Muon1_Energy)
    
        muon2 = TLorentzVector(dimu.Muon2_Px, dimu.Muon2_Py, dimu.Muon2_Pz, dimu.Muon2_Energy)
    
        InvariantMass = (muon1 + muon2).M() 
        
        if InvariantMass > 8.0 and InvariantMass < 12.0:

            upsilon.Fill(InvariantMass)
            


# In[2]:


from ROOT import TF1, kGreen, kRed

Canvas = TCanvas()

Gaussian = TF1("Gaussian","gaus", 9.3, 9.6)
Gaussian2 = TF1("Gaussian2","gaus", 9.84, 10.16)
Gaussian3 = TF1("Gaussian3","gaus", 10.2, 10.5)
Background  = TF1("Background","pol1", 8, 12)

Gaussian.SetLineColor(kRed)
Gaussian2.SetLineColor(kRed)
Gaussian3.SetLineColor(kRed)
Background.SetLineColor(kGreen)

upsilon.Fit(Gaussian,"R")
upsilon.Fit(Gaussian2,"R2+")
upsilon.Fit(Gaussian3,"R3+")
upsilon.Fit(Background,"R+")

Canvas.Draw()


# In[ ]: