low order Gram-Chandler / Edgeworth expansions, coefficients hardcoded

gistfile1.py

import numpy as np
from numpy.polynomial.hermite_e import HermiteE

from scipy.stats import norm, chi2
from scipy.misc import factorial
import matplotlib.pyplot as plt


def gram_chandler_series(cum):
    """Construct the 4th order Gram-Chandler expansion given cumulants.

    cf http://en.wikipedia.org/wiki/Edgeworth_series
    """
    assert len(cum) == 4
    mu, sigma = cum[0], np.sqrt(cum[1])
    coef = [1.,
            0.,
            0.,
            cum[2] / factorial(3) / sigma**3,
            cum[3] / factorial(4) / sigma**4]
    h = HermiteE(coef)
    def f(x):
        y = (x - mu) / sigma
        return h(y) * norm.pdf(y) / sigma
    return f


def edgeworth_series(cum):
    """Construct the 4th order Edgeworth expansion given cumulants.

    cf http://en.wikipedia.org/wiki/Edgeworth_series
    """
    assert len(cum) == 4
    mu, sigma = cum[0], np.sqrt(cum[1])
    coef = [1.,
            0.,
            0.,
            cum[2] / factorial(3) / sigma**3,
            cum[3] / factorial(4) / sigma**4,
            0,
            (cum[2] / factorial(3) / sigma**3)**2 / 2.]
    h = HermiteE(coef)
    def f(x):
        y = (x - mu) / sigma
        return h(y) * norm.pdf(y) / sigma
    return f


def cum_chi2(df, n=4):
    """Cumulants of the chi2 distribution."""
    assert n == 4
    return [2**(j-1) * factorial(j-1) * df for j in range(1, n+1)]


def main():
    df = 12
    c = cum_chi2(df)
    mu, sigma = c[0], np.sqrt(c[1])

    GC = gram_chandler_series(c)
    Ew = edgeworth_series(c)

    x = np.linspace(c[0] - 4*sigma, c[0] + 4*sigma, 100)
    plt.plot(x, chi2.pdf(x, df=df), 'y-', lw=6, 
            label=r'$\chi^2_{%s}$' % df)
    plt.plot(x, norm.pdf(x, loc=mu, scale=sigma), 'g--', lw=2, label='normal')
    plt.plot(x, GC(x), 'r-', lw=3, alpha=0.5, label='Gram-Chandler')
    plt.plot(x, Ew(x), 'b-' ,lw=3, alpha=0.7, label='Edgeworth')

    plt.text(0.1, 0.9, 'df = %s' % df, transform=plt.gca().transAxes,
            fontsize=18)
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()