Создание сюжета Matplotlib с циклом

Question

У меня есть рабочий код, который выдает точный график, который я ищу.Тем не менее, он жестко запрограммирован (особенно легенда), и я предполагаю, что есть способ сделать тот же график, но в цикле.

Вот мой код:

import pandas as pd
import matplotlib.pyplot as plt
raw_data = {'Time':       [1,    2,    3,    4,    5,    6 ,  7,    8,   9,      10],
        'Drug 1':     [23.4, 32.5, 45.6, 46.1, 47.8, 50.1, 51.2, 53.2, 54.5, 55.0],
        'Drug 2':    [10.4,  12.5, 13.7, 13.8, 14.0, 15.6, 17.7, 23.2, 20.4, 19.5], 
        'Drug 3':    [0.4,   1.5,  2.6,  3.7,  4.8,  5.9,  6.2,  8.7,  12.8, 13],
        'Drug 4':    [45,    47,   48,   50,   51,     52,   55,  60,   61,   67],
        'Drug 5':    [17,    21,    20, 20,    20,     24,    26,  28, 29,   30]}


df = pd.DataFrame(raw_data)

plt.errorbar(x = df['Time'], y = df['Drug 1'], yerr=None, linestyle = "--")
plt.errorbar(x = df['Time'], y = df['Drug 2'], yerr=None, linestyle = "--")
plt.errorbar(x = df['Time'], y = df['Drug 3'], yerr=None, linestyle = "--")
plt.errorbar(x = df['Time'], y = df['Drug 4'], yerr=None, linestyle = "--")
plt.errorbar(x = df['Time'], y = df['Drug 5'], yerr=None, linestyle = "--")

plt.legend(['Drug 1', 'Drug 2', 'Drug 3', 'Drug 4', 'Drug 5'], loc = 2)

plt.ylabel('Tumor Size')
plt.xlabel('Time in Years')

plt.title('Effect on Treatment')
plt.grid()

plt.show()

Я знаю, что пять строк, все вызовы plt.errorbar не оптимальны, также жесткое кодирование легенды, как я, также неоптимально.

Я попытался начать с нуля, используя:

for x in df:
    print(x)

... но это включает столбец $ Time $, поэтому я не уверен, как всегда делать времяось х, а затем итерации по остальным столбцам для у.

Answer 1

Вы можете сделать это:

for col in df.columns[1:]:
    plt.errorbar(x=df['Time'], y=df[col], linestyle='--',label=col)

plt.legend(loc=2)

Или, если у вас нет ошибок для передачи на errorbar:

df.plot(x='Time', linestyle='--')

, вывод:

Answer 2

Если я правильно понимаю ваш вопрос, этот код для вас:

drug_features = list(raw_data.keys())[1:]

for drug in drug_features:
    plt.errorbar(x = df['Time'], y = df[drug], yerr=None, linestyle = "--")

plt.legend(drug_features, loc = 2)

plt.ylabel('Tumor Size')
plt.xlabel('Time in Years')

plt.title('Effect on Treatment')
plt.grid()
plt.show()

Я просто извлек имена drug_freautes из словаря необработанных данных, а затем перебрал значенияиз извлеченных особенностей.Точно так же я добавил список возможностей при создании легенды, чтобы не пришлось вручную писать названия лекарств.