Mike
Mike

Reputation: 155

How to change color of graph and draw vertical shaded rectangle in matplotlib pyplot?

I have this dataframe:

            A  a_red  green_bar
date                           
2020-01-01  5  False      False
2020-01-02  6   True      False
2020-01-03  7   True      False
2020-01-04  4  False       True
2020-01-05  5  False       True
2020-01-06  5  False      False

From column A, I draw a line into the chart: Chart

I would like to add two features:

  1. Only on days when the value in column a_red is True, I want to change the color of the blue line (from column A) to red. On all other days, it should remain blue.
  2. Only on days when the value in column green_bar is True, I want to draw a translucent green vertical rectangle that spans the width of the day and the full height of the chart. Essentially, marking this day as special. Multiple consecutive days should appear as one continuous green block.

Here is my code:

import matplotlib.pyplot as plt
import pandas as pd

# create sample dataframe
import io
data = io.StringIO("""
date|A|a_red|green_bar
2020-01-01|5|False|False
2020-01-02|6|True|False
2020-01-03|7|True|False
2020-01-04|4|False|True
2020-01-05|5|False|True
2020-01-06|5|False|False""")
df = pd.read_csv(data, sep='|', parse_dates=['date'])
df = df.set_index('date')

# draw plot
with plt.style.context('seaborn'):
    plt.plot(df['A'], color='blue')
    plt.show()

Thank you so much for your help, I really appreciate it!

Upvotes: 0

Views: 614

Answers (1)

JohanC
JohanC

Reputation: 80329

The desired plot can be created by looping through the dataframe row by row. Note that the dates are points on the plot, while the lines go between two successive dates. So, there won't be a line nor a bar for the last date (to mitigate, you could append a dummy last date).

import matplotlib.pyplot as plt
import pandas as pd
import io

# create sample dataframe
data = io.StringIO("""
date|A|a_red|green_bar
2020-01-01|5|False|False
2020-01-02|6|True|False
2020-01-03|7|True|False
2020-01-04|4|False|True
2020-01-05|5|False|True
2020-01-06|5|False|False""")
df = pd.read_csv(data, sep='|', parse_dates=['date'])
df = df.set_index('date')

# draw plot
with plt.style.context('seaborn'):
    for i, (a_red, green_bar) in enumerate(zip(df['a_red'][:-1], df['green_bar'][:-1])):
        plt.plot(df['A'][i:i + 2], color='red' if a_red else 'blue')
        if green_bar:
            plt.axvspan(df.index[i], df.index[i + 1], color='green', alpha=0.2)
plt.show()

example plot

Upvotes: 1

Related Questions