def decay_vertex(self): nu_tau = None # use production vertex of nu_tau last_tau = self.init.last_self for child in last_tau.iter_children(): if abs(child.pdgId) == pdg.nu_tau: nu_tau = child break if nu_tau is None: return Vector3(0, 0, 0) return Vector3(nu_tau.vx_x, nu_tau.vx_y, nu_tau.vx_z)
def decay_vertex(self): nu_tau = None nu_tau_vertex = None for nu in self.neutrinos: if nu.absPdgId() == pdg.nu_tau: nu_tau = nu nu_tau_vertex = nu.decayVtx() break if nu_tau is None: return Vector3(0, 0, 0) return Vector3(nu_tau_vertex.x(), nu_tau_vertex.y(), nu_tau_vertex.z())
def secvtx(self): if self.obj.nTracks() < 2: return self.privtx else: return Vector3( self.obj.secondaryVertex().x(), self.obj.secondaryVertex().y(), self.obj.secondaryVertex().z())
def write_to_root(tracks, laser): """ Writes tracks and laser data to a root file which is readable by the reconstruction algorithm """ from rootpy.vector import Vector3 import rootpy.stl as stl from rootpy.tree import Tree from rootpy.io import root_open laser_entry, laser_exit = laser Vec = stl.vector(Vector3) track = Vec() laserentry = Vector3(laser_entry[0], laser_entry[1], laser_entry[2]) laserexit = Vector3(laser_exit[0], laser_exit[1], laser_exit[2]) f = root_open("test.root", "recreate") track_tree = Tree('tracks') laser_tree = Tree('lasers') track_tree.create_branches({'track': stl.vector(Vector3)}) laser_tree.create_branches({'entry': Vector3, 'exit': Vector3}) for k in range(10): print(k) for i in range(1000): track.push_back(Vector3(i, k, k * i)) track_tree.track = track track.clear() laser_tree.entry = Vector3(0, 0, 0) laser_tree.exit = Vector3(k, k, k) track_tree.fill() laser_tree.fill() track_tree.write() laser_tree.write() f.close()
def prod_vertex(self): return Vector3(self.init.vx_x, self.init.vx_y, self.init.vx_z)
def transformVars(df): ''' modifies the variables to create the ones that mv2 uses, inserts default values when needed, saves new variables in the dataframe Args: ----- df: pandas dataframe containing all the interesting variables as extracted from the .root file Returns: -------- modified mv2-compliant dataframe ''' from rootpy.vector import LorentzVector, Vector3 import pandautils as pup # -- modify features and set default values df['abs(jet_eta)'] = abs(df['jet_eta']) # -- create new IPxD features for (pu, pb, pc) in zip(df['jet_ip2d_pu'], df['jet_ip2d_pb'], df['jet_ip2d_pc']): pu[np.logical_or(pu >= 10, pu < -1)] = -1 pb[np.logical_or(pu >= 10, pu < -1)] = -1 pc[np.logical_or(pu >= 10, pu < -1)] = -1 for (pu, pb, pc) in zip(df['jet_ip3d_pu'], df['jet_ip3d_pb'], df['jet_ip3d_pc']): pu[pu >= 10] = -1 pb[pu >= 10] = -1 pc[pu >= 10] = -1 df['jet_ip2'] = (df['jet_ip2d_pb'] / df['jet_ip2d_pu']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip2_c'] = (df['jet_ip2d_pb'] / df['jet_ip2d_pc']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip2_cu'] = (df['jet_ip2d_pc'] / df['jet_ip2d_pu']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip3'] = (df['jet_ip3d_pb'] / df['jet_ip3d_pu']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip3_c'] = (df['jet_ip3d_pb'] / df['jet_ip3d_pc']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip3_cu'] = (df['jet_ip3d_pc'] / df['jet_ip3d_pu']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) # -- create new IPMP features for (pu, pb, pc) in zip(df['jet_ipmp_pu'], df['jet_ipmp_pb'], df['jet_ipmp_pc']): pu[pu >= 10] = -1 pb[pu >= 10] = -1 pc[pu >= 10] = -1 df['jet_ip'] = (df['jet_ipmp_pb'] / df['jet_ipmp_pu']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip_c'] = (df['jet_ipmp_pb'] / df['jet_ipmp_pc']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) df['jet_ip_cu'] = (df['jet_ipmp_pc'] / df['jet_ipmp_pu']).apply( lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20)) # -- SV1 features dx = df['jet_sv1_vtx_x'] - df['PVx'] dy = df['jet_sv1_vtx_y'] - df['PVy'] dz = df['jet_sv1_vtx_z'] - df['PVz'] v_jet = LorentzVector() pv2sv = Vector3() sv1_L3d = [] sv1_Lxy = [] dR = [] for index, dxi in enumerate(dx): # loop thru events sv1_L3d_ev = [] sv1L_ev = [] dR_ev = [] for jet in xrange(len(dxi)): # loop thru jets v_jet.SetPtEtaPhiM(df['jet_pt'][index][jet], df['jet_eta'][index][jet], df['jet_phi'][index][jet], df['jet_m'][index][jet]) if (dxi[jet].size != 0): sv1_L3d_ev.append( np.sqrt( pow(dx[index][jet], 2) + pow(dy[index][jet], 2) + pow(dz[index][jet], 2))[0]) sv1L_ev.append(math.hypot(dx[index][jet], dy[index][jet])) pv2sv.SetXYZ(dx[index][jet], dy[index][jet], dz[index][jet]) jetAxis = Vector3(v_jet.Px(), v_jet.Py(), v_jet.Pz()) dR_ev.append(pv2sv.DeltaR(jetAxis)) else: dR_ev.append(-1) sv1L_ev.append(-100) sv1_L3d_ev.append(-100) sv1_Lxy.append(sv1L_ev) dR.append(dR_ev) sv1_L3d.append(sv1_L3d_ev) df['jet_sv1_dR'] = dR df['jet_sv1_Lxy'] = sv1_Lxy df['jet_sv1_L3d'] = sv1_L3d # -- add more default values for sv1 variables sv1_vtx_ok = pup.match_shape( np.asarray([len(el) for event in df['jet_sv1_vtx_x'] for el in event]), df['jet_pt']) for (ok4event, sv1_ntkv4event, sv1_n2t4event, sv1_mass4event, sv1_efrc4event, sv1_sig34event) in zip(sv1_vtx_ok, df['jet_sv1_ntrkv'], df['jet_sv1_n2t'], df['jet_sv1_m'], df['jet_sv1_efc'], df['jet_sv1_sig3d']): sv1_ntkv4event[np.asarray(ok4event) == 0] = -1 sv1_n2t4event[np.asarray(ok4event) == 0] = -1 sv1_mass4event[np.asarray(ok4event) == 0] = -1000 sv1_efrc4event[np.asarray(ok4event) == 0] = -1 sv1_sig34event[np.asarray(ok4event) == 0] = -100 # -- JF features jf_dR = [] for eventN, (etas, phis, masses) in enumerate( zip(df['jet_jf_deta'], df['jet_jf_dphi'], df['jet_jf_m'])): # loop thru events jf_dR_ev = [] for m in xrange(len(masses)): # loop thru jets if (masses[m] > 0): jf_dR_ev.append(np.sqrt(etas[m] * etas[m] + phis[m] * phis[m])) else: jf_dR_ev.append(-10) jf_dR.append(jf_dR_ev) df['jet_jf_dR'] = jf_dR # -- add more default values for jf variables for (jf_mass, jf_n2tv, jf_ntrkv, jf_nvtx, jf_nvtx1t, jf_efrc, jf_sig3) in zip(df['jet_jf_m'], df['jet_jf_n2t'], df['jet_jf_ntrkAtVx'], df['jet_jf_nvtx'], df['jet_jf_nvtx1t'], df['jet_jf_efc'], df['jet_jf_sig3d']): jf_n2tv[jf_mass <= 0] = -1 jf_ntrkv[jf_mass <= 0] = -1 jf_nvtx[jf_mass <= 0] = -1 jf_nvtx1t[jf_mass <= 0] = -1 jf_mass[jf_mass <= 0] = -1e3 jf_efrc[jf_mass <= 0] = -1 jf_sig3[jf_mass <= 0] = -100 return df
laser_entry = np.rec.array([ laser_data[event][1], laser_data[event][2], laser_data[event][3] ], dtype=[('x', 'f'), ('y', 'f'), ('z', 'f')]) laser_exit = np.rec.array([ laser_data[event][4], laser_data[event][5], laser_data[event][6] ], dtype=[('x', 'f'), ('y', 'f'), ('z', 'f')]) for idx in range(int(len(x) / downsample)): track.push_back( Vector3(x[idx * downsample], y[idx * downsample], z[idx * downsample])) track_tree.track = track track_tree.event = int(evt) track.clear() # create some stupid laser positions laser_tree.entry = Vector3(laser_entry.x, laser_entry.y, laser_entry.z) laser_tree.exit = Vector3(laser_exit.x, laser_exit.y, laser_exit.z) laser_tree.side = sides[side] track_tree.fill() laser_tree.fill() track_tree.write()
def privtx(self): return Vector3( self.obj.vertex().x(), self.obj.vertex().y(), self.obj.vertex().z())