#
# Bayesian Treed Gaussian Process Example
# For STAT8810
# Fall, 2017
# M.T. Pratola



# Mixture prior on correlation parameters
x=seq(0,2,length=1000)
da=dgamma(x,shape=1,rate=20)
db=dgamma(x,shape=10,rate=10)
d=0.5*da+0.5*db
par(mfrow=c(1,2))
plot(x,d,type='l',lwd=2,xlab=expression(d[nu]),
     ylab="Density")
lines(x,da,lwd=0.5,col="grey")
lines(x,db,lwd=0.5,col="grey")
abline(v=1/20,lty=2,col="grey")
abline(v=10/10,lty=2,col="grey")


# Show some draws at the means of both mixture components
set.seed(99)
x=seq(0,1,length=100)
D=abs(outer(x,x,"-"))
Ra=exp(-D^2/(1/20))  # like rho=2e-9
Rb=exp(-D^2/(10/10)) # like rho=0.37
La=t(chol(Ra+diag(100)*1e-10))
Lb=t(chol(Ra+diag(100)*1e-10))
Za=La%*%rnorm(100)
Zb=Lb%*%rnorm(100)
plot(x,Za,type='l',lwd=2,col="blue",
     ylim=range(c(Za,Zb)),ylab="Response")
lines(x,Zb,lwd=2,col="red")




# Example using TGP
library(tgp)

# See all the demo's available within the package.
demo(package="tgp")

# You can run a built-in demo using
# demo("demoname",package="tgp")
# But we'll directly look at the moto dataset.




# Example - Moto data
set.seed(88)
library(MASS)
X=data.frame(times=mcycle[,1])
Z=data.frame(accel=mcycle[,2])
fit.gp=bgp(X=X,Z=Z,verb=0) # Regular GP fit (no tree)
fit.tgp=btgp(X=X,Z=Z,bprior="b0",verb=0) # Treed GP

# Plot both fits (posterior mean predictions) side by side
par(mfrow=c(1,2))
plot(fit.gp,layout='surf')
plot(fit.tgp,layout='surf')

# Model details:
str(fit.tgp)





# By default samples are not saved from the posterior, 
# only the posterior quantities we want are recorded.
# Use trace=TRUE to save more information.
# However, storage may be an issue.
fit2.tgp=btgp(X=X,Z=Z,bprior="b0",verb=0,trace=TRUE)
str(fit2.tgp)

# Plot some of the posterior realizations
par(mfrow=c(1,2))
plot(fit2.tgp$trace$hier$s2.a0,type='l')
plot(fit2.tgp$trace$preds$Zp.ks2$XX1,type='l')





# Example - run Moto for more iterations
# BTE=(burn,total,every)
# Default is BTE=(2000,7000,2)
fit3.tgp=btgp(X=X,Z=Z,bprior="b0",verb=0,trace=TRUE,BTE=c(7000,12000,2))
par(mfrow=c(1,2))
plot(fit3.tgp$trace$hier$s2.a0,type='l')
plot(fit3.tgp$trace$preds$Zp.ks2$XX1,type='l')






# Example - good old stationary GP realization - will we just get 1 GP? 

# Realization
set.seed(88)
x=seq(0,1,length=10)
D=abs(outer(x,x,"-"))
R=0.001^(D^2)
L=t(chol(R))
Z=L%*%rnorm(10)
X=data.frame(x)
Z=data.frame(Z)
plot(X,Z,pch=20,col="red",xlab="X",ylab="Response")

# Fit models
fit.gp=bgp(X=X,Z=Z,verb=0) # Regular GP fit (no tree)
fit.tgp=btgp(X=X,Z=Z,bprior="b0",verb=0) # Treed GP

# Plot both fits (posterior mean predictions) side by side
par(mfrow=c(1,2))
plot(fit.gp,layout='surf')
plot(fit.tgp,layout='surf')

# They look identical.