# -*- coding: utf-8 -*-
# <nbformat>2</nbformat>

# <markdowncell>

# Seeds: Random effect logistic regression 
# ==============================
# 
# Ported to PyMC by Abraham Flaxman and Kyle Foreman, 2/15/2012 from http://www.openbugs.info/Examples/Seeds.html
# 
# This example is taken from Table 3 of Crowder (1978), and concerns the proportion of seeds that germinated on each of 21 plates arranged according to a 2 by 2 factorial layout by seed and type of root extract. The data are shown below, where $r_i$ and $n_i$ are the number of germinated and the total number of seeds on the $i$-th plate, $i =1,...,N$. These data are also analysed by, for example, Breslow and Clayton (1993). 
# 
# 
# The model is essentially a random effects logistic, allowing for over-dispersion. If $p_i$ is the probability of germination on the $i$-th plate, we assume
# 
# $$
# r_i \sim \text{Binomial}(p_i , n_i )
# $$
# $$
# \text{logit}(p_i ) = \alpha_0 + \alpha_1 x_{1,i} + \alpha_2 x_{2,i} + \alpha_{1,2} x_{1,i} x_{2,i} + b_i
# $$ 
# $$
#    b_i \sim \text{Normal}(0, \tau ) 
# $$
# 
# where $x_{1,i}, x_{2,i}$ are the seed type and root extract of the $i$-th plate, and an interaction term $a_{1,2} x_{1,i} x_{2,i}$ is included.
# 
# $\alpha_0 , \alpha_1 , \alpha_2 , \alpha_{1,2} , \tau$ are given independent weakly informative priors. 

# <codecell>

import pylab as pl
import pymc as mc

# <codecell>

### data   
# germinated seeds
r =  pl.array([10, 23, 23, 26, 17, 5, 53, 55, 32, 46, 10, 8, 10, 8, 23, 0, 3, 22, 15, 32, 3])

# total seeds
n =  pl.array([39, 62, 81, 51, 39, 6, 74, 72, 51, 79, 13, 16, 30, 28, 45, 4, 12, 41, 30, 51, 7])

# seed type
x1 = pl.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

# root type
x2 = pl.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1])

# number of plates
N =  x1.shape[0]

# <markdowncell>

# Here is the model in BUGS:
# 
#       model
#       {
#          for( i in 1 : N ) {
#             r[i] ~ dbin(p[i],n[i])
#             b[i] ~ dnorm(0.0,tau)
#             logit(p[i]) <- alpha0 + alpha1 * x1[i] + alpha2 * x2[i] + 
#                alpha12 * x1[i] * x2[i] + b[i]
#          }
#          alpha0 ~ dnorm(0.0,1.0E-6)
#          alpha1 ~ dnorm(0.0,1.0E-6)
#          alpha2 ~ dnorm(0.0,1.0E-6)
#          alpha12 ~ dnorm(0.0,1.0E-6)
#          tau ~ dgamma(0.001,0.001)
#          sigma <- 1 / sqrt(tau)
#       }

# <codecell>

### hyperpriors
tau = mc.Gamma('tau', 1.e-3, 1.e-3, value=10.)
sigma = mc.Lambda('sigma', lambda tau=tau: tau**-.5)

### parameters
# fixed effects
alpha_0 =  mc.Normal('alpha_0',  0., 1e-6, value=0.)
alpha_1 =  mc.Normal('alpha_1',  0., 1e-6, value=0.) 
alpha_2 =  mc.Normal('alpha_2',  0., 1e-6, value=0.) 
alpha_12 = mc.Normal('alpha_12', 0., 1e-6, value=0.)  

# random effect
b =        mc.Normal('b',        0., tau,  value=np.zeros(N))

# expected parameter
logit_p =  (alpha_0 + alpha_1*x1 + alpha_2*x2 + alpha_12*x1*x2 + b)


### likelihood
@mc.observed
def obs(value=r, n=n, logit_p=logit_p):
    return mc.binomial_like(r, n, mc.invlogit(logit_p))

# <markdowncell>

# BUGS uses Gibbs steps automatically, so it only takes 10000 steps of MCMC after a 1000 step burn in for this model in their example.
# 
# PyMC only uses Gibbs steps if you set them up yourself, and it uses Metropolis steps by default.  So 10000 steps
# go by quickly, but the chain has not converged to the stationary distribution in this time, so predictions based
# on the samples will be inaccurate.

