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