Python Pandas Matplotlib Plot Colored by type value defined in single column

Question

I have data of the following format:

import pandas as ps
table={'time':[1,2,3,4,5,1,2,3,4,5,1,2,3,4,5],\
    'data':[1,1,2,2,2,1,2,3,4,5,1,2,2,2,3],\
    'type':['a','a','a','a','a','b','b','b','b','b','c','c','c','c','c']}
df=ps.DataFrame(table,columns=['time','data','type']

I would like to plot data as a function of time connected as a line, but I would like each line to be a separate color for unique types. In this example, the result would be three lines: a data(time) line for each type a, b, and, c. Any guidance is appreciated.

I have been unable to produce a line with this data--pandas.scatter will produce a plot, while pandas.plot will not. I have been messing with loops to produce a plot for each type, but I have not found a straight forward way to do this. My data typically has an unknown number of unique 'type's. Does pandas and/or matpltlib have a way to create this type of plot?

Answer 1

Pandas plotting capabilities will allow you to do this if everything is indexed properly. However, sometimes it's easier to just use matplotlib directly:

import pandas as pd
import matplotlib.pyplot as plt

table={'time':[1,2,3,4,5,1,2,3,4,5,1,2,3,4,5],
       'data':[1,1,2,2,2,1,2,3,4,5,1,2,2,2,3],
       'type':['a','a','a','a','a','b','b','b','b','b','c','c','c','c','c']}
df=pd.DataFrame(table, columns=['time','data','type'])

groups = df.groupby('type')

fig, ax = plt.subplots()
for name, group in groups:
    ax.plot(group['time'], group['data'], label=name)
ax.legend(loc='best')

plt.show()

If you'd prefer to use the pandas plotting wrapper, you'll need to override the legend labels:

import pandas as pd
import matplotlib.pyplot as plt

table={'time':[1,2,3,4,5,1,2,3,4,5,1,2,3,4,5],
       'data':[1,1,2,2,2,1,2,3,4,5,1,2,2,2,3],
       'type':['a','a','a','a','a','b','b','b','b','b','c','c','c','c','c']}
df=pd.DataFrame(table, columns=['time','data','type'])

df.index = df['time']
groups = df[['data', 'type']].groupby('type')

fig, ax = plt.subplots()
groups.plot(ax=ax, legend=False)
names = [item[0] for item in groups]
ax.legend(ax.lines, names, loc='best')

plt.show()