Simulated data

This notebook runs bartz on simulated data. It is meant to be run on a GPU. Use the following link to try it out on colab: link

The next cell installs bartz:

%pip install bartz

The next cell defines configuration parameters for the script:

n_train = 100_000  # number of training points
p = 1000           # number of predictors/features
sigma = 0.1        # error standard deviation
n_test = 1000      # number of test points
n_tree = 10_000    # number of trees used by bartz

The next cell generates simulated data from a linear + quadratic test model that comes packaged with bartz:

from jax import random

from bartz.testing import gen_data

# simulate data with bartz's built-in testing data generating process
data = gen_data(
    random.key(2024_04_16_18_53),
    n=n_train + n_test,
    p=p,
    q=2,  # number of interactions, each predictor interacts with other q predictors in the quadratic term
    sigma2_eps=sigma**2,  # error variance
    sigma2_lin=0.5,  # linear term variance
    sigma2_quad=0.5,  # quadratic term variance
)

# split data in train-test
train, test = data.split(n_train)

The next cell runs BART:

from bartz import Bart

# run BART; single-chain to fit on the small free colab gpus
bart = Bart(train.x, train.y, num_trees=n_tree, num_chains=None, seed=2026_07_23_15_06)

How to read the information in the progress bar:

• 100sa: the statistics that follow are averages over the last 100 MCMC iterations.

• acc 21%: the fraction of proposed tree changes that were accepted; if this is too low, it means the MCMC is stuck.

• leaves 2.6/32: the average number of leaves per tree is 2.5, out of a maximum of 32. The number of leaves should stay well below the maximum, otherwise it means BART would like to grow deeper trees and is stuck at the constraint.

The next cell computes predictions on the test set.

from jax import numpy as jnp

# compute predictions
yhat_test = bart.predict(test.x, kind='mean_samples')  # posterior samples, n_samples x n_test
yhat_test_mean = jnp.mean(yhat_test, axis=0)  # posterior mean point-by-point
yhat_test_var = jnp.var(yhat_test, axis=0)  # posterior variance point-by-point

# RMSE
rmse = jnp.sqrt(jnp.mean(jnp.square(yhat_test_mean - test.y)))
avg_sigma = bart.get_error_sdev(mean=True)
expected_rmse = jnp.sqrt(jnp.mean(yhat_test_var + avg_sigma ** 2))

print(f'total sdev: {jnp.std(train.y):#.2g}')
print(f'error sdev: {sigma:#.2g}')
print(f'RMSE: {rmse:#.2g}')
print(f'expected RMSE: {expected_rmse:#.2g}')
print(f'model error sdev: {avg_sigma:#.2g}')

total sdev: 1.0
error sdev: 0.10
RMSE: 0.35
expected RMSE: 0.35
model error sdev: 0.28

The RMSE can at best be as low as the error standard deviation used to generate the data.