# <codecell>

m = mc.MCMC([alpha_0, alpha_1, alpha_2, alpha_12, b, tau, sigma, logit_p, obs])

# not long enough to mix
%time m.sample(10000, 1000, progress_bar=False)

# <markdowncell>

# PyMC has a very sophisticated "Adaptive Metropolis" method, which adjusts the proposal distribution for the Metropolis step
# adaptively.  Using this and running for longer provide more reliable results.

# <codecell>

m = mc.MCMC([alpha_0, alpha_1, alpha_2, alpha_12, b, tau, sigma, logit_p, obs])
m.use_step_method(mc.AdaptiveMetropolis, b)

# about long enough, but maybe better initial values would help more
%time m.sample(20000, 10000, 10, progress_bar=False)

# <markdowncell>

# Starting from carefully chosen initial values and running the chain for even longer yields results that are completely
# mixed.  This is going to give the same answer every time (approximately).

# <codecell>

tau.value=1.
map = mc.MAP([alpha_0, alpha_1, alpha_2, alpha_12, b, logit_p, obs])
map.fit(method='fmin_powell')


m = mc.MCMC([alpha_0, alpha_1, alpha_2, alpha_12, b, tau, sigma, logit_p, obs])
m.use_step_method(mc.AdaptiveMetropolis, b)

# a little longer, good initial values, but not Adaptive Metropolis.  Does not converge.
%time m.sample(200000, 100000, 100, progress_bar=False)

# <markdowncell>

# Now that I've seen that it can work, how necessary is the Adaptive Metropolis?  Necessary.

# <codecell>

tau.value=1.
map = mc.MAP([alpha_0, alpha_1, alpha_2, alpha_12, b, logit_p, obs])
map.fit(method='fmin_powell')


m = mc.MCMC([alpha_0, alpha_1, alpha_2, alpha_12, b, tau, sigma, logit_p, obs])

%time m.sample(200000, 100000, 100, progress_bar=False)

# <markdowncell>

# And how necessary are the carefully chosen initial values?

# <codecell>

### hyperpriors
tau = mc.Gamma('tau', 1.e-3, 1.e-3, value=10.)
sigma = mc.Lambda('sigma', lambda tau=tau: tau**-.5)

### parameters
# fixed effects
alpha_0 =  mc.Normal('alpha_0',  0., 1e-6, value=0.)
alpha_1 =  mc.Normal('alpha_1',  0., 1e-6, value=0.) 
alpha_2 =  mc.Normal('alpha_2',  0., 1e-6, value=0.) 
alpha_12 = mc.Normal('alpha_12', 0., 1e-6, value=0.)  

# random effect
b =        mc.Normal('b',        0., tau,  value=np.zeros(N))

# expected parameter
logit_p =  (alpha_0 + alpha_1*x1 + alpha_2*x2 + alpha_12*x1*x2 + b)


### likelihood
@mc.observed
def obs(value=r, n=n, logit_p=logit_p):
    return mc.binomial_like(r, n, mc.invlogit(logit_p))

m = mc.MCMC([alpha_0, alpha_1, alpha_2, alpha_12, b, tau, sigma, logit_p, obs])
m.use_step_method(mc.AdaptiveMetropolis, b)

# a little longer, but not special initial values, does it converge?
%time m.sample(200000, 100000, 100, progress_bar=False)

# <markdowncell>

# Results
# =======
# 
# BUGS results:
# 
# A burn in of 1000 updates followed by a further 10000 updates gave the following parameter estimates:
# 
#                mean    sd    
#      alpha_0   -0.55   0.19  
#      alpha_1    0.08   0.30  
#      alpha_12  -0.82   0.41  
#      alpha_2    1.35   0.26  
#      sigma      0.27   0.15  

# <codecell>

for node in [alpha_0, alpha_1, alpha_12, alpha_2, sigma]:
    stats = node.stats()
    print '%10s\t%1.2f \t%.2f' % (node.__name__, stats['mean'], stats['standard deviation'])

# <codecell>

for node in [alpha_0, alpha_1, alpha_12, alpha_2, sigma]:
    stats = node.stats()
    print '%10s\t%1.2f \t%.2f' % (node.__name__, stats['mean'], stats['standard deviation'])

# <codecell>

mc.Matplot.plot(m)

# <codecell>

mc.Matplot.plot(b)

# <codecell>


# <codecell>


