Help translating multistate capture-recapture model from JAGS to STAN

Modeling
, ,
2

For anyone seeing this in the future I figured it out. I used the HMM framework found for most of the models here https://github.com/stan-dev/example-models/tree/master/BPA and was able to get the model to sample.

// ---------------------------------
// States (S):
// 1 alive and present
// 2 alive and absent
// 3 dead
// Observations (O):
// 1 seen
// 2 not seen
// ---------------------------------

functions {
  /**
   * Return an integer value denoting occasion of first capture.
   * This function is derived from Stan Modeling Language
   * User's Guide and Reference Manual.
   *
   * @param y         Observed values
   * @return Occasion of first capture
   */
  int first_capture(array[] int y_i) {
    for (k in 1 : size(y_i)) {
      if (y_i[k] == 1) {
        return k;
      }
    }
    return 0;
  }
}

data {
  int<lower=0> nind; //number of individuals
  int<lower=1> nprimary;       // Number of years or primary occasions in this case
  vector<lower=1>[nprimary] nsecondary; //number of secondary occassions
  int<lower=1> nsecondary_max; //max number of secondary occasions
  array[nind, nprimary] int<lower=1,upper=2> ch; //individual capture histories for each primary occasion
  array[nprimary, nsecondary_max] int<lower=0> yes; //fish recaptured by primary and secondary occasion
  array[nprimary, nsecondary_max] int<lower=0> total; //total fish caught during primary and secondary occasion
}

transformed data {
  int<lower=0,upper=nprimary> first[nind];
  vector<lower=0, upper=nind>[nprimary] n_captured;

  for (i in 1:nind)
    first[i] = first_capture(ch[i]);
    
     n_captured = rep_vector(0, nprimary);
  for (t in 1:nprimary) {
    for (i in 1:nind) {
      if (ch[i, t]) {
        n_captured[t] = n_captured[t] + 1;
      }
    }
  }
}

parameters {
  matrix<lower=0,upper=1>[nprimary, nsecondary_max] p; // Recapture probability
  real<lower=0,upper=1> phi; //survival probability
  real<lower=0,upper=1> gP; // the probability of being off the study area, unavailable for capture during primary trapping session (i) given that the animal was not present on the study area during primary trapping session (8 ??? 1), and survives to trapping session (8).
  real<lower=0,upper=1> gPP; // the probability of being off the study area, unavailable for capture during the primary trapping session (i) given that the animal was present during primary trapping session (8 ??? 1), and survives to trapping session (8).
}

transformed parameters {
  vector<lower=0,upper=1>[nprimary] pstar; //primary occasion p's or pooled detection probability
  array[3] simplex[3] ps;
  array[3, nprimary] simplex[2] po; 
    
    //Primary occasions p's or pooled detection probability
    for (t in 1:nprimary){
      pstar[t] = 1 - prod(1 - p[t,1:nsecondary_max]);
    }

 // Define state-transition and observation matrices
    // Define probabilities of state S(t+1) given S(t)
      ps[1, 1] = phi * gPP;
      ps[1, 2] = phi * (1 - gPP);
      ps[1, 3] = 1.0 - phi;
      ps[2, 1] = phi * gP;
      ps[2, 2] = phi * (1 - gP);
      ps[2, 3] = 1.0 - phi;
      ps[3, 1] = 0;
      ps[3, 2] = 0;
      ps[3, 3] = 1.0;

    for (t in 1 : nprimary) {
      // Define probabilities of O(t) given S(t)
      po[1, t, 1] = pstar[t];
      po[1, t, 2] = 1 - pstar[t];
      po[2, t, 1] = 0.0;
      po[2, t, 2] = 1.0;
      po[3, t, 1] = 0.0;
      po[3, t, 2] = 1.0;
 }
  }

model {
  array[3] real acc;
  array[nprimary] vector[3] gamma;

    // priors
    // phi ~ uniform(0,1);     # survival 
    // gP ~ uniform(0,1);      # gamma'
    // gPP ~ uniform(0,1);     # gamma'' 
    // p ~ uniform(0,1);  # detection
    
    for (t in 1:nprimary){
      for (j in 1:nsecondary_max){
        yes[t,j] ~ binomial(total[t,j], p[t,j]);
      }
    }   
  
  // Likelihood
  // Forward algorithm derived from Stan Modeling Language
  // User's Guide and Reference Manual.
  for (i in 1 : nind) {
    if (first[i] > 0) {
      for (k in 1 : 3) {
        gamma[first[i], k] = k == ch[i, first[i]];
      }
      
      for (t in (first[i] + 1) : nprimary) {
        for (k in 1 : 3) {
          for (j in 1 : 3) {
            acc[j] = gamma[t - 1, j] * ps[j, k]
                     * po[k, t - 1, ch[i, t]];
          }
          gamma[t, k] = sum(acc);
        }
      }
      target += log(sum(gamma[nprimary]));
    }
  }
}
4 Likes