R/diagPlot.R

In ick003/SpTMixture: Mixture models with spatio-temporal dependencies

#' Diagnostic plot for Gibbs output
#'
#' Diagnostic plot for Gibbs output
#' @param SPTMresobj Output from estimGibbs.
#' @keywords diagnostic plot
#' @export
#' @examples
#' data("wq_analysis_week2")
#' SPTMData(wq.raw.obs, frequency = "quarter")
diagPlot = function(SPTMresobj, par = "theta", plot.cluster = NULL, q = c(0.01,0.99),...){


  theta = SPTMresobj$Result$theta
  p.sites = SPTMresobj$Result$latent
  priors = SPTMresobj$GibbsOut$priors

  theta.hist = SPTMresobj$GibbsOut$theta

  beta.hist = SPTMresobj$GibbsOut$beta

  random.hist = SPTMresobj$GibbsOut$random

  season.sites = SPTMresobj$seasonSites
  z.hist = SPTMresobj$GibbsOut$z

  idx.mean = 1:ncol(p.sites)
  idx.season = seq_len(ncol(p.sites)*dim(season.sites)[2]) + max(idx.mean)
  if(length(dim(season.sites))==2){
    idx.season = seq_len(ncol(p.sites)*1) + max(idx.mean)
  }
  idx.sigma = (length(theta) - ncol(p.sites) + 1):length(theta)

  idx.burn = round(.3*nrow(theta.hist)):nrow(theta.hist)


  if(par=="theta"){
  dim.plot = c(1 + dim(season.sites)[2], ncol(p.sites))

  if(length(dim(season.sites))==2){
    dim.plot = c(2 + 1, ncol(p.sites))
  }
  if(!is.null(plot.cluster)){
      dim.plot = c(1 + dim(season.sites)[2],length(plot.cluster))
    }

  par(mfrow=dim.plot, mar=c(2,2,2,1))

  idx.plot = 1:length(theta)
  if(!is.null(plot.cluster)){
    idx.plot = NULL
    for(j in plot.cluster){
      idx.plot = c(idx.plot,j + 0:(dim(season.sites)[2])*ncol(p.sites))
    }
    idx.plot = sort(idx.plot)
  }

  for(i in idx.plot){

    if(i <= ncol(p.sites)*dim(season.sites)[2]){

      minim = min(quantile(theta.hist[idx.burn,i], q[1]), qnorm(q[1], priors$theta$m0[i], sqrt(priors$theta$s0[i])), na.rm=T)
      maxim = max(quantile(theta.hist[idx.burn,i], q[2]), qnorm(q[2], priors$theta$m0[i], sqrt(priors$theta$s0[i])), na.rm=T)
      xx = seq(minim, maxim, length.out = 100)
      d.prior =  dnorm(xx, priors$theta$m0[i], sqrt(priors$theta$s0[i]))
    }else{
      minim = min(quantile(theta.hist[idx.burn,i],q[1]), qigamma(q[1], priors$sigma2$a0[i-ncol(p.sites)*dim(season.sites)[2]], priors$sigma2$b0[i-ncol(p.sites)*dim(season.sites)[2]]))
      maxim = max(quantile(theta.hist[idx.burn,i], q[2]), qigamma(q[2], priors$sigma2$a0[i-ncol(p.sites)*dim(season.sites)[2]], priors$sigma2$b0[i-ncol(p.sites)*dim(season.sites)[2]]))
      xx = seq(minim, maxim, length.out = 100)
      d.prior =  densigamma(xx, priors$sigma2$a0[i-ncol(p.sites)*dim(season.sites)[2]], priors$sigma2$b0[i-ncol(p.sites)*dim(season.sites)[2]])
    }
    hist(theta.hist[idx.burn,i], col = "lightblue", border = NA, freq=F, xlim = c(minim, maxim),...)
    points(xx, d.prior, type = "l", lwd = 2, col = "darkgreen")
  }

  }
  if(par=="z"){
    n.sites = nrow(p.sites)
    if(floor(sqrt(n.sites)) == sqrt(n.sites)){
      dim.plot = c(sqrt(n.sites), sqrt(n.sites))
    }else{
      dim.plot = c(floor(sqrt(n.sites)+1), floor(sqrt(n.sites)+1))
      reduction = floor((prod(dim.plot) - n.sites) / dim.plot[1])
      dim.plot = dim.plot - c(reduction, 0)
    }

    if(!is.null(plot.cluster)){
      dim.plot = c(1 + dim(season.sites)[2],length(plot.cluster))
    }


    par(mfrow=dim.plot, mar=c(2,2,2,1))
    idx.plot = 1:n.sites
    if(!is.null(plot.cluster)){
      idx.plot = plot.cluster
    }

    for(i in idx.plot){
      z.i = z.hist[i,,]
      xx = 1:ncol(p.sites)
      d.prior =  priors$z$a[i,]
      data= table(apply(z.i[,idx.burn], 2, which.max)) / length(idx.burn)
      plot(data, col = "lightblue", lwd = 20, ylim = c(0,min(2*max(data),1)),xlim = range(xx),...)
      points(xx, d.prior, type = "l", lwd = 2, col = "darkgreen")
    }
  }
  if(par == "paired-theta"){

    dim.plot = c(length(theta)-ncol(p.sites),length(theta)-ncol(p.sites))

    if(!is.null(plot.cluster)){
      dim.plot = c(dim(season.sites)[2]+1,dim(season.sites)[2]+1)
    }

    par(mfrow=dim.plot, mar=c(2,1,1,1))

    idx.theta = 1:(length(theta)-ncol(p.sites))

    if(!is.null(plot.cluster)){
      idx.theta = NULL
      for(j in plot.cluster){
        idx.theta = c(idx.theta,j + 0:(dim(season.sites)[2])*ncol(p.sites))
      }
      idx.theta = sort(idx.theta)
    }


    for(kk in idx.theta){
      for(kkk in idx.theta){
        if(kk < kkk){
          plot(1, type="n", axes=F, xlab="", ylab=""); text(1,1,labels = round(cor(theta.hist[idx.burn,kk], theta.hist[idx.burn,kkk]),2))
        }
        if(kk == kkk){
          hist(theta.hist[idx.burn,kk], col = "lightblue", border = NA)#,...)
        }
        if(kk > kkk){
          sm2d = kde2d(theta.hist[idx.burn,kkk], theta.hist[idx.burn,kk],n=100)
          contour(sm2d, col = "lightblue", lwd = 2);
          abline(v = mean(theta.hist[idx.burn,kkk], na.rm=T), lty = 2, lwd=2);
          abline(h = mean(theta.hist[idx.burn,kk], na.rm=T), lty=2, lwd = 2);
        }
      }
    }


  }
  if(par=="beta"){
    dim.plot = c(1, 2)

    par(mfrow=dim.plot, mar=c(2,2,2,1))

    hist(beta.hist[idx.burn],col = "lightblue", border = NA,
         ...)
    plot(beta.hist, cex=0.2)
    abline(h = mean(beta.hist[idx.burn]), col = "red", lty = 1, lwd=2)
    abline(h = quantile(beta.hist[idx.burn], 0.05), col = "red", lty = 2, lwd=2)
    abline(h = quantile(beta.hist[idx.burn], 0.95), col = "red", lty = 2, lwd=2)
  }
  if(par=="random-effect"){
    n.sites = nrow(p.sites)
    if(floor(sqrt(n.sites)) == sqrt(n.sites)){
      dim.plot = c(sqrt(n.sites), sqrt(n.sites))
    }else{
      dim.plot = c(floor(sqrt(n.sites)+1), floor(sqrt(n.sites)+1))
      reduction = floor((prod(dim.plot) - n.sites) / dim.plot[1])
      dim.plot = dim.plot - c(reduction, 0)
    }

    if(!is.null(plot.cluster)){
      dim.plot = c(1 + dim(season.sites)[2],length(plot.cluster))
    }


    par(mfrow=dim.plot, mar=c(2,2,2,1))
    idx.plot = 1:n.sites
    if(!is.null(plot.cluster)){
      idx.plot = plot.cluster
    }
    dim.s = dim(season.sites)[2]
    color = c("darkgreen","orange")
    if(dim.s > 2){color = brewer.pal(dim.s, "Set2")}
    color = colorRampAlpha(color, n = dim.s, alpha = 0.5)

    for(i in idx.plot){
      idx.col = i + (0:(dim(season.sites)[2]-1))*n.sites

      random.i = random.hist[,idx.col[1]]
      hist(random.i[idx.burn], col = color[1], border = NA, freq=F, xlim = range(random.hist),...)

      if(length(idx.col)>1){
        for(ii in 2:length(idx.col)){
          random.i = random.hist[,idx.col[ii]]
          hist(random.i[idx.burn], col = color[ii], border = NA, freq=F, add=T, ...)
        }
      }
      abline(v = 0, col = "darkgreen", lwd = 2)
    }
  }
  if(par=="tempBasis"){

    season.lu = updateSeason(season.sites, p.sites)
    season.lu.mod = season.lu * matrix(theta[(ncol(p.sites) + 1):(length(theta) - ncol(p.sites))],
                                       nrow = nrow(season.lu), ncol = ncol(season.lu), byrow = TRUE)
    class(season.lu.mod) <- 'SPTMseason'
    plot(season.lu.mod)

  }
  if(par == "spatial"){
      Coords = SPTMresobj$data$coord
      dd = deldir(Coords[,1], Coords[,2],suppressMsge=TRUE)
      tile.dd = tile.list(dd)
      color = brewer.pal(ncol(p.sites), "Set2")
      bor.fill = color[apply(p.sites, 1, which.max)]
      par(mfrow = c(1,1), mar = c(5,4,4,1))
      plot(Coords[,1], Coords[,2], xlab = "x", ylab = "y", col = bor.fill, pch=16, cex=2)
      plot(tile.dd, add=T, fillcol = NA, border = 1,close = FALSE)
      points(Coords[,1], Coords[,2], col = bor.fill, pch=16, cex=2)
  }
  if(par=="covariates"){

    randomCov = SPTMresobj$GibbsOut$randomCov

    nCov =ncol(randomCov) / dim(season.sites)[2]

    dim.plot = c(dim(season.sites)[2], nCov)

    if(length(dim(season.sites))==2){
      dim.plot = c(2 + 1, nCov)
    }
    if(!is.null(plot.cluster)){
      dim.plot = c(1 + dim(season.sites)[2],length(plot.cluster))
    }

    par(mfrow=dim.plot, mar=c(2,2,2,1))

    idx.plot = 1:ncol(randomCov)
    if(!is.null(plot.cluster)){
      idx.plot = NULL
      for(j in plot.cluster){
        idx.plot = c(idx.plot,j + 0:(dim(season.sites)[2])*nCov)
      }
      idx.plot = sort(idx.plot)
    }

    for(i in idx.plot){

        minim = min(quantile(randomCov[idx.burn,i], q[1]), qnorm(q[1], 0, sqrt(priors$cov$s0[i])), na.rm=T)
        maxim = max(quantile(randomCov[idx.burn,i], q[2]), qnorm(q[2], 0, sqrt(priors$cov$s0[i])), na.rm=T)
        xx = seq(minim, maxim, length.out = 100)
        d.prior =  dnorm(xx, 0, sqrt(priors$cov$s0[i]))

      hist(randomCov[idx.burn,i], col = "lightblue", border = NA, freq=F, xlim = c(minim, maxim),...)
      points(xx, d.prior, type = "l", lwd = 2, col = "darkgreen")
    }
  }
}