Chi-squared distribution

chisq_distribution.py

#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import chi2


def main():
    colors = 'bgrcmyk'
    n = 40000
    for i, k in enumerate((1, 2, 3, 4, 6, 9)):
        vs = []
        for j in range(n):
            v = (np.random.normal(size=k)**2).sum()
            vs.append(v)

        hs, es = np.histogram(vs, bins=2000, range=(0, 200), normed=True)
        xs = (es[1:] + es[:-1]) / 2.

        plt.plot(xs, hs, '.', color=colors[i])

        rv = chi2(k)
        plt.plot(xs, rv.pdf(xs), color=colors[i], label='k={}'.format(k))

    plt.xlim(0, 8)
    plt.ylim(0, 0.5)
    plt.legend()
    plt.xlabel('x')
    plt.ylabel('pdf(x)')
    plt.title('chi-squared distribution')
    plt.show()


if __name__ == '__main__':
    main()