vis.py 文件源码

def plot_actions(cue=0):
    mpl.rcParams['axes.labelsize'] = 'large'
    d_map = {3:1, 8:2, 14:3, 23:4}
    df = pd.read_pickle('data.pkl').reset_index()
    df = df.loc[df['cue'] == cue]
    g = sns.FacetGrid(df, col='subject',
                      col_wrap=6, size=1.5, ylim=(0, 5), aspect=1.5)

    g.map(plt.plot, 'action')
    g.set(xticks=[], yticks=[0,1,2,3], yticklabels=['3', '8', '14', '23'])
    g.set(ylim=(-0.5, 4))
    g.set_ylabels('choice')
    g.fig.tight_layout()
    g.fig.subplots_adjust(top=0.93)


    subjects = df['subject'].unique()
    for ax, subject in zip(g.axes, subjects):
        df_subject = df.loc[df['subject'] == subject]
        df_subject.reset_index(inplace=True)
        df_wins = df_subject.loc[df_subject['reward'] > 0]
        df_lose = df_subject.loc[df_subject['reward'] < 0]
        pos_win = df_wins.loc[df_wins['subject'] == subject].index
        pos_lose = df_lose.loc[df_lose['subject'] == subject].index
        ax.eventplot(pos_win, lineoffsets=3.5, linelength=0.75,
                     linewidths=0.4)
        ax.eventplot(pos_lose, lineoffsets=3.5, linelength=0.75,
                     color='r', linewidths=0.4)
    plt.tight_layout()
    plt.savefig('actions_0.pdf')
    plt.show()
    globals().update(locals())