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())