Supervised Learning¶

Load necessary libraries

In [1]:

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.model_selection import train_test_split

from stochtree import BARTModel

import matplotlib.pyplot as plt import numpy as np import pandas as pd import seaborn as sns from sklearn.model_selection import train_test_split from stochtree import BARTModel

Generate sample data

In [2]:

# RNG
random_seed = 1234
rng = np.random.default_rng(random_seed)

# Generate covariates and basis
n = 1000
p_X = 10
p_W = 1
X = rng.uniform(0, 1, (n, p_X))
W = rng.uniform(0, 1, (n, p_W))


# Define the outcome mean function
def outcome_mean(X, W):
    return np.where(
        (X[:, 0] >= 0.0) & (X[:, 0] < 0.25),
        -7.5 * W[:, 0],
        np.where(
            (X[:, 0] >= 0.25) & (X[:, 0] < 0.5),
            -2.5 * W[:, 0],
            np.where((X[:, 0] >= 0.5) & (X[:, 0] < 0.75), 2.5 * W[:, 0], 7.5 * W[:, 0]),
        ),
    )


# Generate outcome
epsilon = rng.normal(0, 1, n)
y = outcome_mean(X, W) + epsilon

# Standardize outcome
y_bar = np.mean(y)
y_std = np.std(y)
resid = (y - y_bar) / y_std

# RNG random_seed = 1234 rng = np.random.default_rng(random_seed) # Generate covariates and basis n = 1000 p_X = 10 p_W = 1 X = rng.uniform(0, 1, (n, p_X)) W = rng.uniform(0, 1, (n, p_W)) # Define the outcome mean function def outcome_mean(X, W): return np.where( (X[:, 0] >= 0.0) & (X[:, 0] < 0.25), -7.5 * W[:, 0], np.where( (X[:, 0] >= 0.25) & (X[:, 0] < 0.5), -2.5 * W[:, 0], np.where((X[:, 0] >= 0.5) & (X[:, 0] < 0.75), 2.5 * W[:, 0], 7.5 * W[:, 0]), ), ) # Generate outcome epsilon = rng.normal(0, 1, n) y = outcome_mean(X, W) + epsilon # Standardize outcome y_bar = np.mean(y) y_std = np.std(y) resid = (y - y_bar) / y_std

Test-train split

In [3]:

sample_inds = np.arange(n)
train_inds, test_inds = train_test_split(sample_inds, test_size=0.5)
X_train = X[train_inds, :]
X_test = X[test_inds, :]
basis_train = W[train_inds, :]
basis_test = W[test_inds, :]
y_train = y[train_inds]
y_test = y[test_inds]

sample_inds = np.arange(n) train_inds, test_inds = train_test_split(sample_inds, test_size=0.5) X_train = X[train_inds, :] X_test = X[test_inds, :] basis_train = W[train_inds, :] basis_test = W[test_inds, :] y_train = y[train_inds] y_test = y[test_inds]

Demo 1: Using W in a linear leaf regression¶

Run BART

In [4]:

bart_model = BARTModel()
general_params = {"num_chains": 3}
bart_model.sample(
    X_train=X_train,
    y_train=y_train,
    leaf_basis_train=basis_train,
    X_test=X_test,
    leaf_basis_test=basis_test,
    num_gfr=10,
    num_mcmc=100,
    general_params=general_params,
)

bart_model = BARTModel() general_params = {"num_chains": 3} bart_model.sample( X_train=X_train, y_train=y_train, leaf_basis_train=basis_train, X_test=X_test, leaf_basis_test=basis_test, num_gfr=10, num_mcmc=100, general_params=general_params, )

Inspect the MCMC (BART) samples

In [5]:

forest_preds_y_mcmc = bart_model.y_hat_test
y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis=1, keepdims=True)
y_df_mcmc = pd.DataFrame(
    np.concatenate((np.expand_dims(y_test, 1), y_avg_mcmc), axis=1),
    columns=["True outcome", "Average estimated outcome"],
)
sns.scatterplot(data=y_df_mcmc, x="Average estimated outcome", y="True outcome")
plt.axline((0, 0), slope=1, color="black", linestyle=(0, (3, 3)))
plt.show()

forest_preds_y_mcmc = bart_model.y_hat_test y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis=1, keepdims=True) y_df_mcmc = pd.DataFrame( np.concatenate((np.expand_dims(y_test, 1), y_avg_mcmc), axis=1), columns=["True outcome", "Average estimated outcome"], ) sns.scatterplot(data=y_df_mcmc, x="Average estimated outcome", y="True outcome") plt.axline((0, 0), slope=1, color="black", linestyle=(0, (3, 3))) plt.show()

In [6]:

sigma_df_mcmc = pd.DataFrame(
    np.concatenate(
        (
            np.expand_dims(np.arange(bart_model.num_samples), axis=1),
            np.expand_dims(bart_model.global_var_samples, axis=1),
        ),
        axis=1,
    ),
    columns=["Sample", "Sigma"],
)
sns.scatterplot(data=sigma_df_mcmc, x="Sample", y="Sigma")
plt.show()

sigma_df_mcmc = pd.DataFrame( np.concatenate( ( np.expand_dims(np.arange(bart_model.num_samples), axis=1), np.expand_dims(bart_model.global_var_samples, axis=1), ), axis=1, ), columns=["Sample", "Sigma"], ) sns.scatterplot(data=sigma_df_mcmc, x="Sample", y="Sigma") plt.show()

Compute the test set RMSE

In [7]:

np.sqrt(np.mean(np.power(y_test - np.squeeze(y_avg_mcmc), 2)))

