Better tree storage

python/bartrs/compile_pymc.py

@@ -18,7 +18,7 @@
     join_nonshared_inputs,
     make_shared_replacements,
 )
-from numba import carray, cfunc, extending, float64, types, njit
+from numba import carray, cfunc, extending, float64, types, njit, int64
 from numba.core import cgutils
 
 

python/bartrs/pgbart.py

@@ -20,8 +20,6 @@
 
 import numpy as np
 
-import psutil
-import psutil
 from pymc.initial_point import PointType
 from pymc.model import Model, modelcontext
 from pymc.pytensorf import inputvars
@@ -32,7 +30,7 @@
 from pymc_bart.bart import BARTRV
 from bartrs.compile_pymc import CompiledPyMCModel
 from bartrs.bartrs import PyBartSettings, PySampler
-from pymc_bart.utils import _encode_vi, _encode_obj
+from pymc_bart.utils import _encode_vi
 
 
 class PGBART(ArrayStepShared):
@@ -58,8 +56,6 @@ class PGBART(ArrayStepShared):
     generates_stats = True
     stats_dtypes_shapes: dict[str, tuple[type, list]] = {
         "variable_inclusion": (object, []),
-        "trees": (object, []),
-        "baseline_forest": (object, []),
         "tune": (bool, []),
         "time": (float, []),
     }
