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/pymc_bart.py

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