#
# Bayesian Calibration
# For STAT8810
# Fall, 2017
# M.T. Pratola


# Bayesian Implementation of Calibration Model
#
# Data:      yf(xi) = eta(xi,theta) + delta(xi) + ef(xi), ef(xi)~N(0,1/lambdaf), i=1,...,nf ; xi=1xk vector
#         yc(xj,tj) = eta(xj,tj) + ez(xj,tj) , ez(xj,tj)~N(0,1/lambdad), j=1,...,nc ; xj=1xK vector, tj=1xL vector
#
#
# Likelihood: L(Y|.) p. 1/(2pi*det(E))exp(-(Y')(Einv)(Y)) ; Y=[Yf | Yc ]
# where,
#    E=1/lambdaz*[  Rff  |  Rfc  ]   [ 1/lambdaf*Inf |       0       ]
#                [  Rcf  |  Rcc  ] + [      0        |       0       ]
#
#                  1/lambdadel*[ Rdel | 0 ]
#                +             [  0   | 0 ]
#
# Priors:
#         pi(lambdaf) p. lambdaf^(af-1)exp(-bf*lambdaf)                           [Gamma prior]
#         pi(lambdaz) p. lambdaz^(az-1)exp(-bz*lambdaz)                           [Gamma prior]
#         pi(lambdadel) p. lambdadel^(adel-1)exp(-bdel*lambdadel)                 [Gamma prior]
#         pi(rhoz_k) p. rhoz_k^4(1-rhoz_k)^0                                      [Beta prior] } correlation matrix
#         pi(psiz_l) p. psiz_l^4(1-psiz_l)^0                                      [Beta prior] }     parameters
#         pi(rhodel_k) p. rhodel_k^4(1-rhodel_k)^0                                [Beta prior] }
#         pi(theta_l) p. 1/(2pi*sigma2theta)exp(-1/(2*sigma2theta)(theta_l-.5)^2) [Normal prior]
#      OR pi(theta_l) p. I(0,1)                                                   [Uniform prior]
#
# Posterior:
# pi(.|Y) p. L(Y|.)pi(lambdaf)pi(lambdad)pi(lambdaz)pi(lambdadel)pi(rhoz_1)*...*pi(rhoz_K)pi(rhodel_1)*...*pi(rhodel_K)pi(psiz_1)*...*pi(psiz_L)pi(theta_1)*...(pi(theta_L)
#
# We fit the model using MCMC with adaptive stepwidths for the uniform proposal distributions.
#
#


# logposterior:
logp<-function(y,lambdaf,lambdaz,lambdadel,rhoz,psiz,rhodel,l.dez,l.del,theta,pi,ninfo)
{
  af=pi$af; bf=pi$bf; az=pi$az; bz=pi$bz; adel=pi$adel; bdel=pi$bdel; rhoa=pi$rhoa; rhob=pi$rhob;
  if(any(rhoz<=0) || any(rhoz>=1) || any(psiz<=0) || any(psiz>=1) || any(rhodel<=0) || any(rhodel>=1) || lambdaz<=0 || lambdaf<=0 || lambdadel<0)
     logposterior=-Inf
  else
  {
    nf=ninfo$nf
    k=length(rhoz)
    l=length(psiz)
    N=length(y)
    n=N-nf #number of computer model runs
    In=diag(N)
    R=rhogeodacecormat(l.dez,c(rhoz,psiz),alpha=2)$R
    Rdel=rhogeodacecormat(l.del,rhodel,alpha=2)$R
    et=c(rep(1/lambdaf,nf),rep(0,n))
    E=1/lambdaz*R+In*et
    E[1:nf,1:nf]=E[1:nf,1:nf]+1/lambdadel*Rdel
    cholE=chol(E)
    Einv=chol2inv(cholE)
    logdetE=2*sum(log(diag(cholE)))
  
    logpz=-1/2*logdetE-1/2*t(y)%*%Einv%*%(y)
    logplambdaf=(af-1)*log(lambdaf)-bf*lambdaf
    logplambdaz=(az-1)*log(lambdaz)-bz*lambdaz
    logplambdadel=(adel-1)*log(lambdadel)-bdel*lambdadel
    logprhoz=logppsiz=logprhodel=logptheta=0
    for(i in 1:k) logprhoz=logprhoz+(pi$rhoa[i]-1)*log(rhoz[i])+(pi$rhob[i]-1)*log(1-rhoz[i])
    for(i in 1:l) logppsiz=logppsiz+(pi$psia[i]-1)*log(psiz[i])+(pi$psib[i]-1)*log(1-psiz[i])
    for(i in 1:k) logprhodel=logprhodel+(pi$rdla[i]-1)*log(rhodel[i])+(pi$rdlb[i]-1)*log(1-rhodel[i])
    for(i in 1:l) logptheta=logptheta+log(1-any(theta<0))+log(1-any(theta>1))  # uniform prior
#    for(i in 1:l) logptheta=logptheta-log(.25)-1/2*((theta[i]-.5)/.25)^2      # normal(.5,sd=.25) prior
#    logprhoz=logppsiz=logprhodel=0 # If I prefer uniform priors for the correlation params, enable this line.

    logposterior=logpz+logplambdaf+logplambdaz+logplambdadel+logprhoz+logppsiz+logprhodel+logptheta
  }

  logposterior
}

