BATTS compares two samples by estimating the ratio p/q
of their densities. It implements the two additive tree methods of
Awaya, Xu, and Ma (2025): a boosting estimator, which gives a point
estimate, and a generalized Bayesian model with additive trees, which
also gives posterior draws and hence credible bands for the ratio.
Both methods estimate the balancing weight
w = sqrt(p/q). The predict() and
fitted() methods transform it, and by default return the
log-density ratio log(p/q) = 2 * log(w). A positive value
means that the first sample is denser than the second at that point.
This vignette works through simulated examples, first in one dimension and then in two. For each example we keep the true log-density ratio so that the estimates can be compared with it.
We start with two pairs of one-dimensional distributions, each with 250 observations per sample. One dimension is where the estimates and their uncertainty are easiest to see.
The first sample comes from N(0, 1) and the second from
N(0.5, 1). The true log-density ratio is a straight
line.
set.seed(12345)
n1 <- 250L
loc <- list(x = rnorm(n1), y = rnorm(n1, mean = 0.5))
loc_truth <- function(t) dnorm(t, log = TRUE) - dnorm(t, mean = 0.5, log = TRUE)The helper below plots the two samples in different colours. In one
dimension it draws a kernel density estimate of each sample with
density() and marks the observations along the axis, the
first sample at the bottom and the second at the top. In two dimensions
it draws a scatter plot with different symbols, in random order so that
neither sample hides the other.
col_x <- adjustcolor("firebrick", alpha.f = 0.6)
col_y <- adjustcolor("steelblue", alpha.f = 0.6)
draw_samples <- function(sample, main) {
if (NCOL(sample$x) == 1L) {
dens_x <- density(sample$x)
dens_y <- density(sample$y)
plot(dens_x, col = "firebrick", lwd = 2, main = main, xlab = "x",
xlim = range(dens_x$x, dens_y$x), ylim = c(0, max(dens_x$y, dens_y$y)))
lines(dens_y, col = "steelblue", lwd = 2, lty = 2)
rug(sample$x, side = 1, col = col_x)
rug(sample$y, side = 3, col = col_y)
legend("topright", c("first sample (p)", "second sample (q)"),
col = c("firebrick", "steelblue"), lty = c(1, 2), lwd = 2,
bty = "n", cex = 0.8)
} else {
x <- rbind(sample$x, sample$y)
first <- rep(c(TRUE, FALSE), c(nrow(sample$x), nrow(sample$y)))
ord <- sample.int(nrow(x))
plot(x[ord, ], col = ifelse(first[ord], col_x, col_y),
pch = ifelse(first[ord], 16, 17), cex = 0.7, asp = 1,
xlab = "x1", ylab = "x2", main = main)
legend("bottom", c("first sample (p)", "second sample (q)"),
col = c(col_x, col_y), pch = c(16, 17), horiz = TRUE, bty = "n",
inset = c(0, 1), xpd = TRUE, cex = 0.8)
}
}Both samples are centred at zero, with standard deviations 1 and 1.5. The true log-density ratio is a downward parabola: the first sample is denser in the middle and the second in the tails.
set.seed(23456)
scl <- list(x = rnorm(n1), y = rnorm(n1, sd = 1.5))
scl_truth <- function(t) dnorm(t, log = TRUE) - dnorm(t, sd = 1.5, log = TRUE)The settings of each method are collected by
boost_control() and bart_control(). Their
defaults are the settings used in the simulations of the paper. Here we
use boosting trees of depth 2, which give a smoother estimate than the
default depth 4 in these small examples, and a shorter Markov chain (100
trees, 300 burn-in iterations and 301 saved draws) so that the vignette
builds quickly. These shorter chains are for illustration and are not a
recommendation for applied work.
boost_ctl <- boost_control(depth = 2)
bart_ctl <- bart_control(num_trees = 100, burnin = 300, draws = 301,
keep_draws = TRUE)The two samples can be passed directly as x and
y; x is then the numerator p.
Setting keep_draws = TRUE saves the posterior forests, so
that the Bayesian fit can later be evaluated at new points.
set.seed(34567)
loc_boost <- batts_boost(loc$x, y = loc$y, control = boost_ctl, quiet = TRUE)
loc_bart <- batts_bart(loc$x, y = loc$y, control = bart_ctl, quiet = TRUE)
scl_boost <- batts_boost(scl$x, y = scl$y, control = boost_ctl, quiet = TRUE)
scl_bart <- batts_bart(scl$x, y = scl$y, control = bart_ctl, quiet = TRUE)Printing a fit summarizes the data, the ensemble and the direction of the ratio.
loc_boost
#> <BATTS boosting fit>
#> observations: 500 (group "x": 250, group "y": 250)
#> variables: 1
#> trees: 93
#> CV: 5 folds, 1000 candidate tree counts
#> estimand: balancing weight w = sqrt(p/q)
#> p: group "x" (numerator), q: group "y"
#> 2 * log(w) is the log-density ratio
loc_bart
#> <BATTS Bayesian additive trees fit>
#> observations: 500 (group "x": 250, group "y": 250)
#> variables: 1
#> saved draws: 301
#> forests kept: yes
#> boosting init trees: 100
#> estimand: balancing weight w = sqrt(p/q)
#> p: group "x" (numerator), q: group "y"
#> 2 * log(w) is the log-density ratiosummary() adds the fitted domain, outside of which the
fit cannot be evaluated, and the range of the fitted weights.
summary(loc_boost)
#> <BATTS boosting fit>
#> observations: 500 (group "x": 250, group "y": 250)
#> variables: 1
#> trees: 93
#> CV: 5 folds, 1000 candidate tree counts
#> training domain (min, max) per variable:
#> [,1] [,2]
#> [1,] -2.89518 3.282682
#> fitted weight w = sqrt(p/q) range: [0.797663, 1.57782]For boosting, the number of trees is chosen by cross-validation.
plot() shows the held-out loss, averaged over the folds,
against the number of trees, and marks the selected number. The loss
first falls and then rises as the ensemble starts to overfit. If the
selected number were at num_trees_max, the maximum should
be raised.
For the Bayesian fit, plot() shows the trace of the loss
temperature across the saved draws.
A more direct check of the Markov chain looks at the quantity of
interest and compares independent chains. We run a second chain from
another seed and compare the draws of log(p/q) at
x = 0, and the posterior means along a grid.
set.seed(45678)
loc_bart2 <- batts_bart(loc$x, y = loc$y, control = bart_ctl, quiet = TRUE)
chains_at_0 <- cbind(chain1 = predict(loc_bart, 0, summary = "none")[1, ],
chain2 = predict(loc_bart2, 0, summary = "none")[1, ])
check_grid <- seq(-2, 2, length.out = 101)
chain_means <- cbind(chain1 = predict(loc_bart, check_grid),
chain2 = predict(loc_bart2, check_grid))old_par <- par(mfrow = c(1, 2), mar = c(4, 4, 2.5, 1))
matplot(chains_at_0, type = "l", lty = 1, col = c("black", "orange"),
xlab = "saved draw", ylab = "log(p/q) at x = 0",
main = "Draws from two chains")
matplot(check_grid, chain_means, type = "l", lty = 1, lwd = 2,
col = c("black", "orange"), xlab = "x", ylab = "log(p/q)",
main = "Posterior means")
legend("topright", c("chain 1", "chain 2"), col = c("black", "orange"),
lty = 1, lwd = 2, bty = "n", cex = 0.8)The two chains move over the same range and give similar posterior means. Agreement between chains is consistent with, but does not establish, convergence; for applied work, run longer chains and more of them.
We evaluate both fits on a grid covering the central 95% of the
pooled data. For the Bayesian fit, interval = TRUE returns
the posterior mean together with a pointwise credible band.
ratio_curves <- function(sample, boost, bart, truth) {
grid <- seq(quantile(c(sample$x, sample$y), 0.025),
quantile(c(sample$x, sample$y), 0.975), length.out = 200)
list(grid = grid, truth = truth(grid), boost = predict(boost, grid),
band = predict(bart, grid, interval = TRUE, level = 0.95))
}
loc_curves <- ratio_curves(loc, loc_boost, loc_bart, loc_truth)
scl_curves <- ratio_curves(scl, scl_boost, scl_bart, scl_truth)
head(loc_curves$band, 3)
#> fit lower upper
#> [1,] 0.8480653 0.1938299 1.451233
#> [2,] 0.8434791 0.1775229 1.387045
#> [3,] 0.8358078 0.1791229 1.404857The helper below draws the band in grey, the posterior mean as a solid line, the boosting estimate as a dotted line and the truth as a dashed line.
draw_ratio <- function(curves, main) {
band <- curves$band
plot(curves$grid, band[, "fit"], type = "n",
ylim = range(band, curves$truth, curves$boost),
xlab = "x", ylab = "log(p/q)", main = main)
polygon(c(curves$grid, rev(curves$grid)),
c(band[, "lower"], rev(band[, "upper"])), col = "grey85", border = NA)
abline(h = 0, col = "grey60")
lines(curves$grid, curves$truth, col = "red", lwd = 2, lty = 2)
lines(curves$grid, curves$boost, col = "darkgreen", lwd = 2, lty = 3)
lines(curves$grid, band[, "fit"], lwd = 2)
}
old_par <- par(mfrow = c(1, 2), mar = c(4, 4, 2.5, 1))
draw_ratio(loc_curves, "Location shift")
legend("topright", c("truth", "boosting", "posterior mean", "95% band"),
col = c("red", "darkgreen", "black", "grey85"), lty = c(2, 3, 1, 1),
lwd = c(2, 2, 2, 8), bty = "n", cex = 0.8)
draw_ratio(scl_curves, "Scale difference")Both fits follow the line and the parabola over most of the range. The credible band is narrow where the data are dense and widens towards the tails, where few observations inform the ratio and the estimates are drawn towards zero.
Other scales are available through type:
density_ratio() returns p/q itself, and
type = "weight" returns the balancing weight
w = sqrt(p/q).
The two-dimensional examples are based on Scenarios 1 and 2 of Awaya, Xu, and Ma (2025), with 1000 observations per sample. The samples are matrices with one column per variable, and are fitted in the same way as above.
Two correlated normal distributions differ by a shift of 0.75 in both coordinates, along the direction of their correlation. The true log-density ratio is linear and changes along the 45-degree diagonal: the first sample is denser towards the lower left and the second towards the upper right.
log_dnorm2 <- function(x, mean, sigma) {
z <- sweep(x, 2, mean) %*% solve(chol(sigma))
-0.5 * rowSums(z^2) - log(2 * pi) - 0.5 * log(det(sigma))
}
rnorm2 <- function(n, mean, sigma) {
sweep(matrix(rnorm(2 * n), n, 2) %*% chol(sigma), 2, mean, "+")
}
sigma_g <- matrix(c(1, 0.5, 0.5, 1), 2, 2)
shift_g <- c(0.75, 0.75)
set.seed(56789)
n2 <- 1000L
glb <- list(x = rnorm2(n2, c(0, 0), sigma_g), y = rnorm2(n2, shift_g, sigma_g))
glb_truth <- function(x) {
log_dnorm2(x, c(0, 0), sigma_g) - log_dnorm2(x, shift_g, sigma_g)
}Both samples are equal mixtures of three normal components with standard deviation 0.3. In the second sample, only the third component (top right) is shifted, by one standard deviation to the left. The two densities therefore agree except near that component.
centres <- rbind(c(-2.5, 1), c(1, -2), c(2, 2.5))
rmix <- function(n, centres) {
k <- sample.int(3L, n, replace = TRUE)
centres[k, ] + matrix(rnorm(2 * n, sd = 0.3), n, 2)
}
dmix <- function(x, centres) {
dens <- sapply(1:3, function(k) {
dnorm(x[, 1], centres[k, 1], 0.3) * dnorm(x[, 2], centres[k, 2], 0.3)
})
rowMeans(dens)
}
shifted <- centres
shifted[3, ] <- centres[3, ] + c(-0.3, 0)
set.seed(67890)
lcl <- list(x = rmix(n2, centres), y = rmix(n2, shifted))
lcl_truth <- function(x) log(dmix(x, centres)) - log(dmix(x, shifted))We use the same settings as in one dimension. Since we only evaluate these fits at the observed points, we do not keep the posterior forests. The diagnostics of the previous section apply in the same way and are omitted here.
bart_ctl_2d <- bart_control(num_trees = 100, burnin = 300, draws = 301)
set.seed(78901)
glb_boost <- batts_boost(glb$x, y = glb$y, control = boost_ctl, quiet = TRUE)
glb_bart <- batts_bart(glb$x, y = glb$y, control = bart_ctl_2d, quiet = TRUE)
lcl_boost <- batts_boost(lcl$x, y = lcl$y, control = boost_ctl, quiet = TRUE)
lcl_bart <- batts_bart(lcl$x, y = lcl$y, control = bart_ctl_2d, quiet = TRUE)As in the paper, we colour the observations by the value of the log-density ratio, using a common scale for the truth and both estimates. Red points mark where the first sample is denser and blue points where the second is. The last panel shows the posterior standard deviation, the two-dimensional counterpart of the width of the credible band.
diverging <- hcl.colors(101, "Blue-Red 3")
sequential <- hcl.colors(101, "YlOrRd", rev = TRUE)
draw_points <- function(x, value, limits, palette, main) {
index <- round((pmin(pmax(value, limits[1]), limits[2]) - limits[1]) /
diff(limits) * 100) + 1
ord <- order(abs(value - mean(limits)))
plot(x[ord, ], col = palette[index[ord]], pch = 16, cex = 0.5, asp = 1,
xlab = "x1", ylab = "x2", main = main)
}
draw_scenario <- function(sample, boost, bart, truth) {
x <- rbind(sample$x, sample$y)
true_value <- truth(x)
limits <- c(-1, 1) * max(abs(true_value))
draws <- fitted(bart, summary = "none")
old_par <- par(mfrow = c(2, 2), mar = c(4, 4, 2.5, 1))
on.exit(par(old_par))
draw_points(x, true_value, limits, diverging, "true log(p/q)")
draw_points(x, fitted(boost), limits, diverging, "boosting")
draw_points(x, rowMeans(draws), limits, diverging, "posterior mean")
sds <- apply(draws, 1, sd)
draw_points(x, sds, c(0, max(sds)), sequential, "posterior SD")
}For the global shift, both estimates reproduce the linear trend along the diagonal, although each tree splits on one coordinate at a time.
For the local shift, the truth is zero on the first two components and changes sign across the third, where the second sample has moved to the left. Both estimates locate this blue and red pair, and show only faint colour on the unchanged components. The posterior standard deviation is largest around the shifted component.
x and y may be numeric matrices, data
frames with numeric columns, or, for one variable, numeric vectors.
Instead of y, a single sample can be given with a
group vector of two labels; the first level (or the smaller
value) is the numerator p. print() always
states the direction.margin_scale (10% by default).
predict() stops for other points unless
outside = "na" or outside = "clamp" is given.
Columns of newdata are matched by name when both sides have
names.bart_control(keep_draws = TRUE). With the default 200 trees
and 1000 draws the saved forests can exceed 100 MB, so keep them only
when needed. fitted() and predict() without
newdata work in either case.boots(), batts() and
eval_balance_weight(), the interface of versions before
0.4.0, still work and give the same results, but are deprecated. See
?`BATTS-deprecated` for the correspondence of the
arguments.Awaya, N., Xu, Y., and Ma, L. (2025). Two-sample comparison through additive tree models for density ratios. arXiv:2508.03059. https://doi.org/10.48550/arXiv.2508.03059