def get_p4(name):
# Given an alias like "els_p4", return a jagged array of LorentzVectors
vx,vy,vz,vt = t.arrays([
alias_lut[name]+".fCoordinates.fX",
alias_lut[name]+".fCoordinates.fY",
alias_lut[name]+".fCoordinates.fZ",
alias_lut[name]+".fCoordinates.fT",
], outputtype=tuple)
return TLorentzVectorArray.from_cartesian(vx,vy,vz,vt)
def _build_graph(self, constits):
constits_p4 = TLorentzVectorArray.from_cartesian(*list(zip(*constits)))
jet_p4 = constits_p4.sum()
spatialCoord = np.stack([
delta_eta_reflect(constits_p4, jet_p4),
constits_p4.delta_phi(jet_p4)
],
axis=1)
energyFeatures = np.log(
np.stack([constits_p4.pt, constits_p4.energy], axis=1))
features = np.concatenate([spatialCoord, energyFeatures], axis=1)
ret = dgl.DGLGraph()
ret.add_nodes(
len(constits), {
'coordinates': torch.tensor(spatialCoord, dtype=torch.float32),
'features': torch.tensor(features, dtype=torch.float32)
})
# print(ret.number_of_nodes())
return ret
def get_baseline(self, check_labels, j, taus, etas, phis, pts):
baselines = np.empty((0, 5))
m = TLorentzVectorArray.from_cartesian(j.fX, j.fY, j.fZ, j.fE).mass
m = np.nan_to_num(m)
taus = np.nan_to_num(taus)
taus = np.where(taus == 0, 10**-6, taus)
tau21 = taus[:, 1] / taus[:, 0]
tau32 = taus[:, 2] / taus[:, 1]
for label in check_labels:
jid = self.constants.settings[label]['id']
cuts_m = self.constants.settings[label]['cut_m']
mask = (m > cuts_m[0]) & (m < cuts_m[1])
scores = tau32 if label == 't' else tau21
for e, p, pt, s in zip(etas[mask],
phis[mask],
pts[mask],
scores[mask]):
if e < self.edges_eta[0] or e > self.edges_eta[-1]:
continue
e = np.argmax(self.edges_eta >= e) - 1
p = np.argmax(self.edges_phi >= p) - 1
baselines = np.vstack((baselines, [jid, e, p, pt, s]))
return baselines
