def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata['dataset']
integratedLuminosity = 137.19*1000 # fb^{-1} to pb^{-1}
ht = events.HT
weights = integratedLuminosity*events.CrossSection[ht > 1200]/len(events)
GenParticles = events.GenParticles
finalParticles = (GenParticles.Status == 1) & (GenParticles.pt > 1) & (abs(GenParticles.eta) < 2.5) & (GenParticles.Charge != 0)
nTracksGEN = ak.sum(finalParticles[ht > 1200], axis=1)
tracks = events.Tracks
tracks_pt = np.sqrt(tracks.x**2 + tracks.y**2)
tracks_eta = np.arcsinh(tracks.z / tracks_pt)
track_cut = (tracks_pt > 1.) & (abs(tracks_eta) < 2.5) & (tracks.fromPV0 >= 2) & tracks.matchedToPFCandidate
nTracksRECO = ak.to_numpy(ak.sum(track_cut[ht > 1200], axis=1))
output["sumw"][dataset] += len(events)
output["nTracksHist"].fill(
dataset="CMSSW GEN",
nTracks=nTracksGEN,
weight=weights
)
output["nTracksHist"].fill(
dataset="CMSSW RECO",
nTracks=nTracksRECO,
weight=weights
)
return output
class CutStats():
def __init__(self, **kwargs):
self.debug = kwargs.get("debug")
self.event_stats = {}
self.object_stats = {}
self.objects = {}
def add_initial_events(self, sample, n_events):
if sample not in self.event_stats.keys():
self.event_stats[sample] = { "n_events_initial" : n_events }
else:
self.event_stats[sample]["n_events_initial"] += n_events
def add_initial_objects(self, sample, object, n_objects):
if object not in self.object_stats.keys():
self.object_stats[object] = {}
if sample not in self.object_stats[object].keys():
self.object_stats[object][sample] = { "n_objects_initial" : n_objects_initial }
else:
self.object_stats[object][sample]["n_objects_initial"] += n_objects_initial
def add_event_cuts(self, events cuts, names, sample):
if sample not in self.event_stats.keys():
self.event_stats[sample] = { "n_events_initial" : n_events }
else:
self.event_stats[sample]["n_events_initial"] += n_events
for cut, name in zip(cuts, names):
if name not in self.event_stats[sample].keys():
self.event_stats[sample][name] = awkward.sum(events[cut])
else:
self.event_stats[sample][name] += awkward.sum(events[cut])
def rescale(accumulator, xsecs=xsecs, lumi=lumi, data="BTagMu"):
#def rescale(accumulator, xsecs=xsecs, data="BTagMu"):
"""Scale by lumi"""
#lumi = 1000*lumi # Convert lumi to pb^-1
from coffea import hist
scale = {}
sumxsecs = ak.sum(xsecs.values())
#N_data = accumulator['nbtagmu_event_level'][data]
#print("N_data =", N_data)
for dataset, N_mc in collections.OrderedDict(sorted(accumulator['sumw'].items())).items():
if dataset in xsecs:
print(" ", dataset, "\t", N_mc, "events\t", xsecs[dataset], "pb")
#scale[dataset] = (xsecs[dataset]/sumxsecs)*(N_data/N_mc)
scale[dataset] = (xsecs[dataset]*lumi)/N_mc
else:
print(" ", "X ", dataset)
scale[dataset] = 0#lumi / N_mc
print(scale)
datasets_mc = [item for item in list(xsecs.keys()) if not 'GluGlu' in item]
for h in accumulator.values():
if isinstance(h, hist.Hist):
h.scale(scale, axis="dataset")
N_data = ak.sum(h[data].values().values())
N_mc = ak.sum(h[datasets_mc].sum('dataset', 'flavor').values().values())
#scaletodata = dict(zip(scale.keys(), len(scale)*[1./N_data]))
scaletodata = dict(zip(scale.keys(), len(scale)*[N_data/N_mc]))
h.scale(scaletodata, axis="dataset")
return accumulator
def ggTauTau_inclusive_preselection(events, photons, electrons, muons, taus, options, debug):
cut_diagnostics = utils.CutDiagnostics(events = events, debug = debug, cut_set = "[analysis_selections.py : ggTauTau_inclusive_preselection]")
# Get number of electrons, muons, taus
selected_electrons = electrons[lepton_selections.select_electrons(events, photons, electrons, options, debug)]
selected_muons = muons[lepton_selections.select_muons(events, photons, muons, options, debug)]
selected_taus = taus[tau_selections.select_taus(events, photons, selected_muons, selected_electrons, taus, options, debug)]
n_electrons = awkward.num(selected_electrons)
n_muons = awkward.num(selected_muons)
n_taus = awkward.num(selected_taus)
# Require >= 1 lep/tau
n_leptons_and_taus = n_electrons + n_muons + n_taus
lep_tau_cut = n_leptons_and_taus >= options["n_leptons_and_taus"]
# Require OS leptons/taus for events with 2 leptons/taus
sum_charge = awkward.sum(selected_electrons.charge, axis=1) + awkward.sum(selected_muons.charge, axis=1) + awkward.sum(selected_taus.charge, axis=1)
charge_cut = sum_charge == 0
two_leptons = n_leptons_and_taus == 2
not_two_leptons = n_leptons_and_taus != 2
os_cut = (two_leptons & charge_cut) | not_two_leptons # only require 2 OS leptons if there are ==2 leptons in the event
all_cuts = lep_tau_cut & os_cut
cut_diagnostics.add_cuts([lep_tau_cut, os_cut, all_cuts], ["N_leptons + N_taus >= 1", "OS dileptons", "all"])
return events[all_cuts], photons[all_cuts], selected_electrons[all_cuts], selected_muons[all_cuts], selected_taus[all_cuts]
def compute_eff_algo(tau_1, tau_2, a, Pt_thr):
eff_presel = ak.sum(ak.num(tau_1, axis=-1))
deepTau_thr_1 = deep_thr(tau_1, a, Pt_thr)
deepTau_thr_2 = deep_thr(tau_2, a, Pt_thr)
deepTau_mask_1 = tau_1.deepTau_VSjet > deepTau_thr_1
deepTau_mask_2 = tau_2.deepTau_VSjet > deepTau_thr_2
eff = ak.sum(ditau_selection(deepTau_mask_1, deepTau_mask_2))
return compute_eff_witherr(eff, eff_presel)
def test_more():
nparray = np.array(
[np.datetime64("2021-06-03T10:00"), np.datetime64("2021-06-03T11:00")]
)
akarray = ak.Array(nparray)
assert (akarray[1:] - akarray[:-1]).tolist() == [np.timedelta64(60, "m")]
assert ak.sum(akarray[1:] - akarray[:-1]) == np.timedelta64(60, "m")
assert ak.sum(akarray[1:] - akarray[:-1], axis=0) == [np.timedelta64(60, "m")]
def test():
array = ak.Array([
[0, 1, None],
[None, 3],
[],
None,
[4, 5, None, 6],
])
assert ak.sum(array, axis=-1).tolist() == [
1, # 0 + 1
3, # 3
0, # list is empty
None, # list is missing
15, # 4 + 5 + 6
]
assert ak.sum(array, axis=-2).tolist() == [
4,
9,
0,
6,
# 0+4
# 1+3+5
# no data
# 6
]
assert ak.min(array, axis=-1).tolist() == [
0, # min([0, 1])
3, # min([3])
None, # list is empty
None, # list is missing
4, # min([4, 5, 6])
]
assert ak.min(array, axis=-2).tolist() == [
0,
1,
None,
6,
# 0,4
# 1,3,5
# no data
# 6
]
# first bug-fix: single '?'
assert str(ak.min(array, axis=-1).type) == "5 * ?int64"
# second bug-fix: correct mask_identity=False behavior
assert ak.ptp(array, axis=-1).tolist() == [1, 0, None, None, 2]
assert ak.ptp(array, axis=-1,
mask_identity=False).tolist() == [1, 0, 0, None, 2]
assert ak.ptp(array, axis=-2).tolist() == [4, 4, None, 0]
assert ak.ptp(array, axis=-2, mask_identity=False).tolist() == [4, 4, 0, 0]
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x` and `y` components"""
return awkward.zip(
{
"x": awkward.sum(self.x, axis=axis),
"y": awkward.sum(self.y, axis=axis),
},
with_name="TwoVector",
behavior=self.behavior,
)
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x` and `y` components"""
out = awkward.zip(
{
"x": awkward.sum(self.x, axis=axis),
"y": awkward.sum(self.y, axis=axis),
},
with_name="TwoVector",
highlevel=False,
)
return awkward.Array(out, behavior=self.behavior)
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x`, `y`, `z`, `t`, and `charge` components"""
return awkward.zip(
{
"x": awkward.sum(self.x, axis=axis),
"y": awkward.sum(self.y, axis=axis),
"z": awkward.sum(self.z, axis=axis),
"t": awkward.sum(self.t, axis=axis),
"charge": awkward.sum(self.charge, axis=axis),
},
with_name="Candidate",
)
def sum(self, axis=-1):
"""Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` components"""
return awkward.zip(
{
"x": awkward.sum(self.x, axis=axis),
"y": awkward.sum(self.y, axis=axis),
"z": awkward.sum(self.z, axis=axis),
"t": awkward.sum(self.t, axis=axis),
},
with_name="LorentzVector",
behavior=self.behavior,
)
def corrected_polar_met(met_pt, met_phi, jet_pt, jet_phi, jet_pt_orig, deltas=None):
sj, cj = numpy.sin(jet_phi), numpy.cos(jet_phi)
x = met_pt * numpy.cos(met_phi) + awkward.sum(
jet_pt * cj - jet_pt_orig * cj, axis=1
)
y = met_pt * numpy.sin(met_phi) + awkward.sum(
jet_pt * sj - jet_pt_orig * sj, axis=1
)
if deltas:
positive, dx, dy = deltas
x = x + dx if positive else x - dx
y = y + dy if positive else y - dy
return awkward.zip({"pt": numpy.hypot(x, y), "phi": numpy.arctan2(y, x)})
def add_object_cuts(self, objects, object_name, cuts, names, sample):
if object_name not in self.object_stats.keys():
self.object_stats[object_name] = {}
if sample not in self.object_stats[object_name].keys():
self.object_stats[object_name][sample] = { "n_objects_initial" : awkward.sum(awkward.num(objects)) }
else:
self.object_stats[object_name][sample]["n_objects_initial"] += awkward.sum(awkward.num(objects))
for cut, name in zip(cuts, names):
if name not in self.object_stats[object_name][sample].keys():
self.object_stats[object_name][sample][name] = awkward.sum(awkward.num(objects[cut]))
else:
self.object_stats[object_name][sample][name] += awkward.sum(awkward.num(objects[cut]))
def Phi_mpi_pi(array_phi):
"""
Periodically (T=2*pi, shift step done only once) bring values of the given array into [-pi, pi] range.
Arguments:
array_phi: awkward array, values assumed to be radian measure of phi angle
Returns:
Awkward array, values of input array brought to [-pi, pi] range
"""
array_phi = ak.where(array_phi <= np.pi, array_phi, array_phi - 2*np.pi)
array_phi = ak.where(array_phi >= -np.pi, array_phi, array_phi + 2*np.pi)
assert ak.sum(array_phi > np.pi) + ak.sum(array_phi < -np.pi) == 0
return array_phi
def get_norm(file, weight):
with uproot.open(file_map[file]) as f:
tree = f["Events"]
events = tree.arrays(["genWeight", "LHEScaleWeight"],
library="ak",
how="zip")
nominal_weights = events.jagged0.LHEScaleWeight[:, 4]
var_weights = events.jagged0.LHEScaleWeight[:, weight_map[weight]]
sum_nominal = awkward.sum(nominal_weights)
sum_weight = awkward.sum(var_weights)
print(sum_nominal, sum_weight)
return sum_nominal / sum_weight
def add_pdf_weight(weights, pdf_weights):
nom = np.ones(len(weights.weight()))
up = np.ones(len(weights.weight()))
down = np.ones(len(weights.weight()))
# NNPDF31_nnlo_hessian_pdfas
# https://lhapdfsets.web.cern.ch/current/NNPDF31_nnlo_hessian_pdfas/NNPDF31_nnlo_hessian_pdfas.info
if pdf_weights is not None and "306000 - 306102" in pdf_weights.__doc__:
# Hessian PDF weights
# Eq. 21 of https://arxiv.org/pdf/1510.03865v1.pdf
arg = pdf_weights[:, 1:-2] - np.ones((len(weights.weight()), 100))
summed = ak.sum(np.square(arg), axis=1)
pdf_unc = np.sqrt((1. / 99.) * summed)
weights.add('PDF_weight', nom, pdf_unc + nom)
# alpha_S weights
# Eq. 27 of same ref
as_unc = 0.5 * (pdf_weights[:, 102] - pdf_weights[:, 101])
weights.add('aS_weight', nom, as_unc + nom)
# PDF + alpha_S weights
# Eq. 28 of same ref
pdfas_unc = np.sqrt(np.square(pdf_unc) + np.square(as_unc))
weights.add('PDFaS_weight', nom, pdfas_unc + nom)
else:
weights.add('aS_weight', nom, up, down)
weights.add('PDF_weight', nom, up, down)
weights.add('PDFaS_weight', nom, up, down)
def process(self, events):
output = self.accumulator.identity()
dataset = events.metadata['dataset']
integratedLuminosity = 137.19 * 1000 # fb^{-1} to pb^{-1}
ht = events.HT
weights = integratedLuminosity * events.CrossSection[ht > 1200] / len(
events)
GenParticles_pt = events.GenParticles.pt
GenParticles_eta = events.GenParticles.eta
GenParticles_Status = events.GenParticles.Status
GenParticles_PdgId = events.GenParticles.PdgId
GenParticles_Charge = events.GenParticles.Charge
finalParticles = (GenParticles_Status == 1) & (GenParticles_pt > 1) & (
abs(GenParticles_eta) < 2.5) & (GenParticles_Charge != 0)
nTracks = ak.sum(finalParticles[ht > 1200], axis=1)
output["sumw"][dataset] += len(events)
output["nTracks"].fill(dataset=dataset,
nTracks=nTracks,
weight=weights)
return output
def compute_rate(tau_1, tau_2, Nev_den, Pt_thr, a, L1rate):
mask_1 = (tau_1.deepTau_VSjet > deep_thr(
tau_1, a, Pt_thr)) & reco_tau_selection(tau_1, minPt=Pt_thr)
mask_2 = (tau_2.deepTau_VSjet > deep_thr(
tau_2, a, Pt_thr)) & reco_tau_selection(tau_2, minPt=Pt_thr)
Nev_num = ak.sum(ditau_selection(mask_1, mask_2))
return Nev_num / Nev_den * L1rate
def select_deltaR(events, obj1, obj2, threshold, debug):
if debug > 1:
start = time.time()
nEvents = len(events)
n_obj1 = numpy.array(awkward.num(obj1), numpy.int64)
n_obj2 = numpy.array(awkward.num(obj2), numpy.int64)
mask_offsets = numpy.zeros(nEvents + 1, numpy.int64)
mask_contents = numpy.empty(awkward.sum(n_obj1), numpy.bool)
mask_offsets, mask_contents = select_deltaR_raw(events, nEvents, obj1,
n_obj1, obj2, n_obj2,
threshold, mask_offsets,
mask_contents, debug)
if debug > 1:
elapsed_time = time.time() - start
print(
"[object_selections.py] PERFORMANCE: time to execute select_deltaR: %.6f"
% (elapsed_time))
mask = awkward.Array(
awkward.layout.ListOffsetArray64(
awkward.layout.Index64(mask_offsets),
awkward.layout.NumpyArray(mask_contents)))
return mask
def generate_epd_hit_matrix(self):
ring_sum = np.zeros((32, len(self.nMip)))
print("Filling array of dimension", ring_sum.shape)
for i in range(32):
ring_i = ak.sum(ak.where(self.row == i+1, self.nMip, 0), axis=-1)
ring_sum[i] = ring_i
return ring_sum
def process(self, events):
# Note: This is required to ensure that behaviors are registered
# when running this code in a remote task.
ak.behavior.update(candidate.behavior)
output = self.accumulator.identity()
dataset = events.metadata['dataset']
muons = ak.zip(
{
"pt": events.Muon_pt,
"eta": events.Muon_eta,
"phi": events.Muon_phi,
"mass": events.Muon_mass,
"charge": events.Muon_charge,
},
with_name="PtEtaPhiMCandidate")
cut = (ak.num(muons) == 2) & (ak.sum(muons.charge) == 0)
# add first and second muon in every event together
dimuon = muons[cut][:, 0] + muons[cut][:, 1]
output["sumw"][dataset] += len(events)
output["mass"].fill(
dataset=dataset,
mass=dimuon.mass,
)
return output
def load_file(self):
"""
Function which reads a root or csv file and removes
interactions and events without energy depositions.
Returns:
interactions: awkward array
n_simulated_events: Total number of simulated events in the
opened file. (This includes events removed by cuts)
"""
if self.file.endswith(".root"):
interactions, n_simulated_events, start, stop = self._load_root_file(
)
elif self.file.endswith(".csv"):
interactions, n_simulated_events, start, stop = self._load_csv_file(
)
else:
raise ValueError(
f'Cannot load events from file "{self.file}": .root or .cvs file needed.'
)
if np.any(interactions['ed'] < 0):
warnings.warn('At least one of the energy deposits is negative!')
# Removing all events with zero energy deposit
m = interactions['ed'] > 0
if self.cut_by_eventid:
# ufunc does not work here...
m2 = (interactions['evtid'] >= start) & (interactions['evtid'] <
stop)
m = m & m2
interactions = interactions[m]
if self.cut_nr_only:
m = ((interactions['type'] == "neutron") &
(interactions['edproc'] == "hadElastic")) | (
interactions['edproc'] == "ionIoni")
e_dep_er = ak.sum(interactions[~m]['ed'], axis=1)
e_dep_nr = ak.sum(interactions[m]['ed'], axis=1)
interactions = interactions[(e_dep_er < 10) & (e_dep_nr > 0)]
# Removing all events with no interactions:
m = ak.num(interactions['ed']) > 0
interactions = interactions[m]
return interactions, n_simulated_events
def process(self, events):
has2jets = ak.sum(events.Jet.pt > 40, axis=1) >= 2
return (hist.Hist.new.Reg(100,
0,
200,
name="met",
label="$E_{T}^{miss}$ [GeV]").Double().fill(
events[has2jets].MET.pt))
def flip_weight(self, electron):
#f_1 = self.evaluator['el'](electron.pt[:,0:1], abs(electron.eta[:,0:1]))
#f_2 = self.evaluator['el'](electron.pt[:,1:2], abs(electron.eta[:,1:2]))
# For custom measurements
f_1 = yahist_2D_lookup(self.ratio, electron.pt[:, 0:1],
abs(electron.eta[:, 0:1]))
f_2 = yahist_2D_lookup(self.ratio, electron.pt[:, 1:2],
abs(electron.eta[:, 1:2]))
# I'm using ak.prod and ak.sum to replace empty arrays by 1 and 0, respectively
weight = ak.sum(f_1 / (1 - f_1), axis=1) * ak.prod(
1 - f_2 / (1 - f_2), axis=1) + ak.sum(
f_2 / (1 - f_2), axis=1) * ak.prod(1 - f_1 / (1 - f_1), axis=1)
return weight
def compute_rate_base(tau_1, tau_2, Nev_den, Pt_thr, L1rate):
mask_1 = iso_tau_selection(tau_1, "mediumIsoAbs",
"mediumIsoRel") & reco_tau_selection(
tau_1, minPt=Pt_thr)
mask_2 = iso_tau_selection(tau_2, "mediumIsoAbs",
"mediumIsoRel") & reco_tau_selection(
tau_2, minPt=Pt_thr)
Nev_num = ak.sum(ditau_selection(mask_1, mask_2))
return Nev_num / Nev_den * L1rate
def process_event(weights, evt_vars): #, bjets):
out_hists = {}
out_hists["MET"] = Hist1D(ak.to_numpy(evt_vars.MET_pt),
bins=bin_met,
label="MET",
weights=ak.to_numpy(genWeight))
#out_hists["njets"] = Hist1D(ak.to_numpy(evt_vars.nJet), bins = bin_njet, label="njet", weights = ak.to_numpy(genWeight))
#out_hists["nbjets"] = Hist1D(ak.num(bjets), bins = bin_bjet)
out_hists["weight"] = ak.sum(weights)
return out_hists
def add_cuts(self, cuts, names):
for cut, name in zip(cuts, names):
if self.debug > 0:
n_objects_cut = awkward.sum(awkward.num(self.objects[cut]))
if self.n_objects_initial == 0:
return
print(
"%s ObjectCutDiagnostics: Applying cut %s would have an eff of %.4f"
% (self.cut_set, name,
float(n_objects_cut) / float(self.n_objects_initial)))
def FR_weight(self, lepton):
f_1 = yahist_2D_lookup(self.ratio, lepton.conePt[:, 0:1],
abs(lepton.eta[:, 0:1]))
#breakpoint()
# I'm using ak.prod and ak.sum to replace empty arrays by 1 and 0, respectively
weight = ak.sum(
f_1 / (1 - f_1), axis=1
) #*ak.prod(1-f_2/(1-f_2), axis=1) + ak.sum(f_2/(1-f_2), axis=1)*ak.prod(1-f_1/(1-f_1), axis=1)
return weight
def set_genZ(events, selection_options, debug):
electron_idxs = abs(events.GenPart_pdgId) == 11
muon_idxs = abs(events.GenPart_pdgId) == 13
tau_idxs = abs(events.GenPart_pdgId) == 15
motherOfElectrons = events.GenPart_genPartIdxMother[electron_idxs]
motherOfMuons = events.GenPart_genPartIdxMother[muon_idxs]
motherOfTaus = events.GenPart_genPartIdxMother[tau_idxs]
ZToEleEvents = ak.sum(
(events.GenPart_pdgId[motherOfElectrons] == 23), axis=1) >= 2
ZToMuEvents = ak.sum(
(events.GenPart_pdgId[motherOfMuons] == 23), axis=1) >= 2
ZToTauEvents = ak.sum(
(events.GenPart_pdgId[motherOfTaus] == 23), axis=1) >= 2
events[
"genZ_decayMode"] = 1 * ZToEleEvents + 2 * ZToMuEvents + 3 * ZToTauEvents
return events
def test_reducers():
# axis=None reducers are implemented in NumPy.
assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == 6 + 6j
assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == -12 + 12j
# axis != None reducers are implemented in libawkward; this should be ReducerSum.
assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
3 + 3j,
0 + 0j,
3 + 3j,
]
# And this is in ReducerProd.
assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
0 + 4j,
1 + 0j,
3 + 3j,
]
# ReducerCount, ReducerCountNonzero, ReducerAny, and ReducerAll work.
assert ak.count(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [2, 0, 1]
assert ak.count_nonzero(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [2, 0, 1]
assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
True,
False,
True,
]
assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
True,
True,
True,
]
assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]),
axis=1).tolist() == [True, False, True]
assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]),
axis=1).tolist() == [False, True, True]
