Draw several polygons on top of matplotlib plot

Question

I want to plot some data where there was a stimulus given every 3000s for 1500s, 5 times. I would like to plot my signal and in the background display when the stimulus was given. For instance like this, here colored in red:

At the moment my code looks like this:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


df=pd.read_csv('file.csv')

fig, ax = plt.subplots(figsize=(15,5))
df.plot(ax=ax)

#draw verticle lines (I can do this in one line)
ax.vlines(np.arange(0, len(df), 30*50), 0, 1, transform=ax.get_xaxis_transform(), colors='r', alpha=0.5)

#draw polygons (I know only how to write one line per polygon)
ax.axvspan(0, 1500, 0, 1, color='red', alpha=0.4)
ax.axvspan(3000, 4500, 0, 1, color='red', alpha=0.4)
ax.axvspan(6000, 7500, 0, 1, color='red', alpha=0.4)
ax.axvspan(9000, 10500, 0, 1, color='red', alpha=0.4)
ax.axvspan(12000, 13500, 0, 1, color='red', alpha=0.4)

ax.set_xlim([0, len(activity)])
ax.set_ylabel('Normalised Intensity (a.u)')
ax.set_xlabel('Time (s)')
plt.show()

Is there a smarter way of drawing the polygons? Ideally similar to ax.vlines(). I could write a for loop for each axvspan() but I would prefer if there is some function that works with np.arange() similar to ax.vlines().

Thank you very much!

Answer 1

It's pretty straightforward.

def plot_stimuli(ax, start, number:int, step, **kwargs):
    for _ in range(number):
        ax.axvspan(start, start+step, 0, 1, **kwargs)
        start += 2*step

plot_stimuli(ax, 0, len(df), 30*50, color='red', alpha=0.4)

It's simple because the width of the polygons is the same as the gap between them. If you need to separate width and gap, we can do that with a few minor changes.