@@ -106,7 +102,10 @@ def __init__(  # noqa: PLR0915
 
         self.shape = 1 if len(shape) == 1 else shape[0]
 
-        self.bart.n_outputs = self.shape
+        type(self.bart).n_outputs = self.shape
+
+        self._non_tune_steps = 0
+        self._synced_all_trees = False
 
         # Updated later in self.setup() by pymc.sample(...)
         self.n_draws = 0
@@ -249,17 +248,28 @@ def __setstate__(self, state):
     def astep(self, _):
         t0 = perf_counter()
         self.compiled_pymc_model.update_shared_arrays()
-        sum_trees, variable_inclusion, new_batch, new_baseline = self.pg_bart.step(self.tune)
+        sum_trees, variable_inclusion = self.pg_bart.step(self.tune)
         t1 = perf_counter()
 
         stats = {
             "variable_inclusion": _encode_vi(variable_inclusion),
-            "trees": _encode_obj(new_batch) if new_batch is not None else None,
-            "baseline_forest": _encode_obj(new_baseline) if new_baseline is not None else None,
             "tune": self.tune,
             "time": t1 - t0,
         }
 
+        if not self.tune:
+            self._non_tune_steps += 1
+
+        if (
+            not self._synced_all_trees
+            and self.n_draws > 0
+            and self._non_tune_steps == self.n_draws
+        ):
+            batches = self.pg_bart.all_batches
+            baseline = self.pg_bart.baseline_forest
+            self.bart.all_trees.append((baseline, batches))
+            self._synced_all_trees = True
+
         return sum_trees, [stats]
 
 
@@ -272,6 +282,7 @@ def competence(var, has_grad):
         return Competence.INCOMPATIBLE
 
     def setup(self, tune: int, draws: int) -> None:
+        self.n_draws = draws
         self.pg_bart.reserve_draws(draws)
 
 def calculate_max_tree_depth(alpha: float, beta: float, probs_leaf: float) -> int:

src/lib.rs

@@ -241,12 +241,7 @@ impl PySampler {
         &mut self,
         py: Python<'py>,
         tune: Option<bool>,
-    ) -> PyResult<(
-        Bound<'py, PyArray2<f64>>,
-        Vec<u32>,
-        Option<(usize, Vec<TreeArrays>)>,
-        Option<Vec<TreeArrays>>,
-    )> {
+    ) -> PyResult<(Bound<'py, PyArray2<f64>>, Vec<u32>)> {
 
         let mut state = self
             .state
@@ -258,20 +253,14 @@ impl PySampler {
         }
         let was_tune = state.tune;
 
-        // `new_baseline` mirrors what gets stored in `self.baseline_forest`, but only on the
-        // step where it's first computed, so callers can stream it back to Python exactly once.
-        let mut new_baseline: Option<Vec<TreeArrays>> = None;
         if !was_tune && self.baseline_forest.is_none() {
             let mut forest = state.forest.clone();
             for t in &mut forest { t.leaf_indices = Vec::new(); }
-            self.baseline_forest = Some(forest.clone());
-            new_baseline = Some(forest);
+            self.baseline_forest = Some(forest);
         }
 
-
         let (new_state, info) = self.kernel.step(&mut self.rng, state);
 
-        let mut new_batch: Option<(usize, Vec<TreeArrays>)> = None;
         if !was_tune {
             let mut trees = Vec::with_capacity(info.batch_size);
             for k in 0..info.batch_size {
@@ -280,22 +269,38 @@ impl PySampler {
                 t.leaf_indices = Vec::new();
                 trees.push(t);
             }
-            self.all_trees.push(TreeBatch { start: info.batch_start, trees: trees.clone() });
-            new_batch = Some((info.batch_start, trees));
+            self.all_trees.push(TreeBatch { start: info.batch_start, trees });
         }
 
-
         // Return predictions as a 2-D array with shape (n_outputs, n_samples)
         let result = numpy::PyArray2::from_owned_array(py, new_state.predictions.clone());
 
         self.state = Some(new_state);
 
         let variable_inclusion = self.state.as_ref().unwrap().get_variable_inclusion();
 
-        Ok((result, variable_inclusion, new_batch, new_baseline))
+        Ok((result, variable_inclusion))
+    }
+
+    #[getter]
+    fn n_draws(&self) -> usize {
+        self.all_trees.len()
+    }
+
+    #[getter]
+    fn all_batches(&self) -> Vec<(usize, Vec<TreeArrays>)> {
+        self.all_trees
+            .iter()
+            .map(|b| (b.start, b.trees.clone()))
+            .collect()
+    }
+
+    #[getter]
+    fn baseline_forest(&self) -> Option<Vec<TreeArrays>> {
+        self.baseline_forest.clone()
     }
 
-    fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult<Py<PyArray3<f64>>>{
+    fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, draw_indices: Vec<usize>, excluded: Option<&Bound<'py, PyList>>) -> PyResult<Py<PyArray3<f64>>>{
 
         let data = x.as_array().to_owned();
 
@@ -306,22 +311,21 @@ impl PySampler {
             None => Vec::new(),
         };
 
-        let preds = self._sample_posterior(&data, samples, &excl).to_pyarray(py).unbind();
+        let preds = self._sample_posterior(&data, &draw_indices, &excl).to_pyarray(py).unbind();
         return Ok(preds)
     }
 }
 
 impl PySampler {
-    fn _sample_posterior(&mut self, x: &Array<f64, Ix2>, samples: usize, excluded: &Vec<usize>) -> Array<f64, Ix3> {
+    fn _sample_posterior(&mut self, x: &Array<f64, Ix2>, draw_indices: &[usize], excluded: &Vec<usize>) -> Array<f64, Ix3> {
         predict_from_history(
             &self.all_trees,
             self.baseline_forest.as_deref(),
             self.n_trees,
             self.n_outputs,
             x,
-            samples,
+            draw_indices,
             excluded,
-            &mut self.rng,
         )
     }
 
@@ -334,23 +338,19 @@ fn predict_from_history(
     n_trees: usize,
     n_outputs: usize,
     x: &Array<f64, Ix2>,
-    samples: usize,
+    draw_indices: &[usize],
     excluded: &Vec<usize>,
-    rng: &mut SmallRng,
 ) -> Array<f64, Ix3> {
     let n_data_samples: usize = x.nrows();
     let n_draws = batches.len();
+    let samples = draw_indices.len();
 
     assert!(n_draws > 0, "No posterior draws stored yet");
 
-    let random_samples: Vec<usize> = (0..samples)
-        .map(|_| rng.random_range(0..n_draws as u32) as usize)
-        .collect();
-
     let mut predictions = Array3::zeros((samples, n_outputs, n_data_samples));
 
     for posterior_sample_idx in 0..samples {
-        let draw_idx = random_samples[posterior_sample_idx];
+        let draw_idx = draw_indices[posterior_sample_idx];
         let mut sample = predictions.slice_mut(s![posterior_sample_idx, .., ..]);
 
         let mut covered = vec![false; n_trees];
@@ -391,7 +391,6 @@ struct PosteriorSampler {
     baseline_forest: Option<Vec<TreeArrays>>,
     n_trees: usize,
     n_outputs: usize,
-    rng: SmallRng,
 }
 
 #[pymethods]
@@ -402,7 +401,6 @@ impl PosteriorSampler {
         baseline_forest: Option<Vec<TreeArrays>>,
         n_trees: usize,
         n_outputs: usize,
-        seed: u64,
     ) -> Self {
         let batches = batches
             .into_iter()
@@ -413,7 +411,6 @@ impl PosteriorSampler {
             baseline_forest,
             n_trees,
             n_outputs,
-            rng: SmallRng::seed_from_u64(seed),
         }
     }
 
@@ -422,7 +419,12 @@ impl PosteriorSampler {
         self.n_outputs
     }
 
-    fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult<Py<PyArray3<f64>>> {
+    #[getter]
+    fn n_draws(&self) -> usize {
+        self.batches.len()
+    }
+
+    fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, draw_indices: Vec<usize>, excluded: Option<&Bound<'py, PyList>>) -> PyResult<Py<PyArray3<f64>>> {
         let data = x.as_array().to_owned();
 
         let excl = match excluded {
@@ -438,9 +440,8 @@ impl PosteriorSampler {
             self.n_trees,
             self.n_outputs,
             &data,
-            samples,
+            &draw_indices,
             &excl,
-            &mut self.rng,
         ).to_pyarray(py).unbind();
 
         Ok(preds)