Matplotlib plot creation with loop

Question

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.

just curious, why are you plotting error bar without errors? — Quang Hoang
– Quang Hoang, Commented Sep 25, 2019 at 14:01

Quang Hoang · Accepted Answer · 2019-09-25 14:03:47Z

1

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:

answered Sep 25, 2019 at 14:03

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

jpf5046 Over a year ago

-- Okay that is smart. The line df.columns[1:] is saying to ignore the column Time when it goes through the loop.

Massifox · Accepted Answer · 2019-09-25 14:04:03Z

0

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.

answered Sep 25, 2019 at 14:04

Massifox

4,5071 gold badge13 silver badges33 bronze badges

1 Comment

jpf5046 Over a year ago

Thank you for this. I chose to go with @Quang's answer because uses the pandas df to iterate over.

Collectives™ on Stack Overflow

Matplotlib plot creation with loop

2 Answers 2

1 Comment

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related