I am trying to plot lines from a GroupbyDataFrame with properties dependant on other column value. So far the task is being solved using conditions and loops on every pair of y-axis values.
Reproducible example:
import pandas as pd
from itertools import cycle
import seaborn as sns
# Setup DataFrame and Groupby object
df = pd.DataFrame({'week': range(1,11)*2,
'object':['a']*10 + ['b']*10,
'param': [100, 95, 90, 100, 110, 90, 90, 110, 110, 120,
110, 110, 110, 100, 100, 110, 110, 100, 100, 110],
'event': [0, 0, 1, 0, 0, 1, 1, 0, 0, 0,
0, 0, 0, 1, 1, 0, 0, 1, 1, 0]}, index=range(1,21))
print df.head()
groups = df.groupby('object')
# Setup Figure, ax and line properties
color = cycle(sns.color_palette('Set1'))
fig, ax = plt.subplots(figsize=(10,6))
base_line={'ls':'--', 'lw':3}
event_line={'ls':'-', 'lw':2}
# Process each group separately:
for name, g in groups:
# Assign color for the object:
label_set = False
props = {'c':next(color)}
# Set X axis values:
x = g['week']
# Process each pair of points to select desired properties:
for i in range(g.shape[0]):
# Make pair of Y values
y = g.iloc[i:i+2].loc[:,'param'].reindex(g.index).values
# Select properties of the pair of dots based on value of 'event' column:
if g.iloc[i:i+2].loc[:,'event'].sum() >= 1:
props.update(event_line)
# Set properties for the legend if it has not been set before:
if not label_set:
props.update({'label':name})
label_set = True
else:
# Remove label property if legend is already set:
if props.has_key('label'):
props.pop('label')
else:
props.update(base_line)
ax.plot(x, y, **props)
_ = ax.set_ylim(bottom=round(ax.get_yaxis().get_data_interval().min() * .9, -1),
top=round(ax.get_yaxis().get_data_interval().max() * 1.1, -1))
plt.legend(bbox_to_anchor=(1.15, 0.5), fontsize='x-large')
for l in ax.get_xticklabels() + ax.get_yticklabels():
l.set_fontsize('x-large')
plt.show()
This produces desired output:
But I am hoping there is a matplotlib
way to achieve the same result,
which will also allow to easily manipulate legend for a line as whole (currently it is quite cumbersome).