State-space models in greta

Hi developers,

thanks a lot for building this amazing tool that is greta and through your posts to keep providing incredibly detailed and interesting answers to each other’s questions. The material in this forum is advanced stuff and it’s for free, thanks to all for giving us the newcomers an opportunity.

I am trying to fit an SSM to a time series consisting of records of whether animals in dyads were observed at all and if observed together or separately. At each time t, one can have 4 possible observed values depending on if 0 no one in the dyad was observed, 1 only one observed, 2 both observed together and 3 both observed separately. The key thing about the model is that it allows accounting for the fact that one may not see some of the animals sometimes, in my opinion, a great improvement in our field. It was developed by Gimenez et al. 2019, in Inferring animal social networks with imperfect detection - ScienceDirect

The model is based on two probability matrices, a transition between the “real” states (together or separated) probability matrix and an observation matrix. The first computes the next state based on the previous while the other translates from the “real” underlying state to the possible observation values.

I adapted their BUGS model :

model{
# State matrix
px[1,1] <- psAA # probability of staying associated
px[1,2] <- 1 - psAA # probability of associated -> non-associated
px[2,1] <- 1 - psBB # probability of non-associated -> associated
px[2,2] <- psBB # probability of staying non-associated

# Observation matrix
## pp is the individual detection probability
for (t in 1:J){
	po[1,1,t] <- (1-pA[t]) * (1-pA[t])
	po[1,2,t] <- 2 * pA[t] * (1-pA[t])
	po[1,3,t] <- pA[t] * pA[t]
	po[1,4,t] <- 0
	po[2,1,t] <- (1-pB[t]) * (1-pB[t])
	po[2,2,t] <- 2 * pB[t] * (1-pB[t])
	po[2,3,t] <- 0
	po[2,4,t] <- pB[t] * pB[t]
	}
			
# Pr(initial states)
px0[1] <- pi # prob. of being in initial state A
px0[2] <- 1-pi # prob. of being in initial state B

# Model likelihood
	for (i in 1:n){
		
		# record states for every sampling occasion
		x1[i] <- x[i,1]
		x2[i] <- x[i,2]
		x3[i] <- x[i,3]
		x4[i] <- x[i,4]
		x5[i] <- x[i,5]
		
		# for t = 1
		x[i,1] ~ dcat(px0[1:2])
		obs[i,1] ~ dcat(po[x[i,1],1:4,1])

		# for t > 1
		for (t in 2:J){
			
			#-- state equation
			# 1 = associated
			# 2 = non-associated
			x[i,t] ~ dcat(px[x[i,t-1],1:2]) 
			
			#-- observation equation
			# 1 = dyad non-observed, 
			# 2 = one of the two individuals non-observed, 
			# 3 = dyad seen and associated
			# 4 = dyad seen and non-associated 
			obs[i,t] ~ dcat(po[x[i,t],1:4,t])
							}
						}

# Priors
for (t in 1:J){
pA[t] ~ dunif(0,1) # detection pr
pB[t] ~ dunif(0,1)
}
psAA ~ dunif(0,1) # pr of staying associated
psBB ~ dunif(0,1) # pr of staying non-associated
pi ~ dunif(0,1) # initial state pr
}
"

I have tried to adapt this to greta based on @nick answers to @cboettig in [Simulations with greta? (including time-series) - greta forum (greta-stats.org)] (Simulations with greta? (including time-series)). Based on these answers I tried to adapt the ar1ify functions to my model by:

  1. computing the states vectors by obtaining each new state based on the state matrix applied to the t-1 value.
  2. obtain the observations using the states vectors.
    How would you set the model likelihood based on the observations in greta?
    How would you adapt the probabilities for several grouping variables?
library(greta)
library(tidyverse)
library(tensorflow)
library(tfautograph)

