Skip to content

Instantly share code, notes, and snippets.

@douglasdavis

douglasdavis/first_histograms.py

Created March 8, 2018 04:04
Show Gist options
  • Save douglasdavis/78ae3935ee73a41efd6e787f8143a3f2 to your computer and use it in GitHub Desktop.
Save douglasdavis/78ae3935ee73a41efd6e787f8143a3f2 to your computer and use it in GitHub Desktop.
Download ZIP
quick and dirty plots!
Raw
first_histograms.py
#!/usr/bin/env python
import uproot
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['axes.labelsize'] = 12
mpl.rcParams['legend.fontsize'] = 12
mpl.rcParams['legend.frameon'] = False
mpl.rcParams['xtick.top'] = True
mpl.rcParams['ytick.right'] = True
mpl.rcParams['xtick.direction'] = 'in'
mpl.rcParams['ytick.direction'] = 'in'
mpl.rcParams['xtick.labelsize'] = 11
mpl.rcParams['ytick.labelsize'] = 11
mpl.rcParams['xtick.minor.visible'] = True
mpl.rcParams['ytick.minor.visible'] = True
mpl.rcParams['xtick.major.width'] = 0.8 # major tick width in points
mpl.rcParams['xtick.minor.width'] = 0.8 # minor tick width in points
mpl.rcParams['xtick.major.size'] = 6.0 # major tick size in points
mpl.rcParams['xtick.minor.size'] = 4.0 # minor tick size in points
mpl.rcParams['xtick.major.pad'] = 1.5 # distance to major tick label in points
mpl.rcParams['xtick.minor.pad'] = 1.4 # distance to the minor tick label in points
mpl.rcParams['ytick.major.width'] = 0.8 # major tick width in points
mpl.rcParams['ytick.minor.width'] = 0.8 # minor tick width in points
mpl.rcParams['ytick.major.size'] = 6.0 # major tick size in points
mpl.rcParams['ytick.minor.size'] = 4.0 # minor tick size in points
mpl.rcParams['ytick.major.pad'] = 1.5 # distance to major tick label in points
mpl.rcParams['ytick.minor.pad'] = 1.4 # distance to the minor tick label in points
from matplotlib.font_manager import FontProperties
font0 = FontProperties()
ATLASfont = font0.copy()
ATLASfont.set_style('italic')
ATLASfont.set_weight('bold')
def extended_hist(data, xmin, xmax, nbins, underflow=True, overflow=True, weights=None,lumi=1.0):
if weights is not None:
if weights.shape != data.shape:
raise ValueError(
'Unequal shapes data: {}; weights: {}'.format(
data.shape, weights.shape
)
)
edges = np.linspace(xmin,xmax,nbins+1)
neginf = np.array([-np.inf],dtype=np.float32)
posinf = np.array([ np.inf],dtype=np.float32)
uselsp = np.concatenate([neginf,edges,posinf])
if weights is not None:
weights *= lumi
hist, bin_edges = np.histogram(data,bins=uselsp,weights=weights)
else:
hist, bin_edges = np.histogram(data,bins=uselsp)
n = hist[1:-1]
if underflow:
n[0] += hist[0]
if overflow:
n[-1] += hist[-1]
if weights is not None:
bin_sumw2 = np.zeros(nbins+2,dtype=np.float32)
digits = np.digitize(data,edges)
for i in range(nbins+2):
bin_sumw2[i] = np.sum(np.power(weights[np.where(digits==i)[0]],2))
w = bin_sumw2[1:-1]
if underflow:
w[0] += bin_sumw2[0]
if overflow:
w[-1] += bin_sumw2[-1]
w = np.sqrt(w)
else:
w = np.sqrt(n)
centers = np.delete(edges,[0])-(np.ediff1d(edges)/2.0)
return n, w, centers, edges
def make_plot(tt_tree,tW_tree,tt_weight,tW_weight,
tt_cut,tW_cut,var_name,
xmin,xmax,nbin,lumi=36.1,xlabel='',yunit='', et='', ext='pdf',fakes=None, datas=None):
tt_var = tt_tree.array(var_name)[tt_cut]
tW_var = tW_tree.array(var_name)[tW_cut]
tt_w = tt_weight[tt_cut]
tW_w = tW_weight[tW_cut]
tt_hist = extended_hist(tt_var,xmin,xmax,nbin,weights=tt_w,lumi=lumi)
tW_hist = extended_hist(tW_var,xmin,xmax,nbin,weights=tW_w*0.706,lumi=lumi)
fig, ax = plt.subplots()
if fakes is not None:
f_ttbar_var = fakes['ttbar_t'].array(var_name)
f_tW_var = fakes['tW_t'].array(var_name)
f_ttbar_w = fakes['ttbar_t'].array('weight_nominal')
f_tW_w = fakes['tW_t'].array('weight_nominal')
f_var = np.concatenate([f_ttbar_var,f_tW_var])
f_w = np.concatenate([f_ttbar_w,f_tW_w])
f_ttbar_cut1 = fakes['ttbar_t'].array('njets')
f_tW_cut1 = fakes['tW_t'].array('njets')
f_cut1 = np.concatenate([f_ttbar_cut1,f_tW_cut1])
f_ttbar_cut2 = fakes['ttbar_t'].array('nbjets')
f_tW_cut2 = fakes['tW_t'].array('nbjets')
f_cut2 = np.concatenate([f_ttbar_cut2,f_tW_cut2])
f_ttbar_cut3 = fakes['ttbar_t'].array('SS')
f_tW_cut3 = fakes['tW_t'].array('SS')
f_cut3 = np.concatenate([f_ttbar_cut3,f_tW_cut3])
f_cut = (f_cut1 == fakes['njets']) & (f_cut2 == fakes['nbjets']) & (f_cut3 == False)
f_var = f_var[f_cut]
f_w = f_w[f_cut]
f_hist = extended_hist(f_var,xmin,xmax,nbin,weights=f_w,lumi=lumi)
ntt = np.round(np.sum(tt_hist[0]))
ntw = np.round(np.sum(tW_hist[0]))
nf = np.round(np.sum(f_hist[0]))
print(ntw, ntt, nf)
ax.hist([f_hist[2],tt_hist[2],tW_hist[2]],bins=tt_hist[3],weights=[f_hist[0],tt_hist[0],tW_hist[0]],
color=['orange','red','dodgerblue'],label=[r'MC Fakes (%d)' % int(nf),
r'$t\bar{t}$ (MC16a) (%d)' % int(ntt),
r'$tW$ (MC16a) (%d)' % int(ntw)],stacked=True,
histtype='stepfilled')
else:
ax.hist([tt_hist[2],tW_hist[2]],bins=tt_hist[3],weights=[tt_hist[0],tW_hist[0]],
color=['red','dodgerblue'],label=[r'$t\bar{t}$',r'$tW$'],stacked=True,
histtype='stepfilled')
#if datas is not None:
# data_tree = datas['data_t']
# data_var = data_tree.array(var_name)
# data_cut3 = data_tree.array('SS')
# data_cut1 = data_tree.array('1j1b')
# data_cut = (data_cut3 == False) & (data_cut1 == True)
# data_var = data_var[data_cut]
# data_hist = extended_hist(data_var,xmin,xmax,nbin,weights=None,lumi=1.0)
# #print(data_hist)
# ax.errorbar(data_hist[2],data_hist[0],yerr=data_hist[1],label='Data',fmt='o',color='black')
ax.set_xlabel(xlabel,horizontalalignment='right', x=1.0)
ax.set_ylabel('Events / {} {}'.format(((xmax-xmin)/nbin),yunit),horizontalalignment='right', y=1.0)
ax.legend()
ax.set_ylim([0,ax.get_ylim()[1]*1.5])
ax.set_xlim([xmin,xmax])
ax.text(.050,.90,'ATLAS',transform=ax.transAxes,fontproperties=ATLASfont,size=14)
ax.text(.200,.90,'Internal',transform=ax.transAxes,size=14)
ax.text(.050,.84,r'$\sqrt{s}=13$ TeV, 36.1 fb$^{-1}$',transform=ax.transAxes,size=12)
ax.text(.050,.78,et,transform=ax.transAxes,size=12)
fig.savefig('outs/{}.{}'.format(var_name,ext))
if __name__ == '__main__':
tt_tree = uproot.open('~/Desktop/tWhist/tt.root')['WtLoop_nominal']
tW_tree = uproot.open('~/Desktop/tWhist/tW.root')['WtLoop_nominal']
tt_cut1 = tt_tree.array('njets') == 1
tt_cut2 = tt_tree.array('nbjets') == 1
tt_cut3 = tt_tree.array('SS') == False
tW_cut1 = tW_tree.array('njets') == 1
tW_cut2 = tW_tree.array('nbjets') == 1
tW_cut3 = tW_tree.array('SS') == False
tt_cut = tt_cut1 & tt_cut2 & tt_cut3
tW_cut = tW_cut1 & tW_cut2 & tW_cut3
tt_weight = tt_tree.array('weight_nominal')
tW_weight = tW_tree.array('weight_nominal')
fake_dict = { 'ttbar_t' : uproot.open('~/Desktop/tWhist/tt.root')['WtLoop_Fake_nominal'],
'tW_t' : uproot.open('~/Desktop/tWhist/tW.root')['WtLoop_Fake_nominal'],
'njets' : 1 ,
'nbjets' : 1
}
data_dict = { 'data_t' : uproot.open('~/Desktop/tWhist/dd.root')['WtLoop_Fake_nominal'],
'nbjets' : 1,
'njets' : 1
}
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'pTsys_lep1lep2jet1met',
0,150,20,
xlabel=r'$p_\mathrm{T}^\mathrm{sys}(\ell_1\ell_2 E_\mathrm{T}^\mathrm{miss} j_1)$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'deltapT_lep1lep2jet1_met',
-100,100,20,
xlabel=r'$\Delta p_\mathrm{T} (\ell_1\ell_2 j_1,E_\mathrm{T}^\mathrm{miss})$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'sumet',
0,800,20,
xlabel=r'$\sum E_\mathrm{T}$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'eta_lep1lep2jet1met',
-5,5,20,
xlabel=r'$\eta(\ell_1 \ell_2 j_1 E_\mathrm{T}^\mathrm{miss}$)',
yunit='',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'deltapT_lep1lep2_met',
-100,100,20,xlabel=r'$\Delta p_\mathrm{T}(\ell_1\ell_2, E_\mathrm{T}^\mathrm{miss})$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'pTsys_lep1lep2jet1',
0,200,20,
xlabel=r'$p_\mathrm{T}^\mathrm{sys}(\ell_1\ell_2 j_1)$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'cent_lep1lep2',
0,1,20,
xlabel=r'Centrality($\ell_1\ell_2$)',
yunit='',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'mass_lep2jet1',
0,400,20,
xlabel=r'$m(\ell_2 j_1)$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'mass_lep1jet1',
0,500,20,
xlabel=r'$m(\ell_1 j_1)$ [GeV]',
yunit='GeV',et=r'Release 21, 1j1b selection',fakes=fake_dict,datas=data_dict)
"""
tt_cut1 = tt_tree.array('njets') == 2
tW_cut1 = tW_tree.array('njets') == 2
tt_cut = tt_cut1 & tt_cut2
tW_cut = tW_cut1 & tW_cut2
fake_dict['njets'] = 2
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'pTsys_lep1lep2',
0,150,20,
xlabel=r'$p_\mathrm{T}^\mathrm{sys}(\ell_1\ell_2)$ [GeV]',
yunit='GeV', et=r'Release 21, 2j1b selection',fakes=fake_dict,datas=data_dict)
make_plot(tt_tree,tW_tree,tt_weight,tW_weight,tt_cut,tW_cut,
'deltaR_lep1lep2_jet1jet2met',
1.5,4.5,20,
xlabel=r'$\Delta R(\ell_1\ell_2, E_\mathrm{T}^\mathrm{miss} j_1 j_2)$',
yunit='', et=r'Release 21, 2j1b selection',fakes=fake_dict,datas=data_dict)
"""
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment