Spatial capture-recapture in PyMC

Good morning PyMC Community!

I’m writing to see if any one has any ideas on how best to port this Stan code for a spatial capture-recapture model to PyMC.

for (i in 1:M) {
  lp_if_present[i] = bernoulli_lpmf(1 | psi)
                 + binomial_logit_lpmf(y[i, ] | n_occasion, logit_p[i, ]);
  if (sum(y[i, ]) > 0) {
    log_lik[i] = lp_if_present[i];
  } else {
    log_lik[i] = log_sum_exp(lp_if_present[i], bernoulli_lpmf(0 | psi));
}

In this case, M is an integer, y is an array with binomial counts with shape (M, n_trap), psi is the probability for the non-zero portion of the mixture, logit_p is an array of logit-scale probabilities for the number of times that the individual is captured in each trap in n_occassions.

At first glance it seems like a straightforward implementation of the handy pm.ZeroInflatedBinomial. I haven’t figured out how to implement that, however, because the zero-inflation happens along the rows, that is, at the individual-level. In other words, if y[0, 4] = 5, we know that psi = 1 for y[0, 5] = 0. But if all(y[100, :]) == 0, then 0 < psi < 1. Hopefully that makes sense.

This was my attempt. It feels dumb, but it seemed to work for a simple, simulated example. I’m not sure that I trust it in the real world

from pymc.distributions.dist_math import binomln, logpow

def logp(value, n_occasion, p, psi):

    binom = binomln(n_occasion, value) + logpow(p, value) + logpow(1 - p, n_occasion - value)
    bin_sum = pm.math.sum(binom, axis=1)
    bin_exp = pm.math.exp(bin_sum)

    res = pm.math.switch(
        value.sum(axis=1) > 0,
        bin_exp * psi,
        bin_exp * psi + (1 - psi)
    )

    return pm.math.log(res)

Anyway, I would appreciate any help with this fool’s errand! Thanks.

Your logp graph doesn’t look wildly wrong, but I wonder if you can write down a generative simulation instead? That is usually a more natural first step to writing a PyMC model, since we work with the generative model instead of the logp model directly (as Stan does).

Something like:

  1. Sample \psi\in (0,1)
  2. For individual i, sample a bernoulli with p=\psi. If True, all observations are zero.
  3. Otherwise, for each t:
    a. Check y_{t-1} \neq 0. If true, y_t=0
    b. Otherwise, draw y_t \sim Bernoulli(n, p) times trapped.

You might even be able to use the new MarginalModel to marginalize away \psi automatically and improve sampling. To me the harder part to handle is the time dependency. Could just scan, or could do something like:

n_trapped = pm.Bernoulli('n_trapped', logit_p, n)
n_trapped = n_trapped[1:] * (n_trapped[:-1] == 0)

I wonder if MarginalModel could handle removing n_trapped as written above.

In general I think your process is an HMM, it kind of reminds me of this post. The states are “trapped” and “free”, if “free” the emission is a zero, and if “trapped” it’s n_trapped. There is probaility \psi to transition from free to trapped, and probability 1 to transition from trapped to free.

Hello! Thank you for the prompt and thoughtful response. I apologize for the massive delay.

It seems I’ve created some confusion by using our jargon. In this case, “capture” should really be “encounter.” Encounters don’t impede the animal, helping alleviate the complexity that you hit on. Encounter data might involve “traps” of velcro that the animal rubs against, from which we can gather hair, i.e., DNA. (Although I should note that when these data include deaths, that they become HMMs!)

As an ecologist, I also find the generative model more intuitive. Personally, I struggle to write these out without using a discrete latent variable z. Please see below for a simplified example of a non-spatial capture-recapture model. I’ve tried several hacks at marginalizing z out, but haven’t been able to without using the logp, particularly when the parameter p varies by animal and occasion.

Anyway, any tips on marginalizing out this variable out would be greatly appreciated!

import numpy as np
import pymc as pm
import pytensor.tensor as pt
import arviz as az

N_true = 200               # abundance
M = 400                    # augmented dataset size
T = 4                      # number of occasions
P = [0.25, 0.2, 0.4, 0.25] # encounter probabilty at each occastion
P_mat = np.broadcast_to(P, (N_true, T)) 

# simulate encounters
Y = np.random.default_rng().binomial(n=1, p=P_mat)
Y = Y[Y.sum(axis=1) > 0]
n_capt = Y.shape[0]

# augment data with all zero encounter histories 
all_zero_history = np.zeros((M - n_capt, T))
augmented = np.row_stack((Y, all_zero_history))

with pm.Model() as mt:
    
    # detection probability 
    p = pm.Uniform('p', 0, 1, shape=T)
    pmat = pt.tile(p, (M, 1))

    # inclusion probability 
    psi = pm.Uniform('psi', 0, 1)
    z = pm.Bernoulli('z', p=psi, shape=M)
    
    # likelihood
    p_eff = (pmat.T * z).T
    pm.Bernoulli('Y_obs', p_eff, observed=augmented)

    # derived quanities
    N = pm.Deterministic('N', z.sum()) 
    
    trace = pm.sample()

az.summary(trace, var_names=['p', 'N'])