Python/Matplotlib - Change subplot background color based on a binary array of strings-CodePudding

The code:

labels_0 = []
labels_1 = []

for i in range(len(ensemble_ss)):
    if ensemble_ss.iloc[i, ensemble_ss.columns.get_loc('ens_state')] == 'Long':
        labels_0.append(ensemble_ss.iloc[i, ensemble_ss.columns.get_loc('Close')])
        labels_1.append(float('nan'))
    else:
        labels_0.append(float('nan'))
        labels_1.append(ensemble_ss.iloc[i, ensemble_ss.columns.get_loc('Close')])


fig, axs = plt.subplots(4, sharex = True, figsize=(12,6),dpi=500, 
                        gridspec_kw={'height_ratios': [3, 0.33, 0.33, 0.33]})

fig.patch.set_facecolor('silver')

axs[0].plot(labels_0, color="black")
axs[0].plot(labels_1, color="red")
axs[0].plot(ensemble_ss['Cum Return'], color="dimgrey",linewidth = 1.5, linestyle = 'solid')
#axs[0].set_yscale('log')
axs[0].margins(x=0)
axs[0].grid(which='both', axis='both', ls='--')

axs[1].plot(ensemble_ss['ens_state'])

axs[2].plot(ensemble_ss['kf_state'])

axs[3].plot(ensemble_ss['hmm_state'])

The DataFrame which I'm plotting from (ensemble_ss) takes the following form:

    Date          Symbol     Interval   Open    Close     kf_state   hmm_state   ens_state  Cum Return  
0   2013-01-03    SPY         1D       4.983    4.981     Long       Flat        Long       4.981
1   2013-01-04    SPY         1D       4.983    4.986     Long       Long        Flat       4.986
2   2013-01-07    SPY         1D       4.982    4.983     Flat       Long        Long       4.983   
3   2013-01-08    SPY         1D       4.985    4.980     Flat       Long        Flat       4.980

The code above plots this:

Focusing on the subplots. The blue line flicks back and forth along the y-axis between binary states of 'Long' or 'Flat'.

however what I want to display would look (somewhat) like this:

Excuse the rough paint job, but the idea is that I want no line, rather the background of the subplots change color (lets say between green and red) along the timeseries with respect to their state being 'Long' or 'Flat'.

Thanks in advance for any help. I have had a good look over the web but cant seem to find a solution.

CodePudding user response：

You can use