def _kExtra(self, kpt, eta, nl, u, s=0, m=0):
# if it is a jagged array, save the offsets then flatten everything
# needed for the ternary conditions later
abseta = abs(eta)
kData = self._kRes[s][m][1](abseta) # type 1 is data
kMC = self._kRes[s][m][0](abseta) # type 0 is MC
mask = kData > kMC
x = awkward.zeros_like(kpt)
sigma = self._sigma(kpt, eta, nl, s, m)
# Rochester cbA = beta, cbN = m, as well as cbM (always 0?) = loc and cbS = scale to transform y = (x-loc)/scale in the pdf method
cbA = self._cbA[s][m](abseta, nl)
cbN = self._cbN[s][m](abseta, nl)
cbS = self._cbS[s][m](abseta, nl)
counts = awkward.num(u)
u_flat = awkward.flatten(u)
loc = awkward.zeros_like(u_flat)
cbA_flat = awkward.flatten(cbA)
cbN_flat = awkward.flatten(cbN)
cbS_flat = awkward.flatten(cbS)
invcdf = awkward.unflatten(
doublecrystalball.ppf(u_flat, cbA_flat, cbA_flat, cbN_flat,
cbN_flat, loc, cbS_flat),
counts,
)
x = awkward.where(
mask,
(numpy.sqrt(kData * kData - kMC * kMC) * sigma * invcdf),
x,
)
result = awkward.where(x > -1, 1.0 / (1.0 + x), awkward.ones_like(kpt))
if isinstance(kpt, numpy.ndarray):
result = numpy.array(result)
return result
def _assignTruthDef(self, tree):
assert self.splitIdx is not None
nonSplitRecHitSimClusIdx = self._createTruthAssociation(tree)
recHitTruthPID = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_pdgId",nonSplitRecHitSimClusIdx)
recHitTruthEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_boundaryEnergy",nonSplitRecHitSimClusIdx)
recHitDepEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_recEnergy",nonSplitRecHitSimClusIdx)
if not self.use_true_muon_momentum:
recHitTruthEnergy = ak1.where(np.abs(recHitTruthPID[:,:,0])==13, recHitDepEnergy, recHitTruthEnergy)
recHitTruthX = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_x",nonSplitRecHitSimClusIdx)
recHitTruthY = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_y",nonSplitRecHitSimClusIdx)
recHitTruthZ = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_z",nonSplitRecHitSimClusIdx)
recHitTruthTime = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_t",nonSplitRecHitSimClusIdx)
fullyContained = ak1.where(np.abs(recHitTruthZ)[:,:,0]<323.,#somehow that seems necessary
ak1.ones_like(recHitTruthZ),
ak1.zeros_like(recHitTruthZ))
recHitEnergy = self._readSplitAndExpand(tree,"RecHitHGC_energy")
recHitTime = self._readSplitAndExpand(tree,"RecHitHGC_time")
recHitX = self._readSplitAndExpand(tree,"RecHitHGC_x")
recHitY = self._readSplitAndExpand(tree,"RecHitHGC_y")
recHitZ = self._readSplitAndExpand(tree,"RecHitHGC_z")
# should not expand here to allow indexing as done below
recHitSimClusIdx = self._splitJaggedArray(nonSplitRecHitSimClusIdx)
# set noise to rec features
recHitTruthEnergy = ak1.where(recHitSimClusIdx<0, recHitEnergy, recHitTruthEnergy)
recHitTruthX = ak1.where(recHitSimClusIdx<0, recHitX, recHitTruthX)
recHitTruthY = ak1.where(recHitSimClusIdx<0, recHitY, recHitTruthY)
recHitTruthZ = ak1.where(recHitSimClusIdx<0, recHitZ, recHitTruthZ)
recHitTruthTime = ak1.where(recHitSimClusIdx<0, recHitTime, recHitTruthTime)
recHitDepEnergy = ak1.where(recHitSimClusIdx<0, recHitEnergy, recHitDepEnergy)
recHitSpectatorFlag = self._createSpectators(tree)
#remove spectator flag for noise
recHitSpectatorFlag = ak1.where(recHitSimClusIdx<0 , ak1.zeros_like(recHitSpectatorFlag), recHitSpectatorFlag)#this doesn't work for some reason!
#DEBUG!!!
#ticlidx = self._readSplitAndExpand(tree,'RecHitHGC_TICLCandIdx')
self.truth={} #DEBUG!!!
self.truth['t_idx'] = self._expand(recHitSimClusIdx)# now expand to a trailing dimension
self.truth['t_energy'] = recHitTruthEnergy
self.truth['t_pos'] = ak1.concatenate([recHitTruthX, recHitTruthY,recHitTruthZ],axis=-1)
self.truth['t_time'] = recHitTruthTime
self.truth['t_pid'] = recHitTruthPID
self.truth['t_spectator'] = recHitSpectatorFlag
self.truth['t_fully_contained'] = fullyContained
self.truth['t_rec_energy'] = recHitDepEnergy
def calculate_selection(self, syst_tag, events):
"""
"""
electrons = events.ele
electrons["label"] = -1 * awkward.ones_like(electrons.pt)
if not self.is_data:
electrons["label"] = awkward.where(
electrons.genPartFlav == 1, awkward.ones_like(electrons.label),
electrons.label)
electrons["label"] = awkward.where(
(electrons.genPartFlav == 3) | (electrons.genPartFlav == 4) |
(electrons.genPartFlav == 5),
awkward.zeros_like(electrons.label), electrons.label)
fields = [x for x in events.fields if x not in ["ele", "Electron"]]
for x in fields:
if x == "ZCand":
electrons[x] = awkward.firsts(events[x])
else:
electrons[x] = events[x]
electrons = awkward.flatten(electrons)
dummy_cut = electrons.pt >= 0
return dummy_cut, electrons
def _assignTruthByIndexAndSplit(self, tree, label, indices, null=0):
sc = label
if type(label) is str:
sc = ak1.from_awkward0(tree[label].array())
vals = sc[indices]
vals = ak1.where(indices<0, ak1.zeros_like(vals)+null, vals)
ja = self._splitJaggedArray(vals)
return self._expand(ja)
def fsr_recovery(df):
mask = (
(df.Muon.fsrPhotonIdx >= 0)
& (df.Muon.matched_fsrPhoton.relIso03 < 1.8)
& (df.Muon.matched_fsrPhoton.dROverEt2 < 0.012)
& (df.Muon.matched_fsrPhoton.pt / df.Muon.pt < 0.4)
& (abs(df.Muon.matched_fsrPhoton.eta) < 2.4)
)
mask = ak.fill_none(mask, False)
px = ak.zeros_like(df.Muon.pt)
py = ak.zeros_like(df.Muon.pt)
pz = ak.zeros_like(df.Muon.pt)
e = ak.zeros_like(df.Muon.pt)
fsr = {
"pt": df.Muon.matched_fsrPhoton.pt,
"eta": df.Muon.matched_fsrPhoton.eta,
"phi": df.Muon.matched_fsrPhoton.phi,
"mass": 0.0,
}
for obj in [df.Muon, fsr]:
px_ = obj["pt"] * np.cos(obj["phi"])
py_ = obj["pt"] * np.sin(obj["phi"])
pz_ = obj["pt"] * np.sinh(obj["eta"])
e_ = np.sqrt(px_**2 + py_**2 + pz_**2 + obj["mass"] ** 2)
px = px + px_
py = py + py_
pz = pz + pz_
e = e + e_
pt = np.sqrt(px**2 + py**2)
eta = np.arcsinh(pz / pt)
phi = np.arctan2(py, px)
mass = np.sqrt(e**2 - px**2 - py**2 - pz**2)
iso = (df.Muon.pfRelIso04_all * df.Muon.pt - df.Muon.matched_fsrPhoton.pt) / pt
df["Muon", "pt_fsr"] = ak.where(mask, pt, df.Muon.pt)
df["Muon", "eta_fsr"] = ak.where(mask, eta, df.Muon.eta)
df["Muon", "phi_fsr"] = ak.where(mask, phi, df.Muon.phi)
df["Muon", "mass_fsr"] = ak.where(mask, mass, df.Muon.mass)
df["Muon", "iso_fsr"] = ak.where(mask, iso, df.Muon.pfRelIso04_all)
return mask
def find_first_parent(particle, maxgen=10):
tmp_mother = particle.parent
out_mother_pdg = tmp_mother.pdgId
found = ak.zeros_like(particle.pdgId)
for i in range(maxgen):
update = ak.fill_none(
(tmp_mother.pdgId != particle.pdgId) * (found == 0), False)
out_mother_pdg = update * ak.fill_none(tmp_mother.pdgId,
0) + (~update) * out_mother_pdg
found = found + update
tmp_mother = tmp_mother.parent
return out_mother_pdg
def _assignTruth(self, tree):
nonSplitTrackSimClusIdx = self._getMatchIdxs(tree)
truthEnergy = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_boundaryEnergy",nonSplitTrackSimClusIdx)
truthPID = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_pdgId",nonSplitTrackSimClusIdx)
truthX = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_x",nonSplitTrackSimClusIdx)
truthY = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_y",nonSplitTrackSimClusIdx)
truthZ = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_z",nonSplitTrackSimClusIdx)
truthTime = self._assignTruthByIndexAndSplit(tree,"MergedSimCluster_impactPoint_t",nonSplitTrackSimClusIdx)
#some manual sets
zeros = ak1.zeros_like(truthEnergy)
splittruthidx = self._splitJaggedArray(nonSplitTrackSimClusIdx)
spectator = ak1.where(splittruthidx<0, zeros+10., zeros)
trackPt = self._readSplitAndExpand(tree,"Track_pt")
trackVertEta = self._readSplitAndExpand(tree,"Track_eta")
trackMom = trackPt * np.cosh(trackVertEta)
impactX = self._readSplitAndExpand(tree,"Track_HGCFront_x")
impactY = self._readSplitAndExpand(tree,"Track_HGCFront_y")
impactZ = self._readSplitAndExpand(tree,"Track_HGCFront_z")
truthX = ak1.where(splittruthidx<0, impactX, truthX)
truthY = ak1.where(splittruthidx<0, impactY, truthY)
truthZ = ak1.where(splittruthidx<0, impactZ, truthZ)
truthEnergy = ak1.where(splittruthidx<0, trackMom, truthEnergy)
truthidx = self._expand(splittruthidx)
self.truth={}
self.truth['t_idx'] = truthidx# for now
self.truth['t_energy'] = truthEnergy
self.truth['t_pos'] = ak1.concatenate([truthX,truthY,truthZ],axis=-1)
self.truth['t_time'] = truthTime
self.truth['t_pid'] = truthPID
self.truth['t_spectator'] = spectator
self.truth['t_fully_contained'] = zeros+1
self.truth['t_rec_energy'] = trackMom
def _readTree(self, tree):
# no truth here! Only features
self._readSplits(tree, splitlabel='RecHitHGC_z')
recHitEnergy = self._readSplitAndExpand(tree,"RecHitHGC_energy")
recHitTime = self._readSplitAndExpand(tree,"RecHitHGC_time")
recHitX = self._readSplitAndExpand(tree,"RecHitHGC_x")
recHitY = self._readSplitAndExpand(tree,"RecHitHGC_y")
recHitZ = self._readSplitAndExpand(tree,"RecHitHGC_z")
recHitHitR = self._readSplitAndExpand(tree,"RecHitHGC_hitr")
recHitR = np.sqrt(recHitX*recHitX+recHitY*recHitY+recHitZ*recHitZ)
recHitTheta = np.arccos(recHitZ/recHitR)
recHitEta = -np.log(np.tan(recHitTheta/2))
zeros = ak1.zeros_like(recHitEta)
self.features = ak1.concatenate([
recHitEnergy,
recHitEta,
zeros, #indicator if it is track or not
recHitTheta,
recHitR,
recHitX,
recHitY,
recHitZ,
recHitTime,
recHitHitR
], axis=-1)
#this is just for bookkeeping
self.featurenames = [
'recHitEnergy',
'recHitEta',
'isTrack',
'recHitTheta',
'recHitR',
'recHitX',
'recHitY',
'recHitZ',
'recHitTime',
'recHitHitR'
]
def jer_smear(
variation,
forceStochastic,
pt_gen,
jetPt,
etaJet,
jet_energy_resolution,
jet_resolution_rand_gauss,
jet_energy_resolution_scale_factor,
):
pt_gen = pt_gen if not forceStochastic else None
if not isinstance(jetPt, awkward.highlevel.Array):
raise Exception("'jetPt' must be an awkward array of some kind!")
if forceStochastic:
pt_gen = awkward.without_parameters(awkward.zeros_like(jetPt))
jersmear = jet_energy_resolution * jet_resolution_rand_gauss
jersf = jet_energy_resolution_scale_factor[:, variation]
deltaPtRel = (jetPt - pt_gen) / jetPt
doHybrid = (pt_gen > 0) & (numpy.abs(deltaPtRel) <
3 * jet_energy_resolution)
detSmear = 1 + (jersf - 1) * deltaPtRel
stochSmear = 1 + numpy.sqrt(numpy.maximum(jersf**2 - 1, 0)) * jersmear
min_jet_pt = _MIN_JET_ENERGY / numpy.cosh(etaJet)
min_jet_pt_corr = min_jet_pt / jetPt
smearfact = awkward.where(doHybrid, detSmear, stochSmear)
smearfact = awkward.where((smearfact * jetPt) < min_jet_pt,
min_jet_pt_corr, smearfact)
def getfunction(layout, depth):
if isinstance(layout, awkward.layout.NumpyArray) or not isinstance(
layout,
(awkward.layout.Content, awkward.partition.PartitionedArray)):
return lambda: awkward.layout.NumpyArray(smearfact)
return None
smearfact = awkward._util.recursively_apply(
awkward.operations.convert.to_layout(jetPt), getfunction)
smearfact = awkward._util.wrap(smearfact, awkward._util.behaviorof(jetPt))
return smearfact
def _readTree(self, tree):
self._readSplits(tree, splitlabel='Track_HGCFront_z')
trackPt = self._readSplitAndExpand(tree,"Track_pt")
trackEta = self._readSplitAndExpand(tree,"Track_HGCFront_eta")
trackVertEta = self._readSplitAndExpand(tree,"Track_eta")
trackMom = trackPt * np.cosh(trackVertEta)
impactX = self._readSplitAndExpand(tree,"Track_HGCFront_x")
impactY = self._readSplitAndExpand(tree,"Track_HGCFront_y")
impactZ = self._readSplitAndExpand(tree,"Track_HGCFront_z")
chi2 = self._readSplitAndExpand(tree,"Track_normChiSq")
impactR = np.sqrt(impactX**2+impactY**2+impactZ**2)+1e-3
impactTheta = np.arccos(impactZ/impactR)
self.features = ak1.concatenate([
trackMom,
trackEta,
ak1.ones_like(trackMom), #indicator if it is track or not
impactTheta,
impactR,
impactX,
impactY,
impactZ,
ak1.zeros_like(trackMom),#no time info (yet,could be from MTD here)
chi2 #this is radius for hits, here chi2 for tracks, since it's kinda realted to the impact points resolution...
], axis=-1)
#this is just for bookkeeping, keep them the same as for hits
self.featurenames = [
'recHitEnergy',
'recHitEta',
'isTrack',
'recHitTheta',
'recHitR',
'recHitX',
'recHitY',
'recHitZ',
'recHitTime',
'recHitHitR'
]
def get_charge_parent(particle):
parent = find_first_parent(particle)
charge = ak.zeros_like(parent)
one = [[-11, -13, -15, -17, 24, 37], 1]
minus_one = [[11, 13, 15, 17, -24, -37], -1]
two_thirds = [[2, 4, 6, 8], 2 / 3]
minus_two_thirds = [[-2, -4, -6, -8], -2 / 3]
minus_one_third = [[1, 3, 5, 7], -1 / 3]
one_third = [[-1, -3, -5, -7], 1 / 3]
zero = [[12, 14, 16, 18, 9, 21, 22, 23, 25], 0]
charge_pairs = [
one, minus_one, two_thirds, minus_two_thirds, minus_one_third, zero
]
for pair in charge_pairs:
for ID in pair[0]:
charge = (parent == ID) * ak.ones_like(parent) * pair[1] + (
~(parent == ID)) * charge
return charge
def test():
array = ak.Array(
[
[{"x": 0.0, "y": []}, {"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}],
[],
[
{"x": 3.3, "y": [1, 2, None, 3]},
False,
False,
True,
{"x": 4.4, "y": [1, 2, None, 3, 4]},
],
]
)
def assert_array_type(new_array, intended_type):
"""helper function to check each part of the array took the intended type"""
assert isinstance(new_array[0][0]["x"], intended_type)
assert isinstance(new_array[0][1]["x"], intended_type)
assert isinstance(new_array[0][1]["y"][0], intended_type)
assert isinstance(new_array[0][2]["x"], intended_type)
assert isinstance(new_array[0][2]["y"][0], intended_type)
assert isinstance(new_array[0][2]["y"][1], intended_type)
assert isinstance(new_array[2][0]["x"], intended_type)
assert isinstance(new_array[2][0]["y"][0], intended_type)
assert isinstance(new_array[2][0]["y"][1], intended_type)
assert isinstance(new_array[2][0]["y"][3], intended_type)
assert isinstance(new_array[2][1], intended_type)
assert isinstance(new_array[2][2], intended_type)
assert isinstance(new_array[2][3], intended_type)
assert isinstance(new_array[2][4]["x"], intended_type)
assert isinstance(new_array[2][4]["y"][0], intended_type)
assert isinstance(new_array[2][4]["y"][1], intended_type)
assert isinstance(new_array[2][4]["y"][3], intended_type)
assert isinstance(new_array[2][4]["y"][4], intended_type)
int_type = ak.full_like(array, 12, dtype=int)
float_type = ak.full_like(array, 12, dtype=float)
assert int_type.tolist() == float_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
bool_type = ak.full_like(array, 12, dtype=bool)
assert_array_type(bool_type, bool)
int_type = ak.full_like(array, -1.2, dtype=int)
float_type = ak.full_like(array, -1.2, dtype=float)
bool_type = ak.full_like(array, -1.2, dtype=bool)
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.zeros_like(array, dtype=int)
float_type = ak.zeros_like(array, dtype=float)
bool_type = ak.zeros_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.ones_like(array, dtype=int)
float_type = ak.ones_like(array, dtype=float)
bool_type = ak.ones_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
array = ak.Array([["one", "two", "three"], [], ["four", "five"]])
def assert_array_type(new_array, intended_type):
"""helper function to check each part of the array took the intended type"""
assert isinstance(new_array[0][0], intended_type)
assert isinstance(new_array[0][1], intended_type)
assert isinstance(new_array[0][2], intended_type)
assert isinstance(new_array[2][0], intended_type)
assert isinstance(new_array[2][1], intended_type)
int_type = ak.full_like(array, 12, dtype=int)
float_type = ak.full_like(array, 12, dtype=float)
assert int_type.tolist() == float_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
bool_type = ak.full_like(array, 12, dtype=bool)
assert_array_type(bool_type, bool)
int_type = ak.full_like(array, -1.2, dtype=int)
float_type = ak.full_like(array, -1.2, dtype=float)
bool_type = ak.full_like(array, -1.2, dtype=bool)
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.zeros_like(array, dtype=int)
float_type = ak.zeros_like(array, dtype=float)
bool_type = ak.zeros_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
int_type = ak.ones_like(array, dtype=int)
float_type = ak.ones_like(array, dtype=float)
bool_type = ak.ones_like(array, dtype=bool)
assert int_type.tolist() == float_type.tolist()
assert int_type.tolist() == bool_type.tolist()
assert_array_type(int_type, int)
assert_array_type(float_type, float)
assert_array_type(bool_type, bool)
def test():
array = ak.Array([
[{
"x": 0.0,
"y": []
}, {
"x": 1.1,
"y": [1]
}, {
"x": 2.2,
"y": [1, 2]
}],
[],
[
{
"x": 3.3,
"y": [1, 2, None, 3]
},
False,
False,
True,
{
"x": 4.4,
"y": [1, 2, None, 3, 4]
},
],
])
assert ak.full_like(array, 12.3).tolist() == [
[{
"x": 12.3,
"y": []
}, {
"x": 12.3,
"y": [12]
}, {
"x": 12.3,
"y": [12, 12]
}],
[],
[
{
"x": 12.3,
"y": [12, 12, None, 12]
},
True,
True,
True,
{
"x": 12.3,
"y": [12, 12, None, 12, 12]
},
],
]
assert ak.zeros_like(array).tolist() == [
[{
"x": 0.0,
"y": []
}, {
"x": 0.0,
"y": [0]
}, {
"x": 0.0,
"y": [0, 0]
}],
[],
[
{
"x": 0.0,
"y": [0, 0, None, 0]
},
False,
False,
False,
{
"x": 0.0,
"y": [0, 0, None, 0, 0]
},
],
]
assert ak.ones_like(array).tolist() == [
[{
"x": 1.0,
"y": []
}, {
"x": 1.0,
"y": [1]
}, {
"x": 1.0,
"y": [1, 1]
}],
[],
[
{
"x": 1.0,
"y": [1, 1, None, 1]
},
True,
True,
True,
{
"x": 1.0,
"y": [1, 1, None, 1, 1]
},
],
]
array = ak.Array([["one", "two", "three"], [], ["four", "five"]])
assert ak.full_like(array, "hello").tolist() == [
["hello", "hello", "hello"],
[],
["hello", "hello"],
]
assert ak.full_like(array, 1).tolist() == [["1", "1", "1"], [], ["1", "1"]]
assert ak.full_like(array, 0).tolist() == [["0", "0", "0"], [], ["0", "0"]]
assert ak.ones_like(array).tolist() == [["1", "1", "1"], [], ["1", "1"]]
assert ak.zeros_like(array).tolist() == [["", "", ""], [], ["", ""]]
array = ak.Array([[b"one", b"two", b"three"], [], [b"four", b"five"]])
assert ak.full_like(array, b"hello").tolist() == [
[b"hello", b"hello", b"hello"],
[],
[b"hello", b"hello"],
]
assert ak.full_like(array, 1).tolist() == [[b"1", b"1", b"1"], [],
[b"1", b"1"]]
assert ak.full_like(array, 0).tolist() == [[b"0", b"0", b"0"], [],
[b"0", b"0"]]
assert ak.ones_like(array).tolist() == [[b"1", b"1", b"1"], [],
[b"1", b"1"]]
assert ak.zeros_like(array).tolist() == [[b"", b"", b""], [], [b"", b""]]
def __init__(self, ev, obj, wp, year=2018, verbose=0):
self.obj = obj
self.wp = wp
if self.wp == None:
self.selection_dict = {}
else:
self.selection_dict = obj_def[self.obj][self.wp]
self.v = verbose
self.year = year
id_level = None
if wp.lower().count('veto') or wp.lower().count('loose'):
id_level = 0
elif wp.lower().count('fake'):
id_level = 1
elif wp.lower().count('tight'):
id_level = 2
if self.obj == "Muon":
# collections are already there, so we just need to calculate missing ones
ev['Muon', 'absMiniIso'] = ev.Muon.miniPFRelIso_all * ev.Muon.pt
ev['Muon', 'ptErrRel'] = ev.Muon.ptErr / ev.Muon.pt
# this is what we are using:
# - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise
# - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise
# - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise
mask_close = (ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet),
99) < 0.4) * 1
mask_far = ~(ak.fill_none(ev.Muon.delta_r(ev.Muon.matched_jet), 99)
< 0.4) * 1
deepJet = ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB,
0) * mask_close + 0 * mask_far
jetRelIsoV2 = ev.Muon.jetRelIso * mask_close + ev.Muon.pfRelIso03_all * mask_far # default to 0 if no match
#conePt = 0.9 * ak.fill_none(ev.Muon.matched_jet.pt,0) * mask_close + ev.Muon.pt*(1 + ev.Muon.miniPFRelIso_all)*mask_far
if self.year == 2017 or self.year == 2018:
I_1 = 0.11
I_2 = 0.74
I_3 = 6.8
elif self.year == 2016:
I_1 = 0.16
I_2 = 0.76
I_3 = 7.2
PF_unflatten = ak.from_regular(
ev.Muon.miniPFRelIso_all[:, :, np.newaxis])
max_miniIso = ak.max(
ak.concatenate(
[PF_unflatten - I_1,
ak.zeros_like(PF_unflatten)], axis=2),
axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1)
muon_pt_unflatten = ak.from_regular(ev.Muon.pt[:, :, np.newaxis])
jet_pt_unflatten = ak.from_regular(
ev.Muon.matched_jet.pt[:, :, np.newaxis])
max_pt = ak.max(
ak.concatenate([muon_pt_unflatten, jet_pt_unflatten * I_2],
axis=2),
axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2)
conePt = (ev.Muon.pt *
(1 + max_miniIso)) * (ev.Muon.jetPtRelv2 > I_3) + (
max_pt * ~(ev.Muon.jetPtRelv2 > I_3))
ev['Muon', 'deepJet'] = ak.copy(deepJet)
ev['Muon', 'jetRelIsoV2'] = jetRelIsoV2
ev['Muon', 'conePt'] = conePt
ev['Muon', 'id'] = ak.ones_like(conePt) * id_level
self.cand = ev.Muon
elif self.obj == "Electron":
# calculate new variables. asignment is awkward, but what can you do.
ev['Electron',
'absMiniIso'] = ev.Electron.miniPFRelIso_all * ev.Electron.pt
ev['Electron', 'etaSC'] = ev.Electron.eta + ev.Electron.deltaEtaSC
if self.year == 2017 or self.year == 2018:
I_1 = 0.07
I_2 = 0.78
I_3 = 8.0
elif self.year == 2016:
I_1 = 0.12
I_2 = 0.80
I_3 = 7.2
# the following line is only needed if we do our own matching.
# right now, we keep using the NanoAOD match, but check the deltaR distance
# jet_index, mask_match, mask_nomatch = self.matchJets(ev.Electron, ev.Jet)
# this is what we are using:
# - jetRelIso if the matched jet is within deltaR<0.4, pfRelIso03_all otherwise
# - btagDeepFlavB discriminator of the matched jet if jet is within deltaR<0.4, 0 otherwise
# - pt_cone = 0.9*pt of matched jet if jet is within deltaR<0.4, pt/(pt+iso) otherwise
mask_close = (ak.fill_none(
ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1
mask_far = ~(ak.fill_none(
ev.Electron.delta_r(ev.Electron.matched_jet), 99) < 0.4) * 1
deepJet = ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB,
0) * mask_close
jetRelIsoV2 = ev.Electron.jetRelIso * mask_close + ev.Electron.pfRelIso03_all * mask_far # default to 0 if no match
#conePt = 0.9 * ak.fill_none(ev.Electron.matched_jet.pt,0) * mask_close + ev.Electron.pt*(1 + ev.Electron.miniPFRelIso_all)*mask_far
PF_unflatten = ak.from_regular(
ev.Electron.miniPFRelIso_all[:, :, np.newaxis])
max_miniIso = ak.max(
ak.concatenate(
[PF_unflatten - I_1,
ak.zeros_like(PF_unflatten)], axis=2),
axis=2) #equivalent to max(0, ev.Muon.miniPFRelIso_all - I_1)
electron_pt_unflatten = ak.from_regular(ev.Electron.pt[:, :,
np.newaxis])
jet_pt_unflatten = ak.from_regular(
ev.Electron.matched_jet.pt[:, :, np.newaxis])
max_pt = ak.max(
ak.concatenate([electron_pt_unflatten, jet_pt_unflatten * I_2],
axis=2),
axis=2) #max(ev.Muon.pt, ev.Muon.matched_jet.pt * I_2)
conePt = (ev.Electron.pt *
(1 + max_miniIso)) * (ev.Electron.jetPtRelv2 > I_3) + (
max_pt * ~(ev.Electron.jetPtRelv2 > I_3))
ev['Electron', 'deepJet'] = ak.copy(deepJet)
ev['Electron', 'jetRelIsoV2'] = jetRelIsoV2
ev['Electron', 'conePt'] = conePt
ev['Electron', 'id'] = ak.ones_like(conePt) * id_level
ev['Electron', 'jetRelIso'] = ev.Electron.jetRelIso
ev['Electron', 'jetPtRelv2'] = ev.Electron.jetPtRelv2
self.cand = ev.Electron
self.getSelection()
if self.obj == "Electron" and self.wp == "tight":
self.selection = self.selection & self.getElectronMVAID(
) & self.getIsolation(0.07, 0.78, 8.0) & self.isTriggerSafeNoIso()
if self.v > 0: print(" - custom ID and multi-isolation")
if self.obj == "Electron" and self.wp == "tightFCNC":
self.selection = self.selection & self.getElectronMVAID(
) & self.getFCNCIsolation(ev.Electron.jetRelIso,
ev.Electron.jetPtRelv2, I_2,
I_3) & (ev.Electron.miniPFRelIso_all <
I_1) & self.isTriggerSafeNoIso()
if self.v > 0: print(" - custom ID and multi-isolation")
if self.obj == "Muon" and self.wp == "tight":
self.selection = self.selection & self.getIsolation(
0.11, 0.74, 6.8)
if self.v > 0: print(" - custom multi-isolation")
#self.selection = self.selection & ak.fill_none(ev.Muon.matched_jet.btagDeepFlavB<0.2770, True)
#self.selection = self.selection & (ev.Muon.matched_jet.btagDeepFlavB<0.2770)
#if self.v>0: print (" - deepJet")
if self.obj == "Muon" and self.wp == "tightFCNC":
self.selection = self.selection & self.getFCNCIsolation(
ev.Muon.jetRelIso, ev.Muon.jetPtRelv2, I_2,
I_3) & (ev.Muon.miniPFRelIso_all < I_1)
if self.v > 0: print(" - custom multi-isolation")
if self.obj == "Electron" and (self.wp == "tightTTH"
or self.wp == 'fakeableTTH'
or self.wp == "tightSSTTH"
or self.wp == 'fakeableSSTTH'):
self.selection = self.selection & self.getSigmaIEtaIEta
if self.v > 0: print(" - SigmaIEtaIEta")
#self.selection = self.selection & ak.fill_none(ev.Electron.matched_jet.btagDeepFlavB<0.2770, True)
#self.selection = self.selection & (ev.Electron.matched_jet.btagDeepFlavB<0.2770)
#self.selection = self.selection & (ev.Jet[ev.Electron.jetIdx].btagDeepFlavB<0.2770)
#if self.v>0: print (" - deepJet")
if self.obj == "Electron" and self.wp == "looseFCNC":
self.selection = self.selection & (ev.Electron.miniPFRelIso_all <
0.4)
if self.obj == 'Muon' and (self.wp == 'fakeableTTH'
or self.wp == 'fakeableSSTTH'):
#self.selection = self.selection & (self.cand.deepJet < self.getThreshold(self.cand.conePt, min_pt=20, max_pt=45, low=0.2770, high=0.0494))
self.selection = self.selection & (ak.fill_none(
ev.Muon.matched_jet.btagDeepFlavB, 0) < self.getThreshold(
self.cand.conePt, min_pt=20, max_pt=45))
if self.v > 0: print(" - interpolated deepJet")
if self.obj == "Muon" and self.wp == "looseFCNC":
self.selection = self.selection & (ev.Muon.miniPFRelIso_all < 0.4)
def mass(self):
return awkward.without_parameters(awkward.zeros_like(self.pt))
def select_normal(genparts, w_decay_momid):
columns = ["pt", "eta", "phi", "mass", "pdgId", "charge", "decaytype"]
# get tops, defined as last copy
gen_tops = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == 6)]
gen_tops['charge'] = ak.ones_like(gen_tops.pt) * (2. / 3.)
gen_tbars = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == -6)]
gen_tbars['charge'] = ak.ones_like(gen_tbars.pt) * (-2. / 3.)
# get direct top decay products (children are "isHardProcess" and "isFirstCopy")
gen_bs = ak.flatten(gen_tops.children[(gen_tops.children.pdgId == 5)],
axis=2)
gen_bs['charge'] = ak.ones_like(gen_bs.pt) * (-1. / 3.)
gen_bbars = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -5)], axis=2)
gen_bbars['charge'] = ak.ones_like(gen_bbars.pt) * (1. / 3.)
gen_wplus = ak.flatten(
gen_tops.children[(gen_tops.children.pdgId == w_decay_momid)], axis=2)
gen_wplus['charge'] = ak.ones_like(gen_wplus.pt)
gen_wminus = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -1 * w_decay_momid)],
axis=2)
gen_wminus['charge'] = ak.ones_like(gen_wminus.pt) * (-1.)
# get w decay products
# get last copy of W bosons
last_gen_wplus = ak.flatten(gen_tops.distinctChildren[
(gen_tops.distinctChildren.pdgId == w_decay_momid)
& (gen_tops.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
last_gen_wminus = ak.flatten(gen_tbars.distinctChildren[
(gen_tbars.distinctChildren.pdgId == -1 * w_decay_momid)
& (gen_tbars.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
gen_wplus_decaytype = np.zeros(len(gen_wplus))
gen_wminus_decaytype = np.zeros(len(gen_wminus))
# up/down partons from last W+
gen_wpartons_up_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (2. / 3.)
gen_wplus_decaytype[ak.num(gen_wpartons_up_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWplus) > 0)) * 2
gen_wpartons_dw_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (-1. / 3.)
# up/down partons from last W-
gen_wpartons_up_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (2. / 3.)
gen_wminus_decaytype[ak.num(gen_wpartons_up_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWminus) > 0)) * 2
gen_wpartons_dw_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (-1. / 3.)
# charged leps from last W+
gen_charged_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 11) |
(np.abs(last_gen_wplus.children.pdgId) == 13) |
(np.abs(last_gen_wplus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWplus['charge'] = ak.ones_like(
gen_charged_leps_fromWplus.pdgId)
gen_wplus_decaytype[ak.num(gen_charged_leps_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWplus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wplus['decaytype'] = ak.unflatten(gen_wplus_decaytype,
ak.num(gen_wplus))
gen_tops['decaytype'] = gen_wplus['decaytype'] # set decaytype for tops
gen_bs['decaytype'] = gen_wplus['decaytype'] # set decaytype for bs
# neutral leps from last W+
gen_neutral_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 12) |
(np.abs(last_gen_wplus.children.pdgId) == 14) |
(np.abs(last_gen_wplus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWplus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWplus.pt)
# charged leps from last W-
gen_charged_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 11) |
(np.abs(last_gen_wminus.children.pdgId) == 13) |
(np.abs(last_gen_wminus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWminus['charge'] = ak.ones_like(
gen_charged_leps_fromWminus.pdgId) * (-1.)
gen_wminus_decaytype[ak.num(gen_charged_leps_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWminus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wminus['decaytype'] = ak.unflatten(gen_wminus_decaytype,
ak.num(gen_wminus))
gen_tbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for tbars
gen_bbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for bbars
# neutral leps from last W-
gen_neutral_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 12) |
(np.abs(last_gen_wminus.children.pdgId) == 14) |
(np.abs(last_gen_wminus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWminus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWminus.pt)
gen_wpartons_up = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_wpartons_dw = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_charged_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_neutral_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
if column == 'decaytype': continue
gen_wpartons_up[column] = ak.flatten(
ak.concatenate([
gen_wpartons_up_fromWplus[column],
gen_wpartons_up_fromWminus[column]
],
axis=1)) # (up-type partons from W+, W-)
gen_wpartons_dw[column] = ak.flatten(
ak.concatenate([
gen_wpartons_dw_fromWplus[column],
gen_wpartons_dw_fromWminus[column]
],
axis=1)) # (dw-type partons from W+, W-)
gen_charged_leps[column] = ak.flatten(
ak.concatenate([
gen_charged_leps_fromWplus[column],
gen_charged_leps_fromWminus[column]
],
axis=1)) # (charged leps from W+, W-)
gen_neutral_leps[column] = ak.flatten(
ak.concatenate([
gen_neutral_leps_fromWplus[column],
gen_neutral_leps_fromWminus[column]
],
axis=1)) # (neutral leps from W+, W-)
gen_wpartons_up = ak.unflatten(
gen_wpartons_up,
ak.num(gen_wpartons_up_fromWplus) + ak.num(gen_wpartons_up_fromWminus))
gen_wpartons_dw = ak.unflatten(
gen_wpartons_dw,
ak.num(gen_wpartons_dw_fromWplus) + ak.num(gen_wpartons_dw_fromWminus))
gen_charged_leps = ak.unflatten(
gen_charged_leps,
ak.num(gen_charged_leps_fromWplus) +
ak.num(gen_charged_leps_fromWminus))
gen_neutral_leps = ak.unflatten(
gen_neutral_leps,
ak.num(gen_neutral_leps_fromWplus) +
ak.num(gen_neutral_leps_fromWminus))
gen_taus = gen_charged_leps[np.abs(gen_charged_leps.pdgId) == 15]
gen_nu_taus = gen_neutral_leps[np.abs(gen_neutral_leps.pdgId) == 16]
# fully hadronic evts
had_evts = (ak.num(gen_charged_leps) == 0) & (
ak.num(gen_neutral_leps)
== 0) & (ak.num(gen_wpartons_up) == 2) & (ak.num(gen_wpartons_dw) == 2)
#set_trace()
# get direct tau decay products from hard processes (subset of gen_taus events above)
tau_decay_prods = genparts[genparts.hasFlags(
['isDirectHardProcessTauDecayProduct'])]
# only need decays to leptons (e/mu, tau nu) for event classification
tau_TO_tau_nu = tau_decay_prods[np.abs(tau_decay_prods.pdgId) == 16]
tau_TO_charged_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 11)
|
(np.abs(tau_decay_prods.pdgId) == 13)]
tau_TO_neutral_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 12)
|
(np.abs(tau_decay_prods.pdgId) == 14)]
# set decaytype for gen taus
charged_lep_decaytype_array = ak.to_numpy(
ak.flatten(ak.zeros_like(gen_charged_leps['pt'])))
# semilep evts
semilep_evts = (ak.num(gen_charged_leps) == 1) & (
ak.num(gen_neutral_leps)
== 1) & (ak.num(gen_wpartons_up) == 1) & (ak.num(gen_wpartons_dw) == 1)
tau_jets_evts = semilep_evts & (ak.num(gen_taus) == 1)
sl_evts_mask = np.repeat(ak.to_numpy(semilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
semilep_decaytype_array = np.zeros(ak.to_numpy(semilep_evts).sum(),
dtype=int)
# tau -> l
semilep_tau_leptonic_decay = (tau_jets_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
semilep_tau_hadronic_decay = (tau_jets_evts) & (
~semilep_tau_leptonic_decay)
semilep_decaytype_array[
semilep_tau_leptonic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_leptonic_decay).sum(), dtype=int)
semilep_decaytype_array[
semilep_tau_hadronic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_hadronic_decay).sum(), dtype=int) * 2
# set charged_lep_decatype_array for semileptonic events
charged_lep_decaytype_array[sl_evts_mask] = semilep_decaytype_array
# dilep evts
dilep_evts = (ak.num(gen_charged_leps) == 2) & (
ak.num(gen_neutral_leps)
== 2) & (ak.num(gen_wpartons_up) == 0) & (ak.num(gen_wpartons_dw) == 0)
lep_lep_evts = dilep_evts & (ak.num(gen_taus) == 0)
lep_tau_evts = dilep_evts & (ak.num(gen_taus) == 1)
tau_tau_evts = dilep_evts & (ak.num(gen_taus) == 2)
dl_evts_mask = np.repeat(ak.to_numpy(dilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
dilep_decaytype_array = np.zeros((ak.to_numpy(dilep_evts).sum(), 2),
dtype=int)
# tau + tau
# tau + tau -> ll
dilep_TauTau_ll_decay = (tau_tau_evts) & (
ak.num(tau_TO_charged_lep) == 2) & (
ak.num(tau_TO_neutral_lep) == 2) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_ll_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_ll_decay).sum(), 2), dtype=int)
# tau + tau -> hh
dilep_TauTau_hh_decay = (tau_tau_evts) & (
(ak.num(tau_TO_charged_lep) + ak.num(tau_TO_neutral_lep))
== 0) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_hh_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_hh_decay).sum(), 2), dtype=int) * 2
# tau + tau -> lh
dilep_TauTau_lh_decay = (
tau_tau_evts) & ~(dilep_TauTau_ll_decay | dilep_TauTau_hh_decay)
# set index corresponding to leptonically decaying tau to 1, default array is set to 2
dl_TauTau_to_lh_decaytype_array = np.ones(
ak.to_numpy(dilep_TauTau_lh_decay).sum() * 2, dtype=int) * 2
lep_tau_mask = (np.repeat(
ak.to_numpy(
ak.flatten(
tau_TO_charged_lep[dilep_TauTau_lh_decay].parent.pdgId)),
2) == ak.flatten(gen_charged_leps[dilep_TauTau_lh_decay].pdgId))
dl_TauTau_to_lh_decaytype_array[lep_tau_mask] = np.ones(lep_tau_mask.sum(),
dtype=int)
dilep_decaytype_array[dilep_TauTau_lh_decay[
dilep_evts]] = dl_TauTau_to_lh_decaytype_array.reshape(
ak.to_numpy(dilep_TauTau_lh_decay).sum(), 2)
# lep + tau
# tau -> l
dilep_LepTau_l_decay = (lep_tau_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
# set index corresponding to tau to 1
dl_LepTau_to_Lep_decaytype_array = np.zeros(
ak.to_numpy(dilep_LepTau_l_decay).sum() * 2, dtype=int)
dl_LepTau_to_Lep_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_l_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_l_decay), dtype=int)
dilep_decaytype_array[dilep_LepTau_l_decay[
dilep_evts]] = dl_LepTau_to_Lep_decaytype_array.reshape(
ak.sum(dilep_LepTau_l_decay), 2)
# tau -> h
dilep_LepTau_h_decay = (lep_tau_evts) & ~(dilep_LepTau_l_decay)
# set index corresponding to tau to 2
dl_LepTau_to_Had_decaytype_array = np.zeros(ak.sum(dilep_LepTau_h_decay) *
2,
dtype=int)
dl_LepTau_to_Had_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_h_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_h_decay), dtype=int) * 2
dilep_decaytype_array[dilep_LepTau_h_decay[
dilep_evts]] = dl_LepTau_to_Had_decaytype_array.reshape(
ak.sum(dilep_LepTau_h_decay), 2)
# set charged_lep_decatype_array for dileptonic events
charged_lep_decaytype_array[dl_evts_mask] = dilep_decaytype_array.flatten()
# set charged lepton decaytype (defined for taus only, e/mu are 0) (1 is LEPTONIC, 2 is HADRONIC)
gen_charged_leps['decaytype'] = ak.unflatten(charged_lep_decaytype_array,
ak.num(gen_charged_leps))
# make awkward arrays of (top decay prods, tbar decay prods)
Gen_Top_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_B_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_W_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_Wparton_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
Gen_Top_Pairs[column] = ak.flatten(
ak.concatenate([gen_tops[column], gen_tbars[column]],
axis=1)) # (top, tbar)
Gen_B_Pairs[column] = ak.flatten(
ak.concatenate([gen_bs[column], gen_bbars[column]],
axis=1)) # (b, bbar)
Gen_W_Pairs[column] = ak.flatten(
ak.concatenate([gen_wplus[column], gen_wminus[column]],
axis=1)) # (W+, W-)
if column is not "decaytype":
Gen_Wparton_Pairs[column] = ak.flatten(
ak.concatenate(
[
ak.pad_none(gen_wpartons_up[column], 1, axis=1),
ak.pad_none(gen_wpartons_dw[column], 1, axis=1)
],
axis=1)) # (up-type wpartons, down-type wpartons)
Gen_Top_Pairs = ak.unflatten(Gen_Top_Pairs,
ak.num(gen_tops) + ak.num(gen_tbars))
Gen_B_Pairs = ak.unflatten(Gen_B_Pairs, ak.num(gen_bs) + ak.num(gen_bbars))
Gen_W_Pairs = ak.unflatten(Gen_W_Pairs,
ak.num(gen_wplus) + ak.num(gen_wminus))
Gen_Wparton_Pairs = ak.unflatten(
Gen_Wparton_Pairs,
ak.num(ak.pad_none(gen_wpartons_up, 1, axis=1)) +
ak.num(ak.pad_none(gen_wpartons_dw, 1, axis=1)))
Gen_Wparton_Pairs = Gen_Wparton_Pairs[ak.argsort(
Gen_Wparton_Pairs["pt"], ascending=False)] # sort by pt
Gen_TTbar = ak.Array(
{
"pt": (gen_tops + gen_tbars).pt,
"eta": (gen_tops + gen_tbars).eta,
"phi": (gen_tops + gen_tbars).phi,
"mass": (gen_tops + gen_tbars).mass,
"decaytype": gen_tops["decaytype"] + gen_tbars[
"decaytype"], # 0 is for INVALID, 2 for DILEP, 3 for SEMILEP, 4 for HADRONIC
},
with_name="PtEtaPhiMLorentzVector")
## make "table" of gen objects for certain decays
## DILEP
DILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[dilep_evts] if ak.any(dilep_evts) else ak.unflatten(
Gen_TTbar[dilep_evts], ak.values_astype(dilep_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"First_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"First_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"Up_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second plus
"Down_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first plus
"Up_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second minus
"Down_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first minus
})
## HAD
HAD_evts = ak.zip({
"TTbar":
Gen_TTbar[had_evts] if ak.any(had_evts) else ak.unflatten(
Gen_TTbar[had_evts], ak.values_astype(had_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts] if ak.any(had_evts) else
ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts] if ak.any(had_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts],
ak.values_astype(had_evts, int)),
"First_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading positively-charged parton
"Second_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading positively-charged parton
"First_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading negatively-charged parton
"Second_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading negatively-charged parton
"Up_plus":
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged up-type parton
"Down_plus":
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged down-type parton
"Up_minus":
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged up-type parton
"Down_minus":
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged down-type parton
})
## SEMILEP
SEMILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(
Gen_TTbar[semilep_evts], ak.values_astype(semilep_evts, int)),
"THad":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"TLep":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BHad":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BLep":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WHad":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WLep":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Lepton":
gen_charged_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_charged_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Nu":
gen_neutral_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_neutral_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJa":
Gen_Wparton_Pairs[:, 0][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 0][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJb":
Gen_Wparton_Pairs[:, 1][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Up_Had":
gen_wpartons_up[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_up[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Down_Had":
gen_wpartons_dw[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_dw[semilep_evts],
ak.values_astype(semilep_evts, int)),
})
# make dictionary to return
GenObjects = dict({
"SL": SEMILEP_evts,
"DL": DILEP_evts,
"Had": HAD_evts,
})
return GenObjects
def process_shift(self, events, shift_name):
dataset = events.metadata['dataset']
isRealData = not hasattr(events, "genWeight")
selection = PackedSelection()
weights = Weights(len(events), storeIndividual=True)
output = self.make_output()
if shift_name is None and not isRealData:
output['sumw'] = ak.sum(events.genWeight)
if isRealData or self._newTrigger:
trigger = np.zeros(len(events), dtype='bool')
for t in self._triggers[self._year]:
if t in events.HLT.fields:
trigger = trigger | events.HLT[t]
selection.add('trigger', trigger)
del trigger
else:
selection.add('trigger', np.ones(len(events), dtype='bool'))
if isRealData:
selection.add(
'lumimask', lumiMasks[self._year](events.run,
events.luminosityBlock))
else:
selection.add('lumimask', np.ones(len(events), dtype='bool'))
if isRealData and self._skipRunB and self._year == '2017':
selection.add('dropB', events.run > 299329)
else:
selection.add('dropB', np.ones(len(events), dtype='bool'))
if isRealData:
trigger = np.zeros(len(events), dtype='bool')
for t in self._muontriggers[self._year]:
if t in events.HLT.fields:
trigger |= np.array(events.HLT[t])
selection.add('muontrigger', trigger)
del trigger
else:
selection.add('muontrigger', np.ones(len(events), dtype='bool'))
metfilter = np.ones(len(events), dtype='bool')
for flag in self._met_filters[
self._year]['data' if isRealData else 'mc']:
metfilter &= np.array(events.Flag[flag])
selection.add('metfilter', metfilter)
del metfilter
fatjets = events.FatJet
fatjets['msdcorr'] = corrected_msoftdrop(fatjets)
fatjets['qcdrho'] = 2 * np.log(fatjets.msdcorr / fatjets.pt)
fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year=self._year)
fatjets['msdcorr_full'] = fatjets['msdcorr'] * self._msdSF[self._year]
candidatejet = fatjets[
# https://github.com/DAZSLE/BaconAnalyzer/blob/master/Analyzer/src/VJetLoader.cc#L269
(fatjets.pt > 200)
& (abs(fatjets.eta) < 2.5)
& fatjets.isTight # this is loose in sampleContainer
]
candidatejet = candidatejet[:, :
2] # Only consider first two to match generators
if self._jet_arbitration == 'pt':
candidatejet = ak.firsts(candidatejet)
elif self._jet_arbitration == 'mass':
candidatejet = ak.firsts(candidatejet[ak.argmax(
candidatejet.msdcorr, axis=1, keepdims=True)])
elif self._jet_arbitration == 'n2':
candidatejet = ak.firsts(candidatejet[ak.argmin(candidatejet.n2ddt,
axis=1,
keepdims=True)])
elif self._jet_arbitration == 'ddb':
candidatejet = ak.firsts(candidatejet[ak.argmax(
candidatejet.btagDDBvLV2, axis=1, keepdims=True)])
elif self._jet_arbitration == 'ddc':
candidatejet = ak.firsts(candidatejet[ak.argmax(
candidatejet.btagDDCvLV2, axis=1, keepdims=True)])
else:
raise RuntimeError("Unknown candidate jet arbitration")
if self._tagger == 'v1':
bvl = candidatejet.btagDDBvL
cvl = candidatejet.btagDDCvL
cvb = candidatejet.btagDDCvB
elif self._tagger == 'v2':
bvl = candidatejet.btagDDBvLV2
cvl = candidatejet.btagDDCvLV2
cvb = candidatejet.btagDDCvBV2
elif self._tagger == 'v3':
bvl = candidatejet.particleNetMD_Xbb
cvl = candidatejet.particleNetMD_Xcc / (
1 - candidatejet.particleNetMD_Xbb)
cvb = candidatejet.particleNetMD_Xcc / (
candidatejet.particleNetMD_Xcc +
candidatejet.particleNetMD_Xbb)
elif self._tagger == 'v4':
bvl = candidatejet.particleNetMD_Xbb
cvl = candidatejet.btagDDCvLV2
cvb = candidatejet.particleNetMD_Xcc / (
candidatejet.particleNetMD_Xcc +
candidatejet.particleNetMD_Xbb)
else:
raise ValueError("Not an option")
selection.add('minjetkin', (candidatejet.pt >= 450)
& (candidatejet.pt < 1200)
& (candidatejet.msdcorr >= 40.)
& (candidatejet.msdcorr < 201.)
& (abs(candidatejet.eta) < 2.5))
selection.add('_strict_mass', (candidatejet.msdcorr > 85) &
(candidatejet.msdcorr < 130))
selection.add('_high_score', cvl > 0.8)
selection.add('minjetkinmu', (candidatejet.pt >= 400)
& (candidatejet.pt < 1200)
& (candidatejet.msdcorr >= 40.)
& (candidatejet.msdcorr < 201.)
& (abs(candidatejet.eta) < 2.5))
selection.add('minjetkinw', (candidatejet.pt >= 200)
& (candidatejet.pt < 1200)
& (candidatejet.msdcorr >= 40.)
& (candidatejet.msdcorr < 201.)
& (abs(candidatejet.eta) < 2.5))
selection.add('jetid', candidatejet.isTight)
selection.add('n2ddt', (candidatejet.n2ddt < 0.))
if not self._tagger == 'v2':
selection.add('ddbpass', (bvl >= 0.89))
selection.add('ddcpass', (cvl >= 0.83))
selection.add('ddcvbpass', (cvb >= 0.2))
else:
selection.add('ddbpass', (bvl >= 0.7))
selection.add('ddcpass', (cvl >= 0.45))
selection.add('ddcvbpass', (cvb >= 0.03))
jets = events.Jet
jets = jets[(jets.pt > 30.) & (abs(jets.eta) < 2.5) & jets.isTight]
# only consider first 4 jets to be consistent with old framework
jets = jets[:, :4]
dphi = abs(jets.delta_phi(candidatejet))
selection.add(
'antiak4btagMediumOppHem',
ak.max(jets[dphi > np.pi / 2][self._ak4tagBranch],
axis=1,
mask_identity=False) <
BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium'])
ak4_away = jets[dphi > 0.8]
selection.add(
'ak4btagMedium08',
ak.max(ak4_away[self._ak4tagBranch], axis=1, mask_identity=False) >
BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium'])
met = events.MET
selection.add('met', met.pt < 140.)
goodmuon = ((events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.4)
& (events.Muon.pfRelIso04_all < 0.25)
& events.Muon.looseId)
nmuons = ak.sum(goodmuon, axis=1)
leadingmuon = ak.firsts(events.Muon[goodmuon])
if self._looseTau:
goodelectron = ((events.Electron.pt > 10)
& (abs(events.Electron.eta) < 2.5)
&
(events.Electron.cutBased >= events.Electron.VETO))
nelectrons = ak.sum(goodelectron, axis=1)
ntaus = ak.sum(
((events.Tau.pt > 20)
& (abs(events.Tau.eta) < 2.3)
& events.Tau.idDecayMode
& ((events.Tau.idMVAoldDM2017v2 & 2) != 0)
& ak.all(events.Tau.metric_table(events.Muon[goodmuon]) > 0.4,
axis=2)
& ak.all(events.Tau.metric_table(
events.Electron[goodelectron]) > 0.4,
axis=2)),
axis=1,
)
else:
goodelectron = (
(events.Electron.pt > 10)
& (abs(events.Electron.eta) < 2.5)
& (events.Electron.cutBased >= events.Electron.LOOSE))
nelectrons = ak.sum(goodelectron, axis=1)
ntaus = ak.sum(
(events.Tau.pt > 20)
&
events.Tau.idDecayMode # bacon iso looser than Nano selection
& ak.all(events.Tau.metric_table(events.Muon[goodmuon]) > 0.4,
axis=2)
& ak.all(events.Tau.metric_table(events.Electron[goodelectron])
> 0.4,
axis=2),
axis=1,
)
selection.add('noleptons',
(nmuons == 0) & (nelectrons == 0) & (ntaus == 0))
selection.add('onemuon',
(nmuons == 1) & (nelectrons == 0) & (ntaus == 0))
selection.add('muonkin',
(leadingmuon.pt > 55.) & (abs(leadingmuon.eta) < 2.1))
selection.add('muonDphiAK8',
abs(leadingmuon.delta_phi(candidatejet)) > 2 * np.pi / 3)
# W-Tag (Tag and Probe)
# tag side
selection.add(
'ak4btagMediumOppHem',
ak.max(jets[dphi > np.pi / 2][self._ak4tagBranch],
axis=1,
mask_identity=False) >
BTagEfficiency.btagWPs[self._ak4tagger][self._year]['medium'])
selection.add('met40p', met.pt > 40.)
selection.add('tightMuon',
(leadingmuon.tightId) & (leadingmuon.pt > 53.))
# selection.add('ptrecoW', (leadingmuon + met).pt > 250.)
selection.add('ptrecoW200', (leadingmuon + met).pt > 200.)
selection.add(
'ak4btagNearMu',
leadingmuon.delta_r(leadingmuon.nearest(ak4_away, axis=None)) <
2.0)
_bjets = jets[self._ak4tagBranch] > BTagEfficiency.btagWPs[
self._ak4tagger][self._year]['medium']
# _nearAK8 = jets.delta_r(candidatejet) < 0.8
# _nearMu = jets.delta_r(ak.firsts(events.Muon)) < 0.3
# selection.add('ak4btagOld', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1)
_nearAK8 = jets.delta_r(candidatejet) < 0.8
_nearMu = jets.delta_r(leadingmuon) < 0.3
selection.add('ak4btagOld',
ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1)
# _nearAK8 = jets.delta_r(candidatejet) < 0.8
# _nearMu = jets.delta_r(candidatejet.nearest(events.Muon[goodmuon], axis=None)) < 0.3
# selection.add('ak4btagNew', ak.sum(_bjets & ~_nearAK8 & ~_nearMu, axis=1) >= 1)
# probe side
selection.add('minWjetpteta',
(candidatejet.pt >= 200) & (abs(candidatejet.eta) < 2.4))
# selection.add('noNearMuon', candidatejet.delta_r(candidatejet.nearest(events.Muon[goodmuon], axis=None)) > 1.0)
selection.add('noNearMuon', candidatejet.delta_r(leadingmuon) > 1.0)
#####
if isRealData:
genflavor = ak.zeros_like(candidatejet.pt)
else:
if 'HToCC' in dataset or 'HToBB' in dataset:
if self._ewkHcorr:
add_HiggsEW_kFactors(weights, events.GenPart, dataset)
weights.add('genweight', events.genWeight)
if "PSWeight" in events.fields:
add_ps_weight(weights, events.PSWeight)
else:
add_ps_weight(weights, None)
if "LHEPdfWeight" in events.fields:
add_pdf_weight(weights, events.LHEPdfWeight)
else:
add_pdf_weight(weights, None)
if "LHEScaleWeight" in events.fields:
add_scalevar_7pt(weights, events.LHEScaleWeight)
add_scalevar_3pt(weights, events.LHEScaleWeight)
else:
add_scalevar_7pt(weights, [])
add_scalevar_3pt(weights, [])
add_pileup_weight(weights, events.Pileup.nPU, self._year, dataset)
bosons = getBosons(events.GenPart)
matchedBoson = candidatejet.nearest(bosons,
axis=None,
threshold=0.8)
if self._tightMatch:
match_mask = (
(candidatejet.pt - matchedBoson.pt) / matchedBoson.pt <
0.5) & ((candidatejet.msdcorr - matchedBoson.mass) /
matchedBoson.mass < 0.3)
selmatchedBoson = ak.mask(matchedBoson, match_mask)
genflavor = bosonFlavor(selmatchedBoson)
else:
genflavor = bosonFlavor(matchedBoson)
genBosonPt = ak.fill_none(ak.firsts(bosons.pt), 0)
if self._newVjetsKfactor:
add_VJets_kFactors(weights, events.GenPart, dataset)
else:
add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
if shift_name is None:
output['btagWeight'].fill(val=self._btagSF.addBtagWeight(
weights, ak4_away, self._ak4tagBranch))
if self._nnlops_rew and dataset in [
'GluGluHToCC_M125_13TeV_powheg_pythia8'
]:
weights.add('minlo_rew',
powheg_to_nnlops(ak.to_numpy(genBosonPt)))
if self._newTrigger:
add_jetTriggerSF(
weights, ak.firsts(fatjets),
self._year if not self._skipRunB else f'{self._year}CDEF',
selection)
else:
add_jetTriggerWeight(weights, candidatejet.msdcorr,
candidatejet.pt, self._year)
add_mutriggerSF(weights, leadingmuon, self._year, selection)
add_mucorrectionsSF(weights, leadingmuon, self._year, selection)
if self._year in ("2016", "2017"):
weights.add("L1Prefiring", events.L1PreFiringWeight.Nom,
events.L1PreFiringWeight.Up,
events.L1PreFiringWeight.Dn)
logger.debug("Weight statistics: %r" % weights.weightStatistics)
msd_matched = candidatejet.msdcorr * self._msdSF[self._year] * (
genflavor > 0) + candidatejet.msdcorr * (genflavor == 0)
regions = {
'signal': [
'noleptons', 'minjetkin', 'met', 'metfilter', 'jetid',
'antiak4btagMediumOppHem', 'n2ddt', 'trigger', 'lumimask'
],
'signal_noddt': [
'noleptons', 'minjetkin', 'met', 'jetid',
'antiak4btagMediumOppHem', 'trigger', 'lumimask', 'metfilter'
],
# 'muoncontrol': ['minjetkinmu', 'jetid', 'n2ddt', 'ak4btagMedium08', 'onemuon', 'muonkin', 'muonDphiAK8', 'muontrigger', 'lumimask', 'metfilter'],
'muoncontrol': [
'onemuon', 'muonkin', 'muonDphiAK8', 'metfilter',
'minjetkinmu', 'jetid', 'ak4btagMedium08', 'n2ddt',
'muontrigger', 'lumimask'
],
'muoncontrol_noddt': [
'onemuon', 'muonkin', 'muonDphiAK8', 'jetid', 'metfilter',
'minjetkinmu', 'jetid', 'ak4btagMedium08', 'muontrigger',
'lumimask'
],
'wtag': [
'onemuon', 'tightMuon', 'minjetkinw', 'jetid', 'met40p',
'metfilter', 'ptrecoW200', 'ak4btagOld', 'muontrigger',
'lumimask'
],
'wtag0': [
'onemuon', 'tightMuon', 'met40p', 'metfilter', 'ptrecoW200',
'ak4btagOld', 'muontrigger', 'lumimask'
],
'wtag2': [
'onemuon', 'tightMuon', 'minjetkinw', 'jetid',
'ak4btagMediumOppHem', 'met40p', 'metfilter', 'ptrecoW200',
'ak4btagOld', 'muontrigger', 'lumimask'
],
'noselection': [],
}
def normalize(val, cut):
if cut is None:
ar = ak.to_numpy(ak.fill_none(val, np.nan))
return ar
else:
ar = ak.to_numpy(ak.fill_none(val[cut], np.nan))
return ar
import time
tic = time.time()
if shift_name is None:
for region, cuts in regions.items():
allcuts = set([])
cut = selection.all(*allcuts)
output['cutflow_msd'].fill(region=region,
genflavor=normalize(
genflavor, None),
cut=0,
weight=weights.weight(),
msd=normalize(msd_matched, None))
output['cutflow_eta'].fill(region=region,
genflavor=normalize(genflavor, cut),
cut=0,
weight=weights.weight()[cut],
eta=normalize(
candidatejet.eta, cut))
output['cutflow_pt'].fill(region=region,
genflavor=normalize(genflavor, cut),
cut=0,
weight=weights.weight()[cut],
pt=normalize(candidatejet.pt, cut))
for i, cut in enumerate(cuts + ['ddcvbpass', 'ddcpass']):
allcuts.add(cut)
cut = selection.all(*allcuts)
output['cutflow_msd'].fill(region=region,
genflavor=normalize(
genflavor, cut),
cut=i + 1,
weight=weights.weight()[cut],
msd=normalize(msd_matched, cut))
output['cutflow_eta'].fill(
region=region,
genflavor=normalize(genflavor, cut),
cut=i + 1,
weight=weights.weight()[cut],
eta=normalize(candidatejet.eta, cut))
output['cutflow_pt'].fill(
region=region,
genflavor=normalize(genflavor, cut),
cut=i + 1,
weight=weights.weight()[cut],
pt=normalize(candidatejet.pt, cut))
if self._evtVizInfo and 'ddcpass' in allcuts and isRealData and region == 'signal':
if 'event' not in events.fields:
continue
_cut = selection.all(*allcuts, '_strict_mass',
'_high_score')
# _cut = selection.all('_strict_mass'')
output['to_check'][
'mass'] += processor.column_accumulator(
normalize(msd_matched, _cut))
nfatjet = ak.sum(
((fatjets.pt > 200) &
(abs(fatjets.eta) < 2.5) & fatjets.isTight),
axis=1)
output['to_check'][
'njet'] += processor.column_accumulator(
normalize(nfatjet, _cut))
output['to_check'][
'fname'] += processor.column_accumulator(
np.array([events.metadata['filename']] *
len(normalize(msd_matched, _cut))))
output['to_check'][
'event'] += processor.column_accumulator(
normalize(events.event, _cut))
output['to_check'][
'luminosityBlock'] += processor.column_accumulator(
normalize(events.luminosityBlock, _cut))
output['to_check'][
'run'] += processor.column_accumulator(
normalize(events.run, _cut))
if shift_name is None:
systematics = [None] + list(weights.variations)
else:
systematics = [shift_name]
def fill(region, systematic, wmod=None):
selections = regions[region]
cut = selection.all(*selections)
sname = 'nominal' if systematic is None else systematic
if wmod is None:
if systematic in weights.variations:
weight = weights.weight(modifier=systematic)[cut]
else:
weight = weights.weight()[cut]
else:
weight = weights.weight()[cut] * wmod[cut]
output['templates'].fill(
region=region,
systematic=sname,
runid=runmap(events.run)[cut],
genflavor=normalize(genflavor, cut),
pt=normalize(candidatejet.pt, cut),
msd=normalize(msd_matched, cut),
ddb=normalize(bvl, cut),
ddc=normalize(cvl, cut),
ddcvb=normalize(cvb, cut),
weight=weight,
)
if region in [
'wtag', 'wtag0', 'wtag2', 'wtag3', 'wtag4', 'wtag5',
'wtag6', 'wtag7', 'noselection'
]: # and sname in ['nominal', 'pileup_weightDown', 'pileup_weightUp', 'jet_triggerDown', 'jet_triggerUp']:
output['wtag'].fill(
region=region,
systematic=sname,
genflavor=normalize(genflavor, cut),
pt=normalize(candidatejet.pt, cut),
msd=normalize(msd_matched, cut),
n2ddt=normalize(candidatejet.n2ddt, cut),
ddc=normalize(cvl, cut),
ddcvb=normalize(cvb, cut),
weight=weight,
)
# if region in ['signal', 'noselection']:
# output['etaphi'].fill(
# region=region,
# systematic=sname,
# runid=runmap(events.run)[cut],
# genflavor=normalize(genflavor, cut),
# pt=normalize(candidatejet.pt, cut),
# eta=normalize(candidatejet.eta, cut),
# phi=normalize(candidatejet.phi, cut),
# ddc=normalize(cvl, cut),
# ddcvb=normalize(cvb, cut),
# ),
if not isRealData:
if wmod is not None:
_custom_weight = events.genWeight[cut] * wmod[cut]
else:
_custom_weight = np.ones_like(weight)
output['genresponse_noweight'].fill(
region=region,
systematic=sname,
pt=normalize(candidatejet.pt, cut),
genpt=normalize(genBosonPt, cut),
weight=_custom_weight,
)
output['genresponse'].fill(
region=region,
systematic=sname,
pt=normalize(candidatejet.pt, cut),
genpt=normalize(genBosonPt, cut),
weight=weight,
)
if systematic is None:
output['signal_opt'].fill(
region=region,
genflavor=normalize(genflavor, cut),
ddc=normalize(cvl, cut),
ddcvb=normalize(cvb, cut),
msd=normalize(msd_matched, cut),
weight=weight,
)
output['signal_optb'].fill(
region=region,
genflavor=normalize(genflavor, cut),
ddb=normalize(bvl, cut),
msd=normalize(msd_matched, cut),
weight=weight,
)
for region in regions:
cut = selection.all(*(set(regions[region]) - {'n2ddt'}))
if shift_name is None:
output['nminus1_n2ddt'].fill(
region=region,
n2ddt=normalize(candidatejet.n2ddt, cut),
weight=weights.weight()[cut],
)
for systematic in systematics:
if isRealData and systematic is not None:
continue
fill(region, systematic)
if shift_name is None and 'GluGluH' in dataset and 'LHEWeight' in events.fields:
for i in range(9):
fill(region, 'LHEScale_%d' % i, events.LHEScaleWeight[:,
i])
for c in events.LHEWeight.fields[1:]:
fill(region, 'LHEWeight_%s' % c, events.LHEWeight[c])
toc = time.time()
output["filltime"] = toc - tic
if shift_name is None:
output["weightStats"] = weights.weightStatistics
return {dataset: output}
def best_match(gen_hyp=None, jets=None, leptons=None, met=None):
if gen_hyp is None:
raise ValueError("Gen Objects gen_hyp needed for matching")
if jets is None:
raise ValueError("Reco jets needed for matching")
if leptons is None:
raise ValueError("Reco leptons needed for matching")
if met is None:
raise ValueError("Reco met needed for matching")
if not ak.all(ak.num(gen_hyp) == 1):
raise ValueError("Not all events for matching are semileptonic")
jets_ak = ak.with_name(jets[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector")
leps_ak = ak.with_name(leptons[["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector")
# init dict of objects
matched_objects = {}
# match jet closest to gen objects
for genobj in ['BHad', 'BLep', 'WJa', 'WJb']:
genobj_ak = ak.with_name(gen_hyp[genobj][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector")
jets_akc, genobj_akc = ak.unzip(ak.cartesian([jets_ak, genobj_ak], nested=False))
deltaRs = jets_akc.delta_r(genobj_akc) # find deltaRs between jets and gen object
indexOfMin = ak.unflatten(ak.argmin(deltaRs, axis=1), ak.num(genobj_ak))
passing_inds = deltaRs[indexOfMin] < 0.4
matched_jets_inds = indexOfMin[passing_inds]
matched_jets = jets[matched_jets_inds]
## add matched perm objects
matched_objects[genobj] = ak.Array({
'pt' : matched_jets.pt,
'eta' : matched_jets.eta,
'phi' : matched_jets.phi,
'mass' : matched_jets.mass,
'jetIdx' : matched_jets_inds, # index of jet that the gen object is matched to in the event
}, with_name="PtEtaPhiMLorentzVector")
# match lepton closest to gen lepton
genlep_ak = ak.with_name(gen_hyp['Lepton'][["pt", "eta", "phi", "mass"]],"PtEtaPhiMLorentzVector")
lep_akc, genlep_akc = ak.unzip(ak.cartesian([leps_ak, genlep_ak], nested=False))
lepDRs = lep_akc.delta_r(genlep_akc)
lepIdxOfMin = ak.unflatten(ak.argmin(lepDRs, axis=1), ak.num(genlep_ak))
passing_inds = lepDRs[lepIdxOfMin] < 0.4
matched_leps_inds = lepIdxOfMin[passing_inds]
matched_leps = leptons[matched_leps_inds]
## add matched perm objects
matched_objects['Lepton'] = ak.Array({key: matched_leps[key] for key in matched_leps.fields}, with_name="PtEtaPhiMLorentzVector")
# solve for neutrino
nu_array = np.zeros((len(ak.num(jets)), 4), dtype='float64')
# convert all inputs into 2d numpy arrays of dtype=float64 (won't work if they're not float64)
blep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].py, 1), -999))).astype('float64'),\
ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['BLep'].energy, 1), -999))).astype('float64')), axis=-1)
lep_inputs = np.stack((ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].px, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].py, 1), -999))).astype('float64'),\
ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].pz, 1), -999))).astype('float64'), ak.to_numpy(ak.flatten(ak.fill_none(ak.pad_none(matched_objects['Lepton'].energy, 1), -999))).astype('float64')), axis=-1)
met_inputs = np.stack((ak.to_numpy(ak.fill_none(met.px, -999)).astype('float64'), ak.to_numpy(ak.fill_none(met.py, -999)).astype('float64')), axis=-1)
nu_array = find_nu(bleps=blep_inputs, leptons=lep_inputs, met=met_inputs, nu_array=nu_array)
valid_nu = ~((nu_array[:, 3] > 1e20) | (nu_array[:, 3] == 0)) # events that have a solution and matched blep
# convert px, py, pz to pt, eta, phi
nu_px, nu_py, nu_pz = nu_array[:, 0][valid_nu], nu_array[:, 1][valid_nu], nu_array[:, 2][valid_nu]
nu_mom, nu_pt = np.sqrt(np.square(nu_px)+np.square(nu_py)+np.square(nu_pz)), np.sqrt(np.square(nu_px)+np.square(nu_py))
nu_phi = np.arctan2(nu_py, nu_px)
nu_eta = np.arcsinh(nu_pz/nu_pt)
matched_objects['Nu'] = ak.Array({
'pt' : ak.unflatten(nu_pt, valid_nu.astype(int)),
'eta' : ak.unflatten(nu_eta, valid_nu.astype(int)),
'phi' : ak.unflatten(nu_phi, valid_nu.astype(int)),
'mass' : ak.zeros_like(ak.unflatten(nu_array[:, 0][valid_nu], valid_nu.astype(int))),
'chi2' : ak.unflatten(nu_array[:, 3][valid_nu], valid_nu.astype(int)),
}, with_name="PtEtaPhiMLorentzVector")
matched_perm = make_perm_table(bhad=matched_objects['BHad'], blep=matched_objects['BLep'], wja=matched_objects['WJa'], wjb=matched_objects['WJb'], lepton=matched_objects['Lepton'], met=met, nu=matched_objects['Nu'])
return matched_perm
