SpTMixture: Mixture models with spatio-temporal dependencies

Documented in plot.sptmRes

#' Update the seasons for LU
#'
#' Update the seasons for LU
#' @param season.sites the season calculated for each site.
#' @param p.sites the composition of each site.
#' @keywords updating season
#' @export
#' @examples
#' data("wq_analysis_week2")
#' SPTMData(wq.raw.obs, frequency = "quarter")
plot.sptmRes = function(SPTMresobj, type = "prediction", color = NULL, cex.axis=NULL, n_x_lab = 7){

  theta = SPTMresobj$Result$theta
  p.sites = SPTMresobj$Result$latent
  df = SPTMresobj$data$data

  n.colors = ncol(p.sites)

  SiteID = unique(df$obs.data$ID)

  n.sites = length(SiteID)

  n.tempFun = 1 * as.numeric(length(dim(SPTMresobj$data$data$season))==2) +
    (1-as.numeric(length(dim(SPTMresobj$data$data$season))==2)) * dim(SPTMresobj$data$data$season)[2]

  if(length(theta) < (ncol(p.sites) * (2 + n.tempFun))){theta = c(theta, rep(theta[length(theta)], ncol(p.sites)-1))}

  random = rep(0, n.sites*n.tempFun)
 if(!is.null(SPTMresobj$GibbsOut$random)){
   random = colMeans(SPTMresobj$GibbsOut$random)
 }
  Xcov = SPTMresobj$data$covariates
  if(is.null(Xcov)){
    Xcov = matrix(0, ncol = 1, nrow = n.sites)
  }
  randomCov = rep(0, ncol(Xcov)*n.tempFun)
  if(!is.null(SPTMresobj$GibbsOut$randomCov)){
    randomCov = colMeans(SPTMresobj$GibbsOut$randomCov)
  }
  season.sites = SPTMresobj$seasonSites

  phi = SPTMresobj$GibbsOut$Phi
  idx.season = SPTMresobj$data$data$list.idx

  if(!is.null(SPTMresobj$data$data$basis)){
    y.p = SPTMresobj$data$data$basis %*% apply(phi, 1:2, mean)
    phiM = apply(phi, 1:2, mean)
    season.lu = NULL
    for(ii in 1:length(idx.season)){
      if(length(idx.season[[ii]])==1){
        season.lu = cbind(season.lu, matrix(matrix(SPTMresobj$data$data$basis[,idx.season[[ii]]], ncol=1) %*% matrix(phiM[idx.season[[ii]],], nrow=1), nrow = nrow(SPTMresobj$data$data$basis)))
      }else{
        season.lu = cbind(season.lu, SPTMresobj$data$data$basis[,idx.season[[ii]]] %*% phiM[idx.season[[ii]],])
      }
    }
  }else{
    y.p = meanFunc(theta, n.mixt=ncol(p.sites), p.sites=p.sites, season.sites = season.sites)
  }


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

  if(!is.null(SPTMresobj$GibbsOut)){
    theta[idx.sigma] = log(sqrt(theta[idx.sigma]))
  }



  n.mixt = ncol(p.sites)

  if(type == "residuals"){

    par(mfrow =c(2,2),mar = c(5,4,4,1))

    x = sort(unique(df$obs.data$date))
    plot(x, rep(0, length(x)), type="n", main = "Residuals over time",
         ylim = c(-max(abs(df$obs.data$obs), na.rm=T),max(abs(df$obs.data$obs), na.rm=T)))
    res=NULL
    for(i in 1:length(SiteID)){
      obs = df$obs.data$obs[df$obs.data$ID %in% SiteID[i]]
      idx.k = which.max(p.sites[i,])
      mat = matrix(p.sites[i,], ncol=1)
      idx.na = which(is.na(colSums(y.p)))
      idx.a = i+n.sites*(0:(n.tempFun-1))
      n.season = length(randomCov) / ncol(Xcov)
      Xcovi = matrix(0, nrow = n.season, ncol = ncol(Xcov)*n.season)
        idx.r = rep((0)*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,]
      idx.season = idx.k + n.colors*(0:(n.tempFun-1))
      pred= y.p %*% mat + season.lu[,idx.season] %*% (random[idx.a] + (Xcovi %*% randomCov))
      if(length(idx.na)>0){pred= y.p[,-idx.na] %*% mat[-idx.na] - season.lu[,idx.season] %*% random[idx.a]}
      dates = df$obs.data$date[df$obs.data$ID %in% SiteID[i]]
      res = rbind(res,cbind(obs-pred, idx.k))
      points(x,obs-pred)
    }
    plot(0, type="n", xlim = range(df$obs.data$obs, na.rm=T),
         ylim  = range(df$obs.data$obs, na.rm=T), main = "Predicted vs Observed")
    abline(a=0,b=1)
    for(i in 1:length(SiteID)){
      obs = df$obs.data$obs[df$obs.data$ID %in% SiteID[i]]
      idx.k = which.max(p.sites[i,])
      mat = matrix(p.sites[i,], ncol=1)
      idx.na = which(is.na(colSums(y.p)))
      idx.a = i+n.sites*(0:(n.tempFun-1))
      n.season = length(randomCov) / ncol(Xcov)
      Xcovi = matrix(0, nrow = n.season, ncol = ncol(Xcov)*n.season)
      idx.r = rep((0)*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,]
    idx.season = idx.k + n.colors*(0:(n.tempFun-1))
    pred= y.p %*% mat - season.lu[,idx.season] %*% (random[idx.a] + Xcovi %*% randomCov)
      if(length(idx.na)>0){pred= y.p[,-idx.na] %*% mat[-idx.na] - season.lu[,idx.season] %*% (random[idx.a] +  Xcovi %*% randomCov)}
      dates = df$obs.data$date[df$obs.data$ID %in% SiteID[i]]
      points(obs,pred)
    }
    hist(res[,1], main = "Histogram of residuals")
    res = as.data.frame(res)
    boxplot(V1~ idx.k, data = res,main="Residuals per cluster")
  }

  if(type == "prediction"){

  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)
  }

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

  for(i in 1:length(SiteID)){

    obs = df$obs.data$obs[df$obs.data$ID %in% SiteID[i]]
    idx.k = which.max(p.sites[i,])
    mat = matrix(p.sites[i,], ncol=1)
    idx.na = which(is.na(colSums(y.p)))
    idx.a = i+n.sites*(0:(n.tempFun-1))
    idx.season = idx.k + n.colors*(0:(n.tempFun-1))
    pred= y.p %*% mat - season.lu[,idx.season] %*% random[idx.a]
    if(length(idx.na)>0){pred= y.p[,-idx.na] %*% mat[-idx.na] - season.lu[,idx.season] %*% random[idx.a]}
    dates = df$obs.data$date[df$obs.data$ID %in% SiteID[i]]
    ci = cbind(pred - 2*exp(theta[idx.sigma][idx.k]),pred + 2*exp(theta[idx.sigma][idx.k]) )
    plot(dates,obs,main = SiteID[i],
         type="l", ylim = range(ci, na.rm=T))
    polygon(x=c(rev(dates), (dates)), y = c(rev(ci[,1]), (ci[,2])),
            border=NA, col = "grey", density = 100)
    points(dates,pred, type="l", col="red")
    points(df$obs.data$date[df$obs.data$ID %in% SiteID[i]],df$obs.data$obs[df$obs.data$ID %in% SiteID[i]])
  }
  }

  if(type == "tempBasis"){

    if(floor(sqrt(n.mixt)) == sqrt(n.mixt)){
      dim.plot = c(sqrt(n.mixt), sqrt(n.mixt))
    }else{
      dim.plot = c(floor(sqrt(n.mixt)+1), floor(sqrt(n.mixt)+1))
      reduction = floor((prod(dim.plot) - n.mixt) / dim.plot[1])
      dim.plot = dim.plot - c(reduction, 0)
    }
    dim.plot = c(n.mixt,2)
    par(mfrow=dim.plot, mar=c(3,3,2,1))

    x = sort(unique(df$obs.data$date))


    if(!is.null(SPTMresobj$data$data$basis)){

      phiM = apply(phi, 1:2, mean)

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

      dim.s = dim(season.sites)
      if(length(dim.s)==3){dim.s = dim.s[2]
      }else{dim.s=1}

      if(is.null(color)){color = brewer.pal(dim.s, "Set2")}

      for(i in 1:ncol(p.sites)){
        if(!is.na(sum(y.p[,i]))){
          plot(x,y.p[,i], type="l", main = colnames(SPTMresobj$data$lu)[i], bty = "n", axes = F, cex.main = 0.9)
          axis(2)
          axis(1,at = x[seq.int(1, length(x),length.out = n_x_lab)],labels =FALSE)
          text(x=x[seq.int(1, length(x),length.out = n_x_lab)], y=par()$usr[3]-0.1*(par()$usr[4]-par()$usr[3]),
               labels=x[seq.int(1, length(x),length.out = n_x_lab)], srt=30, adj=0.75, xpd=TRUE, cex = cex.axis)
          idx = i + n.mixt*(0:(dim.s-1))
          plot(1, type="n", xlab = "Date", ylab = "", xlim = c(0,nrow(season.l)+1), ylim = 1.05*range(season.l[,idx[1:dim.s]], na.rm=T),
               main = "", axes = F,cex.main = 1)
          axis(2)
          axis(1,at = seq.int(1, length(x),length.out = n_x_lab),labels =FALSE)
          text(x=seq.int(1, length(x),length.out = n_x_lab), y=par()$usr[3]-0.1*(par()$usr[4]-par()$usr[3]),
               labels=x[seq.int(1, length(x),length.out = n_x_lab)], srt=30, adj=0.75, xpd=TRUE,cex = cex.axis)
          abline(h  =0, col = "lightgrey")
          for(j in 1:dim.s){
            idx.j = idx[j]
            if(j %in% c(1,2)){points(season.l[,idx.j] * SPTMresobj$Result$theta[idx.j],type = "l", col = color[j], lwd=2)}
            if(j>2){points(season.l[,idx.j] * SPTMresobj$Result$theta[idx.j],type = "l", col = color[j], lwd=2)}
          }
          legend("bottomleft", legend = paste0(colnames(SPTMresobj$seasonSites),c("","","ly seasonality")), bty = 'n',
                 lty = rep(1,dim.s), lwd = rep(2,dim.s), col = color, cex = 0.8)
        }
      }
    }else{

      season.lu = updateSeason(season.sites, p.sites = p.sites)
      dim.s = dim(season.sites)
      if(length(dim.s)==3){dim.s = dim.s[2]
      }else{dim.s=1}

      if(is.null(color)){color = brewer.pal(dim.s, "Set2")}

      for(i in 1:ncol(p.sites)){
        if(!is.na(sum(y.p[,i]))){
        plot(x,y.p[,i], type="l", main = colnames(SPTMresobj$data$lu)[i], bty = "n", axes = F, cex.main = 0.9)
        axis(2)
        axis(1,at = x[seq.int(1, length(x),length.out = n_x_lab)],labels =FALSE)
        text(x=x[seq.int(1, length(x),length.out = n_x_lab)], y=par()$usr[3]-0.1*(par()$usr[4]-par()$usr[3]),
             labels=x[seq.int(1, length(x),length.out = n_x_lab)], srt=30, adj=0.75, xpd=TRUE, cex = cex.axis)
        idx = i + n.mixt*(0:(dim.s-1))
        plot(1, type="n", xlab = "Date", ylab = "", xlim = c(0,nrow(season.lu)+1), ylim = 1.05*range(y.p[,i], na.rm=T),
             main = "", axes = F,cex.main = 1)
        axis(2)
        axis(1,at = seq.int(1, length(x),length.out = n_x_lab),labels =FALSE)
        text(x=seq.int(1, length(x),length.out = n_x_lab), y=par()$usr[3]-0.1*(par()$usr[4]-par()$usr[3]),
             labels=x[seq.int(1, length(x),length.out = n_x_lab)], srt=30, adj=0.75, xpd=TRUE,cex = cex.axis)
        abline(h  =0, col = "lightgrey")
        for(j in 1:dim.s){
          idx.j = idx[j]
          if(j == 1){points(season.lu[,idx.j] * SPTMresobj$Result$theta[idx.j],type = "l", col = color[j], lwd=2)}
          if(j>1){points(season.lu[,idx.j] * SPTMresobj$Result$theta[idx.j] + SPTMresobj$Result$theta[idx[1]],type = "l", col = color[j], lwd=2)}
        }
        legend("bottomleft", legend = paste0(colnames(SPTMresobj$seasonSites),c("","","ly seasonality")), bty = 'n',
               lty = rep(1,dim.s), lwd = rep(2,dim.s), col = color, cex = 0.8)
      }
    }
    }



  }

}