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