df_new.head()
df_no_pf_cands = df.drop(group_list, axis=1).drop_duplicates()
df_new = df_new.join(df_no_pf_cands.set_index(
['event', 'run', 'lumi', 'ak7pfcand_ijet']),
on=['event', 'run', 'lumi', 'ak7pfcand_ijet'])
df_new['pfcand_centered_phi'] = map(lambda x: np.subtract(x, x[0]),
df_new['ak7pfcand_phi'])
df_new['pfcand_centered_eta'] = map(lambda x: np.subtract(x, x[0]),
df_new['ak7pfcand_eta'])
df_new['ak7pfcand_id'] = df_new.apply(
lambda row: map(np.abs, row['ak7pfcand_id']), axis=1)
temp_vec = df_new.apply(
lambda row: rotate_and_reflect(row['pfcand_centered_eta'], row[
'pfcand_centered_phi'], row['ak7pfcand_pt']),
axis=1)
temp_vec = map(list, temp_vec)
temp_vec = map(list, zip(*temp_vec))
print type(temp_vec)
print "temp_vec:", len(temp_vec)
print "df:", len(df_new['pfcand_centered_phi'])
df_new['pfcand_riz'] = pd.Series(temp_vec)[0]
df_new['pfcand_riy'] = pd.Series(temp_vec)[1]
df_new['jet_image'] = df_new.apply(lambda row: [
np.histogram2d(row['pfcand_riz'],
row['pfcand_riy'],
weights=row['ak7pfcand_eta'],
def plotJet(df_dict_jet,
df_dict_cand,
process='TT',
njets_to_plot=-1,
nx=30,
ny=30,
xbins=[],
ybins=[]):
jet_images = {}
# 4D tensor
# 1st dim is jet index
# 2nd dim is pt value (or rgb, etc.)
# 3rd dim is eta bin
# 4th dim is phi bin
list_x = []
list_y = []
list_w = []
njets = 0
if K.image_dim_ordering() == 'tf':
jet_images[process] = np.zeros((len(df_dict_jet[process]), nx, ny, 1))
else:
jet_images[process] = np.zeros((len(df_dict_jet[process]), 1, nx, ny))
njets_ = []
if njets_to_plot == -1:
njets_ = range(0, len(df_dict_jet[process]))
else:
njets_ = range(0, njets_to_plot)
for i in njets_:
njets += 1
# get the ith jet
df_dict_cand_i = {}
df_dict_cand_i[process] = df_dict_cand[process][
(df_dict_cand[process]['ak7pfcand_ijet'] == df_dict_jet[process]
['ak7pfcand_ijet'].iloc[i])
& (df_dict_cand[process]['event'] == df_dict_jet[process]
['event'].iloc[i])]
# relative eta
x = df_dict_cand_i[process]['ak7pfcand_eta'] - df_dict_cand_i[process][
'ak7pfcand_eta'].iloc[0]
# relative phi
y = df_dict_cand_i[process]['ak7pfcand_phi'] - df_dict_cand_i[process][
'ak7pfcand_phi'].iloc[0]
weights = df_dict_cand_i[process][
'ak7pfcand_pt'] # pt of candidate is the weight
x, y = rotate_and_reflect(x, y, weights)
list_x.append(x)
list_y.append(y)
list_w.append(weights)
hist, xedges, yedges = np.histogram2d(x,
y,
weights=weights,
bins=(xbins, ybins))
for ix in range(0, nx):
for iy in range(0, ny):
if K.image_dim_ordering() == 'tf':
jet_images[process][i, ix, iy, 0] = hist[ix, iy]
else:
jet_images[process][i, 0, ix, iy] = hist[ix, iy]
all_x = np.concatenate(list_x)
all_y = np.concatenate(list_y)
all_w = np.concatenate(list_w)
all_w = 1. * all_w / njets # to get average
plt.figure('W')
plt.hist2d(all_x,
all_y,
weights=all_w,
bins=(xbins, ybins),
norm=mpl.colors.LogNorm())
plt.colorbar()
plt.xlabel('eta')
plt.ylabel('phi')
plt.show()