SSMify <- function (innovations,psAA,psBB) {
  op <- greta::.internals$nodes$constructors$op
  op("SSMify", innovations, psAA,psBB,
     tf_operation = "tf_SSMify")
}
tf_SSMify <- function (innovations,psAA,psBB) {
  
  tf <- tensorflow::tf
  tfp<-tensorflow::tf_probability()
  tfd<-tfp$distributions
  
  # tfp <- tensorflow::tf_probability()
  shape <- tensorflow::shape
  #compute next state based on previous and probabilities AA and BB 
  getnextstate<-function(x,probAA,probBB){
    bl<-autograph({
      if (x == 1L) 
        tfd$Binomial(total_count=1L,probs=probAA)$sample(1L)
      else
        tfd$Binomial(total_count=1L,probs=probBB)$sample(1L)})
    
    bl
  }
  
  
  body <- function(epsilon, innovations, psAA, psBB,iter, maxiter) {
    
    # get next row of epsilon
    innovation <- innovations[, iter, , drop = FALSE]
    previous <- epsilon[, iter - 1L, , drop = FALSE]
    # need to pad both (?!)
    innovation <- tf$expand_dims(innovation, 1L)
    previous <- tf$expand_dims(previous, 1L)
    
    # get nex value of epsilon
    #apply rbinom depending on t-1 value and probs AA or BB
    #1-probBB because BB is a staying in value 0
    new_state<-getnextstate(previous,psAA,psBB)
    
    # append to epsilon
    epsilon <- tf$concat(list(epsilon, new_state), axis = 1L)
    
    list(epsilon, innovations, psAA,psBB, iter + 1L, maxiter)
  }
  
  cond <- function(epsilon, innovations, psAA,psBB, iter, maxiter) {
    tf$less(iter, maxiter)
  }
  
  dim <- dim(innovations)
  n_row <- dim[[2]]
  n_col <- dim[[3]]
  
  shapes <- list(tf$TensorShape(shape(NULL, NULL, n_col)),
                 tf$TensorShape(shape(NULL, n_row, n_col)),
                 tf$TensorShape(shape(NULL, 1L, 1L)),
                 #this extra line should be here for PsBB,right?
                 tf$TensorShape(shape(NULL, 1L, 1L)),
                 tf$TensorShape(shape()),
                 tf$TensorShape(shape()))
  
  # create an epsilon with the first row of innovations
  epsilon <- innovations[, 0, , drop = FALSE]
  
  values <- list(epsilon,
                 innovations,
                 psAA,
                 psBB,
                 tf$constant(1L),
                 tf$constant(n_row))
  
  out <- tf$while_loop(cond,
                       body,
                       values,
                       shape_invariants = shapes)
  out[[1]]
  
}



#based on the state matrix compute obs matrix
SSMifyobs <- function (innovations,poA,poB) {
  op <- greta::.internals$nodes$constructors$op
  op("SSMifyobs", innovations, psAA,psBB,
     tf_operation = "tf_SSMifyobs")
}
tf_SSMifyobs <- function (innovations,pp) {
  
  tf <- tensorflow::tf
  tfp <- tensorflow::tf_probability()
  tfd <- tfp$distributions
  shape <- tensorflow::shape
  
  probsA<-tf$convert_to_tensor(c((1-pp^2),2*pp*(1-pp),pp^2,0))
  probsB<-tf$convert_to_tensor(c((1-pp^2),2*pp*(1-pp),0,pp^2))
  #need to adapt this to an apply() like function
  multin<-function(x,probA,probB){
    ml<-autograph({
      if (x == 1L) 
        tf$where(tfd$Multinomial(total_count=1L,probs=probA)$sample(1L) == 1L)
      else
        tf$where(tfd$Multinomial(total_count=1L,probs=probB)$sample(1L) == 1L)})
    ml
  }
  
  getobsfromstate<-function(x,probA,probB){
    obs<-tf$map_fn(multin,x,probA,probB)
    obs}
  
  args<-list(innovations,probsA,probsB)
  out <- tf$vectorized_map(getobsfromstate,
                           args)
  
  out[[1]]
}


#simul some data
ts <- 100
dyads<-4
myobs<-map(1:dyads,~sample.int(4,10,replace=T)-1) %>% do.call(rbind,.)
#how to add groupping variables into the model? 
#e.g. like sex combination of ids
sexs<-c("FF","MM","MF",rep("FF",7))

#priors
#transition prob states A to A, B to B
psAA<-uniform(0,1,dim=1)
psBB<-uniform(0,1,dim=1)
#prob observing given state A or B 
pp<-uniform(0,1,dim=1)

#prob initial state A (1)
pin<-uniform(0,1,dim=1)

#vector of underlying states
states <- zeros(0, dim = ts)
#define initial state
states[1]<-rbinom(1,1,pin)

#compute states based on transition probabilities
state.m <- SSMify(states,psAA,psBB)
#compute observations based on states
obs<-SSMifyobs(state.m,pp)

#model likelihood? based on observed values
#distribution(myobs)<-

#retrieve the posteriors of the probabilities
#and the states recorded during mcmc?
m <- model(psAA,psBB,poA,poB,state.m)
draws <- mcmc(m)

Thanks a lot for any hint on how to proceed further
Best wishes