Jeffreys interval demo

from __future__ import division

from pylab import *
from scipy.stats import beta, binom
from numpy.random import binomial

def estimate_p(x, N):
    dist = beta(x + 0.5, (N - x) + 0.5)
    p = x / float(N)
    lo = dist.ppf(0.025)
    hi = dist.ppf(0.975)

    return p, lo, hi



def plot_box(x, a, b, c, w=0.4):
    xx = array([-1.0, 1.0]) * w * 0.5
    yy = array([1.0, 1.0])
    plot([x-w*0.5, x+w*0.5, x+w*0.5, x-w*0.5, x-w*0.5],
         [a, a, c, c, a],
         'b-')
    plot([x-w*0.5, x+w*0.5],
         [b, b],
         'r-')


def plota_coisa(p, N, do_sort=False, Ntrials=40):
    # Ntrials = 40
    # N = 10
    # p = 0.3
    samples = binomial(N, p, Ntrials)

    subplot(2,1,1)
    title('{} random samples'.format(Ntrials, p, N))
    mm = []
    for i, x in enumerate(samples):
        phat, lo, hi = estimate_p(x, N)
        plot_box(1+i, lo, phat, hi)
        if i%2:
            text(1+i, 0.01+hi, '{}'.format(x), ha='center', fontsize=8)
        else:
            text(1+i, lo - 0.05, '{}'.format(x), ha='center', fontsize=8)
        mm.append(hi-lo)

    text(0, 0.9, '<length>: {:.2f}'.format(mean(mm)))

    print 'p:',p, 'N:', N, '<hi-lo>:', mean(mm)


    plot([-0.5, Ntrials + 0.5], [p, p], 'k--')

    xlim(-1, Ntrials+1)
    ylim(-0.05, 1.05)

    subplot(2,1,2)
    title('Same samples, ordered')
    samples.sort()
    for i, x in enumerate(samples):
        phat, lo, hi = estimate_p(x, N)
        plot_box(1+i, lo, phat, hi)
        if i%2:
            text(1+i, 0.01+hi, '{}'.format(x), ha='center', fontsize=8)
        else:
            text(1+i, lo - 0.05, '{}'.format(x), ha='center', fontsize=8)


    plot([-0.5, Ntrials + 0.5], [p, p], 'k--')

    xlim(-1, Ntrials+1)
    ylim(-0.05, 1.05)


def plot_benny(N, pp):
    for p in pp:
        plot(mgrid[:N+1], binom(N, p).pmf(mgrid[:N+1]), '-o', label='p:{}'.format(p))
    legend()





if __name__ == '__main__':



    ion()

    N=10
    figure()
    title('Binomial, N:{}'.format(N))
    plot_benny(N, [0.1, 0.5])
    savefig('binom-{}.png'.format(N))

    N=100
    figure()
    title('Binomial, N:{}'.format(N))
    plot_benny(N, [0.1, 0.5])
    savefig('binom-{}.png'.format(N))

    N=1000
    figure()
    title('Binomial, N:{}'.format(N))
    plot_benny(N, [0.1, 0.5])
    savefig('binom-{}.png'.format(N))


    p = 0.1
    N = 10
    figure(figsize=(8,8))
    suptitle('Bernoulli parameter estimates, p:{} N:{}'.format(p, N))
    plota_coisa(p, N)
    savefig('bsmp-{}-{}.png'.format(p,N))

    p = 0.5
    N = 10
    figure(figsize=(8,8))
    suptitle('Bernoulli parameter estimates, p:{} N:{}'.format(p, N))
    plota_coisa(p, N)
    savefig('bsmp-{}-{}.png'.format(p,N))

    p = 0.1
    N = 100
    figure(figsize=(8,8))
    suptitle('Bernoulli parameter estimates, p:{} N:{}'.format(p, N))
    plota_coisa(p, N)
    savefig('bsmp-{}-{}.png'.format(p,N))

    p = 0.5
    N = 100
    figure(figsize=(8,8))
    suptitle('Bernoulli parameter estimates, p:{} N:{}'.format(p, N))
    plota_coisa(p, N)
    savefig('bsmp-{}-{}.png'.format(p,N))

    p = 0.1
    N = 1000
    figure(figsize=(8,8))
    suptitle('Bernoulli parameter estimates, p:{} N:{}'.format(p, N))
    plota_coisa(p, N)
    savefig('bsmp-{}-{}.png'.format(p,N))

    p = 0.5
    N = 1000
    figure(figsize=(8,8))
    suptitle('Bernoulli parameter estimates, p:{} N:{}'.format(p, N))
    plota_coisa(p, N)
    savefig('bsmp-{}-{}.png'.format(p,N))