From 3c33ce3705ac9b2900c0a1747ce785f7d6774c4b Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Mon, 22 Jun 2026 16:01:37 +0300 Subject: [PATCH 1/9] Fixes #279 --- .gitignore | 1 + tests/test_bart.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+) diff --git a/.gitignore b/.gitignore index 5e62ce3..beb8348 100644 --- a/.gitignore +++ b/.gitignore @@ -167,3 +167,4 @@ cython_debug/ **.csv tests/test_hmc.py +tests/bench_bart.py diff --git a/tests/test_bart.py b/tests/test_bart.py index 5883fb6..15857b3 100644 --- a/tests/test_bart.py +++ b/tests/test_bart.py @@ -239,3 +239,17 @@ def test_multiple_bart_variables_manual_step(): assert "mu2" in idata.posterior assert idata.posterior["mu1"].shape == (1, 20, 30) assert idata.posterior["mu2"].shape == (1, 20, 30) + + +def test_mutable_named_dim(): + """Test that BART variables can be created with mutable named dimensions and that sampling works.""" + rng = np.random.default_rng(0) + N = 50 + X = rng.normal(size=(N, 2)) + Y = rng.normal(size=N) + + with pm.Model(coords={"obs": np.arange(N), "feature": ["a", "b"]}) as model: + x = pm.Data("x", X, dims=("obs", "feature")) + mu = pmb.BART("mu", X=x, Y=Y, m=10, dims="obs") + pm.Normal("y", mu=mu, sigma=1.0, observed=Y, dims="obs") + idata = pm.sample(tune=20, draws=20, chains=1) \ No newline at end of file From 44612310eaea05df6e8d9ecfb43382db6af32fbc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 13:02:25 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_bart.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_bart.py b/tests/test_bart.py index 15857b3..843f200 100644 --- a/tests/test_bart.py +++ b/tests/test_bart.py @@ -252,4 +252,4 @@ def test_mutable_named_dim(): x = pm.Data("x", X, dims=("obs", "feature")) mu = pmb.BART("mu", X=x, Y=Y, m=10, dims="obs") pm.Normal("y", mu=mu, sigma=1.0, observed=Y, dims="obs") - idata = pm.sample(tune=20, draws=20, chains=1) \ No newline at end of file + idata = pm.sample(tune=20, draws=20, chains=1) From 10b1e54a052488cacae9d168537b50df6db26d85 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Mon, 22 Jun 2026 16:09:52 +0300 Subject: [PATCH 3/9] fix precommit --- tests/test_bart.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_bart.py b/tests/test_bart.py index 15857b3..1c23dce 100644 --- a/tests/test_bart.py +++ b/tests/test_bart.py @@ -242,7 +242,8 @@ def test_multiple_bart_variables_manual_step(): def test_mutable_named_dim(): - """Test that BART variables can be created with mutable named dimensions and that sampling works.""" + """Test that BART variables can be created with + mutable named dimensions and that sampling works.""" rng = np.random.default_rng(0) N = 50 X = rng.normal(size=(N, 2)) From c48ac3627347757c48228693935315ec1ea4b117 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Wed, 8 Jul 2026 13:50:47 +0300 Subject: [PATCH 4/9] use PosteriorSampler --- pymc_bart/__init__.py | 4 +- pymc_bart/pymc_bart.py | 14 +-- pymc_bart/utils.py | 188 ++++++++++++++++++++++++++--------------- 3 files changed, 130 insertions(+), 76 deletions(-) diff --git a/pymc_bart/__init__.py b/pymc_bart/__init__.py index c3fca0e..74b40b4 100644 --- a/pymc_bart/__init__.py +++ b/pymc_bart/__init__.py @@ -12,10 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# If removed, the PGBART step method will not be detected by PyMC. -import bartrs # registers PGBART with PyMC. +import bartrs -# Fix precommit complaining about unused import if bartrs: pass diff --git a/pymc_bart/pymc_bart.py b/pymc_bart/pymc_bart.py index d92a066..405e6e8 100644 --- a/pymc_bart/pymc_bart.py +++ b/pymc_bart/pymc_bart.py @@ -1,11 +1,15 @@ try: - from bartrs.bartrs import PyBartSettings, PySampler, TreeArrays + from bartrs.bartrs import PosteriorSampler, PyBartSettings, PySampler, TreeArrays except Exception as e: - print(f"Warning: Could not import PyBartSettings, PySampler, or TreeArrays due to: {e}") - PyBartSettings = PySampler = TreeArrays = None + print(f"Warning: Could not import PyBartSettings, PySampler, TreeArrays, or PosteriorSampler due to: {e}") + PyBartSettings = PySampler = TreeArrays = PosteriorSampler = None -for cls in (PyBartSettings, PySampler, TreeArrays): +for cls in (PyBartSettings, PySampler, TreeArrays, PosteriorSampler): if cls is not None: cls.__module__ = "pymc_bart.pymc_bart" -__all__ = [name for name in ("PyBartSettings", "PySampler", "TreeArrays") if globals().get(name)] +__all__ = [ + name + for name in ("PyBartSettings", "PySampler", "TreeArrays", "PosteriorSampler") + if globals().get(name) +] diff --git a/pymc_bart/utils.py b/pymc_bart/utils.py index e0f5aa1..b5ffa54 100644 --- a/pymc_bart/utils.py +++ b/pymc_bart/utils.py @@ -4,9 +4,11 @@ from __future__ import annotations import base64 +import pickle import warnings from collections.abc import Callable from typing import TYPE_CHECKING, Any, TypeVar +from weakref import WeakValueDictionary import matplotlib.pyplot as plt import numpy as np @@ -21,13 +23,13 @@ from scipy.signal import savgol_filter if TYPE_CHECKING: - from pymc_bart.pymc_bart import TreeArrays + from pymc_bart.pymc_bart import PosteriorSampler TensorLike = TypeVar("TensorLike", npt.NDArray, pt.TensorVariable) def _sample_posterior( - all_trees: list[list[TreeArrays]] | list[list[list[TreeArrays]]], + sampler, X: TensorLike, rng: np.random.Generator, size: int | tuple[int, ...] | None = None, @@ -38,8 +40,7 @@ def _sample_posterior( Parameters ---------- - all_trees : list - List of all trees sampled from a posterior + sampler : PosteriorSampler or list of PosteriorSampler X : tensor-like A covariate matrix. Use the same used to fit BART for in-sample predictions or a new one for out-of-sample predictions. @@ -49,15 +50,8 @@ def _sample_posterior( excluded : Optional[npt.NDArray[np.int_]] Indexes of the variables to exclude when computing predictions """ - stacked_trees = all_trees - - if isinstance(X, Variable): - X = X.eval() - - X = np.ascontiguousarray(np.asarray(X, dtype=np.float64)) - if size is None: - size_iter: list | tuple = () + size_iter = () elif isinstance(size, int): size_iter = [size] else: @@ -67,33 +61,68 @@ def _sample_posterior( for s in size_iter: flatten_size *= s - if isinstance(all_trees[0][0], list): - n_draws = len(all_trees[0]) - n_outputs = len(all_trees) + X = np.ascontiguousarray(np.asarray(X, dtype=np.float64)) + excl = list(excluded) if excluded is not None else None - pred = np.zeros((flatten_size, n_outputs, X.shape[0])) + if isinstance(sampler, list): + # shape (flatten_size, n_outputs, n_data_samples) each; stack along the outputs axis + pred = np.concatenate([s.sample_posterior(X, flatten_size, excl) for s in sampler], axis=1) + else: + pred = sampler.sample_posterior(X, flatten_size, excl) - idx = rng.integers(0, n_draws, size=flatten_size) + return pred.transpose((0, 2, 1)).reshape((*size_iter, -1, pred.shape[1])) - for out_idx, trees in enumerate(all_trees): - for ind, p in enumerate(pred[:, out_idx, :]): - for tree in trees[idx[ind]]: - pred[ind, out_idx, :] += np.asarray( - tree.predict(x=X, excluded=excluded) - ).squeeze() - return pred.transpose((0, 2, 1)).reshape((*size_iter, -1, n_outputs)) - else: - n_outputs = all_trees[0][0].n_outputs - idx = rng.integers(0, len(stacked_trees), size=flatten_size) +_posterior_sampler_cache: "WeakValueDictionary[tuple[int, str], PosteriorSampler]" = ( + WeakValueDictionary() +) + + +def _get_posterior_sampler( + idata: Any, + bartrv: Variable, + model: "pm.Model | None" = None, + random_seed: int | None = None, +) -> "PosteriorSampler": + """ + Rebuild a PosteriorSampler from tree data streamed through ``idata.sample_stats``. + + """ + cache_key = (id(idata), bartrv.name) + cached = _posterior_sampler_cache.get(cache_key) + if cached is not None: + return cached - pred = np.zeros((flatten_size, n_outputs, X.shape[0])) + from pymc_bart.pymc_bart import PosteriorSampler - for ind, p in enumerate(pred): - for tree in stacked_trees[idx[ind]]: - pred[ind, :, :] += tree.predict(x=X, excluded=excluded) + trees_xarray = idata["sample_stats"]["trees"] + baseline_xarray = idata["sample_stats"]["baseline_forest"] - return pred.transpose((0, 2, 1)).reshape((*size_iter, -1, n_outputs)) + if trees_xarray.trees_dim_0.size > 1: + if model is None: + raise ValueError( + "The InferenceData was generated from a model with multiple BART variables, \n" + "please provide the model so the right BART variable's trees can be located." + ) + index = [var.name for var in model.free_RVs].index(bartrv.name) + trees_vals = trees_xarray.sel({"trees_dim_0": index}).values.ravel() + baseline_vals = baseline_xarray.sel({"baseline_forest_dim_0": index}).values.ravel() + else: + trees_vals = trees_xarray.values.ravel() + baseline_vals = baseline_xarray.values.ravel() + + batches = [_decode_obj(val) for val in trees_vals if val is not None] + baseline_forest = next((_decode_obj(val) for val in baseline_vals if val is not None), None) + + sampler = PosteriorSampler.from_history( + batches, + baseline_forest, + bartrv.owner.op.m, + bartrv.owner.op.n_outputs, + random_seed if random_seed is not None else 0, + ) + _posterior_sampler_cache[cache_key] = sampler + return sampler def plot_convergence( @@ -108,8 +137,8 @@ def plot_convergence( Parameters ---------- - idata : DataTree - DataTree object containing the posterior samples. + idata : InferenceData + InferenceData object containing the posterior samples. var_name : Optional[str] Name of the BART variable to plot. Defaults to None. kind : str @@ -134,6 +163,7 @@ def plot_convergence( def plot_ice( bartrv: Variable, X: npt.NDArray, + idata: Any, Y: npt.NDArray | None = None, var_idx: list[int] | None = None, var_discrete: list[int] | None = None, @@ -151,6 +181,7 @@ def plot_ice( figsize: tuple[float, float] | None = None, smooth_kwargs: dict[str, Any] | None = None, ax: plt.Axes | None = None, + model: "pm.Model | None" = None, ) -> list[plt.Axes]: """ Individual conditional expectation plot. @@ -161,6 +192,8 @@ def plot_ice( BART variable once the model that include it has been fitted. X : npt.NDArray The covariate matrix. + idata : InferenceData + InferenceData returned by ``pm.sample()``. Y : Optional[npt.NDArray], by default None. The response vector. var_idx : Optional[list[int]], by default None. @@ -201,13 +234,16 @@ def plot_ice( See scipy.signal.savgol_filter() for details. ax : axes Matplotlib axes. + model : Optional[pm.Model] + The PyMC model that contains the BART variable. Only needed if the model contains + multiple BART variables. Returns ------- axes: matplotlib axes """ - all_trees = bartrv.owner.op.all_trees rng = np.random.default_rng(random_seed) + sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) if func is None: @@ -227,7 +263,7 @@ def identity(x): _, ) = _prepare_plot_data(X, Y, "linear", None, var_idx, var_discrete) - fig, axes, shape = _create_figure_axes(bartrv, var_idx, grid, sharey, figsize, ax) + fig, axes, shape = _create_figure_axes(bartrv, idata, var_idx, grid, sharey, figsize, ax, model) instances_ary = rng.choice(range(X.shape[0]), replace=False, size=instances) idx_s = list(range(X.shape[0])) @@ -242,7 +278,7 @@ def identity(x): fake_X[:, indices_mi] = X[:, indices_mi][instance] y_pred.append( np.mean( - _sample_posterior(all_trees, X=fake_X, rng=rng, size=samples), + _sample_posterior(sampler, X=fake_X, rng=rng, size=samples), 0, ) ) @@ -278,6 +314,7 @@ def identity(x): def plot_pdp( bartrv: Variable | list[Variable], X: npt.NDArray, + idata: Any, Y: npt.NDArray | None = None, xs_interval: str = "quantiles", xs_values: int | list[float] | None = None, @@ -296,6 +333,7 @@ def plot_pdp( figsize: tuple[float, float] | None = None, smooth_kwargs: dict[str, Any] | None = None, ax: plt.Axes = None, + model: "pm.Model | None" = None, ) -> list[plt.Axes]: """ Partial dependence plot. @@ -306,6 +344,8 @@ def plot_pdp( BART variable once the model that include it has been fitted. X : npt.NDArray The covariate matrix. + idata : InferenceData + InferenceData returned by ``pm.sample()``. Y : Optional[npt.NDArray], by default None. The response vector. xs_interval : str @@ -353,6 +393,9 @@ def plot_pdp( See scipy.signal.savgol_filter() for details. ax : axes Matplotlib axes. + model : Optional[pm.Model] + The PyMC model that contains the BART variable. Only needed if the model contains + multiple BART variables. Returns ------- @@ -361,9 +404,12 @@ def plot_pdp( if isinstance(bartrv, list): if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") - all_trees = [rv.owner.op.all_trees for rv in bartrv] + sampler = [ + _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) + for rv in bartrv + ] else: - all_trees = bartrv.owner.op.all_trees + sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) rng = np.random.default_rng(random_seed) @@ -385,7 +431,7 @@ def identity(x): xs_values, ) = _prepare_plot_data(X, Y, xs_interval, xs_values, var_idx, var_discrete) - fig, axes, shape = _create_figure_axes(bartrv, var_idx, grid, sharey, figsize, ax) + fig, axes, shape = _create_figure_axes(bartrv, idata, var_idx, grid, sharey, figsize, ax, model) count = 0 fake_X = _create_pdp_data(X, xs_interval, xs_values) @@ -395,7 +441,7 @@ def identity(x): excluded.remove(var) p_d = func( _sample_posterior( - all_trees, + sampler, X=fake_X, rng=rng, size=samples, @@ -452,11 +498,13 @@ def identity(x): def _create_figure_axes( bartrv: Variable | list[Variable], + idata: Any, var_idx: list[int], grid: str = "long", sharey: bool = True, figsize: tuple[float, float] | None = None, ax: plt.Axes | None = None, + model: "pm.Model | None" = None, ) -> tuple[plt.Figure, list[plt.Axes], int]: """ Create and return the figure and axes objects for plotting the variables. @@ -467,6 +515,7 @@ def _create_figure_axes( ---------- bartrv : BART Random Variable BART variable once the model that include it has been fitted. + idata : InferenceData var_idx : Optional[list[int]], by default None. List of the indices of the covariate for which to compute the pdp or ice. var_discrete : Optional[list[int]], by default None. @@ -480,28 +529,24 @@ def _create_figure_axes( Figure size. If None it will be defined automatically. ax : axes Matplotlib axes. - + model : Optional[pm.Model] + The PyMC model that contains the BART variable. Only needed if the model contains + multiple BART variables. Returns ------- tuple[plt.Figure, list[plt.Axes], int] A tuple containing the figure object, list of axes objects, and the shape value. """ - from pymc_bart.pymc_bart import TreeArrays - if not isinstance(bartrv, list): if bartrv.ndim == 1: # type: ignore shape = 1 - elif isinstance(bartrv.owner_op.all_trees[0][0], TreeArrays): - shape = bartrv.owner_op.all_trees[0][0].n_outputs else: - shape = bartrv.eval().shape[0] + shape = _get_posterior_sampler(idata, bartrv, model=model).n_outputs elif all(rv.ndim == 1 for rv in bartrv): shape = len(bartrv) - elif isinstance(bartrv[0].owner_op.all_trees[0][0], TreeArrays): - shape = bartrv[0].owner_op.all_trees[0][0].n_outputs else: - shape = bartrv[0].eval().shape[0] + shape = _get_posterior_sampler(idata, bartrv[0], model=model).n_outputs n_plots = len(var_idx) * shape @@ -722,8 +767,8 @@ def get_variable_inclusion(idata, X, model=None, bart_var_name=None, labels=None Parameters ---------- - idata : DataTree - DataTree with a variable "variable_inclusion" in ``sample_stats`` group + idata : InferenceData + InferenceData with a variable "variable_inclusion" in ``sample_stats`` group X : npt.NDArray The covariate matrix. model : Optional[pm.Model] @@ -751,7 +796,7 @@ def get_variable_inclusion(idata, X, model=None, bart_var_name=None, labels=None if vi_xarray.variable_inclusion_dim_0.size > 1: if model is None or bart_var_name is None: raise ValueError( - "The InfereceData was generated from a model with multiple BART variables, \n" + "The InferenceData was generated from a model with multiple BART variables, \n" "please provide the model and also the name of the BART variable \n" "for which you want to compute the variable inclusion." ) @@ -784,8 +829,8 @@ def plot_variable_inclusion(idata, X, labels=None, figsize=None, plot_kwargs=Non Parameters ---------- - idata : DataTree - DataTree containing a collection of BART_trees in sample_stats group + idata : InferenceData + InferenceData containing a collection of BART_trees in sample_stats group X : npt.NDArray The covariate matrix. labels : Optional[list[str]] @@ -852,8 +897,8 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 Parameters ---------- - idata : DataTree - DataTree containing a "variable_inclusion" variable in the sample_stats group. + idata : InferenceData + InferenceData containing a "variable_inclusion" variable in the sample_stats group. bartrv : BART Random Variable BART variable once the model that include it has been fitted. X : npt.NDArray @@ -883,8 +928,6 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 ------- vi_results: dictionary """ - from pymc_bart.pymc_bart import TreeArrays - if method not in ["VI", "backward", "backward_VI"]: raise ValueError("method must be 'VI', 'backward' or 'backward_VI'") @@ -893,18 +936,19 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 if isinstance(bartrv, list): if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") - all_trees = [rv.owner.op.all_trees for rv in bartrv] + sampler = [ + _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) + for rv in bartrv + ] bart_var_name = [rv.name for rv in bartrv] shape = len(bartrv) else: - all_trees = bartrv.owner.op.all_trees + sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) bart_var_name = bartrv.name if bartrv.ndim == 1: # type: ignore shape = 1 - elif isinstance(all_trees[0][0], TreeArrays): - shape = all_trees[0][0].n_outputs else: - shape = bartrv.eval().shape[0] + shape = sampler.n_outputs n_vars = X.shape[1] @@ -931,14 +975,14 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 else: fixed = 0 init = 0 - predicted_all = _sample_posterior(all_trees, X=X, rng=rng, size=samples, excluded=None) + predicted_all = _sample_posterior(sampler, X=X, rng=rng, size=samples, excluded=None) if method in ["VI", "backward_VI"]: vi_xarray = idata["sample_stats"]["variable_inclusion"] if vi_xarray.variable_inclusion_dim_0.size > 1: if model is None: raise ValueError( - "The InfereceData was generated from a model with multiple BART variables, \n" + "The InferenceData was generated from a model with multiple BART variables, \n" "please provide the model and also the name of the BART variable \n" "for which you want to compute the variable inclusion." ) @@ -969,7 +1013,7 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 for idx, subset in enumerate(subsets): predicted_subset = _sample_posterior( - all_trees=all_trees, + sampler, X=X, rng=rng, size=samples, @@ -1007,7 +1051,7 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 for subset in subsets: # Sample posterior predictions excluding a subset of variables predicted_subset = _sample_posterior( - all_trees=all_trees, X=X, rng=rng, size=samples, excluded=subset + sampler, X=X, rng=rng, size=samples, excluded=subset ) # Calculate Pearson correlation for each sample and find the mean r_2 = np.zeros(samples) @@ -1371,3 +1415,11 @@ def _encode_vi(vec: list[int]) -> str: n >>= 7 result.append(n & 0x7F) return base64.b64encode(bytes(result)).decode("ascii") + + +def _encode_obj(obj: Any) -> str: + return base64.b64encode(pickle.dumps(obj)).decode("ascii") + + +def _decode_obj(s: str) -> Any: + return pickle.loads(base64.b64decode(s)) From dc0fdfae79a585b5ff90f2748354d1a601dc6950 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:53:57 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pymc_bart/pymc_bart.py | 4 +++- pymc_bart/utils.py | 18 ++++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pymc_bart/pymc_bart.py b/pymc_bart/pymc_bart.py index 405e6e8..4642461 100644 --- a/pymc_bart/pymc_bart.py +++ b/pymc_bart/pymc_bart.py @@ -1,7 +1,9 @@ try: from bartrs.bartrs import PosteriorSampler, PyBartSettings, PySampler, TreeArrays except Exception as e: - print(f"Warning: Could not import PyBartSettings, PySampler, TreeArrays, or PosteriorSampler due to: {e}") + print( + f"Warning: Could not import PyBartSettings, PySampler, TreeArrays, or PosteriorSampler due to: {e}" + ) PyBartSettings = PySampler = TreeArrays = PosteriorSampler = None for cls in (PyBartSettings, PySampler, TreeArrays, PosteriorSampler): diff --git a/pymc_bart/utils.py b/pymc_bart/utils.py index b5ffa54..7a8e6cd 100644 --- a/pymc_bart/utils.py +++ b/pymc_bart/utils.py @@ -73,7 +73,7 @@ def _sample_posterior( return pred.transpose((0, 2, 1)).reshape((*size_iter, -1, pred.shape[1])) -_posterior_sampler_cache: "WeakValueDictionary[tuple[int, str], PosteriorSampler]" = ( +_posterior_sampler_cache: WeakValueDictionary[tuple[int, str], PosteriorSampler] = ( WeakValueDictionary() ) @@ -81,9 +81,9 @@ def _sample_posterior( def _get_posterior_sampler( idata: Any, bartrv: Variable, - model: "pm.Model | None" = None, + model: pm.Model | None = None, random_seed: int | None = None, -) -> "PosteriorSampler": +) -> PosteriorSampler: """ Rebuild a PosteriorSampler from tree data streamed through ``idata.sample_stats``. @@ -181,7 +181,7 @@ def plot_ice( figsize: tuple[float, float] | None = None, smooth_kwargs: dict[str, Any] | None = None, ax: plt.Axes | None = None, - model: "pm.Model | None" = None, + model: pm.Model | None = None, ) -> list[plt.Axes]: """ Individual conditional expectation plot. @@ -333,7 +333,7 @@ def plot_pdp( figsize: tuple[float, float] | None = None, smooth_kwargs: dict[str, Any] | None = None, ax: plt.Axes = None, - model: "pm.Model | None" = None, + model: pm.Model | None = None, ) -> list[plt.Axes]: """ Partial dependence plot. @@ -405,8 +405,7 @@ def plot_pdp( if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") sampler = [ - _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) - for rv in bartrv + _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) for rv in bartrv ] else: sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) @@ -504,7 +503,7 @@ def _create_figure_axes( sharey: bool = True, figsize: tuple[float, float] | None = None, ax: plt.Axes | None = None, - model: "pm.Model | None" = None, + model: pm.Model | None = None, ) -> tuple[plt.Figure, list[plt.Axes], int]: """ Create and return the figure and axes objects for plotting the variables. @@ -937,8 +936,7 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") sampler = [ - _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) - for rv in bartrv + _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) for rv in bartrv ] bart_var_name = [rv.name for rv in bartrv] shape = len(bartrv) From 9c50ea4e6ff90ae7ce05cf89e72e64c59f090024 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Wed, 8 Jul 2026 16:00:24 +0300 Subject: [PATCH 6/9] return trees once at the end --- pymc_bart/bart.py | 7 +- pymc_bart/pymc_bart.py | 3 +- pymc_bart/utils.py | 160 +++++++++++++++++------------------------ tests/test_utils.py | 7 +- 4 files changed, 77 insertions(+), 100 deletions(-) diff --git a/pymc_bart/bart.py b/pymc_bart/bart.py index 74c7ef6..27b8e28 100644 --- a/pymc_bart/bart.py +++ b/pymc_bart/bart.py @@ -27,7 +27,7 @@ from pytensor.tensor.sharedvar import TensorSharedVariable from pytensor.tensor.variable import TensorVariable -from .utils import TensorLike, _sample_posterior +from .utils import TensorLike, _get_posterior_sampler, _sample_posterior __all__ = ["BART"] @@ -62,7 +62,10 @@ def rng_fn( # pylint: disable=W0237 else: return np.full(Y.shape[0], Y.mean()) else: - return _sample_posterior(cls.all_trees, cls.X, rng=rng).squeeze().T + shape = size[0] if size is not None else 1 + sampler = _get_posterior_sampler(cls) + pred = _sample_posterior(sampler, X, rng=rng, size=shape) + return pred.squeeze().T bart = BARTRV() diff --git a/pymc_bart/pymc_bart.py b/pymc_bart/pymc_bart.py index 405e6e8..09f2d6b 100644 --- a/pymc_bart/pymc_bart.py +++ b/pymc_bart/pymc_bart.py @@ -1,7 +1,8 @@ try: from bartrs.bartrs import PosteriorSampler, PyBartSettings, PySampler, TreeArrays except Exception as e: - print(f"Warning: Could not import PyBartSettings, PySampler, TreeArrays, or PosteriorSampler due to: {e}") + print(f"""Warning: Could not import PyBartSettings, + PySampler, TreeArrays, or PosteriorSampler due to: {e}""") PyBartSettings = PySampler = TreeArrays = PosteriorSampler = None for cls in (PyBartSettings, PySampler, TreeArrays, PosteriorSampler): diff --git a/pymc_bart/utils.py b/pymc_bart/utils.py index b5ffa54..273e267 100644 --- a/pymc_bart/utils.py +++ b/pymc_bart/utils.py @@ -4,11 +4,9 @@ from __future__ import annotations import base64 -import pickle import warnings from collections.abc import Callable -from typing import TYPE_CHECKING, Any, TypeVar -from weakref import WeakValueDictionary +from typing import Any, TypeVar import matplotlib.pyplot as plt import numpy as np @@ -22,9 +20,6 @@ from scipy.interpolate import griddata from scipy.signal import savgol_filter -if TYPE_CHECKING: - from pymc_bart.pymc_bart import PosteriorSampler - TensorLike = TypeVar("TensorLike", npt.NDArray, pt.TensorVariable) @@ -51,77 +46,88 @@ def _sample_posterior( Indexes of the variables to exclude when computing predictions """ if size is None: - size_iter = () + size_iter: tuple[int, ...] | list[int] = () elif isinstance(size, int): size_iter = [size] else: size_iter = size flatten_size = 1 + s: int for s in size_iter: flatten_size *= s X = np.ascontiguousarray(np.asarray(X, dtype=np.float64)) excl = list(excluded) if excluded is not None else None + first_sampler = sampler[0] if isinstance(sampler, list) else sampler + draw_indices = rng.integers(0, first_sampler.n_draws, size=flatten_size).tolist() if isinstance(sampler, list): # shape (flatten_size, n_outputs, n_data_samples) each; stack along the outputs axis - pred = np.concatenate([s.sample_posterior(X, flatten_size, excl) for s in sampler], axis=1) + pred = np.concatenate( + [s.sample_posterior(X, draw_indices, excl) for s in sampler], axis=1 + ) else: - pred = sampler.sample_posterior(X, flatten_size, excl) + pred = sampler.sample_posterior(X, draw_indices, excl) return pred.transpose((0, 2, 1)).reshape((*size_iter, -1, pred.shape[1])) -_posterior_sampler_cache: "WeakValueDictionary[tuple[int, str], PosteriorSampler]" = ( - WeakValueDictionary() -) +class _MultiChainSampler: + """ + Dispatches each requested draw to the PosteriorSampler for the chain it came from. + """ + def __init__(self, chain_samplers: list): + if not chain_samplers: + raise ValueError("No posterior draws available yet: run pm.sample() first.") + self._chain_samplers = chain_samplers + self._offsets = np.cumsum([0] + [s.n_draws for s in chain_samplers]) + @property + def n_draws(self) -> int: + return int(self._offsets[-1]) -def _get_posterior_sampler( - idata: Any, - bartrv: Variable, - model: "pm.Model | None" = None, - random_seed: int | None = None, -) -> "PosteriorSampler": - """ - Rebuild a PosteriorSampler from tree data streamed through ``idata.sample_stats``. + @property + def n_outputs(self) -> int: + return self._chain_samplers[0].n_outputs + def sample_posterior(self, X, draw_indices, excluded): + draw_indices = np.asarray(draw_indices) + chain_of_draw = np.searchsorted(self._offsets, draw_indices, side="right") - 1 + + out = None + for chain_idx, sampler in enumerate(self._chain_samplers): + mask = chain_of_draw == chain_idx + if not np.any(mask): + continue + local_indices = (draw_indices[mask] - self._offsets[chain_idx]).tolist() + preds = sampler.sample_posterior(X, local_indices, excluded) + if out is None: + out = np.empty((len(draw_indices), *preds.shape[1:]), dtype=preds.dtype) + out[mask] = preds + return out + + +_posterior_sampler_cache: dict[int, tuple[int, "_MultiChainSampler"]] = {} + + +def _get_posterior_sampler(op) -> "_MultiChainSampler": + """ + Rebuild a prediction-only sampler from tree data accumulated on the BART op itself. """ - cache_key = (id(idata), bartrv.name) - cached = _posterior_sampler_cache.get(cache_key) - if cached is not None: - return cached + n_chains = len(op.all_trees) + cached = _posterior_sampler_cache.get(id(op)) + if cached is not None and cached[0] == n_chains: + return cached[1] from pymc_bart.pymc_bart import PosteriorSampler - trees_xarray = idata["sample_stats"]["trees"] - baseline_xarray = idata["sample_stats"]["baseline_forest"] - - if trees_xarray.trees_dim_0.size > 1: - if model is None: - raise ValueError( - "The InferenceData was generated from a model with multiple BART variables, \n" - "please provide the model so the right BART variable's trees can be located." - ) - index = [var.name for var in model.free_RVs].index(bartrv.name) - trees_vals = trees_xarray.sel({"trees_dim_0": index}).values.ravel() - baseline_vals = baseline_xarray.sel({"baseline_forest_dim_0": index}).values.ravel() - else: - trees_vals = trees_xarray.values.ravel() - baseline_vals = baseline_xarray.values.ravel() - - batches = [_decode_obj(val) for val in trees_vals if val is not None] - baseline_forest = next((_decode_obj(val) for val in baseline_vals if val is not None), None) - - sampler = PosteriorSampler.from_history( - batches, - baseline_forest, - bartrv.owner.op.m, - bartrv.owner.op.n_outputs, - random_seed if random_seed is not None else 0, - ) - _posterior_sampler_cache[cache_key] = sampler + chain_samplers = [ + PosteriorSampler.from_history(batches, baseline_forest, op.m, op.n_outputs) + for baseline_forest, batches in op.all_trees + ] + sampler = _MultiChainSampler(chain_samplers) + _posterior_sampler_cache[id(op)] = (n_chains, sampler) return sampler @@ -163,7 +169,6 @@ def plot_convergence( def plot_ice( bartrv: Variable, X: npt.NDArray, - idata: Any, Y: npt.NDArray | None = None, var_idx: list[int] | None = None, var_discrete: list[int] | None = None, @@ -181,7 +186,6 @@ def plot_ice( figsize: tuple[float, float] | None = None, smooth_kwargs: dict[str, Any] | None = None, ax: plt.Axes | None = None, - model: "pm.Model | None" = None, ) -> list[plt.Axes]: """ Individual conditional expectation plot. @@ -192,8 +196,6 @@ def plot_ice( BART variable once the model that include it has been fitted. X : npt.NDArray The covariate matrix. - idata : InferenceData - InferenceData returned by ``pm.sample()``. Y : Optional[npt.NDArray], by default None. The response vector. var_idx : Optional[list[int]], by default None. @@ -234,16 +236,13 @@ def plot_ice( See scipy.signal.savgol_filter() for details. ax : axes Matplotlib axes. - model : Optional[pm.Model] - The PyMC model that contains the BART variable. Only needed if the model contains - multiple BART variables. Returns ------- axes: matplotlib axes """ rng = np.random.default_rng(random_seed) - sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) + sampler = _get_posterior_sampler(bartrv.owner.op) if func is None: @@ -263,7 +262,7 @@ def identity(x): _, ) = _prepare_plot_data(X, Y, "linear", None, var_idx, var_discrete) - fig, axes, shape = _create_figure_axes(bartrv, idata, var_idx, grid, sharey, figsize, ax, model) + fig, axes, shape = _create_figure_axes(bartrv, var_idx, grid, sharey, figsize, ax) instances_ary = rng.choice(range(X.shape[0]), replace=False, size=instances) idx_s = list(range(X.shape[0])) @@ -314,7 +313,6 @@ def identity(x): def plot_pdp( bartrv: Variable | list[Variable], X: npt.NDArray, - idata: Any, Y: npt.NDArray | None = None, xs_interval: str = "quantiles", xs_values: int | list[float] | None = None, @@ -333,7 +331,6 @@ def plot_pdp( figsize: tuple[float, float] | None = None, smooth_kwargs: dict[str, Any] | None = None, ax: plt.Axes = None, - model: "pm.Model | None" = None, ) -> list[plt.Axes]: """ Partial dependence plot. @@ -344,8 +341,6 @@ def plot_pdp( BART variable once the model that include it has been fitted. X : npt.NDArray The covariate matrix. - idata : InferenceData - InferenceData returned by ``pm.sample()``. Y : Optional[npt.NDArray], by default None. The response vector. xs_interval : str @@ -393,9 +388,6 @@ def plot_pdp( See scipy.signal.savgol_filter() for details. ax : axes Matplotlib axes. - model : Optional[pm.Model] - The PyMC model that contains the BART variable. Only needed if the model contains - multiple BART variables. Returns ------- @@ -404,12 +396,9 @@ def plot_pdp( if isinstance(bartrv, list): if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") - sampler = [ - _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) - for rv in bartrv - ] + sampler = [_get_posterior_sampler(rv.owner.op) for rv in bartrv] else: - sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) + sampler = _get_posterior_sampler(bartrv.owner.op) rng = np.random.default_rng(random_seed) @@ -431,7 +420,7 @@ def identity(x): xs_values, ) = _prepare_plot_data(X, Y, xs_interval, xs_values, var_idx, var_discrete) - fig, axes, shape = _create_figure_axes(bartrv, idata, var_idx, grid, sharey, figsize, ax, model) + fig, axes, shape = _create_figure_axes(bartrv, var_idx, grid, sharey, figsize, ax) count = 0 fake_X = _create_pdp_data(X, xs_interval, xs_values) @@ -498,13 +487,11 @@ def identity(x): def _create_figure_axes( bartrv: Variable | list[Variable], - idata: Any, var_idx: list[int], grid: str = "long", sharey: bool = True, figsize: tuple[float, float] | None = None, ax: plt.Axes | None = None, - model: "pm.Model | None" = None, ) -> tuple[plt.Figure, list[plt.Axes], int]: """ Create and return the figure and axes objects for plotting the variables. @@ -515,7 +502,6 @@ def _create_figure_axes( ---------- bartrv : BART Random Variable BART variable once the model that include it has been fitted. - idata : InferenceData var_idx : Optional[list[int]], by default None. List of the indices of the covariate for which to compute the pdp or ice. var_discrete : Optional[list[int]], by default None. @@ -529,9 +515,6 @@ def _create_figure_axes( Figure size. If None it will be defined automatically. ax : axes Matplotlib axes. - model : Optional[pm.Model] - The PyMC model that contains the BART variable. Only needed if the model contains - multiple BART variables. Returns ------- @@ -542,11 +525,11 @@ def _create_figure_axes( if bartrv.ndim == 1: # type: ignore shape = 1 else: - shape = _get_posterior_sampler(idata, bartrv, model=model).n_outputs + shape = _get_posterior_sampler(bartrv.owner.op).n_outputs elif all(rv.ndim == 1 for rv in bartrv): shape = len(bartrv) else: - shape = _get_posterior_sampler(idata, bartrv[0], model=model).n_outputs + shape = _get_posterior_sampler(bartrv[0].owner.op).n_outputs n_plots = len(var_idx) * shape @@ -936,14 +919,11 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 if isinstance(bartrv, list): if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") - sampler = [ - _get_posterior_sampler(idata, rv, model=model, random_seed=random_seed) - for rv in bartrv - ] + sampler = [_get_posterior_sampler(rv.owner.op) for rv in bartrv] bart_var_name = [rv.name for rv in bartrv] shape = len(bartrv) else: - sampler = _get_posterior_sampler(idata, bartrv, model=model, random_seed=random_seed) + sampler = _get_posterior_sampler(bartrv.owner.op) bart_var_name = bartrv.name if bartrv.ndim == 1: # type: ignore shape = 1 @@ -1415,11 +1395,3 @@ def _encode_vi(vec: list[int]) -> str: n >>= 7 result.append(n & 0x7F) return base64.b64encode(bytes(result)).decode("ascii") - - -def _encode_obj(obj: Any) -> str: - return base64.b64encode(pickle.dumps(obj)).decode("ascii") - - -def _decode_obj(s: str) -> Any: - return pickle.loads(base64.b64decode(s)) diff --git a/tests/test_utils.py b/tests/test_utils.py index e3a42e1..1500e0b 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -18,12 +18,13 @@ class TestUtils: y = pm.Normal("y", mu, sigma, observed=Y) idata = pm.sample(tune=200, draws=200, random_seed=3415) + sampler = pmb.utils._get_posterior_sampler(mu.owner.op) + def test_sample_posterior(self): - all_trees = self.mu.owner.op.all_trees rng = np.random.default_rng(3) - pred_all = pmb.utils._sample_posterior(all_trees, X=self.X, rng=rng, size=2) + pred_all = pmb.utils._sample_posterior(self.sampler, X=self.X, rng=rng, size=2) rng = np.random.default_rng(3) - pred_first = pmb.utils._sample_posterior(all_trees, X=self.X[:10], rng=rng) + pred_first = pmb.utils._sample_posterior(self.sampler, X=self.X[:10], rng=rng) assert_almost_equal(pred_first, pred_all[0, :10], decimal=4) assert pred_all.shape == (2, 50, 1) From eebf6b900dfbaf6dbdb5e70f8e5a4847e94f5682 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 13:03:52 +0000 Subject: [PATCH 7/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pymc_bart/utils.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/pymc_bart/utils.py b/pymc_bart/utils.py index 273e267..bfe579e 100644 --- a/pymc_bart/utils.py +++ b/pymc_bart/utils.py @@ -64,9 +64,7 @@ def _sample_posterior( if isinstance(sampler, list): # shape (flatten_size, n_outputs, n_data_samples) each; stack along the outputs axis - pred = np.concatenate( - [s.sample_posterior(X, draw_indices, excl) for s in sampler], axis=1 - ) + pred = np.concatenate([s.sample_posterior(X, draw_indices, excl) for s in sampler], axis=1) else: pred = sampler.sample_posterior(X, draw_indices, excl) @@ -77,6 +75,7 @@ class _MultiChainSampler: """ Dispatches each requested draw to the PosteriorSampler for the chain it came from. """ + def __init__(self, chain_samplers: list): if not chain_samplers: raise ValueError("No posterior draws available yet: run pm.sample() first.") @@ -108,10 +107,10 @@ def sample_posterior(self, X, draw_indices, excluded): return out -_posterior_sampler_cache: dict[int, tuple[int, "_MultiChainSampler"]] = {} +_posterior_sampler_cache: dict[int, tuple[int, _MultiChainSampler]] = {} -def _get_posterior_sampler(op) -> "_MultiChainSampler": +def _get_posterior_sampler(op) -> _MultiChainSampler: """ Rebuild a prediction-only sampler from tree data accumulated on the BART op itself. """ From 6122b89ac55c7e4ffe349e244ce5f80e1318fd31 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Wed, 8 Jul 2026 16:11:23 +0300 Subject: [PATCH 8/9] pre-commit --- pymc_bart/utils.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/pymc_bart/utils.py b/pymc_bart/utils.py index 273e267..10fd83e 100644 --- a/pymc_bart/utils.py +++ b/pymc_bart/utils.py @@ -64,9 +64,7 @@ def _sample_posterior( if isinstance(sampler, list): # shape (flatten_size, n_outputs, n_data_samples) each; stack along the outputs axis - pred = np.concatenate( - [s.sample_posterior(X, draw_indices, excl) for s in sampler], axis=1 - ) + pred = np.concatenate([s.sample_posterior(X, draw_indices, excl) for s in sampler], axis=1) else: pred = sampler.sample_posterior(X, draw_indices, excl) @@ -77,6 +75,7 @@ class _MultiChainSampler: """ Dispatches each requested draw to the PosteriorSampler for the chain it came from. """ + def __init__(self, chain_samplers: list): if not chain_samplers: raise ValueError("No posterior draws available yet: run pm.sample() first.") @@ -108,10 +107,10 @@ def sample_posterior(self, X, draw_indices, excluded): return out -_posterior_sampler_cache: dict[int, tuple[int, "_MultiChainSampler"]] = {} +_posterior_sampler_cache: dict[int, tuple[int, _MultiChainSampler]] = {} -def _get_posterior_sampler(op) -> "_MultiChainSampler": +def _get_posterior_sampler(op) -> _MultiChainSampler: """ Rebuild a prediction-only sampler from tree data accumulated on the BART op itself. """ @@ -393,6 +392,7 @@ def plot_pdp( ------- axes: matplotlib axes """ + sampler: _MultiChainSampler | list[_MultiChainSampler] if isinstance(bartrv, list): if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") @@ -916,6 +916,7 @@ def compute_variable_importance( # noqa: PLR0915 PLR0912 rng = np.random.default_rng(random_seed) + sampler: _MultiChainSampler | list[_MultiChainSampler] if isinstance(bartrv, list): if not all(rv.ndim == 1 for rv in bartrv): raise ValueError("List inputs must contain only 1D BART variables") From 23241b0d98a8217fa5abd19fe37442578be94a11 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Thu, 9 Jul 2026 10:24:32 +0300 Subject: [PATCH 9/9] seed test --- tests/test_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 1500e0b..cc798e4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -7,10 +7,11 @@ class TestUtils: - X_norm = np.random.normal(0, 1, size=(50, 2)) - X_binom = np.random.binomial(1, 0.5, size=(50, 1)) + _rng = np.random.default_rng(3415) + X_norm = _rng.normal(0, 1, size=(50, 2)) + X_binom = _rng.binomial(1, 0.5, size=(50, 1)) X = np.hstack([X_norm, X_binom]) - Y = np.random.normal(0, 1, size=50) + Y = _rng.normal(0, 1, size=50) with pm.Model() as model: mu = pmb.BART("mu", X, Y, m=10)