From 28377ffdb3413d1782f548c51b83590684aead97 Mon Sep 17 00:00:00 2001 From: Drew Herren Date: Tue, 16 Jun 2026 22:05:20 -0500 Subject: [PATCH 1/2] Fix IndexError in mean-only Python BART with categorical covariates The "zero out excluded variable weights" step in BARTModel.sample ran outside the include_*_forest guards, while the weight arrays are only expanded to processed (post-preprocessing, e.g. one-hot) length inside those guards. A mean-only model with categorical covariates therefore indexed an unexpanded variable_weights_variance array with a processed-length boolean mask and raised IndexError. Guard each zero-out with its include_*_forest flag, matching the R implementation in R/bart.R. Adds a regression test (mean-only BART on a categorical DataFrame) and a NEWS entry. Co-Authored-By: Claude Opus 4.8 --- NEWS.md | 1 + stochtree/bart.py | 20 +++++++++++++------- test/python/test_bart.py | 24 ++++++++++++++++++++++++ 3 files changed, 38 insertions(+), 7 deletions(-) diff --git a/NEWS.md b/NEWS.md index c3c2f6bd..ba9bc6db 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,6 +2,7 @@ ## Bug Fixes +* Fix `IndexError` when sampling a mean-only Python BART model with categorical covariates; the excluded-variable weight zero-out is now guarded by the forest-inclusion flags, matching the R implementation * Fix ordinal class prediction bug in R BART [#399](https://github.com/StochasticTree/stochtree/issues/399) # stochtree 0.4.4 diff --git a/stochtree/bart.py b/stochtree/bart.py index 6b57029d..2c352a11 100644 --- a/stochtree/bart.py +++ b/stochtree/bart.py @@ -990,13 +990,19 @@ def sample( * variable_weights_adj ) - # Zero out weights for excluded variables - variable_weights_mean[ - [variable_subset_mean.count(i) == 0 for i in original_var_indices] - ] = 0 - variable_weights_variance[ - [variable_subset_variance.count(i) == 0 for i in original_var_indices] - ] = 0 + # Zero out weights for excluded variables. The weight arrays are only + # expanded to processed (post-preprocessing) length inside the + # include_*_forest guards above, so the zero-out must be guarded the same + # way -- otherwise a mean-only (or variance-only) model with categorical + # covariates indexes an unexpanded array and raises. (Matches R's logic.) + if self.include_mean_forest: + variable_weights_mean[ + [variable_subset_mean.count(i) == 0 for i in original_var_indices] + ] = 0 + if self.include_variance_forest: + variable_weights_variance[ + [variable_subset_variance.count(i) == 0 for i in original_var_indices] + ] = 0 # Set num_features_subsample to default, ncol(X_train), if not already set if num_features_subsample_mean is None: diff --git a/test/python/test_bart.py b/test/python/test_bart.py index 4cab6707..50864f51 100644 --- a/test/python/test_bart.py +++ b/test/python/test_bart.py @@ -1,4 +1,5 @@ import numpy as np +import pandas as pd import pytest from sklearn.model_selection import train_test_split @@ -1578,6 +1579,29 @@ def test_cloglog_ordinal_bart_with_gfr(self): assert bart_model.y_hat_test.shape == (n_test, num_mcmc) assert bart_model.cloglog_cutpoint_samples.shape == (2, num_mcmc) + def test_categorical_covariates_mean_only(self): + """A mean-only BART model with categorical (one-hot expanded) covariates + must sample and predict without error. + + Regression test: the "zero out excluded variable weights" step ran + outside the include_*_forest guards, so variable_weights_variance was + never expanded to the processed (one-hot) length for a model without a + variance forest, and indexing it raised an IndexError. + """ + rng = np.random.default_rng(0) + n = 100 + X_num = rng.uniform(0, 1, (n, 3)) + X = pd.DataFrame(X_num, columns=["a", "b", "c"]) + X["cat"] = pd.Categorical(rng.choice(["x", "y", "z"], size=n)) + y = X_num[:, 0] + rng.normal(scale=0.5, size=n) + + model = BARTModel() + # Mean forest only (no variance forest) is the failing configuration. + model.sample(X_train=X, y_train=y, num_gfr=0, num_burnin=0, num_mcmc=5) + + preds = model.predict(X) + assert preds["y_hat"].shape[0] == n + class TestBARTFloat32: """Tests that float32 inputs are accepted and produce valid results (GH #389).""" From 13d475ff8597bf7000e57cb7d0718423055e90c3 Mon Sep 17 00:00:00 2001 From: Drew Herren Date: Wed, 17 Jun 2026 00:16:29 -0500 Subject: [PATCH 2/2] Update NEWS.md --- NEWS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NEWS.md b/NEWS.md index ba9bc6db..26ed76df 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,7 +2,7 @@ ## Bug Fixes -* Fix `IndexError` when sampling a mean-only Python BART model with categorical covariates; the excluded-variable weight zero-out is now guarded by the forest-inclusion flags, matching the R implementation +* Fix `IndexError` when sampling a mean-only Python BART model with categorical covariates; the excluded-variable weight zero-out is now guarded by the forest-inclusion flags, matching the R implementation [#411](https://github.com/StochasticTree/stochtree/pull/411) * Fix ordinal class prediction bug in R BART [#399](https://github.com/StochasticTree/stochtree/issues/399) # stochtree 0.4.4