CODEWITHSUNDEEP
pythonpandasmatplotlib
jpf5046
jpf5046

Reputation: 797

Matplotlib plot creation with loop

I have working code that produces the exact graph I'm looking for. However, It's hard coded (especially the legend) and I assume there is a way to do the same graph but in a loop.

Here is my code:

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

I know having five lines all calling plt.errorbar is not optimal, also hard coding the legend like I am, also is sub-optimal.

I have tried to start from scratch a little by using: 

for x in df:
    print(x)

...but that includes the $Time$ column, so I'm not sure how to always make time the x axis and then iterate over the rest of the columns for the y.

Upvotes: 0

Views: 59

Answers (2)

Massifox
Massifox

Reputation: 4487

If I understand your question correctly, this code is for you:

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

I extracted simply extracted the names of the drug_freautes from the raw data dictionary, and then iterated over the values ​​of the extracted features. Similarly, I added the list of features in creating the legend, so as not to have to manually write the names of the drugs.

Upvotes: 0

Quang Hoang
Quang Hoang

Reputation: 150735

You can do this:

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

plt.legend(loc=2)

Or if you don't have errors to pass to errorbar:

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

output:

enter image description here

Upvotes: 1

Related Questions