python - 在派生自熊猫数据框的图中自定义图例

Question

我正在使用“mesa”框架（在 Github 中可用）来实现基于代理的模型的 python 实现。在该模型中，网格上的每个“代理”都与其邻居进行囚徒困境游戏。每个代理都有一个策略来确定其移动与其他移动。收益较高的策略取代收益较低的策略。此外，策略会通过突变而演变，因此随着模型的运行，会出现新的和更长的策略。该应用程序会生成一个 pandas 数据框，该数据框在每一步后都会更新。例如，经过 106 步后，df 可能如下所示：

    step strategy count  score
0      0       CC    34   2.08
1      0       DD  1143   2.18
2      0       CD  1261   2.24
3      0       DC    62   2.07
4      1       CC     6   1.88
..   ...      ...   ...    ...
485  106     DDCC    56   0.99
486  106       DD   765   1.00
487  106       DC  1665   1.31
488  106     DCDC    23   1.60
489  106     DDDD    47   0.98

Pandas/matplotlib 为这些数据创建了一个非常好的绘图，调用这个简单的绘图函数：

def plot_counts(df):
    df1 = df.set_index('step')
    df1.groupby('strategy')['count'].plot()
    plt.ylabel('count')
    plt.xlabel('step')
    plt.title('Count of all strategies by step')
    plt.legend(loc='best')
    plt.show()

我得到这个情节：

不错，但这是我无法弄清楚的。自动图例很快变得太长，低频策略没什么兴趣，所以我希望图例 (1) 仅包括上述图例中列出的前 4 个策略，(2) 按它们的顺序列出这些策略根据他们的计数出现在模型的最后一步。例如，查看 df 中步骤 106 中的策略，我希望图例按 DC、DD、DDCC 和 DDDD 的顺序显示前 4 个策略，但不包括 DCDC（或任何其他可能是积极的）。

我已经搜索了大量的 pandas 和 matplotlib 绘图示例，但无法找到解决这个特定问题的方法。很明显，这些图是非常可定制的，所以我怀疑有办法做到这一点。任何帮助将不胜感激。

标签： pythonpandasmatplotliblegendmesa

Answer 1

谢谢你，JohanC，你真的帮助我了解了这个问题的幕后情况。（另外，关于count作为 col 名称的好点。我将其更改为ncount。）

我找到了你的陈述：

strategies_last_step = df.strategy[df['count'][df.step == last_step].nlargest(4).index]

对我不起作用（nlargest 对 dtypes 感到困惑）所以我制定了一种稍微不同的方法。我通过这种方式获得了正确排序的策略名称列表：

def plot_counts(df):
    # to customize plot legend, first get the last step in the df
    last_step = df.step.max()
    # next, make new df_last_step, reverse sorted by 'count' & limited to 4 items  
    df_last_step = df[df['step'] == last_step].sort_values(by='ncount', ascending=False)[0:4]
    # put selected and reordered strategies in a list
    top_strategies = list(df_last_step.strategy)

然后，在索引和分组我的原始 df 并添加我的其他绘图参数之后......

    dfi = df.set_index('step')
    dfi.groupby('strategy')['ncount'].plot()
    plt.ylabel('ncount')
    plt.xlabel('step')
    plt.title('Count of all strategies by step')

我能够从默认句柄列表中挑选出正确的句柄并以这种方式重新排序：

    handles, labels = plt.gca().get_legend_handles_labels()
    # get handles for top_strategies, in order, and replace default handles
    selected_handles = []
    for i in range(len(top_strategies)):
        # get the index of the labels object that matches this strategy
        ix = labels.index(top_strategies[i])
        # get matching handle w the same index, append it to a new handles list in right order
        selected_handles.append(handles[ix])

然后用新的绘图selected_handles：

plt.legend(handles=selected_handles, loc='best')
plt.show()

结果完全符合预期。这是 300 多步后的图。图例的顺序正确，仅限于前 4 种策略：