paramnormal_blog_plots.py

import numpy
from scipy import stats
from matplotlib import pyplot
import seaborn


def plot_norm(mu, sigma):
    example_norm = stats.norm(mu, sigma)
    fig, ax = pyplot.subplots()

    quantiles = numpy.linspace(-5, 15, 250)
    ax.plot(quantiles, example_norm.pdf(quantiles))
    ax.set_xlabel('$X$')
    ax.set_ylabel('Normal PDF')

    sigma_y = sigma * 0.047
    sq = numpy.linspace(mu - 0.5*sigma, mu + 0.5*sigma, 50)
    ax.fill_between(sq, example_norm.pdf(sq))
    arrow = ax.annotate('', xy=(mu - 0.5*sigma, sigma_y), xycoords='data',
                        xytext=(mu + 0.5*sigma, sigma_y), textcoords='data',
                        arrowprops=dict(arrowstyle="<->", ec="w", lw=2))
    text = ax.annotate('$\sigma$', xy=(mu, sigma_y),
                       fontsize=20, ha='center', va='top',
                       color='w', weight='heavy')

    arrow = ax.annotate('', xy=(mu, 0), xycoords='data',
                        xytext=(mu, 0.25*sigma_y), textcoords='data',
                        arrowprops=dict(arrowstyle="->", ec="w", lw=2))
    text = ax.annotate('$\mu$', xy=(mu, 0.25*sigma_y),
                       fontsize=20, ha='center', va='bottom',
                       color='w', weight='heavy')
    seaborn.despine(ax=ax)
    return fig


def format_ax(ax, xlabel, ylabel, legloc):
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.legend(loc=legloc)
    seaborn.despine(ax=ax, offset=5, trim=True)