# Generate dataset according to Friedman function
#--------------------------------------------------
if(1) { #setup and simulate in-sample (x,y) and construct out-of-sample (x)
set.seed(99)
n=5000 #train (in-sample) sample size
p=5 #number of x's
sigma=1 #sd for observational error

X=matrix(runif(p*n),ncol=p) #in-sample x
#out-of-sample (or prediction) x:
sp=10 #number of prediction points per dimension to generate the prediction grid
Xp=as.matrix(expand.grid(seq(0,1,length=sp),seq(0,1,length=sp),seq(0,1,length=sp),seq(0,1,length=sp),seq(0,1,length=sp)))

f=10*sin(2*pi*X[,1]*X[,2])+(X[,3]-.5)^2+X[,4]+X[,5]
y=f+rnorm(n,sd=sigma)
fp=10*sin(2*pi*Xp[,1]*Xp[,2])+(Xp[,3]-.5)^2+Xp[,4]+Xp[,5]

npred=length(fp)
numnodes=8

pv=rep(0,numnodes-1) #number of in-sample observations for each slave node
pv[1:(numnodes-2)] = floor(n/(numnodes-1)) #number of slave = numnodes-1
pv[numnodes-1]=n-sum(pv)
cat("Generating ",numnodes-1," datasets of sizes ",pv,"\n")

# write out the dataset
for(i in 1:(numnodes-1)) #c indices of slaves
{
	if(i==1)
	{
		ii=1:pv[1] 
	} else {
		ii=(sum(pv[1:(i-1)])+1):sum(pv[1:i])
	}

	xx=X[ii,]
	yy=y[ii]
	write(yy,ncol=1,file=paste('y',i,'.txt',sep=""))
	write(t(xx),ncol=p,file=paste('x',i,'.txt',sep=""))
}

# write out the prediction points
pv=rep(0,numnodes) #number of out-of-sample observations for each slave node
pv[1:(numnodes-1)] = floor(nrow(Xp)/numnodes) #number of slaves = numnodes
pv[numnodes]=nrow(Xp)-sum(pv)
cat("Generating ",numnodes," predictive X datasets of sizes ",pv,"\n")

for(i in 1:numnodes) #note xpi.txt i=1,2,..numnodes
{
	if(i==1)
	{
		ii=1:pv[1]
	} else {
		ii=(sum(pv[1:(i-1)])+1):sum(pv[1:i])
	}
	
	xx=Xp[ii,]
	yy=fp[ii]
	write(t(xx),ncol=p,file=paste('xp',i,'.txt',sep=""))
}
}
#--------------------------------------------------
if(0) { # Fit Bart
nu=3
sighat = sd(y)

if(1) { #use least squares for sighat
df = data.frame(X,y)
lmf = lm(y~.,df)
sighat = summary(lmf)$sigma
}

sigq = .90
qchi = qchisq(1.0-sigq,nu)
lambda = (sighat*sighat*qchi)/nu #lambda parameter for sigma prior

nd=3500  #number of MCMC draws to "keep" (eg go into posterior means)
thin=100 #save draws ever "thin" MCMC draws after burn-in
burn=500 #number of initial draws called burn-in	
m= 200   #number of trees in bart sum
kfac=2   #bart prior parameter for determining prior stan dev of mu's

cmd = paste('mpirun -np',numnodes, ' ./mpibart x y ',thin,nd,lambda,burn,m,nu,kfac)
cat('cmd:\n')
cat(cmd)
cat('\n')
system(cmd)
}
#--------------------------------------------------
if(0) { # Plot in-sample fits
par(mfrow=c(1,2))

bfit=c()
for(i in 1:(numnodes-1)) {
fnm = paste('bart-fit',i,'.txt',sep='')
bfit = c(bfit,scan(fnm))
}
plot(f,bfit)
abline(0,1)
title('f(x) vs fhat(x), in-sample')

fnm = "bart-sd.txt"
bsd = c()
bsd = c(bsd,scan(fnm))
plot(bsd,type='l')
abline(h=sigma)
title('sigma draws, including burn-in')
dev.copy2pdf(file='in-sample.pdf',height=6,width=12)
}
#--------------------------------------------------
if(0){ # Do the out-of-sample predictions
cmd = paste('mpirun -np',numnodes, './bartpred xp p')
cat('cmd:\n')
cat(cmd)
cat('\n')
system(cmd)
}
#--------------------------------------------------
if(0) { # Plot out-of-sample predictions
par(mfrow=c(1,1))
bfit=c()
bsd=c()
for(i in 1:numnodes) {
fnm = paste('yp',i-1,'.txt',sep='')
fsd = paste('ysd',i-1,'.txt',sep='')
temp = scan(fnm)
print(length(temp))
bfit = c(bfit,scan(fnm))
temp = scan(fsd)
print(length(temp))
bsd = c(bsd,scan(fsd))
}
plot(fp,bfit-2*bsd,pch=20,col="grey",xlab="true friedman",ylab="est. friedman")
points(fp,bfit+2*bsd,pch=20,col="grey")
points(fp,bfit,pch=20)
abline(0,1)
title('out-of-sample, f(x) vs. fhat(x)')
dev.copy2pdf(file='out-of-sample.pdf',height=10,width=10)
}
#--------------------------------------------------
if(0) { # Do sensitivity analysis
nmc = 100000 #number of Monte Carlo integration points for the sensitivity analysis
ngp = 10 #number of gridpoints to calculate the 1-way sensitivity plots
cmd = paste('mpirun -np',as.integer(numnodes), './bartsa',as.integer(nmc),as.integer(ngp))
cat('cmd:\n')
cat(cmd)
cat('\n')
system(cmd)
}
#--------------------------------------------------
if(0) { # Plot main effect sensitivities
par(mfrow=c(floor(sqrt(p)),ceiling(sqrt(p))))
saf = as.matrix(read.table("saoneway.txt",header=F))
xtmp = seq(0,1,length=100)
fsf=matrix(0,nrow=5,ncol=100)
fsf[1,]=-10/(2*pi*xtmp)*cos(2*pi*xtmp)+10/(2*pi*xtmp)+13/12
fsf[2,]=-10/(2*pi*xtmp)*cos(2*pi*xtmp)+10/(2*pi*xtmp)+13/12
fsf[3,]=3.87964+(xtmp-.5)^2+1
fsf[4,]=3.87964+7/12+xtmp
fsf[5,]=3.87964+7/12+xtmp

for(i in 1:p)
{
ix=which(saf[,1]==(i-1))
plot(xtmp,fsf[i,],col="blue",type='l',xlab=paste("x",i,sep=""),ylab=paste("E[f|x",i,"]",sep=""),cex.lab=.5,cex.axis=.5,ylim=c(min(saf[ix,3]-2*saf[ix,4]),max(saf[ix,3]+2*saf[ix,4])))
points(saf[ix,2],saf[ix,3]+2*saf[ix,4],pch=20,col="grey")
points(saf[ix,2],saf[ix,3]-2*saf[ix,4],pch=20,col="grey")
points(saf[ix,2],saf[ix,3],pch=20)
}
dev.copy2pdf(file='one-way-sa.pdf',height=10*floor(sqrt(p))/ceiling(sqrt(p)),width=10)
}

