Color code a pandas plot based on column values

Question

I have a large pandas dataframe that I want to create a plot - here is a simplified example:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, 4)))


sample = pd.DataFrame({'X': [1,2,3,1,2,3,1,2,1,2,3],
                   'Y': [1,1,1,2,2,2,3,3,4,4,4]})

I want to create a color coded line plot, with color changing based on X column values (X values are always repeats the same numbers, but the length is not always the same) - the default plot is :

plt.plot(sample['X'], sample['Y'], linestyle = '-')

but I want to automate the process, so everytime X column restarts to have a new color - this is the result that I want to have

plt.plot(sample['X'][0:3], sample['Y'][0:3], linestyle = '-', color = next(colors))
plt.plot(sample['X'][3:6], sample['Y'][3:6], linestyle = '-', color = next(colors))
plt.plot(sample['X'][6:8], sample['Y'][6:8], linestyle = '-', color = next(colors))
plt.plot(sample['X'][8:], sample['Y'][8:], linestyle = '-', color = next(colors))

any suggestion on how to achieve this?

Pablo · Accepted Answer · 2017-06-23 15:16:12Z

1

I would build on what you already proposed:

slices = [slice(0, 3), slice(3, 6), slice(6, 8), slice(8, None)]
for _slice, color in zip(slices, colors):
    plt.plot(sample['X'][_slice], sample['Y'][_slice], c=color)
plt.show()

alternatively, if you add an extra column to your data:

sample2 = pd.DataFrame({'X': [1,2,3,1,2,3,1,2,1,2,3],
                        'Y': [1,1,1,2,2,2,3,3,4,4,4],
                        'G': [0,0,0,1,1,1,2,2,3,3,3]})

colors = cm.rainbow(np.linspace(0, 1, 4))
for name, group in sample2.groupby(['G']):
    plt.plot(group['X'], group['Y'], c=colors[name])
plt.show()

answered Jun 23, 2017 at 15:16

Pablo

1,00313 silver badges26 bronze badges

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

1 Comment

Dimitris Over a year ago

thanks - your alternate solution worked for me (I have time stamps so I can make a unique identifier column for each series of numbers)

Collectives™ on Stack Overflow

Color code a pandas plot based on column values

1 Answer 1

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related