np.sqrt(np.mean(np.power(y_test - np.squeeze(y_avg_mcmc), 2)))

Out[7]:

np.float64(1.1162604669336753)

Demo 2: Including W as a covariate in the standard "constant leaf" BART model¶

Run BART

In [8]:

bart_model = BARTModel()
X_train_aug = np.c_[X_train, basis_train]
X_test_aug = np.c_[X_test, basis_test]
bart_model.sample(
    X_train=X_train_aug, y_train=y_train, X_test=X_test_aug, num_gfr=10, num_mcmc=100
)

bart_model = BARTModel() X_train_aug = np.c_[X_train, basis_train] X_test_aug = np.c_[X_test, basis_test] bart_model.sample( X_train=X_train_aug, y_train=y_train, X_test=X_test_aug, num_gfr=10, num_mcmc=100 )

Inspect the MCMC (BART) samples

In [9]:

forest_preds_y_mcmc = bart_model.y_hat_test
y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis=1, keepdims=True)
y_df_mcmc = pd.DataFrame(
    np.concatenate((np.expand_dims(y_test, 1), y_avg_mcmc), axis=1),
    columns=["True outcome", "Average estimated outcome"],
)
sns.scatterplot(data=y_df_mcmc, x="Average estimated outcome", y="True outcome")
plt.axline((0, 0), slope=1, color="black", linestyle=(0, (3, 3)))
plt.show()

forest_preds_y_mcmc = bart_model.y_hat_test y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis=1, keepdims=True) y_df_mcmc = pd.DataFrame( np.concatenate((np.expand_dims(y_test, 1), y_avg_mcmc), axis=1), columns=["True outcome", "Average estimated outcome"], ) sns.scatterplot(data=y_df_mcmc, x="Average estimated outcome", y="True outcome") plt.axline((0, 0), slope=1, color="black", linestyle=(0, (3, 3))) plt.show()

In [10]:

sigma_df_mcmc = pd.DataFrame(
    np.concatenate(
        (
            np.expand_dims(np.arange(bart_model.num_samples), axis=1),
            np.expand_dims(bart_model.global_var_samples, axis=1),
        ),
        axis=1,
    ),
    columns=["Sample", "Sigma"],
)
sns.scatterplot(data=sigma_df_mcmc, x="Sample", y="Sigma")
plt.show()

sigma_df_mcmc = pd.DataFrame( np.concatenate( ( np.expand_dims(np.arange(bart_model.num_samples), axis=1), np.expand_dims(bart_model.global_var_samples, axis=1), ), axis=1, ), columns=["Sample", "Sigma"], ) sns.scatterplot(data=sigma_df_mcmc, x="Sample", y="Sigma") plt.show()

Compute the test set RMSE

In [11]:

np.sqrt(np.mean(np.power(y_test - np.squeeze(y_avg_mcmc), 2)))

np.sqrt(np.mean(np.power(y_test - np.squeeze(y_avg_mcmc), 2)))

Out[11]:

np.float64(1.271813073032276)

Demo 3: Omitting W entirely¶

Run BART

In [12]:

bart_model = BARTModel()
bart_model.sample(
    X_train=X_train, y_train=y_train, X_test=X_test, num_gfr=10, num_mcmc=100
)

bart_model = BARTModel() bart_model.sample( X_train=X_train, y_train=y_train, X_test=X_test, num_gfr=10, num_mcmc=100 )

Inspect the MCMC (BART) samples

In [13]:

forest_preds_y_mcmc = bart_model.y_hat_test
y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis=1, keepdims=True)
y_df_mcmc = pd.DataFrame(
    np.concatenate((np.expand_dims(y_test, 1), y_avg_mcmc), axis=1),
    columns=["True outcome", "Average estimated outcome"],
)
sns.scatterplot(data=y_df_mcmc, x="Average estimated outcome", y="True outcome")
plt.axline((0, 0), slope=1, color="black", linestyle=(0, (3, 3)))
plt.show()

forest_preds_y_mcmc = bart_model.y_hat_test y_avg_mcmc = np.squeeze(forest_preds_y_mcmc).mean(axis=1, keepdims=True) y_df_mcmc = pd.DataFrame( np.concatenate((np.expand_dims(y_test, 1), y_avg_mcmc), axis=1), columns=["True outcome", "Average estimated outcome"], ) sns.scatterplot(data=y_df_mcmc, x="Average estimated outcome", y="True outcome") plt.axline((0, 0), slope=1, color="black", linestyle=(0, (3, 3))) plt.show()

In [14]:

sigma_df_mcmc = pd.DataFrame(
    np.concatenate(
        (
            np.expand_dims(np.arange(bart_model.num_samples), axis=1),
            np.expand_dims(bart_model.global_var_samples, axis=1),
        ),
        axis=1,
    ),
    columns=["Sample", "Sigma"],
)
sns.scatterplot(data=sigma_df_mcmc, x="Sample", y="Sigma")
plt.show()

sigma_df_mcmc = pd.DataFrame( np.concatenate( ( np.expand_dims(np.arange(bart_model.num_samples), axis=1), np.expand_dims(bart_model.global_var_samples, axis=1), ), axis=1, ), columns=["Sample", "Sigma"], ) sns.scatterplot(data=sigma_df_mcmc, x="Sample", y="Sigma") plt.show()

Compute the test set RMSE

In [15]:

np.sqrt(np.mean(np.power(y_test - np.squeeze(y_avg_mcmc), 2)))

np.sqrt(np.mean(np.power(y_test - np.squeeze(y_avg_mcmc), 2)))

Out[15]:

np.float64(1.9887555693328827)