From 32527410992990134b63e2e2b241dbf398a1ec33 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Mon, 29 Jun 2026 11:53:25 +0300 Subject: [PATCH 01/11] store trees in rust to avoid python list appending --- python/bartrs/pgbart.py | 23 ++++++++++++++++++++--- src/lib.rs | 29 ++++++++++++++++++++++++++--- 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 104cb78..1e9ec5b 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -102,6 +102,14 @@ def __init__( # noqa: PLR0915 self.shape = 1 if len(shape) == 1 else shape[0] + self.n_draws = self.bart.n_draws + self.n_tune = self.bart.n_tune + if self.n_draws is None: + self.n_draws = -1 # Use -1 to indicate that n_draws is not specified + self.n_tune = -1 + + self.n_steps = 0 + # Set trees_shape (dim for separate tree structures) # and leaves_shape (dim for leaf node values) # One of the two is always one, the other equal to self.shape @@ -165,6 +173,7 @@ def __init__( # noqa: PLR0915 resampling_rule="systematic", batch_tune=batch[0], batch_post=batch[1], + n_draws=self.n_draws, ) # child processes spawned with spawn need to be able @@ -187,6 +196,7 @@ def __init__( # noqa: PLR0915 "resampling_rule": "systematic", "batch_tune": batch[0], "batch_post": batch[1], + "n_draws": self.n_draws, } # INFO: Only at the end do we return the State structure back to Python to avoid @@ -242,9 +252,16 @@ def astep(self, _): self.compiled_pymc_model.update_shared_arrays() # sum_trees, variable_inclusion = step(self.state, self.tune) - sum_trees, trees, variable_inclusion = self.pg_bart.step(self.tune) - if not self.tune: - self.bart.all_trees.append(trees) # this doubles runtime + sum_trees, variable_inclusion = self.pg_bart.step(self.tune) + + + # if not self.tune: + # self.bart.all_trees.append(trees) # this doubles runtime + + if self.n_steps == (self.n_draws + self.n_tune): + chain_trees = self.pg_bart.get_results() + self.bart.all_trees.extend(chain_trees) + t1 = perf_counter() stats = { diff --git a/src/lib.rs b/src/lib.rs index 4ffe28e..c22a12e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -53,6 +53,7 @@ pub struct PyBartSettings { batch_tune: f64, batch_post: f64, seed: u64, + n_draws: usize, } #[pymethods] @@ -73,6 +74,7 @@ impl PyBartSettings { batch_tune = 0.1, batch_post = 0.1, seed = 0, + n_draws = 0, ))] #[allow(clippy::too_many_arguments)] fn new( @@ -90,6 +92,7 @@ impl PyBartSettings { batch_tune: f64, batch_post: f64, seed: u64, + n_draws: usize, ) -> Self { Self { n_trees, @@ -106,6 +109,7 @@ impl PyBartSettings { batch_tune, batch_post, seed, + n_draws, } } } @@ -115,6 +119,7 @@ struct PySampler { kernel: Box, state: Option, rng: SmallRng, + all_trees: Vec, } #[pymethods] @@ -193,7 +198,13 @@ impl PySampler { }; let data = OwnedData::new(x_data, y_data); - + + let all_trees: Vec = if settings.n_draws == 0 { + Vec::new() + } else { + Vec::with_capacity(config.n_trees * settings.n_draws) + }; + let kernel = BartKernel { split_rules, resampling: SystematicResampling, @@ -208,6 +219,7 @@ impl PySampler { kernel: Box::new(kernel), state: Some(state), rng, + all_trees, }) } @@ -216,7 +228,7 @@ impl PySampler { &mut self, py: Python<'py>, tune: Option, - ) -> PyResult<(Bound<'py, PyArray2>, Vec, Vec)> { + ) -> PyResult<(Bound<'py, PyArray2>, Vec)> { let mut state = self .state @@ -227,12 +239,18 @@ impl PySampler { state.tune = t; } + let tune = state.tune; + let (mut new_state, _info) = self.kernel.step(&mut self.rng, state); // Need to return TreeArrays let trees = std::mem::take(&mut new_state.forest); new_state.forest = trees.clone(); + if !tune { + self.all_trees.extend(trees.clone()); + } + // 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()); @@ -240,8 +258,13 @@ impl PySampler { let variable_inclusion = self.state.as_ref().unwrap().get_variable_inclusion(); - Ok((result, trees, variable_inclusion)) + Ok((result, variable_inclusion)) } + + fn get_results(&mut self) -> Vec { + std::mem::take(&mut self.all_trees) + } + } #[pymodule] From 558b660b18518c33e7f079a24406a4fb5fdddb85 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Tue, 30 Jun 2026 11:34:44 +0300 Subject: [PATCH 02/11] not sure --- python/bartrs/pgbart.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 1e9ec5b..1608b20 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -254,13 +254,17 @@ def astep(self, _): sum_trees, variable_inclusion = self.pg_bart.step(self.tune) - # if not self.tune: # self.bart.all_trees.append(trees) # this doubles runtime if self.n_steps == (self.n_draws + self.n_tune): chain_trees = self.pg_bart.get_results() self.bart.all_trees.extend(chain_trees) + + # Log lik + # Pymc n draws + # n trees + # examples for gradients t1 = perf_counter() From 651e33c57919f1b43e54a0658d1d17173b073b12 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Tue, 30 Jun 2026 11:35:05 +0300 Subject: [PATCH 03/11] remove --- python/bartrs/pgbart.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 1608b20..7c8e01a 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -261,10 +261,6 @@ def astep(self, _): chain_trees = self.pg_bart.get_results() self.bart.all_trees.extend(chain_trees) - # Log lik - # Pymc n draws - # n trees - # examples for gradients t1 = perf_counter() From 080134b6e3d83b67ff37af3e4bee40f8b649f1ef Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Thu, 2 Jul 2026 16:05:30 +0300 Subject: [PATCH 04/11] update ndarray, numpy and pyo3 versions; also expand preliminary tree support in Rust --- Cargo.lock | 66 ++++++++++++++++------------------------- Cargo.toml | 4 +-- python/bartrs/pgbart.py | 33 +++++++++------------ src/lib.rs | 63 ++++++++++++++++++++++++++++++++------- 4 files changed, 94 insertions(+), 72 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index beab318..763bfda 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -36,7 +36,7 @@ dependencies = [ "criterion", "numpy", "pyo3", - "pyo3-build-config", + "pyo3-build-config 0.25.1", "rand", "rand_distr", ] @@ -215,12 +215,6 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" -[[package]] -name = "indoc" -version = "2.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" - [[package]] name = "itertools" version = "0.10.5" @@ -289,15 +283,6 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "ndarray" version = "0.16.1" @@ -343,9 +328,9 @@ dependencies = [ [[package]] name = "numpy" -version = "0.25.0" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29f1dee9aa8d3f6f8e8b9af3803006101bb3653866ef056d530d53ae68587191" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" dependencies = [ "libc", "ndarray", @@ -353,7 +338,7 @@ dependencies = [ "num-integer", "num-traits", "pyo3", - "pyo3-build-config", + "pyo3-build-config 0.29.0", "rustc-hash", ] @@ -432,19 +417,16 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.25.1" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8970a78afe0628a3e3430376fc5fd76b6b45c4d43360ffd6cdd40bdde72b682a" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" dependencies = [ - "indoc", "libc", - "memoffset", "once_cell", "portable-atomic", - "pyo3-build-config", + "pyo3-build-config 0.29.0", "pyo3-ffi", "pyo3-macros", - "unindent", ] [[package]] @@ -457,21 +439,30 @@ dependencies = [ "target-lexicon", ] +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + [[package]] name = "pyo3-ffi" -version = "0.25.1" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7114fe5457c61b276ab77c5055f206295b812608083644a5c5b2640c3102565c" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" dependencies = [ "libc", - "pyo3-build-config", + "pyo3-build-config 0.29.0", ] [[package]] name = "pyo3-macros" -version = "0.25.1" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8725c0a622b374d6cb051d11a0983786448f7785336139c3c94f5aa6bef7e50" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -481,13 +472,12 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.25.1" +version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4109984c22491085343c05b0dbc54ddc405c3cf7b4374fc533f5c3313a572ccc" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" dependencies = [ "heck", "proc-macro2", - "pyo3-build-config", "quote", "syn", ] @@ -673,9 +663,9 @@ dependencies = [ [[package]] name = "target-lexicon" -version = "0.13.2" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] name = "tinytemplate" @@ -693,12 +683,6 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" -[[package]] -name = "unindent" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" - [[package]] name = "walkdir" version = "2.5.0" diff --git a/Cargo.toml b/Cargo.toml index 57ae468..b635f19 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,8 +22,8 @@ name = "bartrs" crate-type = ["cdylib", "rlib"] [dependencies] -pyo3 = { version = "0.25.1", features = ["macros"] } -numpy = { version = "0.25.0" } +pyo3 = { version = "0.29.0", features = ["macros"] } +numpy = { version = "0.29.0" } rand = { version = "0.9.1" } rand_distr = { version = "0.5.0" } diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 7c8e01a..6fb8ccb 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -102,13 +102,10 @@ def __init__( # noqa: PLR0915 self.shape = 1 if len(shape) == 1 else shape[0] - self.n_draws = self.bart.n_draws - self.n_tune = self.bart.n_tune - if self.n_draws is None: - self.n_draws = -1 # Use -1 to indicate that n_draws is not specified - self.n_tune = -1 - - self.n_steps = 0 + # Updated later in self.setup() by pymc.sample(...) + self.n_draws = 0 + self.n_tune = 0 + self.n_total = 0 # Set trees_shape (dim for separate tree structures) # and leaves_shape (dim for leaf node values) @@ -215,6 +212,8 @@ def __init__( # noqa: PLR0915 ) self.tune = True + + self.check = True super().__init__(vars, self.compiled_pymc_model.shared) def __getstate__(self): @@ -247,23 +246,15 @@ def __setstate__(self, state): def astep(self, _): - # # Record time to quantify performance improvements t0 = perf_counter() self.compiled_pymc_model.update_shared_arrays() - # sum_trees, variable_inclusion = step(self.state, self.tune) - sum_trees, variable_inclusion = self.pg_bart.step(self.tune) - - # if not self.tune: - # self.bart.all_trees.append(trees) # this doubles runtime - - if self.n_steps == (self.n_draws + self.n_tune): - chain_trees = self.pg_bart.get_results() - self.bart.all_trees.extend(chain_trees) - - t1 = perf_counter() + if self.check == True: + print(self.n_draws, self.n_tune, self.n_total) + self.check = False + stats = { "variable_inclusion": _encode_vi(variable_inclusion), "tune": self.tune, @@ -281,6 +272,10 @@ def competence(var, has_grad): return Competence.IDEAL return Competence.INCOMPATIBLE + def setup(self, tune: int, draws: int) -> None: + self.n_draws = draws + self.n_tune = tune + self.n_total = draws + tune def calculate_max_tree_depth(alpha: float, beta: float, probs_leaf: float) -> int: """Calculates the maximum tree depth for which the probability of a node diff --git a/src/lib.rs b/src/lib.rs index c22a12e..d778726 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,11 +23,13 @@ use crate::splitting::{ContinuousSplit, SplitRules}; use crate::weight::PyMCWeightFn; use numpy::{ - PyReadonlyArray, PyReadonlyArrayDyn, PyArray2, PyUntypedArrayMethods, - ndarray::{Array1, Array2, ArrayView1, ArrayViewMut1, Ix2}, + PyReadonlyArray, PyReadonlyArrayDyn, PyArray2, PyArray3, PyUntypedArrayMethods, PyReadonlyArray2, + ToPyArray, + ndarray::{Array, Array1, Array2, Array3, ArrayView1, ArrayViewMut1, Ix2, Ix3, s}, }; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::types::{PyList}; use rand::Rng; use rand::SeedableRng; use rand::rngs::SmallRng; @@ -245,7 +247,7 @@ impl PySampler { // Need to return TreeArrays let trees = std::mem::take(&mut new_state.forest); - new_state.forest = trees.clone(); + new_state.forest = trees.clone(); // Vec if !tune { self.all_trees.extend(trees.clone()); @@ -261,8 +263,49 @@ impl PySampler { Ok((result, variable_inclusion)) } - fn get_results(&mut self) -> Vec { - std::mem::take(&mut self.all_trees) + + fn sample_posterior<'py>(&self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult>>{ + + let data = x.as_array().to_owned(); + + let excl = match excluded { + Some(list) => list.iter() + .map(|item| item.extract::()) + .collect::>>()?, + None => Vec::new(), + }; + + let preds = self._sample_posterior(&data, samples, &excl).to_pyarray(py).unbind(); + return Ok(preds) + } +} + +impl PySampler { + + fn _sample_posterior(&self, x: &Array, samples: usize, excluded: &Vec) -> Array { + + let n_outputs = self.kernel.config.n_outputs; + + let n_data_samples: usize = x.nrows(); + let n_forests: u32 = self.all_trees.len() as u32; + + let mut random_samples = vec![0usize; samples]; + + for x in random_samples { + x = self.rng.random_range(0..n_forests) as usize; + } + + let predictions = Array3::zeros((samples, n_outputs, n_data_samples)); + + for posterior_sample_idx in 0..samples { + + let draw_forest_idx = random_samples[posterior_sample_idx]; + + for tree in self.all_trees[draw_forest_idx].iter() { + predictions.slice_mut(s![posterior_sample_idx as i32, :, :]) += &tree.predict_batch_test(x, excluded); + } + } + predictions } } @@ -306,18 +349,18 @@ fn are_whole_number(col: ArrayView1<'_, f64>) -> bool { } fn nanstd(col: ArrayView1<'_, f64>) -> f64 { - let mut count = 0usize; - let mut mean = 0.0f64; - let mut m2 = 0.0f64; + let mut count: usize = 0usize; + let mut mean: f64 = 0.0f64; + let mut m2: f64 = 0.0f64; for &value in col.iter() { if value.is_nan() { continue; } count += 1; - let delta = value - mean; + let delta: f64 = value - mean; mean += delta / count as f64; - let delta2 = value - mean; + let delta2: f64 = value - mean; m2 += delta * delta2; } From e77d21cecf15b98aaf66e6effcd8562f0186699c Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Thu, 2 Jul 2026 20:21:10 +0300 Subject: [PATCH 05/11] change tree storage to nested Vecs --- src/lib.rs | 27 ++++++++++++--------------- src/tree.rs | 8 ++++---- 2 files changed, 16 insertions(+), 19 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index d778726..84fd78c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,7 +37,7 @@ use rand_distr::StandardNormal; type LogpFunc = unsafe extern "C" fn(*const f64, usize) -> c_double; -#[pyclass] +#[pyclass(from_py_object)] #[derive(Clone, Debug)] #[allow(dead_code)] pub struct PyBartSettings { @@ -121,7 +121,7 @@ struct PySampler { kernel: Box, state: Option, rng: SmallRng, - all_trees: Vec, + all_trees: Vec>, } #[pymethods] @@ -201,10 +201,10 @@ impl PySampler { let data = OwnedData::new(x_data, y_data); - let all_trees: Vec = if settings.n_draws == 0 { + let all_trees: Vec> = if settings.n_draws == 0 { Vec::new() } else { - Vec::with_capacity(config.n_trees * settings.n_draws) + Vec::with_capacity(settings.n_draws) }; let kernel = BartKernel { @@ -250,7 +250,7 @@ impl PySampler { new_state.forest = trees.clone(); // Vec if !tune { - self.all_trees.extend(trees.clone()); + self.all_trees.push(trees); } // Return predictions as a 2-D array with shape (n_outputs, n_samples) @@ -264,7 +264,7 @@ impl PySampler { } - fn sample_posterior<'py>(&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>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult>>{ let data = x.as_array().to_owned(); @@ -282,27 +282,24 @@ impl PySampler { impl PySampler { - fn _sample_posterior(&self, x: &Array, samples: usize, excluded: &Vec) -> Array { + fn _sample_posterior(&mut self, x: &Array, samples: usize, excluded: &Vec) -> Array { - let n_outputs = self.kernel.config.n_outputs; + let n_outputs = self.state.as_ref().expect("Sampler state is missing").predictions.nrows(); let n_data_samples: usize = x.nrows(); let n_forests: u32 = self.all_trees.len() as u32; - let mut random_samples = vec![0usize; samples]; + let random_samples: Vec = (0..samples).map( | _ | self.rng.random_range(0..n_forests) as usize).collect(); - for x in random_samples { - x = self.rng.random_range(0..n_forests) as usize; - } - - let predictions = Array3::zeros((samples, n_outputs, n_data_samples)); + let mut predictions = Array3::zeros((samples, n_outputs, n_data_samples)); for posterior_sample_idx in 0..samples { let draw_forest_idx = random_samples[posterior_sample_idx]; for tree in self.all_trees[draw_forest_idx].iter() { - predictions.slice_mut(s![posterior_sample_idx as i32, :, :]) += &tree.predict_batch_test(x, excluded); + let mut sample = predictions.slice_mut(s![posterior_sample_idx, .., ..]); + sample += &tree.predict_batch_test(x, excluded); } } predictions diff --git a/src/tree.rs b/src/tree.rs index a9e9072..f6e731e 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -12,7 +12,7 @@ use crate::data::NotNan; /// Internal nodes store split variable and threshold. Leaf nodes store /// predicted values. The `leaf_indices` vector maps each training sample /// to its assigned leaf node. -#[pyclass(module = "bartrs.bartrs", get_all)] +#[pyclass(module = "bartrs.bartrs", get_all, from_py_object)] #[derive(Clone, Debug)] pub struct TreeArrays { /// Split variable per node (u32::MAX = leaf sentinel) @@ -113,7 +113,7 @@ impl TreeArrays { } pub fn __setstate__(&mut self, state: &Bound<'_, PyAny>) -> PyResult<()> { - let dict: &Bound<'_, PyDict> = state.downcast::()?; + let dict: &Bound<'_, PyDict> = state.cast::()?; self.split_var = dict.get_item("split_var").unwrap().expect("split_var not found").extract()?; self.split_val = dict.get_item("split_val").unwrap().expect("split_val not found").extract()?; self.leaf_val = dict.get_item("leaf_val").unwrap().expect("leaf_val not found").extract()?; @@ -129,12 +129,12 @@ impl TreeArrays { Ok(()) } - pub fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(PyObject, (f64, usize, u8, usize), Py)> { + pub fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<(Py, (f64, usize, u8, usize), Py)> { let cls = py.get_type::(); // args are what `new(init_leaf_value, n_samples, max_depth, n_outputs)` expects let args = (self.leaf_val.get(0).copied().unwrap_or(0.0), self.leaf_indices.len(), self.max_depth, self.n_outputs); let state = self.to_state(py)?; - Ok((cls.into(), args, state)) + Ok((cls.as_any().clone().unbind(), args, state)) } From 0003c888cbe33da4bae9b8ae75f98ce0ca6dd985 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Tue, 7 Jul 2026 11:31:56 +0300 Subject: [PATCH 06/11] store less trees --- src/kernel.rs | 4 +++ src/lib.rs | 85 ++++++++++++++++++++++++++++++++++++--------------- src/state.rs | 4 +++ src/tree.rs | 5 +++ 4 files changed, 74 insertions(+), 24 deletions(-) diff --git a/src/kernel.rs b/src/kernel.rs index 1ce9f4a..e0f656d 100644 --- a/src/kernel.rs +++ b/src/kernel.rs @@ -127,6 +127,7 @@ where let batch_size = ((batch_frac * n_trees as f64).round() as usize) .max(1) .min(n_trees); + let batch_start = state.next_tree_idx; let mut acceptance_count = 0; let mut tree_depths = Vec::with_capacity(batch_size); @@ -200,6 +201,9 @@ where log_likelihood: total_log_likelihood, acceptance_count, tree_depths, + batch_start, + batch_size, + }; (state, info) diff --git a/src/lib.rs b/src/lib.rs index 84fd78c..61a4f68 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,7 +14,7 @@ pub mod tree; pub mod update; pub mod weight; -use crate::tree::TreeArrays; +use crate::tree::{TreeArrays, TreeBatch}; use crate::config::BartConfig; use crate::data::OwnedData; use crate::kernel::{BartKernel, ErasedKernel, SamplingAlgorithm}; @@ -121,7 +121,9 @@ struct PySampler { kernel: Box, state: Option, rng: SmallRng, - all_trees: Vec>, + all_trees: Vec, + baseline_forest: Option>, + n_trees: usize, } #[pymethods] @@ -201,7 +203,7 @@ impl PySampler { let data = OwnedData::new(x_data, y_data); - let all_trees: Vec> = if settings.n_draws == 0 { + let all_trees: Vec = if settings.n_draws == 0 { Vec::new() } else { Vec::with_capacity(settings.n_draws) @@ -222,6 +224,8 @@ impl PySampler { state: Some(state), rng, all_trees, + baseline_forest: None, + n_trees: settings.n_trees, }) } @@ -240,17 +244,21 @@ impl PySampler { if let Some(t) = tune { state.tune = t; } + let was_tune = state.tune; - let tune = state.tune; - - let (mut new_state, _info) = self.kernel.step(&mut self.rng, state); + if !was_tune && self.baseline_forest.is_none() { + self.baseline_forest = Some(state.forest.clone()); + } - // Need to return TreeArrays - let trees = std::mem::take(&mut new_state.forest); - new_state.forest = trees.clone(); // Vec + let (new_state, info) = self.kernel.step(&mut self.rng, state); - if !tune { - self.all_trees.push(trees); + if !was_tune { + let mut trees = Vec::with_capacity(info.batch_size); + for k in 0..info.batch_size { + let idx = (info.batch_start + k) % self.n_trees; + trees.push(new_state.forest[idx].clone()); + } + self.all_trees.push(TreeBatch {start: info.batch_start, trees}); } // Return predictions as a 2-D array with shape (n_outputs, n_samples) @@ -281,28 +289,57 @@ impl PySampler { } impl PySampler { - fn _sample_posterior(&mut self, x: &Array, samples: usize, excluded: &Vec) -> Array { - - let n_outputs = self.state.as_ref().expect("Sampler state is missing").predictions.nrows(); - + let n_outputs = self.state.as_ref().expect("Sampler state is missing").predictions.nrows(); let n_data_samples: usize = x.nrows(); - let n_forests: u32 = self.all_trees.len() as u32; - - let random_samples: Vec = (0..samples).map( | _ | self.rng.random_range(0..n_forests) as usize).collect(); + let n_trees = self.n_trees; + let n_draws = self.all_trees.len(); + + assert!(n_draws > 0, "No posterior draws stored yet"); + + let random_samples: Vec = (0..samples) + .map(|_| self.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 mut sample = predictions.slice_mut(s![posterior_sample_idx, .., ..]); + + let mut covered = vec![false; n_trees]; + let mut remaining = n_trees; + + // Walk backward from draw_idx, taking the most recent version of + // each tree slot until every slot is accounted for. + let mut b = draw_idx as isize; + while remaining > 0 && b >= 0 { + let batch = &self.all_trees[b as usize]; + for (k, tree) in batch.trees.iter().enumerate() { + let tree_idx = (batch.start + k) % n_trees; + if !covered[tree_idx] { + covered[tree_idx] = true; + remaining -= 1; + sample += &tree.predict_batch_test(x, excluded); + } + } + b -= 1; + } - let draw_forest_idx = random_samples[posterior_sample_idx]; - - for tree in self.all_trees[draw_forest_idx].iter() { - let mut sample = predictions.slice_mut(s![posterior_sample_idx, .., ..]); - sample += &tree.predict_batch_test(x, excluded); + // Ran out of stored history before covering all slots: fill the rest + // from the end-of-tuning snapshot. + if remaining > 0 { + if let Some(baseline) = &self.baseline_forest { + for tree_idx in 0..n_trees { + if !covered[tree_idx] { + sample += &baseline[tree_idx].predict_batch_test(x, excluded); + } + } + } } } - predictions + + predictions } } diff --git a/src/state.rs b/src/state.rs index 998c498..07f971c 100644 --- a/src/state.rs +++ b/src/state.rs @@ -34,4 +34,8 @@ pub struct BartInfo { pub log_likelihood: f64, pub acceptance_count: usize, pub tree_depths: Vec, + /// Index of the start of trees updated this step + pub batch_start: usize, + /// Number of trees updated per step + pub batch_size: usize, } diff --git a/src/tree.rs b/src/tree.rs index f6e731e..fa4fa7b 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -6,6 +6,11 @@ use pyo3::types::{PyDict, PyList}; use crate::response::{LeafKind, LeafPayload, LeafProposal}; use crate::data::NotNan; +pub struct TreeBatch { + pub start: usize, + pub trees: Vec, +} + /// Bartz-style heap-indexed tree with separate internal/leaf arrays. /// /// Uses heap convention: root=0, left=2i+1, right=2i+2. From 35278413e6b878243164242ae7bef68616ddf1f2 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Tue, 7 Jul 2026 15:06:54 +0300 Subject: [PATCH 07/11] remove redundant checks --- python/bartrs/pgbart.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 6fb8ccb..15d6cb0 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -101,7 +101,6 @@ def __init__( # noqa: PLR0915 self.shape = 1 if len(shape) == 1 else shape[0] - # Updated later in self.setup() by pymc.sample(...) self.n_draws = 0 self.n_tune = 0 @@ -203,7 +202,6 @@ def __init__( # noqa: PLR0915 # pg = PySampler(...) # state = pg.init(...) # new_state, info = pg.step(rng, state) - self.pg_bart = PySampler.init( x=np.asfortranarray(self.X), y=self.bart.Y, @@ -212,8 +210,6 @@ def __init__( # noqa: PLR0915 ) self.tune = True - - self.check = True super().__init__(vars, self.compiled_pymc_model.shared) def __getstate__(self): @@ -251,10 +247,6 @@ def astep(self, _): sum_trees, variable_inclusion = self.pg_bart.step(self.tune) t1 = perf_counter() - if self.check == True: - print(self.n_draws, self.n_tune, self.n_total) - self.check = False - stats = { "variable_inclusion": _encode_vi(variable_inclusion), "tune": self.tune, From 4c6121b25d5b419728405e6197bac041869591c9 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Tue, 7 Jul 2026 16:17:56 +0300 Subject: [PATCH 08/11] diagnoses --- Cargo.lock | 41 +++++++++++++++ Cargo.toml | 1 + python/bartrs/pgbart.py | 28 +++++++++++ src/lib.rs | 109 ++++++++++++++++++++++++++++++++++++++-- 4 files changed, 174 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 763bfda..0a9d344 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,6 +34,7 @@ name = "bartrs" version = "0.3.0" dependencies = [ "criterion", + "mimalloc", "numpy", "pyo3", "pyo3-build-config 0.25.1", @@ -59,6 +60,16 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "cc" +version = "1.2.66" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.0" @@ -187,6 +198,12 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + [[package]] name = "getrandom" version = "0.3.3" @@ -261,6 +278,15 @@ version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "log" version = "0.4.27" @@ -283,6 +309,15 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "ndarray" version = "0.16.1" @@ -650,6 +685,12 @@ dependencies = [ "serde", ] +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + [[package]] name = "syn" version = "2.0.90" diff --git a/Cargo.toml b/Cargo.toml index b635f19..a2ae9c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ pyo3 = { version = "0.29.0", features = ["macros"] } numpy = { version = "0.29.0" } rand = { version = "0.9.1" } rand_distr = { version = "0.5.0" } +mimalloc = "0.1" [dev-dependencies] criterion = { version = "0.6.0", features = ["html_reports"] } diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index 15d6cb0..d95fd5a 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -14,11 +14,14 @@ import math +import os from time import perf_counter from typing import Optional, Tuple 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 @@ -210,6 +213,7 @@ def __init__( # noqa: PLR0915 ) self.tune = True + self._astep_calls = 0 super().__init__(vars, self.compiled_pymc_model.shared) def __getstate__(self): @@ -246,6 +250,29 @@ def astep(self, _): self.compiled_pymc_model.update_shared_arrays() sum_trees, variable_inclusion = self.pg_bart.step(self.tune) t1 = perf_counter() + + import psutil + _PROC = psutil.Process() + + def cur_rss_mb(): + return _PROC.memory_info().rss / 1e6 + + + import resource, os + def rss_mb(): + r = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return r / (1024*1024 if os.uname().sysname == "Darwin" else 1024) + + # in astep, print RSS at the tune->post transition and at the end + self._astep_calls += 1 + if self._astep_calls == self.n_tune: + print(f"[end of tune] cur_rss={cur_rss_mb():.1f}MB calls={self._astep_calls}", flush=True) + if self._astep_calls == self.n_total: + print(f"[end of sampling] cur_rss={cur_rss_mb():.1f}MB", flush=True) + self.pg_bart.report_stored_bytes() + + if self._astep_calls % 200 == 0: + print(f"call={self._astep_calls} cur_rss={cur_rss_mb():.1f}MB", flush=True) stats = { "variable_inclusion": _encode_vi(variable_inclusion), @@ -253,6 +280,7 @@ def astep(self, _): "time": t1 - t0, } + return sum_trees, [stats] diff --git a/src/lib.rs b/src/lib.rs index 61a4f68..ac6a625 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,6 +37,11 @@ use rand_distr::StandardNormal; type LogpFunc = unsafe extern "C" fn(*const f64, usize) -> c_double; + +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + + #[pyclass(from_py_object)] #[derive(Clone, Debug)] #[allow(dead_code)] @@ -247,20 +252,26 @@ impl PySampler { let was_tune = state.tune; if !was_tune && self.baseline_forest.is_none() { - self.baseline_forest = Some(state.forest.clone()); + let mut forest = state.forest.clone(); + for t in &mut forest { t.leaf_indices = Vec::new(); } + self.baseline_forest = Some(forest); } + let (new_state, info) = self.kernel.step(&mut self.rng, state); if !was_tune { let mut trees = Vec::with_capacity(info.batch_size); - for k in 0..info.batch_size { + for k in 0..info.batch_size { let idx = (info.batch_start + k) % self.n_trees; - trees.push(new_state.forest[idx].clone()); + let mut t = new_state.forest[idx].clone(); + t.leaf_indices = Vec::new(); // only the STORED copy loses it + trees.push(t); } - self.all_trees.push(TreeBatch {start: 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()); @@ -271,6 +282,75 @@ impl PySampler { Ok((result, variable_inclusion)) } + /// Diagnostic: per-field byte breakdown of everything currently retained in + /// `all_trees` + `baseline_forest`. Reports both `len`-based bytes (actual + /// data) and `capacity`-based bytes (allocated), so fixed-width buffers and + /// capacity slack are both visible. Call once after sampling finishes. + fn report_stored_bytes(&self) { + use std::collections::BTreeMap; + + let mut len_tot: BTreeMap<&str, usize> = BTreeMap::new(); + let mut cap_tot: BTreeMap<&str, usize> = BTreeMap::new(); + let mut n_trees_counted = 0usize; + let mut node_count_sum = 0usize; // sum of split_var.len() across trees + let mut max_node_count = 0usize; + + let all = self + .all_trees + .iter() + .flat_map(|b| b.trees.iter()) + .chain(self.baseline_forest.iter().flatten()); + + for t in all { + n_trees_counted += 1; + node_count_sum += t.split_var.len(); + max_node_count = max_node_count.max(t.split_var.len()); + for (name, len_b, cap_b) in tree_field_report(t) { + *len_tot.entry(name).or_default() += len_b; + *cap_tot.entry(name).or_default() += cap_b; + } + } + + eprintln!( + "--- stored tree field breakdown ({} trees, {} batches, baseline={}) ---", + n_trees_counted, + self.all_trees.len(), + self.baseline_forest.is_some() + ); + + let mut len_sum = 0usize; + let mut cap_sum = 0usize; + for (name, &lb) in &len_tot { + let cb = cap_tot[name]; + len_sum += lb; + cap_sum += cb; + eprintln!( + "{:<16} len={:>9.2}MB cap={:>9.2}MB slack={:>8.2}MB", + name, + lb as f64 / 1e6, + cb as f64 / 1e6, + (cb - lb) as f64 / 1e6 + ); + } + eprintln!( + "{:<16} len={:>9.2}MB cap={:>9.2}MB slack={:>8.2}MB", + "TOTAL", + len_sum as f64 / 1e6, + cap_sum as f64 / 1e6, + (cap_sum - len_sum) as f64 / 1e6 + ); + + if n_trees_counted > 0 { + eprintln!( + "avg per tree: len={:.0}B cap={:.0}B | avg nodes/tree={:.1} max nodes={}", + len_sum as f64 / n_trees_counted as f64, + cap_sum as f64 / n_trees_counted as f64, + node_count_sum as f64 / n_trees_counted as f64, + max_node_count + ); + } + } + fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult>>{ @@ -344,6 +424,25 @@ impl PySampler { } +/// Per-field byte accounting for a single stored tree. +/// Returns (field name, bytes from len, bytes from capacity) for each field. +fn tree_field_report(t: &TreeArrays) -> [(&'static str, usize, usize); 10] { + [ + ("split_var", t.split_var.len() * 4, t.split_var.capacity() * 4), + ("split_val", t.split_val.len() * 8, t.split_val.capacity() * 8), + ("leaf_val", t.leaf_val.len() * 8, t.leaf_val.capacity() * 8), + ("leaf_indices", t.leaf_indices.len() * 4, t.leaf_indices.capacity() * 4), + ("node_nvalue", t.node_nvalue.len() * 8, t.node_nvalue.capacity() * 8), + ("leaf_kind", t.leaf_kind.len(), t.leaf_kind.capacity()), + ("leaf_param_idx", t.leaf_param_idx.len() * 8, t.leaf_param_idx.capacity() * 8), + ("linear_intercept", t.linear_intercept.iter().map(|v| v.len() * 8).sum(), + t.linear_intercept.iter().map(|v| v.capacity() * 8).sum()), + ("linear_slope", t.linear_slope.iter().map(|v| v.len() * 8).sum(), + t.linear_slope.iter().map(|v| v.capacity() * 8).sum()), + ("linear_var", t.linear_var.len() * 4, t.linear_var.capacity() * 4), + ] +} + #[pymodule] fn bartrs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; @@ -403,4 +502,4 @@ fn nanstd(col: ArrayView1<'_, f64>) -> f64 { } (m2 / count as f64).sqrt() -} +} \ No newline at end of file From de44794ffa168a2730b9f5a9a15bfc39196df85a Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Wed, 8 Jul 2026 13:37:05 +0300 Subject: [PATCH 09/11] add PosteriorSampler --- Cargo.lock | 41 ------ Cargo.toml | 1 - python/bartrs/pgbart.py | 41 ++---- src/lib.rs | 284 ++++++++++++++++++++-------------------- 4 files changed, 154 insertions(+), 213 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0a9d344..763bfda 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,7 +34,6 @@ name = "bartrs" version = "0.3.0" dependencies = [ "criterion", - "mimalloc", "numpy", "pyo3", "pyo3-build-config 0.25.1", @@ -60,16 +59,6 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" -[[package]] -name = "cc" -version = "1.2.66" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" -dependencies = [ - "find-msvc-tools", - "shlex", -] - [[package]] name = "cfg-if" version = "1.0.0" @@ -198,12 +187,6 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" -[[package]] -name = "find-msvc-tools" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" - [[package]] name = "getrandom" version = "0.3.3" @@ -278,15 +261,6 @@ version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" -[[package]] -name = "libmimalloc-sys" -version = "0.1.49" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" -dependencies = [ - "cc", -] - [[package]] name = "log" version = "0.4.27" @@ -309,15 +283,6 @@ version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" -[[package]] -name = "mimalloc" -version = "0.1.52" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" -dependencies = [ - "libmimalloc-sys", -] - [[package]] name = "ndarray" version = "0.16.1" @@ -685,12 +650,6 @@ dependencies = [ "serde", ] -[[package]] -name = "shlex" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" - [[package]] name = "syn" version = "2.0.90" diff --git a/Cargo.toml b/Cargo.toml index a2ae9c4..b635f19 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,7 +26,6 @@ pyo3 = { version = "0.29.0", features = ["macros"] } numpy = { version = "0.29.0" } rand = { version = "0.9.1" } rand_distr = { version = "0.5.0" } -mimalloc = "0.1" [dev-dependencies] criterion = { version = "0.6.0", features = ["html_reports"] } diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index d95fd5a..07ddcd4 100644 --- a/python/bartrs/pgbart.py +++ b/python/bartrs/pgbart.py @@ -32,7 +32,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 +from pymc_bart.utils import _encode_vi, _encode_obj class PGBART(ArrayStepShared): @@ -58,6 +58,8 @@ 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, []), } @@ -104,14 +106,13 @@ def __init__( # noqa: PLR0915 self.shape = 1 if len(shape) == 1 else shape[0] + self.bart.n_outputs = self.shape + # Updated later in self.setup() by pymc.sample(...) self.n_draws = 0 self.n_tune = 0 self.n_total = 0 - # Set trees_shape (dim for separate tree structures) - # and leaves_shape (dim for leaf node values) - # One of the two is always one, the other equal to self.shape if self.bart.split_prior.size == 0: self.alpha_vec = np.ones(self.X.shape[1]) @@ -248,39 +249,17 @@ def __setstate__(self, state): def astep(self, _): t0 = perf_counter() self.compiled_pymc_model.update_shared_arrays() - sum_trees, variable_inclusion = self.pg_bart.step(self.tune) + sum_trees, variable_inclusion, new_batch, new_baseline = self.pg_bart.step(self.tune) t1 = perf_counter() - - import psutil - _PROC = psutil.Process() - - def cur_rss_mb(): - return _PROC.memory_info().rss / 1e6 - - - import resource, os - def rss_mb(): - r = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - return r / (1024*1024 if os.uname().sysname == "Darwin" else 1024) - - # in astep, print RSS at the tune->post transition and at the end - self._astep_calls += 1 - if self._astep_calls == self.n_tune: - print(f"[end of tune] cur_rss={cur_rss_mb():.1f}MB calls={self._astep_calls}", flush=True) - if self._astep_calls == self.n_total: - print(f"[end of sampling] cur_rss={cur_rss_mb():.1f}MB", flush=True) - self.pg_bart.report_stored_bytes() - - if self._astep_calls % 200 == 0: - print(f"call={self._astep_calls} cur_rss={cur_rss_mb():.1f}MB", flush=True) 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, } - return sum_trees, [stats] @@ -293,9 +272,7 @@ def competence(var, has_grad): return Competence.INCOMPATIBLE def setup(self, tune: int, draws: int) -> None: - self.n_draws = draws - self.n_tune = tune - self.n_total = draws + tune + self.pg_bart.reserve_draws(draws) def calculate_max_tree_depth(alpha: float, beta: float, probs_leaf: float) -> int: """Calculates the maximum tree depth for which the probability of a node diff --git a/src/lib.rs b/src/lib.rs index ac6a625..2d5cd59 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,11 +37,6 @@ use rand_distr::StandardNormal; type LogpFunc = unsafe extern "C" fn(*const f64, usize) -> c_double; - -#[global_allocator] -static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; - - #[pyclass(from_py_object)] #[derive(Clone, Debug)] #[allow(dead_code)] @@ -129,6 +124,7 @@ struct PySampler { all_trees: Vec, baseline_forest: Option>, n_trees: usize, + n_outputs: usize, } #[pymethods] @@ -231,15 +227,26 @@ impl PySampler { all_trees, baseline_forest: None, n_trees: settings.n_trees, + n_outputs: settings.n_outputs, }) } + fn reserve_draws(&mut self, draws: usize) { + self.all_trees = Vec::with_capacity(draws); + } + #[pyo3(signature = (tune = None))] + #[allow(clippy::type_complexity)] fn step<'py>( &mut self, py: Python<'py>, tune: Option, - ) -> PyResult<(Bound<'py, PyArray2>, Vec)> { + ) -> PyResult<( + Bound<'py, PyArray2>, + Vec, + Option<(usize, Vec)>, + Option>, + )> { let mut state = self .state @@ -251,24 +258,30 @@ 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); + self.baseline_forest = Some(forest.clone()); + new_baseline = 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 { let idx = (info.batch_start + k) % self.n_trees; let mut t = new_state.forest[idx].clone(); - t.leaf_indices = Vec::new(); // only the STORED copy loses it + t.leaf_indices = Vec::new(); trees.push(t); } - self.all_trees.push(TreeBatch { start: info.batch_start, trees }); + self.all_trees.push(TreeBatch { start: info.batch_start, trees: trees.clone() }); + new_batch = Some((info.batch_start, trees)); } @@ -279,79 +292,9 @@ impl PySampler { let variable_inclusion = self.state.as_ref().unwrap().get_variable_inclusion(); - Ok((result, variable_inclusion)) - } - - /// Diagnostic: per-field byte breakdown of everything currently retained in - /// `all_trees` + `baseline_forest`. Reports both `len`-based bytes (actual - /// data) and `capacity`-based bytes (allocated), so fixed-width buffers and - /// capacity slack are both visible. Call once after sampling finishes. - fn report_stored_bytes(&self) { - use std::collections::BTreeMap; - - let mut len_tot: BTreeMap<&str, usize> = BTreeMap::new(); - let mut cap_tot: BTreeMap<&str, usize> = BTreeMap::new(); - let mut n_trees_counted = 0usize; - let mut node_count_sum = 0usize; // sum of split_var.len() across trees - let mut max_node_count = 0usize; - - let all = self - .all_trees - .iter() - .flat_map(|b| b.trees.iter()) - .chain(self.baseline_forest.iter().flatten()); - - for t in all { - n_trees_counted += 1; - node_count_sum += t.split_var.len(); - max_node_count = max_node_count.max(t.split_var.len()); - for (name, len_b, cap_b) in tree_field_report(t) { - *len_tot.entry(name).or_default() += len_b; - *cap_tot.entry(name).or_default() += cap_b; - } - } - - eprintln!( - "--- stored tree field breakdown ({} trees, {} batches, baseline={}) ---", - n_trees_counted, - self.all_trees.len(), - self.baseline_forest.is_some() - ); - - let mut len_sum = 0usize; - let mut cap_sum = 0usize; - for (name, &lb) in &len_tot { - let cb = cap_tot[name]; - len_sum += lb; - cap_sum += cb; - eprintln!( - "{:<16} len={:>9.2}MB cap={:>9.2}MB slack={:>8.2}MB", - name, - lb as f64 / 1e6, - cb as f64 / 1e6, - (cb - lb) as f64 / 1e6 - ); - } - eprintln!( - "{:<16} len={:>9.2}MB cap={:>9.2}MB slack={:>8.2}MB", - "TOTAL", - len_sum as f64 / 1e6, - cap_sum as f64 / 1e6, - (cap_sum - len_sum) as f64 / 1e6 - ); - - if n_trees_counted > 0 { - eprintln!( - "avg per tree: len={:.0}B cap={:.0}B | avg nodes/tree={:.1} max nodes={}", - len_sum as f64 / n_trees_counted as f64, - cap_sum as f64 / n_trees_counted as f64, - node_count_sum as f64 / n_trees_counted as f64, - max_node_count - ); - } + Ok((result, variable_inclusion, new_batch, new_baseline)) } - fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult>>{ let data = x.as_array().to_owned(); @@ -370,83 +313,146 @@ impl PySampler { impl PySampler { fn _sample_posterior(&mut self, x: &Array, samples: usize, excluded: &Vec) -> Array { - let n_outputs = self.state.as_ref().expect("Sampler state is missing").predictions.nrows(); - let n_data_samples: usize = x.nrows(); - let n_trees = self.n_trees; - let n_draws = self.all_trees.len(); - - assert!(n_draws > 0, "No posterior draws stored yet"); - - let random_samples: Vec = (0..samples) - .map(|_| self.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 mut sample = predictions.slice_mut(s![posterior_sample_idx, .., ..]); + predict_from_history( + &self.all_trees, + self.baseline_forest.as_deref(), + self.n_trees, + self.n_outputs, + x, + samples, + excluded, + &mut self.rng, + ) + } - let mut covered = vec![false; n_trees]; - let mut remaining = n_trees; +} - // Walk backward from draw_idx, taking the most recent version of - // each tree slot until every slot is accounted for. - let mut b = draw_idx as isize; - while remaining > 0 && b >= 0 { - let batch = &self.all_trees[b as usize]; - for (k, tree) in batch.trees.iter().enumerate() { - let tree_idx = (batch.start + k) % n_trees; - if !covered[tree_idx] { - covered[tree_idx] = true; - remaining -= 1; - sample += &tree.predict_batch_test(x, excluded); - } +#[allow(clippy::too_many_arguments)] +fn predict_from_history( + batches: &[TreeBatch], + baseline_forest: Option<&[TreeArrays]>, + n_trees: usize, + n_outputs: usize, + x: &Array, + samples: usize, + excluded: &Vec, + rng: &mut SmallRng, +) -> Array { + let n_data_samples: usize = x.nrows(); + let n_draws = batches.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 mut sample = predictions.slice_mut(s![posterior_sample_idx, .., ..]); + + let mut covered = vec![false; n_trees]; + let mut remaining = n_trees; + + let mut b = draw_idx as isize; + while remaining > 0 && b >= 0 { + let batch = &batches[b as usize]; + for (k, tree) in batch.trees.iter().enumerate() { + let tree_idx = (batch.start + k) % n_trees; + if !covered[tree_idx] { + covered[tree_idx] = true; + remaining -= 1; + sample += &tree.predict_batch_test(x, excluded); } - b -= 1; } + b -= 1; + } - // Ran out of stored history before covering all slots: fill the rest - // from the end-of-tuning snapshot. - if remaining > 0 { - if let Some(baseline) = &self.baseline_forest { - for tree_idx in 0..n_trees { - if !covered[tree_idx] { - sample += &baseline[tree_idx].predict_batch_test(x, excluded); - } + if remaining > 0 { + if let Some(baseline) = baseline_forest { + for tree_idx in 0..n_trees { + if !covered[tree_idx] { + sample += &baseline[tree_idx].predict_batch_test(x, excluded); } } } } - - predictions } + predictions +} + + +#[pyclass(module = "bartrs.bartrs", weakref)] +struct PosteriorSampler { + batches: Vec, + baseline_forest: Option>, + n_trees: usize, + n_outputs: usize, + rng: SmallRng, } -/// Per-field byte accounting for a single stored tree. -/// Returns (field name, bytes from len, bytes from capacity) for each field. -fn tree_field_report(t: &TreeArrays) -> [(&'static str, usize, usize); 10] { - [ - ("split_var", t.split_var.len() * 4, t.split_var.capacity() * 4), - ("split_val", t.split_val.len() * 8, t.split_val.capacity() * 8), - ("leaf_val", t.leaf_val.len() * 8, t.leaf_val.capacity() * 8), - ("leaf_indices", t.leaf_indices.len() * 4, t.leaf_indices.capacity() * 4), - ("node_nvalue", t.node_nvalue.len() * 8, t.node_nvalue.capacity() * 8), - ("leaf_kind", t.leaf_kind.len(), t.leaf_kind.capacity()), - ("leaf_param_idx", t.leaf_param_idx.len() * 8, t.leaf_param_idx.capacity() * 8), - ("linear_intercept", t.linear_intercept.iter().map(|v| v.len() * 8).sum(), - t.linear_intercept.iter().map(|v| v.capacity() * 8).sum()), - ("linear_slope", t.linear_slope.iter().map(|v| v.len() * 8).sum(), - t.linear_slope.iter().map(|v| v.capacity() * 8).sum()), - ("linear_var", t.linear_var.len() * 4, t.linear_var.capacity() * 4), - ] +#[pymethods] +impl PosteriorSampler { + #[staticmethod] + fn from_history( + batches: Vec<(usize, Vec)>, + baseline_forest: Option>, + n_trees: usize, + n_outputs: usize, + seed: u64, + ) -> Self { + let batches = batches + .into_iter() + .map(|(start, trees)| TreeBatch { start, trees }) + .collect(); + Self { + batches, + baseline_forest, + n_trees, + n_outputs, + rng: SmallRng::seed_from_u64(seed), + } + } + + #[getter] + fn n_outputs(&self) -> usize { + self.n_outputs + } + + fn sample_posterior<'py>(&mut self, py: Python<'py>, x: PyReadonlyArray2<'py, f64>, samples: usize, excluded: Option<&Bound<'py, PyList>>) -> PyResult>> { + let data = x.as_array().to_owned(); + + let excl = match excluded { + Some(list) => list.iter() + .map(|item| item.extract::()) + .collect::>>()?, + None => Vec::new(), + }; + + let preds = predict_from_history( + &self.batches, + self.baseline_forest.as_deref(), + self.n_trees, + self.n_outputs, + &data, + samples, + &excl, + &mut self.rng, + ).to_pyarray(py).unbind(); + + Ok(preds) + } } + #[pymodule] fn bartrs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } From e1550b9da3d9e813bc13e349940eac49be4f447f Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Thu, 9 Jul 2026 09:36:12 +0300 Subject: [PATCH 10/11] return all batches --- python/bartrs/compile_pymc.py | 3 +- python/bartrs/pgbart.py | 29 +++++++++----- src/lib.rs | 72 +++++++++++++++++------------------ 3 files changed, 57 insertions(+), 47 deletions(-) diff --git a/python/bartrs/compile_pymc.py b/python/bartrs/compile_pymc.py index f94246a..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 @@ -232,7 +232,6 @@ def _make_persistent_arrays(self): the compiled function interface. """ arrays = [item.storage[0].copy() for item in self.logp_fn_ptr.input_storage[1:]] - assert all(arr.dtype == np.float64 for arr in arrays) return arrays # TODO: fast update for shared arrays 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..31874f3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -236,17 +236,11 @@ impl PySampler { } #[pyo3(signature = (tune = None))] - #[allow(clippy::type_complexity)] fn step<'py>( &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 +252,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 +268,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 +278,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 +310,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 +337,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 +390,6 @@ struct PosteriorSampler { baseline_forest: Option>, n_trees: usize, n_outputs: usize, - rng: SmallRng, } #[pymethods] @@ -402,7 +400,6 @@ impl PosteriorSampler { baseline_forest: Option>, n_trees: usize, n_outputs: usize, - seed: u64, ) -> Self { let batches = batches .into_iter() @@ -413,7 +410,6 @@ impl PosteriorSampler { baseline_forest, n_trees, n_outputs, - rng: SmallRng::seed_from_u64(seed), } } @@ -422,7 +418,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 +439,8 @@ impl PosteriorSampler { self.n_trees, self.n_outputs, &data, - samples, + &draw_indices, &excl, - &mut self.rng, ).to_pyarray(py).unbind(); Ok(preds) From 8d600b7aab63c389be4cd1e8c1bddbedf4b7dc08 Mon Sep 17 00:00:00 2001 From: OttoVintola Date: Thu, 9 Jul 2026 10:36:49 +0300 Subject: [PATCH 11/11] remove psutil --- python/bartrs/pgbart.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/python/bartrs/pgbart.py b/python/bartrs/pgbart.py index cada385..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, []), } @@ -257,8 +253,6 @@ def astep(self, _): 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, }