Can this plt.fill_between() approach be modified in a loop?

Question

I am trying to shade an (x,y) region in a specific region of the plot. As a simplified example, consider a normal distribution with confidence intervals. I would like to shade the confidence intervals such that the region within one standard deviation (or one sigma) is darkest, the region within two standard deviations (or 2 sigmas) is a little lighter, etc. I have a way of doing this, but I am trying to make my script more flexible. Code is below.

## imports
import numpy as np
import matplotlib.pyplot as plt
from math import pi

## y = f(x)
def get_f(x, mu, sigma):
    """ Normal Distribution Probability Density Function """
    norm_constant = (sigma* (2*pi)**(1/2))
    return [norm_constant * np.exp((-1) * (x[idx] - mu)**2 / (2* sigma**2)) for idx in range(len(x))]

x = np.linspace(0, 100, 5000)

Now that we have x and a function f(x), we can make a plot. I left in the part of the code that works, and commented out my attempt at a solution. I would prefer my solution method if it worked because it is more convenient to shade based on the number of desired intervals and the code isn't as repetitive.

## generate plot
def get_plot(x, num_intervals=None, line_color='g', shade_color='b', mu=48, sigma=7):
    """ Returns (x,y) plot; confidence intervals shading is optional """
    y = get_f(x, mu, sigma)
    plt.plot(x, y, line_color)
    if num_intervals is not None:

        ## THIS CODE SEGMENT BELOW WORKS BUT I WOULD LIKE TO MAKE IT BETTER

        plt.fill_between(x, y, where=(mu - sigma <= x), alpha=0.18, color=shade_color)
        plt.fill_between(x, y, where=(x <= mu + sigma), alpha=0.18, color=shade_color)
        plt.fill_between(x, y, where=(mu - 2*sigma <= x), alpha=0.11, color=shade_color)
        plt.fill_between(x, y, where=(x <= mu + 2*sigma), alpha=0.11, color=shade_color)
        plt.fill_between(x, y, where=(mu - 3*sigma <= x), alpha=0.02, color=shade_color)
        plt.fill_between(x, y, where=(x <= mu + 3*sigma), alpha=0.02, color=shade_color)

        ## THIS CODE SEGMENT BELOW DOES NOT WORK AS I WOULD LIKE
        ## IT WILL SHADE THE REGIONS IN THE WRONG SHADE/DARKNESS 

        ## choose shading level via dictionary
        # alpha_keys = [idx+1 for idx in range(num_intervals)]
        # alpha_vals = [0.18, 0.11, 0.02]
        # alpha_dict = dict(zip(alpha_keys, alpha_vals))
        # for idx in range(num_intervals):
            # print("\nidx & stdev = %d & %d, \nmu - (stdev * sigma) = %.2f, \nmu + (stdev * sigma) = %.2f, alpha = %.2f" %(idx, stdev, mu - stdev*sigma, mu + stdev*sigma, alpha_dict[stdev]), "\n")
            # stdev = idx + 1 ## number of standard deviations away from mu
            # plt.fill_between(x, y, where=(mu - stdev * sigma <= x), alpha=alpha_dict[stdev], color=shade_color)
            # plt.fill_between(x, y, where=(x >= mu + stdev * sigma), alpha=alpha_dict[stdev], color=shade_color)


    plt.show()

Running the correct code produces this plot. My attempt at a more convenient solution produces this plot and produces the output below (via the print statement), though I can't find the source of my mistake.

idx & stdev = 0 & 1, 
mu - (stdev * sigma) = 41.00, 
mu + (stdev * sigma) = 55.00, alpha = 0.18 


idx & stdev = 1 & 2, 
mu - (stdev * sigma) = 34.00, 
mu + (stdev * sigma) = 62.00, alpha = 0.11 


idx & stdev = 2 & 3, 
mu - (stdev * sigma) = 27.00, 
mu + (stdev * sigma) = 69.00, alpha = 0.02

Is my approach for a more convenient solution adaptable?

Answer 1

Here I offer my version of the normal distribution plot that is more compact than yours. I use the Normal distribution function from Scipy package rather than reinvent the wheel.

from scipy.stats import norm   # import normal dist.
import matplotlib.pyplot as plt
import numpy as np

# mean and standard deviation
mu,sigma = 48,7  

# normal_dist(mu,sigma)
anorm = norm(loc=mu, scale=sigma)
factors = [1,2,3]            # multiple of sigma
alphas = [0.18, 0.11, 0.08]  # level of alpha

fig, ax = plt.subplots(1, 1)
fig.set_size_inches(10,8)

# plot full normal curve
segs = 100
x = np.linspace(anorm.ppf(0.0005), anorm.ppf(0.9995), segs)
ax.plot(x, anorm.pdf(x), 'b-', lw=0.5, alpha=0.6)

# plot color-filled portions
for fac, alp in zip(factors, alphas):
    # print(mu-fac*sigma, mu+fac*sigma, alp)
    lo = mu-fac*sigma
    hi = mu+fac*sigma
    xs = np.linspace(lo, hi, fac*segs/4)  # prep array of x's
    plt.fill_between(xs, anorm.pdf(xs), y2=0, where= xs >= lo , \
                     interpolate=False, \
                     color='blue', alpha=alp)

plt.ylim(0, 0.06)

plt.show()

The resulting plot: