Ejemplo n.º 1
0
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()