Coverage for src/bartz/stochtree/_stochtree.py: 87%
295 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 11:03 +0000
1# bartz/src/bartz/stochtree/_stochtree.py
2#
3# Copyright (c) 2026, The Bartz Contributors
4#
5# This file is part of bartz.
6#
7# Permission is hereby granted, free of charge, to any person obtaining a copy
8# of this software and associated documentation files (the "Software"), to deal
9# in the Software without restriction, including without limitation the rights
10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
11# copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in all
15# copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
25"""Implement class `BARTModel` that mimics the Python package stochtree."""
27from collections.abc import Mapping, Sequence
28from dataclasses import dataclass, field, fields
29from functools import partial
31# WORKAROUND(python<3.15): use frozendict instead of MappingProxyType
32from types import MappingProxyType
33from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload
35from jax import numpy as jnp
36from jax.scipy.special import ndtr, ndtri
37from jaxtyping import Array, Float, Float32, Key, Real, Shaped
39from bartz._interface import Bart, DataFrame, PredictKind, Series
40from bartz.mcmcstep._state import ArrayLike, FloatLike
41from bartz.prepcovars import RangeEvenBinner
42from bartz.stochtree._preprocess import _PreprocessorBase, make_preprocessor
44if TYPE_CHECKING:
45 from _typeshed import DataclassInstance
47_MAX_DEPTH_LIMIT = 16
50@dataclass(frozen=True)
51class OutcomeModel:
52 """Outcome model specification, matching `stochtree.OutcomeModel`.
54 Only ``('continuous', 'identity')`` and ``('binary', 'probit')`` are
55 supported.
56 """
58 outcome: Literal['continuous', 'binary'] = 'continuous'
59 """Outcome family."""
61 link: Literal['identity', 'probit'] | None = None
62 """Link function. If `None`, defaults to ``'identity'`` for ``'continuous'`` and ``'probit'`` for ``'binary'``."""
64 def __post_init__(self) -> None:
65 if self.link is None:
66 default_link = {'continuous': 'identity', 'binary': 'probit'}.get(
67 self.outcome
68 )
69 object.__setattr__(self, 'link', default_link)
70 if (self.outcome, self.link) not in (
71 ('continuous', 'identity'),
72 ('binary', 'probit'),
73 ):
74 msg = (
75 f'unsupported outcome_model (outcome={self.outcome!r}, '
76 f"link={self.link!r}); only ('continuous', 'identity') "
77 "and ('binary', 'probit') are supported."
78 )
79 raise NotImplementedError(msg)
82class NotSampledError(ValueError, AttributeError):
83 """Raised when calling a method that requires `BARTModel.sample` to have been called."""
86@dataclass(frozen=True, kw_only=True)
87class GeneralParams:
88 """Mirror of stochtree's ``general_params`` dict, with the keys bartz handles."""
90 standardize: bool = True
91 """Whether to standardize the outcome before fitting. Ignored for probit binary."""
93 sigma2_init: FloatLike | None = None
94 """Starting value of the global error variance. If `None` (default), uses ``var(resid_train)`` for continuous and ``1.0`` for probit."""
96 sigma2_global_shape: FloatLike = 0.0
97 """Shape parameter of the inverse-gamma prior on the global error variance. The default ``0`` is mapped to a near-improper prior, since bartz's scaled-inv-chi² cannot represent ``IG(0, 0)`` exactly."""
99 sigma2_global_scale: FloatLike = 0.0
100 """Scale parameter of the inverse-gamma prior on the global error variance. The default ``0`` is mapped to a near-improper prior, since bartz's scaled-inv-chi² cannot represent ``IG(0, 0)`` exactly."""
102 variable_weights: Float[ArrayLike, ' p'] | None = None
103 """Per-predictor sampling weights. Must be strictly positive; pass a small positive value to suppress a variable."""
105 random_seed: int | Key[Array, ''] | None = None
106 """Seed for the random number generator. Unlike stochtree, the default
107 `None` is deterministic (equivalent to seed ``0``) rather than drawing a
108 random seed, so repeated fits reproduce by default."""
110 keep_every: int = 1
111 """Thinning factor for retained MCMC samples."""
113 num_chains: int = 1
114 """Number of independent MCMC chains."""
116 outcome_model: OutcomeModel = field(default_factory=OutcomeModel)
117 """Outcome family and link specification. Defaults to continuous with
118 identity link."""
121@dataclass(frozen=True, kw_only=True)
122class MeanForestParams:
123 """Mirror of stochtree's ``mean_forest_params`` dict, restricted to the keys bartz handles."""
125 num_trees: int = 200
126 """Number of trees in the conditional mean ensemble."""
128 alpha: FloatLike = 0.95
129 """Tree split prior base."""
131 beta: FloatLike = 2.0
132 """Tree split prior decay."""
134 min_samples_leaf: int = 5
135 """Minimum number of training samples at a leaf."""
137 max_depth: int = 10
138 """Maximum tree depth. Must be a non-negative integer at most ``16``."""
140 sample_sigma2_leaf: bool = True
141 """Whether to sample the leaf-variance prior. Must be set to ``False``."""
143 sigma2_leaf_init: FloatLike | None = None
144 """Initial leaf-variance prior (held fixed since ``sample_sigma2_leaf=False``). If `None`, matches stochtree's defaults of ``var(resid_train) / num_trees`` for continuous and ``2 / num_trees`` for probit."""
146 def __post_init__(self) -> None:
147 if self.sample_sigma2_leaf:
148 msg = (
149 'sample_sigma2_leaf=True is not supported (bartz uses a fixed'
150 " leaf-variance prior); pass mean_forest_params={'sample_sigma2_leaf':"
151 ' False} to acknowledge this.'
152 )
153 raise NotImplementedError(msg)
154 if self.max_depth < 0: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true
155 msg = (
156 f'max_depth={self.max_depth} is not supported; bartz stores trees'
157 ' as heap arrays of size 2**max_depth, so the stochtree'
158 ' convention max_depth=-1 (unbounded) is rejected. Pass a'
159 f' non-negative integer at most {_MAX_DEPTH_LIMIT}.'
160 )
161 raise NotImplementedError(msg)
162 if self.max_depth > _MAX_DEPTH_LIMIT: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true
163 msg = (
164 f'max_depth={self.max_depth} exceeds {_MAX_DEPTH_LIMIT}; bartz'
165 ' stores trees as heap arrays of size 2**max_depth, so memory'
166 ' grows exponentially with depth.'
167 )
168 raise ValueError(msg)
171if TYPE_CHECKING:
172 # static only, beartype does not support type[DataclassInstance]
173 T = TypeVar('T', bound='DataclassInstance')
174else:
175 T = TypeVar('T')
178def build_dataclass(cls: type[T], params: Mapping[str, Any] | None, name: str) -> T:
179 """Convert a user-supplied dict to a dataclass, with friendly errors."""
180 if params is None:
181 params = {}
182 allowed = {f.name for f in fields(cls)}
183 extra = set(params) - allowed
184 if extra:
185 msg = (
186 f'{name} contains unsupported key(s) {sorted(extra)}; valid keys'
187 f' are {sorted(allowed)}'
188 )
189 raise ValueError(msg)
190 return cls(**params)
193class BARTModel:
194 R"""
195 BART model with a `stochtree`-compatible interface, powered by bartz.
197 This class mimics `stochtree.BARTModel` so that bartz can be used as a
198 drop-in reference implementation for testing. The intersection of features
199 is targeted: continuous regression (Gaussian outcome, identity link) and
200 binary classification (probit link) on tabular covariates.
202 Use the same idiomatic pattern as `stochtree.BARTModel`::
204 m = BARTModel()
205 m.sample(
206 X_train=X, y_train=y, X_test=X_test,
207 num_gfr=0, num_mcmc=200,
208 mean_forest_params={'sample_sigma2_leaf': False},
209 )
210 yhat = m.predict(X_new, terms='y_hat', type='mean')
212 See `GeneralParams` and `MeanForestParams` for the supported keys in the
213 ``general_params`` and ``mean_forest_params`` dicts.
215 Notes
216 -----
217 Differences from `stochtree`, by design:
219 - ``num_gfr`` has no default and must be set explicitly to ``0``.
220 - ``mean_forest_params['sample_sigma2_leaf']`` must be ``False``.
221 - ``mean_forest_params['max_depth']`` must be a non-negative integer at
222 most ``16``; stochtree's ``-1`` (unbounded depth) sentinel is not
223 accepted.
224 - The deprecated ``general_params['probit_outcome_model']`` flag is not
225 accepted; pass ``outcome_model=OutcomeModel('binary', 'probit')``
226 instead.
227 - ``general_params['cutpoint_grid_size']`` is not accepted; bartz uses a
228 fixed grid of 256 evenly-spaced bins per predictor. stochtree only
229 uses this parameter for the GFR sampler, which bartz does not support.
230 - Leaf-basis regression, random effects, heteroskedastic variance
231 forests, and warm-starting from a previous model are not supported.
232 - bartz uses single-precision floats, so outputs differ from stochtree
233 at the float32 precision level.
234 - ``general_params['random_seed']`` defaults to deterministic behavior
235 (seed ``0``) when unset, whereas stochtree draws a random seed. This is
236 intentional, to make repeated fits reproducible by default.
238 References
239 ----------
240 Herren, A., Hahn, P. R., Murray, J., Carvalho, C. (2026). "StochTree:
241 BART-based modeling in R and Python". arXiv:2512.12051.
242 """
244 # public, set by sample()
245 sampled: bool
246 """Whether `sample` has been called."""
248 standardize: bool
249 """Whether the outcome was standardized before fitting."""
251 sample_sigma2_global: bool
252 """Whether the global error variance is sampled (always ``True``)."""
254 probit_outcome_model: bool
255 """Whether the model uses a binary outcome with probit link."""
257 outcome_model: OutcomeModel
258 """Outcome family and link specification used during fitting."""
260 num_gfr: int
261 """Number of grow-from-root iterations (always ``0``)."""
263 num_burnin: int
264 """Number of MCMC burn-in iterations."""
266 num_mcmc: int
267 """Number of retained MCMC iterations per chain."""
269 num_chains: int
270 """Number of independent MCMC chains."""
272 num_samples: int
273 """Total number of retained posterior samples (``num_mcmc * num_chains``)."""
275 sigma2_init: FloatLike
276 """Starting value of the global error variance actually used to seed the chain."""
278 y_bar: Float32[Array, '']
279 """Mean used to standardize the outcome (``0`` if not standardized)."""
281 y_std: Float32[Array, '']
282 """Standard deviation used to standardize the outcome (``1`` if not standardized)."""
284 has_rfx: bool
285 """Whether the model includes random effects (always ``False``)."""
287 include_mean_forest: bool
288 """Whether the model includes a conditional mean forest (always ``True``)."""
290 include_variance_forest: bool
291 """Whether the model includes a variance forest (always ``False``)."""
293 y_hat_train: Float32[Array, 'n num_samples']
294 """Posterior predictions at the training covariates, in the original outcome scale."""
296 global_var_samples: Float32[Array, ' num_samples']
297 """Posterior samples of the global error variance. For probit binary regression, an array of ones."""
299 y_hat_test: Float32[Array, 'm num_samples'] | None
300 """Posterior predictions at `X_test` if it was supplied to `sample`, else `None`."""
302 _bart: Bart
303 _preprocessor: _PreprocessorBase | None
305 def __init__(self) -> None:
306 self.sampled = False
307 self._preprocessor = None
309 def is_sampled(self) -> bool:
310 """Return whether `sample` has been called."""
311 return self.sampled
313 def _prepare_training_inputs(
314 self,
315 X_train: Real[ArrayLike, 'n p'] | DataFrame,
316 y_train: Real[ArrayLike, ' n'] | Series,
317 gp: GeneralParams,
318 ) -> tuple[Real[Array, 'n p'], Real[Array, ' n'], Float32[Array, ' p'] | None]:
319 """Coerce inputs and build variable weights, fitting the DataFrame preprocessor if any."""
320 y_train_arr = _coerce_response(y_train, name='y_train')
322 self._preprocessor = make_preprocessor(X_train)
323 if self._preprocessor is None:
324 X_train_arr = check_X(X_train, name='X_train')
325 _, p = X_train_arr.shape
326 varprob = check_variable_weights(gp.variable_weights, p)
327 else:
328 # The preprocessor decides the default weights: uniform over the
329 # *original* columns split across each one-hot expansion (so every
330 # original variable keeps an equal splitting budget), or `None` when
331 # nothing expands (deferring to bartz's native uniform fast-path).
332 weights_np = self._preprocessor.fit(
333 X_train, variable_weights=gp.variable_weights
334 )
335 X_train_np = self._preprocessor.transform(X_train)
336 if X_train_np.shape[1] == 0: 336 ↛ 337line 336 didn't jump to line 337 because the condition on line 336 was never true
337 msg = 'X_train has no usable columns after preprocessing'
338 raise ValueError(msg)
339 X_train_arr = jnp.asarray(X_train_np)
340 varprob = None if weights_np is None else jnp.asarray(weights_np)
342 n, _ = X_train_arr.shape
343 if y_train_arr.shape[0] != n: 343 ↛ 344line 343 didn't jump to line 344 because the condition on line 343 was never true
344 msg = (
345 f'X_train and y_train length mismatch: X_train has {n} rows,'
346 f' y_train has {y_train_arr.shape[0]} entries'
347 )
348 raise ValueError(msg)
349 return X_train_arr, y_train_arr, varprob
351 def sample(
352 self,
353 X_train: Real[ArrayLike, 'n p'] | DataFrame,
354 y_train: Real[ArrayLike, ' n'] | Series,
355 X_test: Real[ArrayLike, 'm p'] | DataFrame | None = None,
356 observation_weights: Float[ArrayLike, ' n'] | Series | None = None,
357 *,
358 num_gfr: int,
359 num_burnin: int = 0,
360 num_mcmc: int = 100,
361 general_params: Mapping[str, Any] | None = None,
362 mean_forest_params: Mapping[str, Any] | None = None,
363 bart_kwargs: Mapping[str, Any] = MappingProxyType({}),
364 ) -> None:
365 """Fit the model.
367 The signature mirrors `stochtree.BARTModel.sample`, restricted to the
368 keyword arguments bartz supports.
370 Parameters
371 ----------
372 X_train
373 Training covariates with shape ``(n, p)``.
374 y_train
375 Training outcomes of length ``n``.
376 X_test
377 Optional test covariates; if given, predictions are cached on
378 them in `y_hat_test`.
379 observation_weights
380 Optional positive per-observation weights scaling the residual
381 variance (``y_i | - ~ N(mu(X_i), sigma^2 / w_i)``). Not supported
382 with a probit outcome model.
383 num_gfr
384 Number of grow-from-root iterations. Must be ``0``.
385 num_burnin
386 Number of MCMC burn-in iterations.
387 num_mcmc
388 Number of retained MCMC iterations per chain.
389 general_params
390 Optional override for the keys of `GeneralParams`.
391 mean_forest_params
392 Override for the keys of `MeanForestParams`. Must explicitly
393 disable ``sample_sigma2_leaf``.
394 bart_kwargs
395 Additional arguments forwarded to `bartz.Bart`. Use this to set
396 ``devices`` and ``rm_const=False`` when wrapping `sample` in
397 `jax.jit`.
399 Raises
400 ------
401 NotImplementedError
402 If ``num_gfr`` is non-zero.
403 ValueError
404 If `observation_weights` are set with a probit outcome model.
405 """
406 if num_gfr != 0:
407 msg = (
408 'num_gfr must be 0; the grow-from-root sampler is not available'
409 ' in bartz.'
410 )
411 raise NotImplementedError(msg)
413 gp = build_dataclass(GeneralParams, general_params, 'general_params')
414 mfp = build_dataclass(
415 MeanForestParams, mean_forest_params, 'mean_forest_params'
416 )
418 is_probit = gp.outcome_model.outcome == 'binary'
420 # stochtree does not support heteroskedastic probit
421 if is_probit and observation_weights is not None:
422 msg = (
423 'observation_weights are not supported with a probit outcome'
424 ' model; stochtree does not support heteroskedastic probit.'
425 )
426 raise ValueError(msg)
428 X_train_arr, y_train_arr, variable_weights = self._prepare_training_inputs(
429 X_train, y_train, gp
430 )
432 y_bar, y_std, y_for_bartz = standardize_y(
433 y_train_arr, is_probit, gp.standardize
434 )
436 bart_num_chains = None if gp.num_chains == 1 else gp.num_chains
438 # variance of the standardized residual, matching stochtree
439 # (np.var(resid_train) with ddof=0). For standardize=True it is exactly
440 # 1.0; we hardcode that so the value stays trace-time concrete.
441 if is_probit:
442 var_resid_train: FloatLike = 1.0 # bartz ignores σ² for binary
443 elif gp.standardize: 443 ↛ 446line 443 didn't jump to line 446 because the condition on line 443 was always true
444 var_resid_train = 1.0
445 else:
446 var_resid_train = jnp.var(y_for_bartz)
448 # leaf-prior: bartz uses sigma_mu = tau_num / (k * sqrt(num_trees));
449 # stochtree's sigma2_leaf is the leaf-variance prior. Hold k=2 and solve
450 # for tau_num so that the two parameterizations agree.
451 bartz_k = 2.0
452 sigma2_leaf_init = resolve_sigma2_leaf_init(
453 mfp.sigma2_leaf_init, mfp.num_trees, is_probit, var_resid_train
454 )
455 tau_num_arg = bartz_k * jnp.sqrt(mfp.num_trees * sigma2_leaf_init)
457 if is_probit:
458 # stochtree pins σ²=1 for probit; bartz binary branch ignores the
459 # variance prior, so we leave the scale/init at their 'auto'
460 # defaults (bartz rejects explicit values for binary outcomes).
461 sigma_df_arg: FloatLike = 3.0
462 sigma_scale_arg: FloatLike | Literal['auto'] = 'auto'
463 sigma_init_arg: FloatLike | Literal['auto'] = 'auto'
464 sigma2_init_stored: FloatLike = 1.0
465 else:
466 sigma_df_arg, sigma_scale_arg, sigma_init_arg, sigma2_init_stored = (
467 resolve_variance_prior(
468 gp.sigma2_global_shape,
469 gp.sigma2_global_scale,
470 gp.sigma2_init,
471 var_resid_train,
472 )
473 )
475 binner = partial(RangeEvenBinner, max_bins=256)
477 seed = 0 if gp.random_seed is None else gp.random_seed
479 kwargs: dict = dict(
480 x_train=X_train_arr.T,
481 y_train=y_for_bartz,
482 outcome_type='binary' if is_probit else 'continuous',
483 binner=binner,
484 varprob=variable_weights,
485 sigma_df=sigma_df_arg,
486 sigma_scale=sigma_scale_arg,
487 sigma_init=sigma_init_arg,
488 k=bartz_k,
489 power=mfp.beta,
490 base=mfp.alpha,
491 tau_num=tau_num_arg,
492 error_scale=observation_weights,
493 num_trees=mfp.num_trees,
494 n_save=num_mcmc,
495 n_burn=num_burnin,
496 n_skip=gp.keep_every,
497 printevery=None,
498 num_chains=bart_num_chains,
499 seed=seed,
500 maxdepth=mfp.max_depth + 1,
501 precompute_predict_train=True,
502 )
503 kwargs.update(bart_kwargs)
504 # match stochtree's gating: only acceptance-time veto on
505 # min_samples_leaf, no per-leaf affluence filter (stochtree picks
506 # leaves uniformly over all of them). User-supplied init_kw values
507 # win on conflicts.
508 kwargs = dict(
509 kwargs,
510 init_kw=dict(
511 {
512 'min_points_per_leaf': mfp.min_samples_leaf,
513 'min_points_per_decision_node': None,
514 },
515 **kwargs.get('init_kw', {}),
516 ),
517 )
518 self._bart = Bart(**kwargs)
519 self._finalize_sample(
520 outcome_model=gp.outcome_model,
521 num_burnin=num_burnin,
522 num_mcmc=num_mcmc,
523 num_chains=gp.num_chains,
524 sigma2_init=sigma2_init_stored,
525 y_bar=y_bar,
526 y_std=y_std,
527 standardize=gp.standardize,
528 X_test=X_test,
529 )
531 def _finalize_sample(
532 self,
533 *,
534 outcome_model: OutcomeModel,
535 num_burnin: int,
536 num_mcmc: int,
537 num_chains: int,
538 sigma2_init: FloatLike,
539 y_bar: Float32[Array, ''],
540 y_std: Float32[Array, ''],
541 standardize: bool,
542 X_test: Real[ArrayLike, 'm p'] | DataFrame | None,
543 ) -> None:
544 """Populate the public attributes after `_bart` has been constructed."""
545 is_probit = outcome_model.outcome == 'binary'
546 self.sampled = True
547 self.standardize = standardize
548 self.sample_sigma2_global = True
549 self.probit_outcome_model = is_probit
550 self.outcome_model = outcome_model
551 self.num_gfr = 0
552 self.num_burnin = num_burnin
553 self.num_mcmc = num_mcmc
554 self.num_chains = num_chains
555 self.num_samples = num_mcmc * num_chains
556 self.sigma2_init = sigma2_init
557 self.y_bar = y_bar
558 self.y_std = y_std
559 self.has_rfx = False
560 self.include_mean_forest = True
561 self.include_variance_forest = False
563 # cached outputs in stochtree's (n, num_samples) layout, original scale
564 self.y_hat_train = self._predict_y_hat_internal('train')
565 if X_test is not None:
566 self.y_hat_test = self._predict_y_hat_internal(self._prepare_x(X_test).T)
567 else:
568 self.y_hat_test = None
570 if is_probit:
571 self.global_var_samples = jnp.ones((self.num_samples,))
572 else:
573 sigma = self._bart.get_error_sdev()
574 self.global_var_samples = (sigma * y_std) ** 2
576 @overload
577 def predict(
578 self,
579 X: Real[ArrayLike, 'm p'] | DataFrame,
580 *,
581 type: Literal['posterior', 'mean'] = 'posterior',
582 terms: Literal['y_hat', 'mean_forest'],
583 scale: Literal['linear', 'probability', 'class'] = 'linear',
584 ) -> Shaped[Array, 'm num_samples'] | Shaped[Array, ' m']: ...
586 @overload
587 def predict(
588 self,
589 X: Real[ArrayLike, 'm p'] | DataFrame,
590 *,
591 type: Literal['posterior', 'mean'] = 'posterior',
592 terms: Literal['all'] = 'all',
593 scale: Literal['linear', 'probability', 'class'] = 'linear',
594 ) -> dict[str, Shaped[Array, 'm num_samples']] | dict[str, Shaped[Array, ' m']]: ...
596 @overload
597 def predict(
598 self,
599 X: Real[ArrayLike, 'm p'] | DataFrame,
600 *,
601 type: Literal['posterior', 'mean'] = 'posterior',
602 terms: Sequence[Literal['y_hat', 'mean_forest', 'all']],
603 scale: Literal['linear', 'probability', 'class'] = 'linear',
604 ) -> (
605 Shaped[Array, 'm num_samples']
606 | Shaped[Array, ' m']
607 | dict[str, Shaped[Array, 'm num_samples']]
608 | dict[str, Shaped[Array, ' m']]
609 ): ...
611 def predict(
612 self,
613 X: Real[ArrayLike, 'm p'] | DataFrame,
614 *,
615 type: Literal['posterior', 'mean'] = 'posterior', # noqa: A002
616 terms: Literal['y_hat', 'mean_forest', 'all']
617 | Sequence[Literal['y_hat', 'mean_forest', 'all']] = 'all',
618 scale: Literal['linear', 'probability', 'class'] = 'linear',
619 ) -> (
620 Shaped[Array, 'm num_samples']
621 | Shaped[Array, ' m']
622 | dict[str, Shaped[Array, 'm num_samples']]
623 | dict[str, Shaped[Array, ' m']]
624 ):
625 """Predict at new covariates.
627 Parameters
628 ----------
629 X
630 New covariates with shape ``(m, p)``.
631 type
632 ``'posterior'`` returns one prediction per posterior sample, with
633 shape ``(m, num_samples)``. ``'mean'`` averages the posterior
634 samples, returning a vector of shape ``(m,)``.
635 terms
636 One of ``'y_hat'``, ``'mean_forest'``, ``'all'``, or a list. Since
637 random effects and a variance forest are not supported, ``'y_hat'``
638 and ``'mean_forest'`` produce the same result.
639 scale
640 For probit binary regression: ``'linear'`` returns the eta values,
641 ``'probability'`` returns ``Phi(eta)``, ``'class'`` returns 0 / 1.
642 Only ``'linear'`` is valid for continuous outcomes.
644 Returns
645 -------
646 Either a single jax array (for a single requested term) or a dict keyed by term name.
648 Raises
649 ------
650 NotSampledError
651 If `sample` has not been called yet.
652 """
653 if not self.sampled:
654 msg = (
655 "This BARTModel instance is not fitted yet. Call 'sample' before"
656 ' using this model.'
657 )
658 raise NotSampledError(msg)
659 terms_tuple = check_predict_args(type, scale, terms, self.probit_outcome_model)
661 pred = self._predict_y_hat_internal(self._prepare_x(X).T)
663 if self.probit_outcome_model and scale in ('probability', 'class'):
664 prob = ndtr(pred)
665 pred_out = jnp.where(prob < 0.5, 0, 1) if scale == 'class' else prob
666 else:
667 pred_out = pred
669 if type == 'mean':
670 pred_out = jnp.mean(pred_out, axis=1)
672 wants_y_hat = ('y_hat' in terms_tuple) or ('all' in terms_tuple)
673 wants_mean_forest = ('mean_forest' in terms_tuple) or ('all' in terms_tuple)
674 single = sum([wants_y_hat, wants_mean_forest]) == 1
675 if single:
676 return pred_out
677 result: dict[str, Shaped[Array, '...']] = {}
678 if wants_y_hat: 678 ↛ 680line 678 didn't jump to line 680 because the condition on line 678 was always true
679 result['y_hat'] = pred_out
680 if wants_mean_forest: 680 ↛ 682line 680 didn't jump to line 682 because the condition on line 680 was always true
681 result['mean_forest_predictions'] = pred_out
682 return result
684 def _prepare_x(self, X: Real[ArrayLike, 'm p'] | DataFrame) -> Real[Array, 'm p']:
685 """Convert covariates to a 2-D jax array, replaying the fitted preprocessor if any."""
686 if self._preprocessor is None:
687 return check_X(X)
688 if make_preprocessor(X) is None:
689 msg = (
690 'this model was fit on a DataFrame, so prediction covariates must'
691 ' also be a pandas/polars DataFrame with the same columns; got a'
692 ' non-DataFrame. Passing a raw array would bypass the fitted'
693 ' preprocessing (e.g. one-hot encoding) and silently misalign the'
694 ' features.'
695 )
696 raise TypeError(msg)
697 return jnp.asarray(self._preprocessor.transform(X))
699 def _predict_y_hat_internal(
700 self, x: Real[ArrayLike, 'p m'] | Literal['train']
701 ) -> Float32[Array, 'm num_samples']:
702 """Return predictions on the original outcome scale, layout ``(m, num_samples)``."""
703 latent = self._bart.predict(x, kind=PredictKind.latent_samples)
704 if self.probit_outcome_model:
705 # bartz integrates the binary offset into latent; result already on probit scale.
706 return latent.T
707 if self.standardize: 707 ↛ 709line 707 didn't jump to line 709 because the condition on line 707 was always true
708 return (latent * self.y_std + self.y_bar).T
709 return latent.T
712def standardize_y(
713 y_train: Real[ArrayLike, ' n'], is_probit: bool, standardize: bool
714) -> tuple[Float32[Array, ''], Float32[Array, ''], Float32[Array, ' n']]:
715 """Return ``(y_bar, y_std, y_for_bartz)`` matching stochtree's standardization."""
716 y = jnp.asarray(y_train, jnp.float32)
717 if is_probit:
718 return ndtri(y.mean()), jnp.float32(1.0), (y != 0).astype(jnp.float32)
719 if standardize: 719 ↛ 724line 719 didn't jump to line 724 because the condition on line 719 was always true
720 y_bar = y.mean()
721 y_std_val = y.std()
722 y_std = jnp.where(y_std_val > 0, y_std_val, 1.0)
723 return y_bar, y_std, (y - y_bar) / y_std
724 return jnp.float32(0.0), jnp.float32(1.0), y
727def resolve_sigma2_leaf_init(
728 sigma2_leaf_init: FloatLike | None,
729 num_trees: int,
730 is_probit: bool,
731 var_resid_train: FloatLike,
732) -> FloatLike:
733 """Default `sigma2_leaf_init` per stochtree: probit→2/num_trees, continuous→var(resid)/num_trees."""
734 if sigma2_leaf_init is not None:
735 return sigma2_leaf_init
736 if is_probit:
737 return 2.0 / num_trees
738 return var_resid_train / num_trees
741def resolve_variance_prior(
742 shape: FloatLike,
743 scale: FloatLike,
744 sigma2_init: FloatLike | None,
745 var_resid_train: FloatLike,
746) -> tuple[Float32[Array, ''], Float32[Array, ''], Float32[Array, ''], FloatLike]:
747 """Translate stochtree's IG(shape, scale) prior to bartz's error variance prior.
749 The IG(shape, scale) prior on σ² is the scaled-inverse-χ² with
750 ``sigma_df = 2*shape`` and prior harmonic mean ``square(sigma_scale) =
751 scale/shape``; the chain starts at `sigma2_init` (default
752 ``var(resid_train)``), decoupled from the prior. The mapping is branchless so
753 `shape` / `scale` may be traced; the unrepresentable IG(0, scale>0) (positive
754 rate, zero df) yields a NaN that surfaces downstream rather than an error.
756 Parameters
757 ----------
758 shape
759 Stochtree's ``sigma2_global_shape``.
760 scale
761 Stochtree's ``sigma2_global_scale``.
762 sigma2_init
763 Stochtree's ``sigma2_init``. If `None`, defaults to `var_resid_train`.
764 var_resid_train
765 Variance of the residual, the default chain start for σ².
767 Returns
768 -------
769 sigma_df : Float32[Array, '']
770 Degrees of freedom of bartz's error variance prior.
771 sigma_scale : Float32[Array, '']
772 Scale of bartz's prior (sqrt of the prior harmonic mean of the variance).
773 sigma_init : Float32[Array, '']
774 Initial error standard deviation seeding the chain.
775 sigma2_init_stored : FloatLike
776 The chain starting value of σ², suitable for ``BARTModel.sigma2_init``.
777 """
778 shape = jnp.asarray(shape, jnp.float32)
779 scale = jnp.asarray(scale, jnp.float32)
780 sigma2_start = sigma2_init if sigma2_init is not None else var_resid_train
781 sigma_init = jnp.sqrt(jnp.asarray(sigma2_start, jnp.float32))
782 # IG(shape, scale) <=> scaled-inv-chi2(df=2*shape, harmonic mean=scale/shape).
783 # The `scale > 0` guard keeps IG(0, 0) at harmonic mean 0 (avoiding 0/0) while
784 # letting IG(0, scale>0) overflow to inf -> NaN rate, flagging it as invalid.
785 harmonic_mean = jnp.where(scale > 0, scale / shape, 0.0)
786 return 2.0 * shape, jnp.sqrt(harmonic_mean), sigma_init, sigma2_start
789def check_variable_weights(
790 variable_weights: Float[ArrayLike, ' p'] | None, p: int
791) -> Float32[Array, ' p'] | None:
792 """Validate `variable_weights`, returning the jax array (or None)."""
793 if variable_weights is None:
794 return None
795 arr = jnp.asarray(variable_weights, jnp.float32)
796 if arr.shape != (p,): 796 ↛ 797line 796 didn't jump to line 797 because the condition on line 796 was never true
797 msg = f'variable_weights must have shape (p,)=({p},), got {arr.shape}'
798 raise ValueError(msg)
799 return arr
802def check_predict_args(
803 type_: Literal['posterior', 'mean'],
804 scale: Literal['linear', 'probability', 'class'],
805 terms: Literal['y_hat', 'mean_forest', 'all']
806 | Sequence[Literal['y_hat', 'mean_forest', 'all']],
807 probit_outcome_model: bool,
808) -> tuple[str, ...]:
809 """Validate `BARTModel.predict` arguments, returning the normalized terms tuple."""
810 if scale not in ('linear', 'probability', 'class'): 810 ↛ 811line 810 didn't jump to line 811 because the condition on line 810 was never true
811 msg = f"scale must be 'linear', 'probability', or 'class'; got {scale!r}"
812 raise ValueError(msg)
813 if type_ not in ('posterior', 'mean'): 813 ↛ 814line 813 didn't jump to line 814 because the condition on line 813 was never true
814 msg = f"type must be 'posterior' or 'mean'; got {type_!r}"
815 raise ValueError(msg)
816 if not probit_outcome_model and scale != 'linear': 816 ↛ 817line 816 didn't jump to line 817 because the condition on line 816 was never true
817 msg = (
818 "scale must be 'linear' for non-probit (continuous) regression;"
819 f' got {scale!r}'
820 )
821 raise ValueError(msg)
822 if type_ == 'mean' and scale == 'class': 822 ↛ 823line 822 didn't jump to line 823 because the condition on line 822 was never true
823 msg = "scale='class' is incompatible with type='mean'"
824 raise ValueError(msg)
825 terms_tuple = (terms,) if isinstance(terms, str) else tuple(terms)
826 for t in terms_tuple:
827 if t not in ('y_hat', 'mean_forest', 'all'): 827 ↛ 828line 827 didn't jump to line 828 because the condition on line 827 was never true
828 msg = f'unknown term {t!r}; valid terms are y_hat, mean_forest, all'
829 raise ValueError(msg)
830 if scale == 'class' and set(terms_tuple) != {'y_hat'}:
831 # match stochtree: 'class' converts only the single 'y_hat' term, so it
832 # rejects 'mean_forest' and 'all' (the latter also pulls in mean_forest)
833 msg = "scale='class' is only supported when requesting a single 'y_hat' term"
834 raise ValueError(msg)
835 return terms_tuple
838def check_X(
839 X: Real[ArrayLike, 'n p'] | DataFrame, *, name: str = 'X'
840) -> Real[Array, 'n p']:
841 """Convert a DataFrame/array-like to a 2-D jax array in ``(n, p)`` layout."""
842 if isinstance(X, DataFrame): 842 ↛ 843line 842 didn't jump to line 843 because the condition on line 842 was never true
843 X = X.to_numpy()
844 arr = jnp.asarray(X)
845 if arr.ndim == 1: 845 ↛ 846line 845 didn't jump to line 846 because the condition on line 845 was never true
846 arr = arr[:, None]
847 if arr.ndim != 2: 847 ↛ 848line 847 didn't jump to line 848 because the condition on line 847 was never true
848 msg = f'{name} must be 2D (n, p); got shape {arr.shape}'
849 raise ValueError(msg)
850 return arr
853def _coerce_response(
854 y: Real[ArrayLike, ' n'] | Series, *, name: str
855) -> Real[Array, ' n']:
856 """Convert a Series/array-like response to a 1-D jax array."""
857 if isinstance(y, Series): 857 ↛ 858line 857 didn't jump to line 858 because the condition on line 857 was never true
858 y = y.to_numpy()
859 arr = jnp.asarray(y)
860 if arr.ndim != 1: 860 ↛ 861line 860 didn't jump to line 861 because the condition on line 860 was never true
861 msg = f'{name} must be 1D (n,); got shape {arr.shape}'
862 raise ValueError(msg)
863 return arr