Implementing a corrected truncated normal sampler (for the case with non...

dbarts has an issue when using a binary offset. It appears that it's sampling from the wrong distribution, and sometimes returns NaN's. The inverse cdf sampler for truncated normals can also be numerically unstable when the mean is var from the truncation level, relative to the SD. I implemented a more stable truncated normal sampler (based on efficient rejection sampling) and checked it against BayesTree and the truth. My new version of sampleProbitLatentVariables includes a (hardcoded) toggle between my new sampler and the original implementation for easy comparison.

There still seems to be an issue where the offset is erroneously added to the predicted values, but I haven't been able to track down where that is happening (it seems to be a postprocessing issue). R code illustrating the bug and correct results after the change follows.

######################################################
# File: bug-ex.R
# Run with old=true in sampleProbitLatentVariables, or equivalently 
# the current master in the dbarts repo 
######################################################
library(dbarts); library(nnet)

set.seed(2)

# With offset = -1.0, bartfit1 returns NaN's
# with offset = -0.5, bartfit1 returns proper values but gives the wrong answer
# in both cases barfit2 matches bartfit3 & bartfit4
offset = -0.5

X <- matrix(rnorm(500*2),500,2)
gam <- 0.25*(sin((X[,1]+X[,2])*pi)+X[,1]) + offset
Y <- as.numeric(gam+rnorm(500) > 0)

bartfit1 <- bart(X,Y,binaryOffset=offset)
bartfit2 <- bart(X,Y)
bartfit3 <- BayesTree::bart(X, Y, binaryOffset=offset)
bartfit4 <- BayesTree::bart(X, Y)

print(bartfit1$yhat.train[1:10])
print(bartfit2$yhat.train[1:10])

# This one is wrong & not by something additive; the truncated normal
# sampler is incorrect or just unstable (I think the mean is incorrect but
# I haven't looked at it very carefully)
plot(gam, colMeans(bartfit1$yhat.train), col=Y+1); abline(0,1)
# The next three are the same up to Monte Carlo error (i.e., correct)
plot(gam, colMeans(bartfit2$yhat.train), col=Y+1); abline(0,1)
plot(gam, colMeans(bartfit3$yhat.train), col=Y+1); abline(0,1)
plot(gam, colMeans(bartfit4$yhat.train), col=Y+1); abline(0,1)

######################################################
# File: fixed-ex.R
# Run with old=false in sampleProbitLatentVariables
######################################################
library(dbarts); library(nnet)

set.seed(2)

offset = -0.5

X <- matrix(rnorm(500*2),500,2)
gam <- 0.25*(sin((X[,1]+X[,2])*pi)+X[,1]) + offset
Y <- as.numeric(gam+rnorm(500) > 0)

bartfit1 <- bart(X,Y,binaryOffset=offset)
bartfit2 <- bart(X,Y)
#bartfit3 <- BayesTree::bart(X, Y, binaryOffset=offset)
#bartfit4 <- BayesTree::bart(X, Y)

# Using the ref interface
dbf <- dbarts(Y~., data.frame(Y=Y, X=X), offset=offset*rep(1, 500))
bartfit5 <- dbf$run()

# OK
print(bartfit1$yhat.train[1:10])
print(bartfit2$yhat.train[1:10])

# Somehow the offset is added twice (3x in the BayesTree-stype interface?)
# The sampler is correct though. After adjusting for multiple 
# additions of the offset the results match the BayesTree results &
# the zero offset results, all of which "match" the truth
plot(gam, colMeans(bartfit1$yhat.train)-2*offset, col=Y+1); abline(0,1)
plot(gam, colMeans(bartfit2$yhat.train), col=Y+1); abline(0,1)
plot(gam, rowMeans(bartfit5$train)-offset, col=Y+1); abline(0,1)