def get_fjparticles(self,
df_tracks,
m,
offset_indices=False,
random_mass=False):
# If offset_indices is true, then offset the user_index by a large negative value
user_index_offset = 0
if offset_indices:
user_index_offset = int(-1e6)
m_array = np.full((df_tracks['ParticlePt'].values.size), m)
# Randomly assign K and p mass for systematic check
if random_mass:
rand_val = np.random.random((len(m_array)))
K_mass = 0.4937 # GeV/c^2
p_mass = 0.938272 # GeV/c^2
# (p + pbar) / (pi+ + pi-) ~ 5.5%
# (K+ + K-) / (pi+ + pi-) ~ 13%
# But these are numbers with respect to the final _unreplaced_ pions, so there is
# an additional factor of 1/(1 + 5.5% + 13%) to get things right
K_factor = 0.13
p_factor = 0.055
K_prob = K_factor / (1 + K_factor + p_factor)
p_prob = 1 - p_factor / (1 + K_factor + p_factor
) # 1- just to look at diff random vals
m_array = np.where(rand_val < K_prob, K_mass, m_array)
m_array = np.where(rand_val > p_prob, p_mass, m_array)
# Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi
fj_particles = fjext.vectorize_pt_eta_phi_m(
df_tracks['ParticlePt'].values, df_tracks['ParticleEta'].values,
df_tracks['ParticlePhi'].values, m_array, user_index_offset)
return fj_particles
def process_event(self, df):
self.pbar.update(1)
if 'ev_id_ext' in list(self.event_df):
_ev_query = "run_number == {} & ev_id == {} & ev_id_ext == {}".format(df['run_number'], df['ev_id'], df['ev_id_ext'])
else:
_ev_query = "run_number == {} & ev_id == {}".format(df['run_number'], df['ev_id'])
_df_d0 = self.d0_df.query(_ev_query)
_df_d0.reset_index(drop=True)
_n_d0s = len(_df_d0.index)
if _n_d0s < 1:
return
_ev_query = ""
_df_tracks = self.track_df.query(_ev_query)
_df_tracks.reset_index(drop=True)
_parts = fjext.vectorize_pt_eta_phi(_df_tracks['ParticlePt'].values, _df_tracks['ParticleEta'].values, _df_tracks['ParticlePhi'].values)
self._user_index_offset = 10000
# _d0s = fjext.vectorize_pt_eta_phi(_df_d0['pt_cand'].values, _df_d0['eta_cand'].values, _df_d0['phi_cand'].values, self._user_index_offset)
_d0s = fjext.vectorize_pt_eta_phi_m(_df_d0['pt_cand'].values, _df_d0['eta_cand'].values, _df_d0['phi_cand'].values, _df_d0['inv_mass'].values, self._user_index_offset)
_d0s_gh = [p * 1.e-6 for p in _d0s]
_parts_and_ds = _parts
_tmp = [_parts_and_ds.push_back(p) for p in _d0s_gh]
# pinfo('n parts = ', len(_parts_and_ds))
ja = jet_analysis.JetAnalysis(jet_R = 0.6, particle_eta_max=0.9, jet_pt_min=5.0)
ja.analyze_event(_parts_and_ds)
_d0_imass_list = _df_d0['inv_mass'].values.tolist()
self.tw.fill_branches(dpsj = _d0s)
self.tw.fill_branches(dpsjgh = _d0s_gh)
self.tw.fill_branches(minv = _d0_imass_list)
self.tw.fill_branches(jets = ja.jets_as_psj_vector())
self.tw.fill_tree()
self.d0_jet_correl(ja.jets, _d0s, _d0_imass_list)
def analyze(self, df):
self.compile_selection(df)
_df = df[self.df_selection]
self.fj_parts = fjext.vectorize_pt_eta_phi(_df['ParticlePt'].values,
_df['ParticleEta'].values,
_df['ParticlePhi'].values)
_df = _df.drop(columns=['ParticlePt', 'ParticleEta', 'ParticlePhi'])
self.fj_Dcands = fjext.vectorize_pt_eta_phi_m(_df['pt_cand'].values,
_df['eta_cand'].values,
_df['phi_cand'].values,
_df['inv_mass'].values)
_Dgh_pt = [1e-5 for v in _df['pt_cand'].values]
self.fj_DcandsGhosts = fjext.vectorize_pt_eta_phi_m(
_Dgh_pt, _df['eta_cand'].values, _df['phi_cand'].values,
_df['inv_mass'].values)
_df.reset_index(drop=True)
_df.drop_duplicates(inplace=True)
self.analysis(_df)
if self.callback is not None:
self.callback(df['ev_id'].values[0])
def get_fjparticles(self, df_tracks):
m = 0.1396
user_index_offset = 0
m_array = np.full((df_tracks['ParticlePt'].values.size), m)
# Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi
fj_particles = fjext.vectorize_pt_eta_phi_m(
df_tracks['ParticlePt'].values, df_tracks['ParticleEta'].values,
df_tracks['ParticlePhi'].values, m_array, user_index_offset)
self.ev_idx.append(df_tracks['ev_id'].values[0])
return fj_particles
def subtract_thermal_4momsub(self, jet, jetR, diagnostic=False):
m = 0.1396
jet_particles = 0.
dummies = 0.
background = ROOT.TLorentzVector()
for constit in jet.constituents():
if constit.pt() < 1e-5:
for th in range(0, len(self.event_bck_df)):
th_eta = self.event_bck_df['ParticleEta'].iloc[th]
th_phi = self.event_bck_df['ParticlePhi'].iloc[th]
th_pt = self.event_bck_df['ParticlePt'].iloc[th]
if math.sqrt(
pow(constit.eta() - th_eta, 2) +
pow(constit.phi() - th_phi, 2)) < 1e-5:
temp_background = ROOT.TLorentzVector()
temp_background.SetPtEtaPhiM(th_pt, th_eta, th_phi, m)
background += temp_background
dummies += 1.
else:
jet_particles += 1.
if jet_particles == 0:
print('Clustered jet only contains dummies. Bailing out!')
exit()
jet_df = pd.DataFrame(
columns=['ParticlePt', 'ParticleEta', 'ParticlePhi', 'm'])
for constit in jet.constituents():
if constit.pt() > 1e-5 and dummies > 0:
jet_df = jet_df.append(
{
'ParticlePt':
constit.pt() - background.Pt() /
((float)(jet_particles)),
'ParticleEta':
constit.eta(),
'ParticlePhi':
constit.phi(),
'm':
m
},
ignore_index=True)
elif constit.pt() > 1e-5:
jet_df = jet_df.append(
{
'ParticlePt': constit.pt(),
'ParticleEta': constit.eta(),
'ParticlePhi': constit.phi(),
'm': m
},
ignore_index=True)
if len(jet_df) == 0:
return None
# Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi
fj_particles = fjext.vectorize_pt_eta_phi_m(
jet_df['ParticlePt'].values, jet_df['ParticleEta'].values,
jet_df['ParticlePhi'].values, jet_df['m'].values, 0)
return fj_particles
def subtract_thermal_gridsub1(self,
jets,
jetR,
gridsize,
diagnostic=False):
m = 0.1396
#-----------------------------------
# Clean up what's leftover from the previous event
self.reset_grids(gridsize) # empty grid from previous event
self.populated_cells = [
] # empty this list, which was populated in the previous event
self.populated_cells_w_constit = [
] # empty this list, which was populated in the previous event
unsubtracted_thermal_pT_fraction = -1000
#-----------------------------------
# Loop over all jets in the event
for jet in jets:
jet_pt = jet.pt()
jet_eta = jet.eta()
jet_phi = jet.phi()
if jet_phi > math.pi:
jet_phi -= 2. * math.pi
# Add jet particles to grid
[
self.populate_grid_with_constituents(constit, gridsize)
for constit in jet.constituents()
]
# Subtract thermal particles from grid
if not diagnostic:
[
self.subtract_thermals_from_grid(th, gridsize, jet_eta,
jet_phi, jetR)
for th in range(0, len(self.event_bck_df))
]
if self.total_thermal_momentum > 0 and not diagnostic:
unsubtracted_thermal_pT_fraction = self.unsubtracted_thermal_momentum / self.total_thermal_momentum
name = 'h_thermal_fraction_not_subtracted_v_pT_R{}_gridsize{}'.format(
jetR, gridsize)
getattr(self, name).Fill(jet_pt,
unsubtracted_thermal_pT_fraction)
# Create dataframe to store the surviving constituents
jet_df = pd.DataFrame(
columns=['ParticlePt', 'ParticleEta', 'ParticlePhi', 'm'])
# Loop over grid
#for k1 in self.grid_dict[gridsize].keys():
# for k2 in self.grid_dict[gridsize][k1].keys():
for pair in self.populated_cells:
grid_cell = self.grid_dict[gridsize][pair[0]][pair[1]]
# If a cell has negative pT, set the 4 momentum to 0
if grid_cell.Pt() < 0:
grid_cell.SetPtEtaPhiM(0, 0, 0, 0)
pt = grid_cell.Pt()
eta = grid_cell.Eta()
phi = grid_cell.Phi()
# Append to dataframe all grid entries with positive pT
if pt > 0:
jet_df = jet_df.append(
{
'ParticlePt': grid_cell.Pt(),
'ParticleEta': grid_cell.Eta(),
'ParticlePhi': grid_cell.Phi(),
'm': m
},
ignore_index=True)
# If no entry was stored in grid at this point, return None
if len(jet_df) == 0:
print(
'This should never happen! The end of the world must be close')
return None
# Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi
fj_particles = fjext.vectorize_pt_eta_phi_m(
jet_df['ParticlePt'].values, jet_df['ParticleEta'].values,
jet_df['ParticlePhi'].values, jet_df['m'].values, 0)
return fj_particles
