def plot(input_files, output_files, title, legend):
handle_lhe = Handle('LHEEventProduct')
label_lhe = ('externalLHEProducer')
bins = array.array('f', [1. * i for i in range(0, 201)])
h = ROOT.TH1F(legend, legend, len(bins) - 1, bins)
for input_file in input_files:
events = Events(input_file)
for event in events:
event.getByLabel(label_lhe, handle_lhe)
lhe = handle_lhe.product()
invmass = []
for status, pdg, mom in zip(lhe.hepeup().ISTUP,
lhe.hepeup().IDUP,
lhe.hepeup().PUP):
if status == 1 and abs(pdg) in [11, 13, 15]:
l = ROOT.TLorentzVector(mom.x[0], mom.x[1], mom.x[2],
mom.x[3])
invmass.append(l)
if len(invmass) == 2:
h.Fill((invmass[0] + invmass[1]).M())
else:
raise RuntimeError(
'Did not find exactly 2 but %d LHE leptons in %s' %
(len(invmass), input_file))
c = ROOT.TCanvas('c', 'c')
c.SetLogy()
c.SetGrid()
h.SetTitle(title)
h.GetXaxis().SetTitle('m_ll [GeV]')
h.GetYaxis().SetTitle('# events')
h.Draw()
for output_file in output_files:
c.SaveAs(output_file)
del h
del c
input_tree.SetBranchStatus('*', 0)
input_tree.SetBranchStatus(collection_sz_name, 1)
input_tree.SetBranchStatus(collection_pt_name, 1)
input_tree.SetBranchStatus(collection_eta_name, 1)
input_tree.SetBranchStatus(collection_phi_name, 1)
input_tree.SetBranchStatus(collection_mass_name, 1)
input_tree.SetBranchStatus(collection_pdgId_name, 1)
if use_genweight:
input_tree.SetBranchStatus(GENWEIGHT_BR_NAME, 1)
for idx in range(nof_events):
input_tree.GetEntry(idx)
higgses = []
for partIdx in range(sz_br[0]):
if pdgId_br[partIdx] == 25:
p4 = ROOT.TLorentzVector()
p4.SetPtEtaPhiM(pt_br[partIdx], eta_br[partIdx],
phi_br[partIdx], mass_br[partIdx])
higgses.append(p4)
if len(higgses) != 2:
logging.warning(
"Event #{} did not contain two LHE parton-level Higgses!".
format(idx))
continue
leading_idx = 0 if higgses[0].E() > higgses[1] else 1
leading_higgs = higgses[leading_idx]
subleading_higgs = higgses[1 - leading_idx]
higgs_sum = leading_higgs + subleading_higgs
mhh = higgs_sum.M()
def dump_yields(fn, fn_out):
print('Dumping event yields from %s to %s' % (fn, fn_out))
f = ROOT.TFile.Open(fn)
t = f.Get('Events')
x_branch = array.array(*get_type(x_var))
y_branch = array.array(*get_type(y_var)) if is2D else None
req_nLHESW = 9
run = array.array('I', [0])
luminosityBlock = array.array('I', [0])
event = array.array('L', [0])
genWeight = array.array('f', [0.])
puWeight = array.array('f', [0.])
puWeightUp = array.array('f', [0.])
puWeightDown = array.array('f', [0.])
nLHEScaleWeight = array.array('I', [0])
LHEScaleWeight = array.array('f', [0.] * req_nLHESW)
max_objs = 32
nLHEPart = array.array('B', [0])
LHEPart_pt = array.array('f', [0.] * max_objs)
LHEPart_eta = array.array('f', [0.] * max_objs)
LHEPart_phi = array.array('f', [0.] * max_objs)
LHEPart_mass = array.array('f', [0.] * max_objs)
LHEPart_pdgId = array.array('i', [0] * max_objs)
t.SetBranchAddress('run', run)
t.SetBranchAddress('luminosityBlock', luminosityBlock)
t.SetBranchAddress('event', event)
t.SetBranchAddress(x_var, x_branch)
if is2D:
t.SetBranchAddress(y_var, y_branch)
t.SetBranchAddress('genWeight', genWeight)
t.SetBranchAddress('puWeight', puWeight)
t.SetBranchAddress('puWeightUp', puWeightUp)
t.SetBranchAddress('puWeightDown', puWeightDown)
t.SetBranchAddress('nLHEScaleWeight', nLHEScaleWeight)
t.SetBranchAddress('LHEScaleWeight', LHEScaleWeight)
if apply_mll_cut:
t.SetBranchAddress('nLHEPart', nLHEPart)
t.SetBranchAddress('LHEPart_pt', LHEPart_pt)
t.SetBranchAddress('LHEPart_eta', LHEPart_eta)
t.SetBranchAddress('LHEPart_phi', LHEPart_phi)
t.SetBranchAddress('LHEPart_mass', LHEPart_mass)
t.SetBranchAddress('LHEPart_pdgId', LHEPart_pdgId)
bins = {
'LHE_Njets': list(map(float, range(6))),
'LHE_HT': [0., 100., 200., 400., 600., 800., 1200., 2500., 1.e5]
}
bins_arr = {key: array.array('d', bins[key]) for key in bins}
mll_bins = [10., 50.]
def get_mll_str(val):
assert (len(mll_bins) == 2)
if val < mll_bins[0]:
return 'lt%s' % str(int(mll_bins[0]))
elif mll_bins[0] <= val < mll_bins[1]:
return '%sto%s' % (str(int(mll_bins[0])), str(int(mll_bins[1])))
elif val >= mll_bins[1]:
return 'gt%s' % (str(int(mll_bins[1])))
else:
raise RuntimeError('Unexpected value: %f (%s)' % (val, mll_bins))
histograms = collections.OrderedDict()
bins_x = bins_arr[x_var]
bins_y = bins_arr[y_var] if is2D else None
def create_histogram(key, title):
if is2D:
histograms[key] = ROOT.TH2D(key, title,
len(bins_x) - 1, bins_x,
len(bins_y) - 1, bins_y)
for bin_idx in range(len(bins_x) - 1):
histograms[key].GetXaxis().SetBinLabel(
bin_idx + 1, '%d <= %s < %d' %
(bins_x[bin_idx], x_var, bins_x[bin_idx + 1]))
for bin_idx in range(len(bins_y) - 1):
histograms[key].GetYaxis().SetBinLabel(
bin_idx + 1, '%d <= %s < %d' %
(bins_y[bin_idx], y_var, bins_y[bin_idx + 1]))
histograms[key].SetXTitle(x_var)
histograms[key].SetYTitle(y_var)
else:
histograms[key] = ROOT.TH1D(key, title, len(bins_x) - 1, bins_x)
for bin_idx in range(len(bins_x) - 1):
histograms[key].GetXaxis().SetBinLabel(
bin_idx + 1, '%d <= %s < %d' %
(bins_x[bin_idx], x_var, bins_x[bin_idx + 1]))
histograms[key].SetXTitle(x_var)
count_keys = collections.OrderedDict([
('Count', {
'nbins': 1,
'title': 'sum(1)',
}),
('CountFullWeighted', {
'nbins': 3,
'title': 'sum(gen * PU(central,up,down))',
}),
('CountWeighted', {
'nbins': 3,
'title': 'sum(sgn(gen) * PU(central,up,down))',
}),
('CountFullWeightedNoPU', {
'nbins': 1,
'title': 'sum(gen)',
}),
('CountPosWeight', {
'nbins': 1,
'title': 'sum(gen > 0)'
}),
('CountNegWeight', {
'nbins': 1,
'title': 'sum(gen < 0)'
}),
('CountWeightedNoPU', {
'nbins': 1,
'title': 'sum(sgn(gen))',
}),
('CountWeightedLHEWeightScale', {
'nbins': req_nLHESW,
'title': 'sum(sgn(gen) * PU(central) * LHE(scale))',
}),
('CountWeightedLHEWeightScaleNoPU', {
'nbins': req_nLHESW,
'title': 'sum(sgn(gen) * LHE(scale))',
}),
('CountFullWeightedLHEWeightScale', {
'nbins': req_nLHESW,
'title': 'sum(gen * PU(central) * LHE(scale))',
}),
('CountFullWeightedLHEWeightScaleNoPU', {
'nbins': req_nLHESW,
'title': 'sum(gen * LHE(scale))',
}),
])
if apply_mll_cut:
for val in [
mll_bins[0] - 1., (mll_bins[0] + mll_bins[1]) / 2.,
mll_bins[1] + 1
]:
for count_key, count_settings in count_keys.items():
for histogram_idx in range(count_settings['nbins']):
mll_str = get_mll_str(val)
key = '%s_%d_%s' % (count_key, histogram_idx, mll_str)
title = count_settings['title']
if count_settings['nbins'] > 1:
title += ' [bin = %d]' % histogram_idx
title += ' (mll %s)' % mll_str
create_histogram(key, title)
else:
for count_key, count_settings in count_keys.items():
for histogram_idx in range(count_settings['nbins']):
key = '%s_%d' % (count_key, histogram_idx)
title = count_settings['title']
if count_settings['nbins'] > 1:
title += ' [bin = %d]' % histogram_idx
create_histogram(key, title)
def clip(value, min_val=-10., max_val=10.):
return min(max(value, min_val), max_val)
def plot2d(histogram, plot_fn_base, width=1200, height=900):
canvas = ROOT.TCanvas('c1', 'c1', width, height)
ROOT.gStyle.SetOptStat(0)
histogram.Draw('col text')
canvas.SetLogy()
canvas.SaveAs('%s.png' % plot_fn_base)
canvas.SaveAs('%s.pdf' % plot_fn_base)
del canvas
def plot1d(histogram, plot_fn_base, width=1200, height=900):
canvas = ROOT.TCanvas('c1', 'c1', width, height)
ROOT.gStyle.SetOptStat(0)
histogram.SetLineWidth(2)
histogram.Draw('hist')
canvas.SetLogx()
canvas.SetLogy()
canvas.SetGrid()
canvas.SaveAs('%s.png' % plot_fn_base)
canvas.SaveAs('%s.pdf' % plot_fn_base)
del canvas
n = t.GetEntries()
printEvery = 100000
for i in range(n):
t.GetEntry(i)
if i % printEvery == 0:
rle = ':'.join(
map(lambda x: str(x[0]), [run, luminosityBlock, event]))
print('Processing event %d: %s' % (i, rle))
if apply_mll_cut:
invmass = []
for j in range(nLHEPart[0]):
if abs(LHEPart_pdgId[j]) in [11, 13, 15]:
lv = ROOT.TLorentzVector()
lv.SetPtEtaPhiM(LHEPart_pt[j], LHEPart_eta[j],
LHEPart_phi[j], LHEPart_mass[j])
invmass.append(lv)
if len(invmass) != 2:
continue
mll = (invmass[0] + invmass[1]).M()
suffix = '_%s' % get_mll_str(mll)
else:
suffix = ''
genWeight_sign = np.sign(genWeight[0])
counts = {
'Count_0': 1.,
'CountWeighted_0': genWeight_sign * puWeight[0],
'CountWeighted_1': genWeight_sign * puWeightUp[0],
'CountWeighted_2': genWeight_sign * puWeightDown[0],
'CountFullWeighted_0': genWeight[0] * puWeight[0],
'CountFullWeighted_1': genWeight[0] * puWeightUp[0],
'CountFullWeighted_2': genWeight[0] * puWeightDown[0],
'CountWeightedNoPU_0': genWeight_sign,
'CountFullWeightedNoPU_0': genWeight[0],
'CountPosWeight_0': genWeight[0] * (genWeight_sign > 0),
'CountNegWeight_0': genWeight[0] * (genWeight_sign < 0),
}
if nLHEScaleWeight[0] != req_nLHESW:
print('Error: event #%d' % i)
continue
for j in range(nLHEScaleWeight[0]):
LHEScaleWeight_clipped = clip(LHEScaleWeight[j])
counts['CountWeightedLHEWeightScale_%d' %
j] = genWeight_sign * puWeight[0] * LHEScaleWeight_clipped
counts['CountWeightedLHEWeightScaleNoPU_%d' %
j] = genWeight_sign * LHEScaleWeight_clipped
counts['CountFullWeightedLHEWeightScale_%d' %
j] = genWeight[0] * puWeight[0] * LHEScaleWeight_clipped
counts['CountFullWeightedLHEWeightScaleNoPU_%d' %
j] = genWeight[0] * LHEScaleWeight_clipped
for count_key in counts:
key = count_key + suffix
evtWeight = counts[count_key]
if is2D:
histograms[key].Fill(x_branch[0], y_branch[0], evtWeight)
else:
histograms[key].Fill(x_branch[0], evtWeight)
f_out = ROOT.TFile.Open(fn_out, 'recreate')
f_out.cd()
for histogram in histograms.values():
histogram.Write()
if plot_dir:
if is2D:
plot2d(histogram, os.path.join(plot_dir, histogram.GetName()))
else:
plot1d(histogram, os.path.join(plot_dir, histogram.GetName()))
f.Close()
f_out.Close()
