ColorPlot.py

import pandas as pd
import Helpers
import matplotlib.pyplot as plt
import seaborn as sns


def ReduceListSeries(a: pd.Series, *, diff=0, rangeOnly=False):
    """Given a pd.Series of integers, this function returns a list of tuples
    where each pair consists of the indexes of the boundaries (1st inclusive, 2nd exclusive [, 1st  value])
    of sublists where the n-th element consists of the n-1 th element + diff
    diff = 1:
    [0,1,2,3,5,6,8] => [(0, 4), (4, 6), (6, 7)]
    """
    first = last = a.index[0]
    prev = a[first]
    for last, value in a[1:].items():
        if value != prev + diff:
            yield (first, last) if rangeOnly else (first, last, a[first])
            first = last
        prev = value
    yield (first, last) if rangeOnly else (first, last, a[first])


 

def ColorPlot(s : pd.Series, pltAxis, colorFunc):
    x = ReduceListSeries(s, diff=0, rangeOnly=False)
    for left,right,state in x:
        color = colorFunc(state)
        pltAxis.axvspan(xmin=left, xmax=right, alpha=0.35, color=color)

        
df = pd.DataFrame(
    {
        "y": [x * x / 100 for x in range(50)],
        "state": [x // 2 for x in range(50)]
    }
)
cmap = plt.get_cmap("viridis")
cfunc = lambda x: cmap((x % 6) / 5)

fig, ax = plt.subplots(1)
df[["y"]].plot(ax=ax)
ColorPlot(df['state'],ax,cfunc)
plt.show()