Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Background shape of plot based on variable

I have a pandas Dataframe which has several columns with data and one column that encodes the state of the process of interest (non-consecutive integers).

Instead of plotting the state column as a line I would like to use it to add a shade to the background of the plot, e.g. as follows:

Example dataframe:

df = pd.DataFrame(
{
    "y": [x * x / 100 for x in range(10)],
    "state": [0 if x < 5 else 1 for x in range(10)],
})
    y   state
0   0.00    0
1   0.01    0
2   0.04    0
3   0.09    0
4   0.16    0
5   0.25    1
6   0.36    1
7   0.49    1
8   0.64    1
9   0.81    1

Desired plot (note that state is included as a line to get the point across, in the final picture I would of course omit it):

enter image description here

like image 596
oerpli Avatar asked Sep 16 '26 07:09

oerpli


1 Answers

You can find blocks where state is constant and then use axvspan to fill these blocks with different colors:

fig, ax = plt.subplots(1)
ax.set_ylim(0,1)

df[['y']].plot(ax=ax)

x = df.loc[df['state'] != df['state'].shift(1), 'state'].reset_index()
x['next_index'] = x['index'].shift(-1).fillna(df.index.max())

for i in x.index:
    c = 'blue' if (x.at[i, 'state']==1) else 'red'
    xa = x.at[i, 'index']
    xb = x.at[i, 'next_index']
    ax.axvspan(xa, xb, alpha=0.15, color=c)

Output:

output

like image 144
perl Avatar answered Sep 18 '26 21:09

perl



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!