Mt

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

# Constants
T=2
Cells=4
# Observations seen by neither observer, both, just the first observer, or just the second
f=(None,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))

# Probabilities, conditional on observation
pi_obs = Lambda('pi_obs', lambda pi=pi: pi[1:]/(1-pi[0]))
# Probability of observation
p_n = Lambda('p_n', lambda pi=pi: 1-pi[0])

# Multinomial likelihood for conditional probabilities
F = Multinomial('F', n=sum(f[1:]), p=pi_obs, value=f[1:], observed=True)

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

# Binomial likelihood for complement of the probability of seeing nothing
n = Binomial('n', n=N, p=p_n, value=sum(f[1:]), observed=True)



"""
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) 
} 

"""