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

Posted 2023-02-16

【中文标题】在派生自熊猫数据框的图中自定义图例

【英文标题】：customizing the legend in a plot derived from a pandas dataframe

【发布时间】：2020-05-07 02:51:08

【问题描述】：

我正在使用“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 绘图示例，但无法找到解决这个特定问题的方法。很明显，这些图是非常可定制的，所以我怀疑有办法做到这一点。任何帮助将不胜感激。

【参考方案1】：

这篇文章和你问的有点相似，我想你应该在这个页面上查看答案：Show only certain items in legend Python Matplotlib。希望这会有所帮助！

【讨论】：

没有一个答案是对那里提出的问题的真正答案。在这里可行的是重命名 df1 中的所有策略，除了 4 个以“_”开头的所需策略。但这仍然无法管理最后一步中所需的出现顺序。

【参考方案2】：

这是一种方法。我没有完整的数据框，所以只测试问题中显示的数据框。

问题的熊猫部分可以通过将最后一步分配给一个变量来解决，然后查询该步骤的策略，然后获得最高计数。

为了找到句柄，我们向 matplotlib 询问它生成的所有句柄和标签。然后我们在标签列表中搜索每个策略，获取其索引以获得对应的句柄。

请注意，“count”是一个令人讨厌的列名称。它也是一个 pandas 函数的名称，它阻止了它在点符号中的使用。

import pandas as pd
from matplotlib import pyplot as plt

df = pd.DataFrame(columns=['step', 'strategy', 'count', 'score'],
                  data=[[0, 'CC', 34, 2.08],
                        [0, 'DD', 1143, 2.18],
                        [0, 'CD', 1261, 2.24],
                        [0, 'DC', 62, 2.07],
                        [1, 'CC', 6, 1.88],
                        [106, 'DDCC', 56, 0.99],
                        [106, 'DD', 765, 1.00],
                        [106, 'DC', 1665, 1.31],
                        [106, 'DCDC', 23, 1.60],
                        [106, 'DDDD', 47, 0.98]])
last_step = df.step.max()
strategies_last_step = df.strategy[df['count'][df.step == last_step].nlargest(4).index]

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

handles, labels = plt.gca().get_legend_handles_labels()
selected_handles = [handles[labels.index(strategy)] for strategy in strategies_last_step]

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

plt.show()

【参考方案3】：

谢谢你，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 种策略：

【讨论】：

请注意，在 Python 中，强烈建议将 for i in range(len(top_strategies)): ix = labels.index(top_strategies[i]) 重写为 for strategy in top_strategies: ix = labels.index(strategy)