Is there a way to loop through pandas dataframe names to plot a figure?

Viewed 569

I have the following code where I plot different dataframes:

plt.scatter(df1['phys_red'],df1['lost_frac_damage'], s = 10, c='DarkBlue')
plt.scatter(df2.phys_red,df2.lost_frac_damage, s = 10, c='Red')
plt.scatter(df3.phys_red,df3.lost_frac_damage, s = 10, c='Green')
plt.scatter(df4.phys_red,df4.lost_frac_damage, s = 10, c='Black')
plt.scatter(df5.phys_red,df5.lost_frac_damage, s = 10, c='LightBlue')

this continues for another 10 dataframes. So in essence, one for loop, which can loop through the different df's would do the trick and reduce the code to one line, as the colors c can be looped through too by:

colors = ['DarkBlue','Red','Green','Black','LightBlue']

and using colors[i]. I was thinking of something like

for i in range(1,6,1):
    df+str(i) = ...

but I could not yet find out how to exactly do it. Is there a way to do this? Thanks.

3 Answers
dfs = [df1, df2, df3, df4, df5]
colours = ['red', 'green', 'blue', 'cyan', 'black']

for df, colour in zip(dfs,colours):
    plt.scatter(df.phys_red, df.lost_frac_damage, s=10, c=colour)

You have two options:

  1. You can store the dataframe in a container (like a dictionary or an array) and iterate through the with a for loop.
  2. You can use the eval function to convert a sting into code.

Method 1

You create an array to hold your frames and colors. This could look like:

dataframes = [df1, df2, df3, ...]
colors = ['DarkBlue','Red','Green','Black','LightBlue']

You can then iterate through the pairs using the zip() function. This will map the nth element in one list to the nth element in another and return a tuple.

In practice, this would look like:

for df, color in zip(dataframes, colors):
    plt.scatter(df.phys_red,df.lost_frac_damage, s = 10, c=color)

Method 2

I don't recommend using eval because it will result in a performance toll. It's also just generally considered a bad practice but I'm not sure why that is.

Essentially, you would just construct the string that represents the current var that you are calling but wrap it with eval so that it is executed as code.

This could look something like this:

for idx, color in enumerate(colors):
    df_red_str = "df%d.phys_red" % (idx+1)
    df_frac_str = "df%d.lost_frac_damage" % (idx+1)
    plt.scatter(eval(df_red_str).eval(df_frac_str), s = 10, c='LightBlue')

This method has the advantage that you don't have to make any existing changes. However, it might not be considered a best practice.

Store your DataFrames in a dict where the keys are some useful piece of information that labels them. Ideally you would create the dict as the DataFrames are created so you don't have all of those variables floating around anyway, but for this example I'll just throw the already existing DataFrames in a dict container.

Then you can loop over the dict and use another dictionary of dictionaries to specify the plot arguments that depend on the labels.

Sample Data

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

df1 = pd.DataFrame({'phys_red': np.random.normal(0,10,100),
                   'lost_frac_damage': np.random.normal(0,10,100)})
df2 = pd.DataFrame({'phys_red': np.random.normal(0,10,100),
                   'lost_frac_damage': np.random.normal(0,10,100)})
df3 = pd.DataFrame({'phys_red': np.random.normal(0,10,100),
                   'lost_frac_damage': np.random.normal(0,10,100)})

Code

d = {'key1': df1, 'key2': df2, 'key3': df3}
kwargs = {'key1': {'s': 10, 'color': 'DarkBlue'}, 
          'key2': {'s': 15, 'color': 'red'},
          'key3': {'s': 4, 'color': 'Green'}}

for label, df in d.items():
    plt.scatter(df['phys_red'], df['lost_frac_damage'], **kwargs[label])
plt.show()

enter image description here

Related