# Gibbs/MCMC Sampler
# y - [field observations | computer model runs ]
# l.dez - distances in my format.  Space/time first, then calib dimensions last.
# nf - the first nf entires of y form the 'field observations' component of the vector
# ncdims - the dimensionality of the calibration parameter.  Correspondingly, the last ncaldims entries in l.dez
#            correspond to the calibraiton parameter dimensions.
# N - number of Gibbs/MCMC iterations for the sampler
#
# with my usual prios, i use: rr=.03,rp=.1,re=.05,rf=50,rz=.5,rl=5e8,rd=10000,rth=.15
# mh=list(rr=.07,rp=.1,re=.05,rf=50,rz=.5,rl=5e8,rth=.15)
calibrate<-function(y,l.dez,ninfo,N,pi,mh,last=1000)#,rr=.03,rp=.1,re=.1,rf=50,rz=.5,rl=.7,rth=.15)
{
  if(N<2000) stop("N>=2000 please\n")

  rr=mh$rr; rp=mh$rp; re=mh$re; rf=mh$rf; rz=mh$rz; rl=mh$rl; rth=mh$rth
  nf=ninfo$nf
  l=ncdims=ninfo$ncdims    #number of calibration parameter dimesions, these are the last l dimensions in l.dez
  derivs=ninfo$derivs  
  last=last-1
  k=length(l.dez)-ncdims   #number of space/time parameter dimensions, these are the first k dimsions in l.dez
  nn=nrow(l.dez$l1$m1)     #total number of locations for comp model runs and field observations
  n=nn-nf                  #number of locations of computer model runs only

  l.del=l.dez[1:k] # for the discrepancy, only need the space-time part of the design
  for(i in 1:k)    # and only at the nf field locations.
    l.del[[i]][[1]]=l.del[[i]][[1]][1:nf,1:nf]
  
  # based on the above, y=[ (nf x 1) vector of field observations | (n x 1) vector of computer model runs]
  # check:
  if(length(y)!=(n+nf)) stop("\nThe length(y) != n+nf !\n")
  ntot=length(y) #should equal n+nf
  ncomp=ntot-nf

  draw.lambdaf=rep(NA,N)
  draw.lambdaz=rep(NA,N)
  draw.lambdadel=rep(NA,N)
  draw.lambdad=rep(NA,N)
  draw.rhoz=matrix(NA,nrow=N,ncol=k)
  draw.psiz=matrix(NA,nrow=N,ncol=l)
  draw.rhodel=matrix(NA,nrow=N,ncol=k)
  draw.th=matrix(NA,nrow=N,ncol=l)
  rr=rep(rr,k)
  rp=rep(rp,l)
  re=rep(re,k)
  rth=rep(rth,l)
  
  #initial guesses
  draw.lambdaf[1]=10 #100
  draw.lambdaz[1]=1
  draw.lambdadel[1]=1 #1e12
  draw.rhoz[1,]=rep(.5,k)
  draw.psiz[1,]=rep(.5,l)
  draw.rhodel[1,]=rep(.5,k) #.001
  draw.th[1,]=rep(.5,l)
  
  #accept/reject ratio trackers
  accept.rhoz=rep(0,k)
  accept.psiz=rep(0,l)
  accept.rhodel=rep(0,k)
  accept.lambdaf=0
  accept.lambdaz=0
  accept.lambdadel=0
  accept.th=rep(0,l)

  alpha=2
  
  I=diag(ntot)
  one=rep(1,ntot)
  rhoz.new=rep(NA,k)
  psiz.new=rep(NA,l)

  lastadapt=0
  
  l.dezprev=l.dez
  l.deznew=l.dez
  for(j in 1:l)
  {
    l.dezprev[[k+j]][[1]][1:nf,]=l.dezprev[[k+j]][[1]][1:nf,]+draw.th[1,j]
    l.dezprev[[k+j]][[1]][,1:nf]=l.dezprev[[k+j]][[1]][,1:nf]-draw.th[1,j]
  }

  
  cat("\n Generalized Bayesian Gaussian Process Response Calibration with Derivatives Model")
  cat("\n The last ",last," samples from the posterior will be reported")
  cat("\n The stepwidth for uniform corr. param proposal distn (sp/tm dim) is rr=",rr)
  cat("\n The stepwidth for uniform corr. param proposal distn (calib dim) is rp=",rp)
  cat("\n The stepwidth for uniform corr. param proposal distn (discrpncy) is re=",re)
  cat("\n The stepwidth for uniform field error proposal distn is rf=",rf)
  cat("\n The stepwidth for uniform variance proposal distn is rz=",rz)
  cat("\n The stepwidth for uniform variance proposal distn (discrpncy) is rl=",rl)
  cat("\n The stepwidth for uniform calibration param proposal distn is rth=",rth)
  cat("\n Prior params: af=",pi$af," bf=",pi$bf," az=",pi$az," bz=",pi$bz, " adel=",pi$adel," bdel=",pi$bdel)
  cat("\n ----------------------------------------------------------------\n\n\n")

  for(i in 2:N)
  {
    draw.rhoz[i,]=draw.rhoz[i-1,]
    for(j in 1:k)
    {
      rhoz.new=draw.rhoz[i,]
      rhoz.new[j]=runif(1,draw.rhoz[i-1,j]-rr[j],draw.rhoz[i-1,j]+rr[j])
      a=min(0,logp(y,draw.lambdaf[i-1],draw.lambdaz[i-1],draw.lambdadel[i-1],rhoz.new,
                   draw.psiz[i-1,],draw.rhodel[i-1,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo)
             -logp(y,draw.lambdaf[i-1],draw.lambdaz[i-1],draw.lambdadel[i-1],draw.rhoz[i,],
                   draw.psiz[i-1,],draw.rhodel[i-1,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo))
      if(log(runif(1))<a)
      {
        draw.rhoz[i,j]=rhoz.new[j]
        accept.rhoz[j]=accept.rhoz[j]+1
      }
    }

    draw.psiz[i,]=draw.psiz[i-1,]
    for(j in 1:l)
    {
      psiz.new=draw.psiz[i,]
      psiz.new[j]=runif(1,draw.psiz[i-1,j]-rp[j],draw.psiz[i-1,j]+rp[j])
      a=min(0,logp(y,draw.lambdaf[i-1],draw.lambdaz[i-1],draw.lambdadel[i-1],
                   draw.rhoz[i,],psiz.new,draw.rhodel[i-1,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo)
             -logp(y,draw.lambdaf[i-1],draw.lambdaz[i-1],draw.lambdadel[i-1],
                   draw.rhoz[i,],draw.psiz[i,],draw.rhodel[i-1,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo))
      if(log(runif(1))<a)
      {
        draw.psiz[i,j]=psiz.new[j]
        accept.psiz[j]=accept.psiz[j]+1
      }
    }

    draw.rhodel[i,]=draw.rhodel[i-1,]
    for(j in 1:k)
    {
      rhodel.new=draw.rhodel[i,]
      rhodel.new[j]=runif(1,draw.rhodel[i-1,j]-re[j],draw.rhodel[i-1,j]+re[j])
      a=min(0,logp(y,draw.lambdaf[i-1],draw.lambdaz[i-1],draw.lambdadel[i-1],draw.rhoz[i,],
                   draw.psiz[i,],rhodel.new,l.dezprev,l.del,draw.th[i-1,],pi,ninfo)
             -logp(y,draw.lambdaf[i-1],draw.lambdaz[i-1],draw.lambdadel[i-1],draw.rhoz[i,],
                   draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo))
      if(log(runif(1))<a)
      {
        draw.rhodel[i,j]=rhodel.new[j]
        accept.rhodel[j]=accept.rhodel[j]+1
      }
    }

    
    draw.lambdaf[i]=draw.lambdaf[i-1]
    lambdaf.new=runif(1,draw.lambdaf[i-1]-rf,draw.lambdaf[i-1]+rf)
    a=min(0,logp(y,lambdaf.new,draw.lambdaz[i-1],draw.lambdadel[i-1],draw.rhoz[i,],
                 draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo)
           -logp(y,draw.lambdaf[i],draw.lambdaz[i-1],draw.lambdadel[i-1],draw.rhoz[i,],
                 draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo))
    if(log(runif(1))<a)
    {
      draw.lambdaf[i]=lambdaf.new
      accept.lambdaf=accept.lambdaf+1
    }

    draw.lambdaz[i]=draw.lambdaz[i-1]
    lambdaz.new=runif(1,draw.lambdaz[i-1]-rz,draw.lambdaz[i-1]+rz)
    a=min(0,logp(y,draw.lambdaf[i],lambdaz.new,draw.lambdadel[i-1],draw.rhoz[i,],
                 draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo)
           -logp(y,draw.lambdaf[i],draw.lambdaz[i],draw.lambdadel[i-1],draw.rhoz[i,],
                 draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo))
    if(log(runif(1))<a)
    {
      draw.lambdaz[i]=lambdaz.new
      accept.lambdaz=accept.lambdaz+1
    }

    draw.lambdadel[i]=draw.lambdadel[i-1]
    lambdadel.new=runif(1,draw.lambdadel[i-1]-rl,draw.lambdadel[i-1]+rl)
    a=min(0,logp(y,draw.lambdaf[i],draw.lambdaz[i],lambdadel.new,draw.rhoz[i,],
                 draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo)
           -logp(y,draw.lambdaf[i],draw.lambdaz[i],draw.lambdadel[i],draw.rhoz[i,],
                 draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i-1,],pi,ninfo))
    if(log(runif(1))<a)
    {
      draw.lambdadel[i]=lambdadel.new
      accept.lambdadel=accept.lambdadel+1
    }
 
    
    draw.th[i,]=draw.th[i-1,]
    for(j in 1:l)
    {
      th.new=draw.th[i,]
      th.new[j]=runif(1,draw.th[i,j]-rth[j],draw.th[i,j]+rth[j])

      l.deznew=l.dez
      l.dezprev=l.dez
      for(jj in 1:l)
      {
        l.dezprev[[k+jj]][[1]][1:nf,]=l.dezprev[[k+jj]][[1]][1:nf,]+draw.th[i,jj]
        l.dezprev[[k+jj]][[1]][,1:nf]=l.dezprev[[k+jj]][[1]][,1:nf]-draw.th[i,jj]
        l.deznew[[k+jj]][[1]][1:nf,]=l.deznew[[k+jj]][[1]][1:nf,]+th.new[jj]
        l.deznew[[k+jj]][[1]][,1:nf]=l.deznew[[k+jj]][[1]][,1:nf]-th.new[jj]
      }
      a=min(0,logp(y,draw.lambdaf[i],draw.lambdaz[i],draw.lambdadel[i],draw.rhoz[i,],
              draw.psiz[i,],draw.rhodel[i,],l.deznew,l.del,th.new,pi,ninfo)
             -logp(y,draw.lambdaf[i],draw.lambdaz[i],draw.lambdadel[i],draw.rhoz[i,],
              draw.psiz[i,],draw.rhodel[i,],l.dezprev,l.del,draw.th[i,],pi,ninfo))
      if(log(runif(1))<a)
      {
        draw.th[i,j]=th.new[j]
        accept.th[j]=accept.th[j]+1
      }
    }


    l.dezprev=l.dez
    for(j in 1:l)
    {
      l.dezprev[[k+j]][[1]][1:nf,]=l.dezprev[[k+j]][[1]][1:nf,]+draw.th[i,j]
      l.dezprev[[k+j]][[1]][,1:nf]=l.dezprev[[k+j]][[1]][,1:nf]-draw.th[i,j]
    }
    
    cat(i/N*100," percent complete\r")
    # Adaptive MCMC stuff:
    # adapt the proposal step every N/20 iterations, but only for the first 50% of the iterations
    if(i%%(N/20)==0 && i<(N*.5+1))
    {
      rate.rhoz=accept.rhoz/(i-lastadapt)
      rate.psiz=accept.psiz/(i-lastadapt)
      rate.rhodel=accept.rhodel/(i-lastadapt)
      rate.lambdaf=accept.lambdaf/(i-lastadapt)
      rate.lambdaz=accept.lambdaz/(i-lastadapt)
      rate.lambdadel=accept.lambdadel/(i-lastadapt)
      rate.theta=accept.th/(i-lastadapt)
      cat("Adapting rates from ",rate.rhoz," ",rate.psiz," ",rate.rhodel," ",rate.lambdaf," ",
          rate.lambdaz," ",rate.lambdadel," ",rate.theta,"\n");
      for(j in 1:k)
        if(rate.rhoz[j]>.49 || rate.rhoz[j]<.39) rr[j]=rr[j]*rate.rhoz[j]/.44
      for(j in 1:l)
        if(rate.psiz[j]>.49 || rate.psiz[j]<.39) rp[j]=rp[j]*rate.psiz[j]/.44
      for(j in 1:k)
        if(rate.rhodel[j]>.49 || rate.rhodel[j]<.39) re[j]=re[j]*rate.rhodel[j]/.44
      if(rate.lambdaf>.49 || rate.lambdaf<.39) rf=rf*rate.lambdaf/.44
      if(rate.lambdaz>.49 || rate.lambdaz<.39) rz=rz*rate.lambdaz/.44
      if(rate.lambdadel>.49 || rate.lambdadel<.39) rl=rl*rate.lambdadel/.44
      for(j in 1:l)
        if(rate.theta[j]>.49 || rate.theta[j]<.39 && rate.theta[j]>0) rth[j]=rth[j]*rate.theta[j]/.44
      lastadapt=i
      accept.rhodel=accept.rhoz=rep(0,k)
      accept.psiz=accept.th=rep(0,l)
      accept.lambdaf=accept.lambdaz=accept.lambdadel=0
    }
  }
  cat("\n Complete.\n")
  rate.rhoz=accept.rhoz/(i-lastadapt)
  rate.psiz=accept.psiz/(i-lastadapt)
  rate.rhodel=accept.rhodel/(i-lastadapt)
  rate.lambdaf=accept.lambdaf/(i-lastadapt)
  rate.lambdaz=accept.lambdaz/(i-lastadapt)
  rate.lambdadel=accept.lambdadel/(i-lastadapt)
  rate.theta=accept.th/(i-lastadapt)
  cat("[rate.rhoz=",rate.rhoz,"]\n")
  cat("[rate.psiz=",rate.psiz,"]\n")
  cat("[rate.rhodel=",rate.rhodel,"]\n")
  cat("[rate.lambdaf=",rate.lambdaf,"]\n")
  cat("[rate.lambdaz=",rate.lambdaz,"]\n")
  cat("[rate.lambdadel=",rate.lambdadel,"]\n")
  cat("[rate.theta=",rate.theta,"]\n\n")
  
  return(list(y=y,lambdaf=draw.lambdaf[(N-last):N],lambdaz=draw.lambdaz[(N-last):N],
              lambdadel=draw.lambdadel[(N-last):N],rhoz=as.matrix(draw.rhoz[(N-last):N,]),
              psiz=as.matrix(draw.psiz[(N-last):N,]),rhodel=as.matrix(draw.rhodel[(N-last):N,]),
              theta=as.matrix(draw.th[(N-last):N,])))
}


predict<-function(fit,l.pred,ninfo)
{
  l=ncol(fit$psiz)
  k=ncol(fit$rhoz)
  Npost=length(fit$lambdaf)
  nf=ninfo$nf  # field data
  nc=ninfo$nc  # computer model runs
  npred=ninfo$npred  # prediction locations
  nn=length(fit$y)  #nn=nf+nc
  N=npred+nf+nc  # npred + nf + nc
  z.pred=matrix(0,nrow=Npost,ncol=npred)
  eta.pred=matrix(0,nrow=Npost,ncol=npred)
  Ntot=npred+nn
  I=diag(Ntot)

  l.delpred=l.pred[1:k]  # for the discrepancy, only need the space-time part of the design
  for(i in 1:k)          # and only at the npred prediction + nf field locations.
    l.delpred[[i]][[1]]=l.delpred[[i]][[1]][1:(npred+nf),1:(npred+nf)]

  
  rix=1:N
  
  for(i in 1:Npost)
  {
    l.dez=l.pred
    for(j in 1:l)
    {
      l.dez[[k+j]][[1]][1:(npred+nf),]=l.dez[[k+j]][[1]][1:(npred+nf),]+fit$theta[i,j]
      l.dez[[k+j]][[1]][,1:(npred+nf)]=l.dez[[k+j]][[1]][,1:(npred+nf)]-fit$theta[i,j]
    }

    # predict field:
    R=rhogeodacecormat(l.dez,c(fit$rhoz[i,],fit$psiz[i,]),alpha=2)$R[rix,rix]
    Rdel=rhogeodacecormat(l.delpred,fit$rhodel[i,],alpha=2)$R
    et=c(rep(0,npred),rep(1/fit$lambdaf[i],nf),rep(0,nc))
    E=1/fit$lambdaz[i]*R+I*et
    E[1:(npred+nf),1:(npred+nf)]=E[1:(npred+nf),1:(npred+nf)]+1/fit$lambdadel[i]*Rdel
    Ed=E[(npred+1):Ntot,(npred+1):Ntot]
    Epd=E[1:npred,(npred+1):Ntot]
    cholEd=chol(Ed)
    Edinv=chol2inv(cholEd)
    z.pred[i,]=Epd%*%Edinv%*%fit$y

    Ep=E[1:npred,1:npred]
    Vp=Ep-Epd%*%Edinv%*%t(Epd)
    e=eigen(Vp,symmetric=TRUE)
    psd=max(which(e$values>0))
    Lp=e$vectors[,1:psd]%*%diag(sqrt(e$values[1:psd]))%*%t(e$vectors[,1:psd])
    u=rnorm(npred)
    z.pred[i,]=z.pred[i,]+Lp%*%u
    
    # predict computer model, eta(theta):
    E=1/fit$lambdaz[i]*R+I*et
    E[(npred+1):(npred+nf),(npred+1):(npred+nf)]=E[(npred+1):(npred+nf),(npred+1):(npred+nf)]+
                                                 1/fit$lambdadel[i]*Rdel[(npred+1):(npred+nf),(npred+1):(npred+nf)]
    Ed=E[(npred+1):Ntot,(npred+1):Ntot]
    Epd=E[1:npred,(npred+1):Ntot]
    cholEd=chol(Ed)
    Edinv=chol2inv(cholEd)
    eta.pred[i,]=Epd%*%Edinv%*%fit$y

    Ep=E[1:npred,1:npred]
    Vp=Ep-Epd%*%Edinv%*%t(Epd)
    e=eigen(Vp,symmetric=TRUE)
    psd=max(which(e$values>0))
    Lp=e$vectors[,1:psd]%*%diag(sqrt(e$values[1:psd]))%*%t(e$vectors[,1:psd])
    u=rnorm(npred)
    eta.pred[i,]=eta.pred[i,]+Lp%*%u
  }
  z.mu=apply(z.pred,2,mean)
  z.sd=apply(z.pred,2,sd)
  eta.mu=apply(eta.pred,2,mean)
  eta.sd=apply(eta.pred,2,sd)

  return(list(y=fit$y,lambdaf=fit$lambdaf,lambdaz=fit$lambdaz,lambdadel=fit$lambdadel,rhoz=fit$rhoz,
              psiz=fit$psiz,rhodel=fit$rhodel,theta=fit$theta,z.pred=z.pred,z.mu=z.mu,z.sd=z.sd,
              eta.pred=eta.pred,eta.mu=eta.mu,eta.sd=eta.sd))
}