def embed_crossref(source, idx_name, dest, dest_name):
"""Embed a cross-reference
Parameters
----------
source : ak.Array
any array with shape N * var * {record}
idx_name : str
A field in the source record
dest : ak.Array
any array with shape N * var * {record}, where:
``ak.max(source[idx_name], axis=1) < ak.num(dest)`` and
``ak.min(source[idx_name], axis=1) >= 0``
"""
print(ak.max(source[idx_name], axis=1))
print(ak.num(dest))
print(ak.all(ak.max(source[idx_name], axis=1) < ak.num(dest)))
assert ak.all(ak.max(source[idx_name], axis=1) < ak.num(dest))
assert ak.all(ak.min(source[idx_name], axis=1) >= 0)
id_global = ak.flatten(
source[idx_name] + np.asarray(dest.layout.starts), axis=None
)
source[dest_name] = ak.Array(
ak.layout.ListOffsetArray64(
source.layout.offsets,
ak.layout.ListOffsetArray64(
source.layout.content.offsets,
ak.flatten(dest)[id_global].layout,
),
)
)
def test_read_nanomc():
factory = NanoEventsFactory.from_file(
os.path.abspath('tests/samples/nano_dy.root'))
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all((abs(events.Electron.matched_gen.pdgId) == 11)
| (events.Electron.matched_gen.pdgId == 22))
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False, True, True, True, False, False, False, False, False
]
def genroundtrips(genpart):
# check genpart roundtrip
assert ak.all(genpart.children.parent.pdgId == genpart.pdgId)
assert ak.all(
ak.any(genpart.parent.children.pdgId == genpart.pdgId,
axis=-1,
mask_identity=True))
# distinctParent should be distinct and it should have a relevant child
assert ak.all(genpart.distinctParent.pdgId != genpart.pdgId)
assert ak.all(
ak.any(genpart.distinctParent.children.pdgId == genpart.pdgId,
axis=-1,
mask_identity=True))
# exercise hasFlags
genpart.hasFlags(['isHardProcess'])
genpart.hasFlags(['isHardProcess', 'isDecayedLeptonHadron'])
genPart = final.GenPart
tops = genPart[abs(genPart.pdgId) == 6]
#The isLastCopy Flag filters out copy Genparticles:
tops = tops[tops.hasFlags('isLastCopy')]
tDecay = tops.distinctChildren
tDecay = tDecay[tDecay.hasFlags('isLastCopy')]
t_Events = tDecay[abs(tDecay.pdgId) == 5]
W = tDecay[abs(tDecay.pdgId) == 24]
W = W[W.hasFlags('isLastCopy')]
WDecay = W.distinctChildren
WDecay = WDecay[WDecay.hasFlags('isLastCopy')]
#t_events is the lone bottom, W_events is the -> two jets
#select the hadronically decaying W
W_Events = ak.flatten(WDecay[ak.all(abs(WDecay.pdgId) <= 8, axis=-1)], axis=3)
#HadW is mask for Quark deacying W boson
hadW = ak.num(W_Events, axis=2) == 2
#filters out t_events that have a hadronically decayign W Boson
hadB = t_Events[hadW]
hadB = ak.flatten(hadB, axis=2)
W_quarks = W_Events[hadW]
W_quarks = ak.flatten(W_quarks, axis=2)
#concatentating these two arrays make an array of events with the correctly decaying GenParticles.
qqb = ak.concatenate([hadB, W_quarks], axis=1)
print("qqb Genparts matched")
#Filtering Out events with extra tops
fig.savefig(f'valplots/{hist_spec[0]}.png')
print("Checking Coffea/NanoEvents compatibility:")
import coffea
import uproot
import awkward1 as ak
import numpy as np
from coffea.nanoevents import NanoEventsFactory
for fn in check_files:
print(f" {fn}")
def nano_evts(fname):
factory = NanoEventsFactory.from_file(
fname,
entry_start=0,
entry_stop=10000,
metadata={"dataset": ""},
)
return factory.events()
evts = nano_evts(fn)
print(" Subjets pt ordered (should be false):",
ak.all(evts.SubJet.pt[:, :-1] - evts.SubJet.pt[:, 1:] >= 0))
print(" SDmass from subjets hist:")
print(
np.histogram(ak.to_numpy(
ak.flatten(evts.FatJet.subjets.sum().mass, None)),
bins=np.linspace(40, 200, 20))[0])
def process(self, events):
# get meta infos
dataset = events.metadata["dataset"]
isRealData = not hasattr(events, "genWeight")
n_events = len(events)
selection = processor.PackedSelection()
weights = processor.Weights(n_events)
output = self.accumulator.identity()
# weights
if not isRealData:
output['sumw'][dataset] += awkward1.sum(events.genWeight)
# trigger
triggers = {}
for channel in ["e","mu"]:
trigger = np.zeros(len(events), dtype='bool')
for t in self._trigger[channel]:
try:
trigger = trigger | events.HLT[t]
except:
warnings.warn("Missing trigger %s" % t, RuntimeWarning)
triggers[channel] = trigger
# met filter
met_filters = ["goodVertices",
"globalSuperTightHalo2016Filter",
"HBHENoiseFilter",
"HBHENoiseIsoFilter",
"EcalDeadCellTriggerPrimitiveFilter",
"BadPFMuonFilter",
]
met_filters_mask = np.ones(len(events), dtype='bool')
for t in met_filters:
met_filters_mask = met_filters_mask & events.Flag[t]
selection.add("met_filter", awkward1.to_numpy(met_filters_mask))
# load objects
muons = events.Muon
electrons = events.Electron
jets = events.Jet
fatjets = events.FatJet
subjets = events.SubJet
fatjetsLS = events.FatJetLS
met = events.MET
# muons
goodmuon = (
(muons.mediumId)
& (muons.miniPFRelIso_all <= 0.2)
& (muons.pt >= 27)
& (abs(muons.eta) <= 2.4)
& (abs(muons.dz) < 0.1)
& (abs(muons.dxy) < 0.05)
& (muons.sip3d < 4)
)
good_muons = muons[goodmuon]
ngood_muons = awkward1.sum(goodmuon, axis=1)
# electrons
goodelectron = (
(electrons.mvaFall17V2noIso_WP90)
& (electrons.pt >= 30)
& (abs(electrons.eta) <= 1.479)
& (abs(electrons.dz) < 0.1)
& (abs(electrons.dxy) < 0.05)
& (electrons.sip3d < 4)
)
good_electrons = electrons[goodelectron]
ngood_electrons = awkward1.sum(goodelectron, axis=1)
# good leptons
good_leptons = awkward1.concatenate([good_muons, good_electrons], axis=1)
good_leptons = good_leptons[awkward1.argsort(good_leptons.pt)]
# lepton candidate
candidatelep = awkward1.firsts(good_leptons)
# lepton channel selection
selection.add("ch_e", awkward1.to_numpy((triggers["e"]) & (ngood_electrons==1) & (ngood_muons==0))) # not sure if need to require 0 muons or 0 electrons in the next line
selection.add("ch_mu", awkward1.to_numpy((triggers["mu"]) & (ngood_electrons==0) & (ngood_muons==1)))
# jets
ht = awkward1.sum(jets[jets.pt > 30].pt,axis=1)
selection.add("ht_400", awkward1.to_numpy(ht>=400))
goodjet = (
(jets.isTight)
& (jets.pt > 30)
& (abs(jets.eta) <= 2.5)
)
good_jets = jets[goodjet]
# fat jets
jID = "isTight"
# TODO: add mass correction
# a way to get the first two subjets
# cart = awkward1.cartesian([fatjets, subjets], nested=True)
# idxes = awkward1.pad_none(awkward1.argsort(cart['0'].delta_r(cart['1'])), 2, axis=2)
# sj1 = subjets[idxes[:,:,0]]
# sj2 = subjets[idxes[:,:,1]]
good_fatjet = (
(getattr(fatjets, jID))
& (abs(fatjets.eta) <= 2.4)
& (fatjets.pt > 50)
& (fatjets.msoftdrop > 30)
& (fatjets.msoftdrop < 210)
#& (fatjets.pt.copy(content=fatjets.subjets.content.counts) == 2) # TODO: require 2 subjets?
# this can probably be done w FatJet_subJetIdx1 or FatJet_subJetIdx2
& (awkward1.all(fatjets.subjets.pt >= 20))
& (awkward1.all(abs(fatjets.subjets.eta) <= 2.4))
)
good_fatjets = fatjets[good_fatjet]
# hbb candidate
mask_hbb = (
(good_fatjets.pt > 200)
& (good_fatjets.delta_r(candidatelep) > 2.0)
)
candidateHbb = awkward1.firsts(good_fatjets[mask_hbb])
# b-tag #& (good_fatjets.particleNetMD_Xbb > 0.9)
selection.add('hbb_btag',awkward1.to_numpy(candidateHbb.deepTagMD_ZHbbvsQCD >= 0.8)) # score would be larger for tight category (0.97)
# No AK4 b-tagged jets away from bb jet
jets_HbbV = jets[good_jets.delta_r(candidateHbb) >= 1.2]
selection.add('hbb_vetobtagaway', awkward1.to_numpy(awkward1.max(jets_HbbV.btagDeepB, axis=1, mask_identity=False) > BTagEfficiency.btagWPs[self._year]['medium']))
# fat jets Lepton Subtracted
# wjj candidate
mask_wjj = (
(fatjetsLS.pt > 50)
& (fatjetsLS.delta_r(candidatelep) > 1.2)
# need to add 2 subjets w pt > 20 & eta<2.4
# need to add ID?
)
candidateWjj = awkward1.firsts(fatjetsLS[mask_wjj][awkward1.argmin(fatjetsLS[mask_wjj].delta_r(candidatelep),axis=1,keepdims=True)])
# add t2/t1 <= 0.75 (0.45 HP)
selection.add('hww_mass', awkward1.to_numpy(candidateWjj.mass >= 10))
print('met ',met)
# wjjlnu info
#HSolverLiInfo hwwInfoLi;
# qqSDmass = candidateWjj.msoftdrop
# hwwLi = hSolverLi->minimize(candidatelep.p4(), met.p4(), wjjcand.p4(), qqSDmass, hwwInfoLi)
#neutrino = hwwInfoLi.neutrino;
#wlnu = hwwInfoLi.wlnu;
#wqq = hwwInfoLi.wqqjet;
#hWW = hwwInfoLi.hWW;
#wwDM = PhysicsUtilities::deltaR( wlnu,wqq) * hWW.pt()/2.0;
# add dlvqq <= 11 (2.5 HP)
# in the meantime let's add the mass
'''
mm = (candidatejet - candidatelep).mass2
jmass = (mm>0)*np.sqrt(np.maximum(0, mm)) + (mm<0)*candidatejet.mass
joffshell = jmass < 62.5
massassumption = 80.*joffshell + (125 - 80.)*~joffshell
x = massassumption**2/(2*candidatelep.pt*met.pt) + np.cos(candidatelep.phi - met.phi)
met_eta = (
(x < 1)*np.arcsinh(x*np.sinh(candidatelep.eta))
+ (x > 1)*(
candidatelep.eta - np.sign(candidatelep.eta)*np.arccosh(candidatelep.eta)
)
)
met_p4 = TLorentzVectorArray.from_ptetaphim(np.array([0.]),np.array([0.]),np.array([0.]),np.array([0.]))
if met.size > 0:
met_p4 = TLorentzVectorArray.from_ptetaphim(met.pt, met_eta.fillna(0.), met.phi, np.zeros(met.size))
# hh system
candidateHH = candidateWjj + met_p4 + candidateHbb
selection.add('hh_mass', candidateHH.mass >= 700)
selection.add('hh_centrality', candidateHH.pt/candidateHH.mass >= 0.3)
'''
channels = {"e": ["met_filter","ch_e","ht_400","hbb_btag","hbb_vetobtagaway","hww_mass"], #,"hh_mass","hh_centrality"],
"mu": ["met_filter","ch_mu","ht_400","hbb_btag","hbb_vetobtagaway","hww_mass"] #,"hh_mass","hh_centrality"],
}
# need to add gen info
if not isRealData:
weights.add('genweight', events.genWeight)
add_pileup_weight(weights, events.Pileup.nPU, self._year, dataset)
for channel, cuts in channels.items():
allcuts = set()
output['cutflow'].fill(dataset=dataset, channel=channel, cut=0, weight=weights.weight())
for i, cut in enumerate(cuts):
allcuts.add(cut)
cut = selection.all(*allcuts)
output['cutflow'].fill(dataset=dataset, channel=channel, cut=i + 1, weight=weights.weight()[cut])
return output
def trilep_baseline(self, omit=[], cutflow=None, tight=False):
'''
give it a cutflow object if you want it to be filed.
cuts in the omit list will not be applied
'''
self.selection = PackedSelection()
is_trilep = ((ak.num(self.ele) + ak.num(self.mu))==3)
pos_charge = ((ak.sum(self.ele.pdgId, axis=1) + ak.sum(self.mu.pdgId, axis=1))<0)
neg_charge = ((ak.sum(self.ele.pdgId, axis=1) + ak.sum(self.mu.pdgId, axis=1))>0)
lep0pt = ((ak.num(self.ele[(self.ele.pt>25)]) + ak.num(self.mu[(self.mu.pt>25)]))>0)
lep1pt = ((ak.num(self.ele[(self.ele.pt>20)]) + ak.num(self.mu[(self.mu.pt>20)]))>1)
lepveto = ((ak.num(self.ele_veto) + ak.num(self.mu_veto))==3)
dimu = choose(self.mu, 2)
diele = choose(self.ele, 2)
dilep = cross(self.mu, self.ele)
OS_dimu = dimu[(dimu['0'].charge*dimu['1'].charge < 0)]
OS_diele = diele[(diele['0'].charge*diele['1'].charge < 0)]
offZ = (ak.all(abs(OS_dimu.mass-91.2)>10, axis=1) & ak.all(abs(OS_diele.mass-91.2)>10, axis=1))
lepton = ak.concatenate([self.ele, self.mu], axis=1)
lepton_pdgId_pt_ordered = ak.fill_none(ak.pad_none(lepton[ak.argsort(lepton.pt, ascending=False)].pdgId, 2, clip=True), 0)
triggers = getTriggers(self.events,
ak.flatten(lepton_pdgId_pt_ordered[:,0:1]),
ak.flatten(lepton_pdgId_pt_ordered[:,1:2]), year=self.year, dataset=self.dataset)
ht = ak.sum(self.jet_all.pt, axis=1)
st = self.met.pt + ht + ak.sum(self.mu.pt, axis=1) + ak.sum(self.ele.pt, axis=1)
self.selection.add('lepveto', lepveto)
self.selection.add('trilep', is_trilep)
self.selection.add('filter', self.filters)
self.selection.add('trigger', triggers)
self.selection.add('p_T(lep0)>25', lep0pt)
self.selection.add('p_T(lep1)>20', lep1pt)
self.selection.add('N_jet>2', (ak.num(self.jet_all)>2) )
self.selection.add('N_jet>3', (ak.num(self.jet_all)>3) )
self.selection.add('N_central>1', (ak.num(self.jet_central)>1) )
self.selection.add('N_central>2', (ak.num(self.jet_central)>2) )
self.selection.add('N_btag>0', (ak.num(self.jet_btag)>0) )
self.selection.add('N_fwd>0', (ak.num(self.jet_fwd)>0) )
self.selection.add('MET>50', (self.met.pt>50) )
self.selection.add('ST>600', (st>600) )
self.selection.add('offZ', offZ )
reqs = [
'filter',
'lepveto',
'trilep',
'p_T(lep0)>25',
'p_T(lep1)>20',
'trigger',
'offZ',
'MET>50',
'N_jet>2',
'N_central>1',
'N_btag>0',
'N_fwd>0',
]
if tight:
reqs += [
'N_jet>3',
'N_central>2',
'ST>600',
#'MET>50',
#'delta_eta',
]
reqs_d = { sel: True for sel in reqs if not sel in omit }
selection = self.selection.require(**reqs_d)
self.reqs = [ sel for sel in reqs if not sel in omit ]
if cutflow:
#
cutflow_reqs_d = {}
for req in reqs:
cutflow_reqs_d.update({req: True})
cutflow.addRow( req, self.selection.require(**cutflow_reqs_d) )
return selection
def equals(self, other):
if isinstance(other, AwkwardSeries):
other = other.data
return ak.all(self.data == other)
def argon_kshell_cuts_3kV(events):
#4) largest_s2_amplitude_pmt < 11600
events = events[ak.all(events.s2.height_mvdc_bot < 11600, axis=1)]
#5)second_largest_s2_area_pmt = 0
events = events[events.s2s_per_waveform == 1]
#6)largest_s1_area_pmt >0
events = events[ak.any(events.s1.area_pe_bot > 0, axis=1)]
#7)largest_s2_area_pmt > largest_s1_area_pmt
event = events[ak.max(events.s2.area_pe_bot, axis=1) > ak.max(
events.s1.area_pe_bot, axis=1)]
#8) 0.20 < s2_frt < 0.34
events["s2", "s2_frt"] = events.s2.area_pe_top / (events.s2.area_pe_top +
events.s2.area_pe_bot)
mask = ak.any(events.s2.s2_frt > 0.16, axis=1) & ak.any(
events.s2.s2_frt < 0.34, axis=1)
events = events[mask]
#9)s2_area+s1_area>1175
mask = ak.max(events.s2.area_pe_top,axis=1)+ak.max(events.s2.area_pe_bot,axis=1)+ak.max(events.s1.area_pe_top,axis=1)+\
ak.max(events.s1.area_pe_bot,axis=1)>1175
events = events[mask]
#print("Length array pos 1")
#print(ak.num(events,axis=0))
if ak.num(events, axis=0) == 0:
return ak.Array([])
#Add the drifttime
mask_s1_before_s2 = events.s1.pos_bot < ak.flatten(events.s2.pos_bot)
#Cut away the event which don't have an s1 befor the s2. these are s2 only.
events = events[ak.any(events.s1.area_pe_bot[mask_s1_before_s2], axis=1)]
#print("Length array pos 2")
#print(ak.num(events,axis=0))
if ak.num(events, axis=0) == 0:
return ak.Array([])
mask_s1_before_s2 = events.s1.pos_bot < ak.flatten(events.s2.pos_bot)
max_s1_before_s2 = ak.argmax(events.s1.area_pe_bot[mask_s1_before_s2],
axis=1,
keepdims=True)
events["mask_max_s1_before_s2"] = max_s1_before_s2
events["drifttime_musec"] = ak.flatten(
(ak.max(events.s2.pos_bot, axis=1) -
events.s1.pos_bot[max_s1_before_s2]) / 100)
#10) S2 width cut
lower_width = 34.5 + 0.22 * events["drifttime_musec"] + np.sqrt(
32.6 * events["drifttime_musec"])
upper_width = 48.2 + 0.19 * events["drifttime_musec"] + np.sqrt(
47.3 * events["drifttime_musec"])
upper_mask = ak.any(events.s2.width_bot < upper_width, axis=1)
lower_mask = ak.any(events.s2.width_bot > lower_width, axis=1)
events = events[upper_mask & lower_mask]
#11)FDV cut
v = 1.9608556663165941
gate_time = 1.537109944449019
events["z"] = -v * (events["drifttime_musec"] - gate_time)
events = events[events.z < -2]
events = events[events.z > -28]
events["r_s2"] = ak.flatten(
np.sqrt(np.square(events.s2.x_corr) + np.square(events.s2.y_corr)))
events = events[events.r_s2 < 10]
return events
def crossref(events):
# check some cross-ref roundtrips (some may not be true always but they are for the test file)
assert ak.all(events.Jet.matched_muons.matched_jet.pt == events.Jet.pt)
assert ak.all(events.Electron.matched_photon.matched_electron.r9 ==
events.Electron.r9)