diff --git a/python/bartrs/compile_pymc.py b/python/bartrs/compile_pymc.py index d8048cf..96157f4 100644 --- a/python/bartrs/compile_pymc.py +++ b/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 diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 07ddcd4..0255374 100644 --- a/python/bartrs/pgbart.py +++ b/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: diff --git a/src/lib.rs b/src/lib.rs index 2d5cd59..1fd8491 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -241,12 +241,7 @@ impl PySampler { &mut self, py: Python<'py>, tune: Option, - ) -> PyResult<( - Bound<'py, PyArray2>, - Vec, - Option<(usize, Vec)>, - Option>, - )> { + ) -> PyResult<(Bound<'py, PyArray2>, Vec)> { 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> = 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)> = None; if !was_tune { let mut trees = Vec::with_capacity(info.batch_size); for k in 0..info.batch_size { @@ -280,11 +269,9 @@ 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()); @@ -292,10 +279,28 @@ impl PySampler { 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)> { + self.all_trees + .iter() + .map(|b| (b.start, b.trees.clone())) + .collect() + } + + #[getter] + fn baseline_forest(&self) -> Option> { + 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>>{ + fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, draw_indices: Vec, excluded: Option<&Bound<'py, PyList>>) -> PyResult>>{ 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, samples: usize, excluded: &Vec) -> Array { + fn _sample_posterior(&mut self, x: &Array, draw_indices: &[usize], excluded: &Vec) -> Array { 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, - samples: usize, + draw_indices: &[usize], excluded: &Vec, - rng: &mut SmallRng, ) -> Array { 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 = (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>, n_trees: usize, n_outputs: usize, - rng: SmallRng, } #[pymethods] @@ -402,7 +401,6 @@ impl PosteriorSampler { baseline_forest: Option>, 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>> { + #[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, excluded: Option<&Bound<'py, PyList>>) -> PyResult>> { 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)