Mt2

# =======================================================================================
# = Estimating the size of a closed population, with time-varying capture probabilities =
# =======================================================================================
from pymc import Lambda, Beta, DiscreteUniform, stochastic, multinomial_like
from numpy import reshape, ma, prod, array, append
import pdb

# Constants
T=2
Cells=4
# Observations seen by both observers, just the first observer, or just the second
f=(22,363,174)
# Inidicators for observer sightings
omegas=reshape((0,0, 1,0, 0,1, 1,1), (Cells, T))

# Capture probabilities
p = Beta('p', alpha=1, beta=1, size=T)

# Cell probabilities
c = Lambda('c', lambda p=p: p**omegas * (1-p)**(1.-omegas))

# Joint probabilities
pi = Lambda('pi', lambda c=c: prod(c, 1))

# Bound total N by the number of observed individuals
N = DiscreteUniform('N', lower=sum(f), upper=10000)

@stochastic(observed=True, dtype=int)
def F(value=f, n=N, p=pi):
    """Complete multinomial for observed and unobserved quatities"""
    return multinomial_like(append(n-sum(value), value), n, p)


"""
Original BUGS model from Link and Barker 2009

# Data 
list(T=2, Cells=4, f=c(NA,22,363,174), 
omegas=structure(.Data=c(0,0, 
1,0, 
0,1, 
1,1),.Dim=c(4,2))) 

model 
{ 
    for (t in 1:T) 
    { 
        p[t] ~ dunif(0,1) 
    }     
    for (j in 1:Cells) 
    { 
        for (t in 1:T) 
        { 
            c[j,t] <- pow(p[t],omegas[j,t]) * pow(1-p[t],1-omegas[j,t]) 
        } 
        pi[j] <- prod(c[j,1:T]) 
    } 
    for (j in 2:Cells) 
    { 
        pi.obs[j] <- pi[j]/(1-pi[1]) 
    } 
    n <- sum(f[2:Cells]) 
    f[2:Cells] ~ dmulti(pi.obs[2:Cells],n) 
    pn <- 1-pi[1] 
    n ~ dbin(pn,N) 
    cN ~ dunif(n,10000) 
    N <- round(cN) 
} 

"""