SpTMixture: Mixture models with spatio-temporal dependencies

Documented in gibbSPTMixed

#' Gibbs sampler with spatial Potts model for neighbouring behaviours and random site effect.
#'
#' Gibbs sampler with sptial Potts model for neighbouring behaviours and random site effect.
#' @param theta parameters to be estimated
#' @param y data
#' @param p.sites the composition of each site.
#' @param season.sites season data for each site.
#' @param Coords Coordinates of the sites.
#' @param Bds Matrix of edges between sites.
#' @param ID ID vector of the sites.
#' @param N.run Number of run of the sampler.
#' @param beta.range Range of beta for the Potts model.
#' @keywords Gibbs sampler
#' @export
#' @examples
#' data("wq_analysis_week2")
#' SPTMData(wq.raw.obs, frequency = "quarter")
gibbSPTMixed = function(priors= list(theta = list(m0 = 0, s0 = 1), sigma2 = list(a0=2, b0=1), cov = list(s0=0.1)), y, p.sites, season.sites, basis, Coords=NULL ,Bds=NULL, ID, N.run = 10000,
                    beta.range = c(0.5,3), meanFunc = NULL, debug = FALSE, print.res = FALSE, Xcov = NULL, ...){




  theta0 = c(priors$theta$m0, priors$sigma2$b0 / (priors$sigma2$a0-1))

  n.sites = nrow(p.sites)

  n.lu = ncol(p.sites)

  n.param = length(theta0)

  n.season = dim(season.sites)[2]

  if(length(dim(season.sites))==2){
    n.season=1
  }

  type = "stl"
  if(!is.null(basis)){type = "splines"}

  beta.max.min = beta.range

  p.hist = array(0, dim = c(n.sites, n.lu, N.run))

  theta.hist = matrix(0, nrow = N.run, ncol = n.param)

  phi.hist = array(0, dim = c(ncol(basis),n.lu, N.run))

  random.hist = matrix(0, nrow = N.run, ncol = n.sites*n.season)

  if(is.null(Xcov)){
    Xcov = matrix(0, ncol = 1, nrow = n.sites)
  }

  priors$cov$s0 = rep(0.1, ncol(Xcov)*n.season)

  randomb.hist = matrix(0, nrow = N.run, ncol = n.season*ncol(Xcov))

  z.hist = array(0, c(n.sites, n.lu, N.run))

  beta.hist = matrix(0, nrow = N.run, ncol = 1)

  p.hist[,,1] =as.matrix(p.sites)

  theta.hist[1,] = theta0

  beta.hist[1] = runif(1,beta.max.min[1],beta.max.min[2])

  idx.max = apply(p.sites, 1,which.max)

  for(li in 1:nrow(p.sites)){
    z.hist[li,idx.max[li],1] = 1
  }

  n.colors = ncol(p.sites)

  b = t(apply(season.sites,1,c))

  zt = NULL
  for(ii in 1:nrow(z.hist[,,1])){
    zt = cbind(zt, kronecker(diag(n.season),z.hist[ii,,1]))
  }

  idx.season = attr(basis, "indexSeason")

  for(t in 2:N.run){
    if(debug){browser()}
    # Init

    Ym = matrix(y$obs, ncol = n.sites, byrow=T)

    z = z.hist[,,t-1]

#    Ym[is.na(Ym)] <- 0

    bY = NULL
    for(iii in 1:n.sites){
      idx.na = which(!is.na(Ym[,iii]))
      temp = ginv(t(basis[idx.na,]) %*% basis[idx.na,]) %*% t(basis[idx.na,]) %*% Ym[idx.na,iii]
      bY = cbind(bY, temp)
    }

#    phiB = ginv(t(basis) %*% basis) %*% t(basis) %*% Ym %*% z %*% ginv(t(z) %*% z)

    phi = bY %*% z %*% ginv(t(z) %*% z)

    phi.hist[,,t] = phi

    zt = NULL
    for(ii in 1:nrow(z.hist[,,t-1])){
      zt = cbind(zt, kronecker(diag(n.season),z.hist[ii,,t-1]))
    }



    season.l = NULL
    for(ii in 1:length(idx.season)){
      if(length(idx.season[[ii]])==1){
        season.l = cbind(season.l, matrix(matrix(basis[,idx.season[[ii]]], ncol=1) %*% matrix(phi[idx.season[[ii]],], nrow=1), nrow = nrow(basis)))
      }else{
      season.l = cbind(season.l, basis[,idx.season[[ii]]] %*% phi[idx.season[[ii]],])
      }
    }

#    Bl = season.lu[,idx.na] %*% zt[idx.na,]
    idx.na = which(!is.na(colSums(season.l)))
    Bl = season.l[,idx.na] %*% zt[idx.na,]


    Bti = matrix(0, nrow = nrow(Bl)*n.sites, ncol = n.sites*n.season)
    for(i in 1:n.sites){
      idx = (i-1)*n.season + 1:n.season
      idx.r = (i-1)*nrow(Bl) + 1:nrow(Bl)
      idx.c = (i-1)*n.season + 1:n.season
      Bti[idx.r, idx.c] = Bl[,idx]
    }

    A = random.hist[t-1,]

    Xcovi = matrix(0, nrow = nrow(Xcov)*n.season, ncol = ncol(Xcov)*n.season)
    for(i in 1:nrow(Xcov)){
      idx.r = rep((i-1)*n.season + 1:n.season, each = ncol(Xcov))
      idx.c = 1:(ncol(Xcov)*n.season)
      Xcovi[cbind(idx.r,idx.c)] = as.matrix(Xcov)[i,]
    }

    A_t = randomb.hist[t-1,]

    Bz = t(zt) %*% theta.hist[t-1,1:(n.season*n.colors)]

    sigma = theta.hist[t-1, (length(theta0) - n.lu + 1):length(theta0)]

    beta = beta.hist[t-1]

#     B = b %*% t(zt) %*% ginv(zt %*% t(zt))
#
#     theta.d = NULL
#     for(j in 1:n.season){
#       theta.d = rbind(theta.d,diag(theta.hist[t-1,((j-1)*n.colors+1):(j*n.colors)]))
#     }
#
#     md = B %*% theta.d
#


#     plot(Ym[,1]);points(Ym[,2]);points(Ym[,6]);points(Ym[,8]);
#     points(basis %*% phi[,2], type="l", col = "red", lwd=2)

#     season.lu = updateSeason(season.sites, p.sites = z.hist[,,t-1])
#
#     m =  meanFunc(theta.hist[t-1,], n.mixt = n.lu, priors$z$a, season.sites)
    m = basis %*% phi
    idx.nan = which(is.nan(colSums(m)))

    # Updating beta

    zz = z.hist[,,t-1]

    zz.temp = apply(zz,1,which.max)

    out = MurraySchemePotts(beta.init = beta, Coords = Coords ,Nb = Bds ,n.colors = n.colors, col.obs = zz.temp, N.run = 100, range.beta = beta.max.min)

    beta.hist[t] = mean(out$beta.sample, na.rm=T)


    # Updating z

    if(!is.null(Coords)){Bds = getBonds(as.matrix(Coords), NN = 4, th = 2)}

    vois = getCanStatF(Bds, zz.temp, n.colors)

    Ys = Bti %*% (Bz + A + Xcovi %*% A_t)

    Ypm = matrix(Bti %*% (A + Xcovi %*% A_t), ncol=n.sites, byrow = FALSE)

    for(i in 1:n.sites){
      l.prob = NULL

      y.p = y[y$ID == ID[i], ]

      for(j in 1:n.lu){

        #y.m = m[,j]
        y.m = m[,j] + Ypm[,i]

        n.na = sum(!is.na(y.p$obs))

        py.cx = sum(dnorm(y.p$obs, y.m , sigma[j], log = T), na.rm=T)+ n.na*(log(sqrt(2*pi)*sigma[j]))

        # l.prob = c(l.prob, py.cx  + beta * vois[j] + max(log(p.hist[i,j,t-1]), -1*10e100, na.rm=T))

        l.prob = c(l.prob, py.cx  + beta * vois[i,j])
      }

      l.prob[idx.nan] <- -Inf

      max.l = max(l.prob)

      logprob = (l.prob - max.l) + log(p.hist[i,,t-1])

      prob = exp(logprob - max(logprob))

      prob = prob / sum(prob)

      z.hist[i,,t] = 0
      z.hist[i,sample(1:n.lu,1, prob = prob),t] = 1
    }


    n.t = colSums(z.hist[,,t])

    # Updating pi

    for(i in 1:n.sites){
      alpha = priors$z$a[i,] + z.hist[i,,t] / n.colors
      p.hist[i,,t] = rdirichlet(c(as.matrix(alpha)))
    }

    # Updating theta_k

    season.lu = updateSeason(season.sites, p.sites = z.hist[,,t])

#    zt = apply(z.hist[,,t],1,function(x) kronecker(diag(n.season),matrix(x, nrow=1)))

    l.j = NULL
    for(j in 1:n.lu){

      index.j = ID[which(z.hist[,j,t]==1)]

      idx.Xlu = j + (0:((length(theta0) - n.lu)/n.lu-1))*n.lu
      idx.jlu = 0:(n.season-1)*n.colors + j
      Blu = season.l[,idx.jlu]

      if(length(index.j)>0){

        y.p = y[y$ID %in% index.j,]

        l.j = c(l.j,length(y.p$obs))

        Bs = kronecker(Blu,matrix(1,ncol = 1, nrow = length(index.j)))

        idx.na = which(!is.na(y.p$obs))
        Bs = Bs[idx.na,]

        Zs = sapply(unique(y.p$ID), function(x) as.numeric(y.p$ID == x))
        Zs = Zs[idx.na,]

        idx.a = outer(match(unique(y.p$ID), ID),n.sites*(0:(n.season-1)),FUN =  '+')

        a = matrix(random.hist[t-1,idx.a], ncol=3, byrow = F)

        est.b = ginv(t(Bs) %*% Bs) %*% t(Bs) %*% (y.p$obs[idx.na] - matrix(diag(Bs %*% t(Zs %*% a)), ncol = 1))

        theta.hist[t,idx.Xlu] = est.b


      }else{

        l.j = c(l.j, 0)

        idx.Xlu = j + (0:((length(theta0) - n.lu)/n.lu-1))*n.lu

        theta.hist[t,idx.Xlu] = theta.hist[t-1,idx.Xlu]
      }


    }

    Y = y[order(y$ID),]$obs

    idx.na = which(!is.na(Y))

    BtiF = Bti[idx.na,]

    est.b = ginv(zt %*% t(zt)) %*% zt %*% (ginv(t(BtiF) %*% BtiF) %*% t(BtiF) %*% Y[idx.na] - A - Xcovi %*% A_t)

    n.j = colSums(z.hist[,,t])

    var.post = (1 / priors$theta$s0 + rep(n.j, n.season) / theta.hist[t-1,(length(theta0) - n.lu + 1):length(theta0)])^(-1)

    mean.post = var.post * (priors$theta$m0 / priors$theta$s0 + n.j * est.b / theta.hist[t-1,(length(theta0) - n.lu + 1):length(theta0)])

    theta.hist[t,1:(n.season*n.colors)] = rnorm(n.season*n.colors,mean.post, sd=sqrt(var.post))

    # Updating a_k

    # if(t == 1000){browser()}

    Y = y[order(y$ID),]$obs

    idx.na = which(!is.na(Y))

    BtiF = Bti[idx.na,]

    Bz = t(zt) %*% theta.hist[t,1:(n.season*n.colors)]

    est.a = (ginv(t(BtiF) %*% BtiF) %*% t(BtiF) %*% Y[idx.na] - Bz - Xcovi %*% A_t)

    n.j = colSums(z.hist[,,t])

    tna = table(y$ID, !is.na(y$obs))

    nn.j = tna[,ncol(tna)]

    index.vv0 = cbind(rep(1:n.season, each = n.sites),rep(apply(z.hist[,,t],1,which.max), n.season))

    vv0 = matrix(priors$theta$s0, ncol = n.lu, byrow = T)[index.vv0]

    vv = rep(theta.hist[t-1,(length(theta0) - n.lu + 1):length(theta0)][apply(z.hist[,,t],1,which.max)], n.season)

    var.post.a = (1 / vv0 + nn.j / vv)^(-1)

    mean.post.a = var.post.a * (nn.j * est.a / vv)

    random.hist[t,] = rnorm(n.sites*n.season,mean.post.a, sd=sqrt(var.post.a))

# Updating cov_k

    Y = y[order(y$ID),]$obs

    idx.na = which(!is.na(Y))

    BtiF = Bti[idx.na,]

    Bz = t(zt) %*% theta.hist[t,1:(n.season*n.colors)]

    A = random.hist[t,]

    est.ab = ginv(t(Xcovi) %*% Xcovi) %*% t(Xcovi) %*% (ginv(t(BtiF) %*% BtiF) %*% t(BtiF) %*% Y[idx.na] - Bz - A)

    n.j = colSums(z.hist[,,t])

    nnn.j = apply(Xcov,2,function(x) sum(!is.na(x)))

    index.vv0 = cbind(rep(1:n.season, each = n.sites),rep(apply(z.hist[,,t],1,which.max), n.season))

    vv0 = priors$cov$s0

    vvb = rep(max(theta.hist[t-1,(length(theta0) - n.lu + 1):length(theta0)]), n.season*ncol(Xcov))

    var.post.ab = (1 / vv0 + nnn.j / vvb)^(-1)

    mean.post.ab = var.post.ab * (nnn.j * est.ab / vvb)

    randomb.hist[t,] = rnorm(length(mean.post.ab),mean.post.ab, sd=sqrt(var.post.ab))

    # Updating sigma_k

    #m =  meanFunc(theta.hist[t,], n.mixt = n.lu, z.hist[,,t], season.sites)

    A_t = randomb.hist[t,]

#    A[8] = random.hist[t,8]-2

    Ys = Bti %*% (Bz + A + Xcovi %*% A_t)

    # Ys1 = Bti %*% (Bz + A + Xcovi %*% A_t)

    Yf = y[order(y$ID),]

#    idx.16 = which(Yf$ID == "16")

    for(j in 1:n.lu){

      idx.obs = which(Yf$ID %in% ID[which(z.hist[,j,t] == 1)])

      mean.se = 0

      if(length(idx.obs)>0){mean.se = mean((Ys[idx.obs] - Yf$obs[idx.obs])^2, na.rm=T)}


      n.obs= length(which(z.hist[,j,t] == 1)) #length(idx.obs) - sum(is.na(y$obs[idx.obs]))

      colSums((matrix(Ys[idx.obs], ncol = n.obs, byrow=FALSE) - matrix(Yf$obs[idx.obs], ncol = n.obs, byrow=FALSE))^2, na.rm=T)/apply(matrix(Yf$obs[idx.obs], ncol = n.obs, byrow=FALSE),2,function(x) sum(!is.na(x)))

      a = priors$sigma2$a0[j] +n.obs/2
      b = n.obs*mean.se /2 + priors$sigma2$b0[j]

      theta.hist[t,(length(theta0) - n.lu + 1):length(theta0)][j] = rigamma(1,a,b)
    }

    #theta.hist[t,ncol(theta.hist)] = sd.t
    if(round(t/100) == (t/100)  & print.res){print(t)}
  }


  RET = list(theta = theta.hist, z= z.hist, p = p.hist, beta = beta.hist, priors = priors, random = random.hist, randomCov = randomb.hist, Phi = phi.hist)

  return(RET)

}