настройка легенды на графике, полученном из pandas кадра данных

Question

Я работаю над реализацией python агентной модели, использующей инфраструктуру 'mesa' (доступна в Github). В этой модели каждый «агент» в сетке играет в игру «Дилемма заключенного» против своих соседей. У каждого агента есть стратегия, которая определяет его ход против других ходов. Стратегии с более высокими выплатами заменяют стратегии с более низкими выплатами. Кроме того, стратегии развиваются в результате мутаций, поэтому в процессе работы модели появляются новые и более длинные стратегии. Приложение создает 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 создает довольно хороший график этих данных, вызывая эту простую функцию plot:

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

Я получаю этот сюжет:

Неплохо, но вот что я не могу понять. Легенда автомата c быстро становится слишком длинной, а низкочастотные стратегии малоинтересны, поэтому я хочу, чтобы легенда (1) включала только 4 верхние стратегии, перечисленные в вышеупомянутой легенде, и (2) перечисляла эти стратегии в порядок, в котором они появляются на последнем шаге модели, исходя из их количества. Глядя на стратегии в шаге 106 в df, например, я хочу, чтобы легенда показывала 4 верхние стратегии в порядке D C, DD, DD CC и D DDD, но не включала DCD C (или любые другие стратегии с меньшим числом, которые могут быть активны).

Я искал тонны pandas и примеров построения графиков matplotlib, но не смог найти решение этой проблемы c проблема. Ясно, что эти графики чрезвычайно настраиваемы, поэтому я подозреваю, что есть способ сделать это. Любая помощь будет принята с благодарностью.

Answer 1

Вот подход. У меня нет полного фрейма данных, поэтому тест проводится только с теми, которые показаны в вопросе.

Часть вопроса pandas может быть решена путем назначения последнего шага переменной, а затем запроса для стратегий этого шага, а затем получить наибольшее число.

Чтобы найти дескрипторы, мы запрашиваем у 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()

Answer 2

Спасибо, Йохан C, вы действительно помогли мне увидеть, что происходит под капотом с этой проблемой. (Кроме того, хорошая отметка о 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+ шагов. Легенда в правильном порядке и ограничена четырьмя основными стратегиями:

Answer 3

Этот пост чем-то похож на то, что вы просили, я думаю, вам стоит проверить ответ на этой странице: Показать только определенные элементы в легенде Python Matplotlib . Надеюсь, это поможет!