Better tree storage

.gitignore

@@ -167,3 +167,4 @@ cython_debug/
 
 **.csv
 tests/test_hmc.py
+tests/bench_bart.py

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
 

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()

pymc_bart/pymc_bart.py

@@ -1,11 +1,16 @@
 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)
+]