library(BayesTree)

if(1) { # simulate simple model with binary y
f = function(x){(x/2)^2-1}
binaryOffset = 0.0
set.seed(99)
n = 2000
x= rnorm(n)
x = sort(x)
y = rbinom(n,1,pnorm(f(x)+binaryOffset))
}

if(1) { 
set.seed(99)
if(1) { # run binary bart

binbart = bart(matrix(x,ncol=1),y)

#plot sequences of f(x) draws
bartfit = apply(binbart$yhat.train,2,mean)
colind = floor(c(.25*n,.5*n,.75*n))
par(mfrow=c(2,2))
matplot(binbart$yhat.train[,colind])

#plot inference for f
plot(f(x),bartfit)
abline(0,1)
qts = apply(binbart$yhat.train,2,quantile,prob=c(.025,.975))
for(i in 1:n) {
   lines(rep(f(x[i]),2),qts[,i],col='grey')
}
points(f(x),bartfit)

#plot inference for p(x)
plot(x,pnorm(f(x)+binaryOffset),type='l',col='blue')
points(x,pnorm(bartfit),type='l',col='red')
pdraws = pnorm(binbart$yhat.train)
points(x,apply(pdraws,2,mean),type='l',col='green')
for(i in 1:n) {
   lines(c(x[i],x[i]),pnorm(qts[,i]),col='grey')
}
points(x,pnorm(f(x)+binaryOffset),type='l',col='blue')
points(x,pnorm(bartfit),type='l',col='red')

#default plot method
plot(binbart)
}

if(1) { #univariate partial dependence plot
par(mfrow=c(1,1))
pdbinbart = pdbart(matrix(x,ncol=1),y,k=2,keepevery=10)
lines(x,f(x))
}
}

if(0) { # binary Friedman (example with 10 x's)
#simulate
f = function(x){
10*sin(pi*x[,1]*x[,2]) + 20*(x[,3]-.5)^2+10*x[,4]+5*x[,5]
}
sigma = 1.0  #y = f(x) + sigma*z , z~N(0,1)
nf = 500      #number of observations
set.seed(99)
xf=matrix(runif(nf*10),nf,10) #10 variables, only first 5 matter
Eyf = scale(f(xf))
yf=Eyf+sigma*rnorm(nf)
yfb = as.factor(yf>0)

#binary bart
if(0) {
rbartfb = bart(xf,yfb)

#plot inference for f
par(mfrow=c(1,1))
plot(Eyf,apply(rbartfb$yhat.train,2,mean),ylim=c(-3,3))
abline(0,1)
qts = apply(rbartfb$yhat.train,2,quantile,prob=c(.05,.95))
for(i in 1:100) {
   lines(rep(Eyf[i],2),qts[,i],col='black')
}
}

if(1) { #2-d partial depedence plot for first two vars in Friedman
frpd2 = pd2bart(xf,yfb,k=2,keepevery=10)
}
}