Coverage for src/bartz/stochtree/_stochtree.py: 88%

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. 

24 

25"""Implement class `BARTModel` that mimics the Python package stochtree.""" 

26 

27from collections.abc import Mapping, Sequence (empty)

28from dataclasses import dataclass, field, fields (empty)

29from functools import partial (empty)

30 

31# WORKAROUND(python<3.15): use frozendict instead of MappingProxyType 

32from types import MappingProxyType (empty)

33from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload (empty)

34 

35from jax import lax (empty)

36from jax import numpy as jnp (empty)

37from jax.scipy.special import ndtr, ndtri (empty)

38from jaxtyping import Array, Float, Float32, Key, Real, Shaped (empty)

39 

40from bartz._interface import Bart, DataFrame, PredictKind, Series (empty)

41from bartz._jaxext import jit (empty)

42from bartz.mcmcloop import MainTrace (empty)

43from bartz.mcmcstep import Wishart (empty)

44from bartz.mcmcstep._axes import chain_to_axis, chain_vmap_axes (empty)

45from bartz.mcmcstep._state import ArrayLike, FloatLike (empty)

46from bartz.prepcovars import RangeEvenBinner (empty)

47from bartz.stochtree._preprocess import _PreprocessorBase, make_preprocessor (empty)

48 

49if TYPE_CHECKING: (empty)

50 from _typeshed import DataclassInstance 

51 

52_MAX_DEPTH_LIMIT = 16 (empty)

53 

54 

55@dataclass(frozen=True) (empty)

56class OutcomeModel: (empty)

57 """Outcome model specification, matching `stochtree.OutcomeModel`. 

58 

59 Only ``('continuous', 'identity')`` and ``('binary', 'probit')`` are 

60 supported. 

61 """ 

62 

63 outcome: Literal['continuous', 'binary'] = 'continuous' (empty)

64 """Outcome family.""" (empty)

65 

66 link: Literal['identity', 'probit'] | None = None (empty)

67 """Link function. If `None`, defaults to ``'identity'`` for ``'continuous'`` and ``'probit'`` for ``'binary'``.""" (empty)

68 

69 def __post_init__(self) -> None: (empty)

70 if self.link is None: (empty)

71 default_link = {'continuous': 'identity', 'binary': 'probit'}.get( (empty)

72 self.outcome 

73 ) 

74 object.__setattr__(self, 'link', default_link) (empty)

75 if (self.outcome, self.link) not in ( (empty)

76 ('continuous', 'identity'), 

77 ('binary', 'probit'), 

78 ): 

79 msg = ( (empty)

80 f'unsupported outcome_model (outcome={self.outcome!r}, ' 

81 f"link={self.link!r}); only ('continuous', 'identity') " 

82 "and ('binary', 'probit') are supported." 

83 ) 

84 raise NotImplementedError(msg) (empty)

85 

86 

87class NotSampledError(ValueError, AttributeError): (empty)

88 """Raised when calling a method that requires `BARTModel.sample` to have been called.""" 

89 

90 

91@dataclass(frozen=True, kw_only=True) (empty)

92class GeneralParams: (empty)

93 """Mirror of stochtree's ``general_params`` dict, with the keys bartz handles.""" 

94 

95 standardize: bool = True (empty)

96 """Whether to standardize the outcome before fitting. Ignored for probit binary.""" (empty)

97 

98 sigma2_init: FloatLike | None = None (empty)

99 """Starting value of the global error variance. If `None` (default), uses ``var(resid_train)`` for continuous and ``1.0`` for probit.""" (empty)

100 

101 sigma2_global_shape: FloatLike = 0.0 (empty)

102 """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.""" (empty)

103 

104 sigma2_global_scale: FloatLike = 0.0 (empty)

105 """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.""" (empty)

106 

107 variable_weights: Float[ArrayLike, ' p'] | None = None (empty)

108 """Per-predictor sampling weights. Must be strictly positive; pass a small positive value to suppress a variable.""" (empty)

109 

110 random_seed: int | Key[Array, ''] | None = None (empty)

111 """Seed for the random number generator. Unlike stochtree, the default (empty)

112 `None` is deterministic (equivalent to seed ``0``) rather than drawing a 

113 random seed, so repeated fits reproduce by default.""" 

114 

115 keep_every: int = 1 (empty)

116 """Thinning factor for retained MCMC samples.""" (empty)

117 

118 num_chains: int = 1 (empty)

119 """Number of independent MCMC chains.""" (empty)

120 

121 outcome_model: OutcomeModel = field(default_factory=OutcomeModel) (empty)

122 """Outcome family and link specification. Defaults to continuous with (empty)

123 identity link.""" 

124 

125 

126@dataclass(frozen=True, kw_only=True) (empty)

127class MeanForestParams: (empty)

128 """Mirror of stochtree's ``mean_forest_params`` dict, restricted to the keys bartz handles.""" 

129 

130 num_trees: int = 200 (empty)

131 """Number of trees in the conditional mean ensemble.""" (empty)

132 

133 alpha: FloatLike = 0.95 (empty)

134 """Tree split prior base.""" (empty)

135 

136 beta: FloatLike = 2.0 (empty)

137 """Tree split prior decay.""" (empty)

138 

139 min_samples_leaf: int = 5 (empty)

140 """Minimum number of training samples at a leaf.""" (empty)

141 

142 max_depth: int = 10 (empty)

143 """Maximum tree depth. Must be a non-negative integer at most ``16``.""" (empty)

144 

145 sample_sigma2_leaf: bool = True (empty)

146 """Whether to sample the leaf variance with an inverse-gamma prior, or hold it fixed at `sigma2_leaf_init`.""" (empty)

147 

148 sigma2_leaf_init: FloatLike | None = None (empty)

149 """Initial leaf variance. If `None`, matches stochtree's defaults of ``var(resid_train) / num_trees`` for continuous and ``2 / num_trees`` for probit.""" (empty)

150 

151 sigma2_leaf_shape: FloatLike = 3.0 (empty)

152 """Twice the shape parameter of the inverse-gamma prior on the leaf variance, i.e., the prior is ``IG(sigma2_leaf_shape / 2, sigma2_leaf_scale / 2)``.""" (empty)

153 

154 sigma2_leaf_scale: FloatLike | None = None (empty)

155 """Twice the scale parameter of the inverse-gamma prior on the leaf variance. If `None`, defaults to ``var(resid_train) / num_trees`` for continuous and ``1 / num_trees`` for probit.""" (empty)

156 

157 def __post_init__(self) -> None: (empty)

158 if self.max_depth < 0: 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true(empty)

159 msg = ( 

160 f'max_depth={self.max_depth} is not supported; bartz stores trees' 

161 ' as heap arrays of size 2**max_depth, so the stochtree' 

162 ' convention max_depth=-1 (unbounded) is rejected. Pass a' 

163 f' non-negative integer at most {_MAX_DEPTH_LIMIT}.' 

164 ) 

165 raise NotImplementedError(msg) 

166 if self.max_depth > _MAX_DEPTH_LIMIT: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true(empty)

167 msg = ( 

168 f'max_depth={self.max_depth} exceeds {_MAX_DEPTH_LIMIT}; bartz' 

169 ' stores trees as heap arrays of size 2**max_depth, so memory' 

170 ' grows exponentially with depth.' 

171 ) 

172 raise ValueError(msg) 

173 

174 

175if TYPE_CHECKING: (empty)

176 # static only, beartype does not support type[DataclassInstance] 

177 T = TypeVar('T', bound='DataclassInstance') 

178else: 

179 T = TypeVar('T') (empty)

180 

181 

182def build_dataclass(cls: type[T], params: Mapping[str, Any] | None, name: str) -> T: (empty)

183 """Convert a user-supplied dict to a dataclass, with friendly errors.""" 

184 if params is None: (empty)

185 params = {} (empty)

186 allowed = {f.name for f in fields(cls)} (empty)

187 extra = set(params) - allowed (empty)

188 if extra: (empty)

189 msg = ( (empty)

190 f'{name} contains unsupported key(s) {sorted(extra)}; valid keys' 

191 f' are {sorted(allowed)}' 

192 ) 

193 raise ValueError(msg) (empty)

194 return cls(**params) (empty)

195 

196 

197class BARTModel: (empty)

198 R""" 

199 BART model with a `stochtree`-compatible interface, powered by bartz. 

200 

201 This class mimics `stochtree.BARTModel` so that bartz can be used as a 

202 drop-in reference implementation for testing. The intersection of features 

203 is targeted: continuous regression (Gaussian outcome, identity link) and 

204 binary classification (probit link) on tabular covariates. 

205 

206 Use the same idiomatic pattern as `stochtree.BARTModel`:: 

207 

208 m = BARTModel() 

209 m.sample(X_train=X, y_train=y, X_test=X_test, num_gfr=0, num_mcmc=200) 

210 yhat = m.predict(X_new, terms='y_hat', type='mean') 

211 

212 See `GeneralParams` and `MeanForestParams` for the supported keys in the 

213 ``general_params`` and ``mean_forest_params`` dicts. 

214 

215 Notes 

216 ----- 

217 Differences from `stochtree`, by design: 

218 

219 - ``num_gfr`` has no default and must be set explicitly to ``0``. 

220 - ``mean_forest_params['max_depth']`` must be a non-negative integer at 

221 most ``16``; stochtree's ``-1`` (unbounded depth) sentinel is not 

222 accepted. 

223 - The deprecated ``general_params['probit_outcome_model']`` flag is not 

224 accepted; pass ``outcome_model=OutcomeModel('binary', 'probit')`` 

225 instead. 

226 - ``general_params['cutpoint_grid_size']`` is not accepted; bartz uses a 

227 fixed grid of 256 evenly-spaced bins per predictor. stochtree only 

228 uses this parameter for the GFR sampler, which bartz does not support. 

229 - Leaf-basis regression, random effects, heteroskedastic variance 

230 forests, and warm-starting from a previous model are not supported. 

231 - bartz uses single-precision floats, so outputs differ from stochtree 

232 at the float32 precision level. 

233 - ``general_params['random_seed']`` defaults to deterministic behavior 

234 (seed ``0``) when unset, whereas stochtree draws a random seed. This is 

235 intentional, to make repeated fits reproducible by default. 

236 

237 References 

238 ---------- 

239 Herren, A., Hahn, P. R., Murray, J., Carvalho, C. (2026). "StochTree: 

240 BART-based modeling in R and Python". arXiv:2512.12051. 

241 """ 

242 

243 # public, set by sample() 

244 sampled: bool 

245 """Whether `sample` has been called.""" (empty)

246 

247 standardize: bool 

248 """Whether the outcome was standardized before fitting.""" (empty)

249 

250 sample_sigma2_global: bool 

251 """Whether the global error variance is sampled (always ``True``).""" (empty)

252 

253 probit_outcome_model: bool 

254 """Whether the model uses a binary outcome with probit link.""" (empty)

255 

256 outcome_model: OutcomeModel 

257 """Outcome family and link specification used during fitting.""" (empty)

258 

259 num_gfr: int 

260 """Number of grow-from-root iterations (always ``0``).""" (empty)

261 

262 num_burnin: int 

263 """Number of MCMC burn-in iterations.""" (empty)

264 

265 num_mcmc: int 

266 """Number of retained MCMC iterations per chain.""" (empty)

267 

268 num_chains: int 

269 """Number of independent MCMC chains.""" (empty)

270 

271 num_samples: int 

272 """Total number of retained posterior samples (``num_mcmc * num_chains``).""" (empty)

273 

274 sigma2_init: FloatLike 

275 """Starting value of the global error variance actually used to seed the chain.""" (empty)

276 

277 y_bar: Float32[Array, ''] 

278 """Mean used to standardize the outcome (``0`` if not standardized).""" (empty)

279 

280 y_std: Float32[Array, ''] 

281 """Standard deviation used to standardize the outcome (``1`` if not standardized).""" (empty)

282 

283 has_rfx: bool 

284 """Whether the model includes random effects (always ``False``).""" (empty)

285 

286 include_mean_forest: bool 

287 """Whether the model includes a conditional mean forest (always ``True``).""" (empty)

288 

289 include_variance_forest: bool 

290 """Whether the model includes a variance forest (always ``False``).""" (empty)

291 

292 y_hat_train: Float32[Array, 'n num_samples'] 

293 """Posterior predictions at the training covariates, in the original outcome scale.""" (empty)

294 

295 global_var_samples: Float32[Array, ' num_samples'] 

296 """Posterior samples of the global error variance. For probit binary regression, an array of ones.""" (empty)

297 

298 sample_sigma2_leaf: bool 

299 """Whether the leaf variance is sampled.""" (empty)

300 

301 leaf_scale_samples: Float32[Array, ' num_samples'] | None 

302 """Posterior samples of the leaf variance, in the standardized outcome scale, or `None` if not sampled.""" (empty)

303 

304 y_hat_test: Float32[Array, 'm num_samples'] | None 

305 """Posterior predictions at `X_test` if it was supplied to `sample`, else `None`.""" (empty)

306 

307 _bart: Bart 

308 _preprocessor: _PreprocessorBase | None 

309 

310 def __init__(self) -> None: (empty)

311 self.sampled = False (empty)

312 self._preprocessor = None (empty)

313 

314 def is_sampled(self) -> bool: (empty)

315 """Return whether `sample` has been called.""" 

316 return self.sampled (empty)

317 

318 def _prepare_training_inputs( (empty)

319 self, 

320 X_train: Real[ArrayLike, 'n p'] | DataFrame, 

321 y_train: Real[ArrayLike, ' n'] | Series, 

322 gp: GeneralParams, 

323 ) -> tuple[Real[Array, 'n p'], Real[Array, ' n'], Float32[Array, ' p'] | None]: 

324 """Coerce inputs and build variable weights, fitting the DataFrame preprocessor if any.""" 

325 y_train_arr = _coerce_response(y_train, name='y_train') (empty)

326 

327 self._preprocessor = make_preprocessor(X_train) (empty)

328 if self._preprocessor is None: (empty)

329 X_train_arr = check_X(X_train, name='X_train') (empty)

330 _, p = X_train_arr.shape (empty)

331 varprob = check_variable_weights(gp.variable_weights, p) (empty)

332 else: 

333 # The preprocessor decides the default weights: uniform over the 

334 # *original* columns split across each one-hot expansion (so every 

335 # original variable keeps an equal splitting budget), or `None` when 

336 # nothing expands (deferring to bartz's native uniform fast-path). 

337 weights_np = self._preprocessor.fit( (empty)

338 X_train, variable_weights=gp.variable_weights 

339 ) 

340 X_train_np = self._preprocessor.transform(X_train) (empty)

341 if X_train_np.shape[1] == 0: 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true(empty)

342 msg = 'X_train has no usable columns after preprocessing' 

343 raise ValueError(msg) 

344 X_train_arr = jnp.asarray(X_train_np) (empty)

345 varprob = None if weights_np is None else jnp.asarray(weights_np) (empty)

346 

347 n, _ = X_train_arr.shape (empty)

348 if y_train_arr.shape[0] != n: 348 ↛ 349line 348 didn't jump to line 349 because the condition on line 348 was never true(empty)

349 msg = ( 

350 f'X_train and y_train length mismatch: X_train has {n} rows,' 

351 f' y_train has {y_train_arr.shape[0]} entries' 

352 ) 

353 raise ValueError(msg) 

354 return X_train_arr, y_train_arr, varprob (empty)

355 

356 def sample( (empty)

357 self, 

358 X_train: Real[ArrayLike, 'n p'] | DataFrame, 

359 y_train: Real[ArrayLike, ' n'] | Series, 

360 X_test: Real[ArrayLike, 'm p'] | DataFrame | None = None, 

361 observation_weights: Float[ArrayLike, ' n'] | Series | None = None, 

362 *, 

363 num_gfr: int, 

364 num_burnin: int = 0, 

365 num_mcmc: int = 100, 

366 general_params: Mapping[str, Any] | None = None, 

367 mean_forest_params: Mapping[str, Any] | None = None, 

368 bart_kwargs: Mapping[str, Any] = MappingProxyType({}), 

369 ) -> None: 

370 """Fit the model. 

371 

372 The signature mirrors `stochtree.BARTModel.sample`, restricted to the 

373 keyword arguments bartz supports. 

374 

375 Parameters 

376 ---------- 

377 X_train 

378 Training covariates with shape ``(n, p)``. 

379 y_train 

380 Training outcomes of length ``n``. 

381 X_test 

382 Optional test covariates; if given, predictions are cached on 

383 them in `y_hat_test`. 

384 observation_weights 

385 Optional positive per-observation weights scaling the residual 

386 variance (``y_i | - ~ N(mu(X_i), sigma^2 / w_i)``). Not supported 

387 with a probit outcome model. 

388 num_gfr 

389 Number of grow-from-root iterations. Must be ``0``. 

390 num_burnin 

391 Number of MCMC burn-in iterations. 

392 num_mcmc 

393 Number of retained MCMC iterations per chain. 

394 general_params 

395 Optional override for the keys of `GeneralParams`. 

396 mean_forest_params 

397 Optional override for the keys of `MeanForestParams`. 

398 bart_kwargs 

399 Additional arguments forwarded to `bartz.Bart`. Use this to set 

400 ``devices`` and ``rm_const=False`` when wrapping `sample` in 

401 `jax.jit`. 

402 

403 Raises 

404 ------ 

405 NotImplementedError 

406 If ``num_gfr`` is non-zero. 

407 ValueError 

408 If `observation_weights` are set with a probit outcome model. 

409 """ 

410 if num_gfr != 0: (empty)

411 msg = ( (empty)

412 'num_gfr must be 0; the grow-from-root sampler is not available' 

413 ' in bartz.' 

414 ) 

415 raise NotImplementedError(msg) (empty)

416 

417 gp = build_dataclass(GeneralParams, general_params, 'general_params') (empty)

418 mfp = build_dataclass( (empty)

419 MeanForestParams, mean_forest_params, 'mean_forest_params' 

420 ) 

421 

422 is_probit = gp.outcome_model.outcome == 'binary' (empty)

423 

424 # stochtree does not support heteroskedastic probit 

425 if is_probit and observation_weights is not None: (empty)

426 msg = ( (empty)

427 'observation_weights are not supported with a probit outcome' 

428 ' model; stochtree does not support heteroskedastic probit.' 

429 ) 

430 raise ValueError(msg) (empty)

431 

432 X_train_arr, y_train_arr, variable_weights = self._prepare_training_inputs( (empty)

433 X_train, y_train, gp 

434 ) 

435 

436 y_bar, y_std, y_for_bartz = standardize_y( (empty)

437 y_train_arr, is_probit, gp.standardize 

438 ) 

439 

440 bart_num_chains = None if gp.num_chains == 1 else gp.num_chains (empty)

441 

442 # variance of the standardized residual, matching stochtree 

443 # (np.var(resid_train) with ddof=0). For standardize=True it is exactly 

444 # 1.0; we hardcode that so the value stays trace-time concrete. 

445 if is_probit: (empty)

446 var_resid_train: FloatLike = 1.0 # bartz ignores σ² for binary (empty)

447 elif gp.standardize: 447 ↛ 450line 447 didn't jump to line 450 because the condition on line 447 was always true(empty)

448 var_resid_train = 1.0 (empty)

449 else: 

450 var_resid_train = jnp.var(y_for_bartz) 

451 

452 # leaf-prior: bartz uses sigma_mu = tau_num / (k * sqrt(num_trees)); 

453 # stochtree's sigma2_leaf is the leaf-variance prior. Hold k=2 and solve 

454 # for tau_num so that the two parameterizations agree. 

455 bartz_k = 2.0 (empty)

456 sigma2_leaf_init = resolve_sigma2_leaf_init( (empty)

457 mfp.sigma2_leaf_init, mfp.num_trees, is_probit, var_resid_train 

458 ) 

459 tau_num_arg = bartz_k * jnp.sqrt(mfp.num_trees * sigma2_leaf_init) (empty)

460 

461 if mfp.sample_sigma2_leaf: (empty)

462 leaf_init_kw: dict = { (empty)

463 'leaf_prior_cov_inv': resolve_leaf_variance_prior( 

464 mfp.sigma2_leaf_shape, 

465 mfp.sigma2_leaf_scale, 

466 sigma2_leaf_init, 

467 num_trees=mfp.num_trees, 

468 is_probit=is_probit, 

469 var_resid_train=var_resid_train, 

470 ) 

471 } 

472 else: 

473 leaf_init_kw = {} (empty)

474 

475 if is_probit: (empty)

476 # stochtree pins σ²=1 for probit; bartz binary branch ignores the 

477 # variance prior, so we leave the scale/init at their 'auto' 

478 # defaults (bartz rejects explicit values for binary outcomes). 

479 sigma_df_arg: FloatLike = 3.0 (empty)

480 sigma_scale_arg: FloatLike | Literal['auto'] = 'auto' (empty)

481 sigma_init_arg: FloatLike | Literal['auto'] = 'auto' (empty)

482 sigma2_init_stored: FloatLike = 1.0 (empty)

483 else: 

484 sigma_df_arg, sigma_scale_arg, sigma_init_arg, sigma2_init_stored = ( (empty)

485 resolve_variance_prior( 

486 gp.sigma2_global_shape, 

487 gp.sigma2_global_scale, 

488 gp.sigma2_init, 

489 var_resid_train, 

490 ) 

491 ) 

492 

493 binner = partial(RangeEvenBinner, max_bins=256) (empty)

494 

495 seed = 0 if gp.random_seed is None else gp.random_seed (empty)

496 

497 kwargs: dict = dict( (empty)

498 x_train=X_train_arr.T, 

499 y_train=y_for_bartz, 

500 outcome_type='binary' if is_probit else 'continuous', 

501 binner=binner, 

502 varprob=variable_weights, 

503 sigma_df=sigma_df_arg, 

504 sigma_scale=sigma_scale_arg, 

505 sigma_init=sigma_init_arg, 

506 k=bartz_k, 

507 power=mfp.beta, 

508 base=mfp.alpha, 

509 tau_num=tau_num_arg, 

510 error_scale=observation_weights, 

511 num_trees=mfp.num_trees, 

512 n_save=num_mcmc, 

513 n_burn=num_burnin, 

514 n_skip=gp.keep_every, 

515 printevery=None, 

516 num_chains=bart_num_chains, 

517 seed=seed, 

518 maxdepth=mfp.max_depth + 1, 

519 precompute_predict_train=True, 

520 ) 

521 kwargs.update(bart_kwargs) (empty)

522 # min_samples_leaf maps to bartz's acceptance-time veto. 

523 # min_points_per_decision_node = 2 * min_samples_leaf makes the 

524 # proposal skip leaves too small to split, without changing the target 

525 # posterior, which thus matches stochtree. The leaf variance prior goes 

526 # through init_kw because Bart's sigma_mu_df ties the prior scale to 

527 # the initial value, while stochtree sets them independently. 

528 # User-supplied init_kw values win on conflicts. 

529 kwargs = dict( (empty)

530 kwargs, 

531 init_kw=dict( 

532 { 

533 'min_points_per_leaf': mfp.min_samples_leaf, 

534 'min_points_per_decision_node': 2 * mfp.min_samples_leaf, 

535 **leaf_init_kw, 

536 }, 

537 **kwargs.get('init_kw', {}), 

538 ), 

539 ) 

540 self._bart = Bart(**kwargs) (empty)

541 self._finalize_sample( (empty)

542 outcome_model=gp.outcome_model, 

543 num_burnin=num_burnin, 

544 num_mcmc=num_mcmc, 

545 num_chains=gp.num_chains, 

546 sample_sigma2_leaf=mfp.sample_sigma2_leaf, 

547 sigma2_init=sigma2_init_stored, 

548 y_bar=y_bar, 

549 y_std=y_std, 

550 standardize=gp.standardize, 

551 X_test=X_test, 

552 ) 

553 

554 def _finalize_sample( (empty)

555 self, 

556 *, 

557 outcome_model: OutcomeModel, 

558 num_burnin: int, 

559 num_mcmc: int, 

560 num_chains: int, 

561 sample_sigma2_leaf: bool, 

562 sigma2_init: FloatLike, 

563 y_bar: Float32[Array, ''], 

564 y_std: Float32[Array, ''], 

565 standardize: bool, 

566 X_test: Real[ArrayLike, 'm p'] | DataFrame | None, 

567 ) -> None: 

568 """Populate the public attributes after `_bart` has been constructed.""" 

569 is_probit = outcome_model.outcome == 'binary' (empty)

570 self.sampled = True (empty)

571 self.standardize = standardize (empty)

572 self.sample_sigma2_global = True (empty)

573 self.sample_sigma2_leaf = sample_sigma2_leaf (empty)

574 self.probit_outcome_model = is_probit (empty)

575 self.outcome_model = outcome_model (empty)

576 self.num_gfr = 0 (empty)

577 self.num_burnin = num_burnin (empty)

578 self.num_mcmc = num_mcmc (empty)

579 self.num_chains = num_chains (empty)

580 self.num_samples = num_mcmc * num_chains (empty)

581 self.sigma2_init = sigma2_init (empty)

582 self.y_bar = y_bar (empty)

583 self.y_std = y_std (empty)

584 self.has_rfx = False (empty)

585 self.include_mean_forest = True (empty)

586 self.include_variance_forest = False (empty)

587 

588 # cached outputs in stochtree's (n, num_samples) layout, original scale 

589 self.y_hat_train = self._predict_y_hat_internal('train') (empty)

590 if X_test is not None: (empty)

591 self.y_hat_test = self._predict_y_hat_internal(self._prepare_x(X_test).T) (empty)

592 else: 

593 self.y_hat_test = None (empty)

594 

595 if is_probit: (empty)

596 self.global_var_samples = jnp.ones((self.num_samples,)) (empty)

597 else: 

598 sigma = self._bart.get_error_sdev() (empty)

599 self.global_var_samples = (sigma * y_std) ** 2 (empty)

600 

601 if sample_sigma2_leaf: (empty)

602 self.leaf_scale_samples = get_leaf_scale_samples( (empty)

603 self._bart._main_trace # noqa: SLF001 

604 ) 

605 else: 

606 self.leaf_scale_samples = None (empty)

607 

608 @overload (empty)

609 def predict( (empty)

610 self, 

611 X: Real[ArrayLike, 'm p'] | DataFrame, (empty)

612 *, 

613 type: Literal['posterior', 'mean'] = 'posterior', (empty)

614 terms: Literal['y_hat', 'mean_forest'], (empty)

615 scale: Literal['linear', 'probability', 'class'] = 'linear', (empty)

616 ) -> Shaped[Array, 'm num_samples'] | Shaped[Array, ' m']: ... (empty)

617 

618 @overload (empty)

619 def predict( (empty)

620 self, 

621 X: Real[ArrayLike, 'm p'] | DataFrame, (empty)

622 *, 

623 type: Literal['posterior', 'mean'] = 'posterior', (empty)

624 terms: Literal['all'] = 'all', (empty)

625 scale: Literal['linear', 'probability', 'class'] = 'linear', (empty)

626 ) -> dict[str, Shaped[Array, 'm num_samples']] | dict[str, Shaped[Array, ' m']]: ... (empty)

627 

628 @overload (empty)

629 def predict( (empty)

630 self, 

631 X: Real[ArrayLike, 'm p'] | DataFrame, (empty)

632 *, 

633 type: Literal['posterior', 'mean'] = 'posterior', (empty)

634 terms: Sequence[Literal['y_hat', 'mean_forest', 'all']], (empty)

635 scale: Literal['linear', 'probability', 'class'] = 'linear', (empty)

636 ) -> ( 

637 Shaped[Array, 'm num_samples'] (empty)

638 | Shaped[Array, ' m'] (empty)

639 | dict[str, Shaped[Array, 'm num_samples']] (empty)

640 | dict[str, Shaped[Array, ' m']] (empty)

641 ): ... 

642 

643 def predict( (empty)

644 self, 

645 X: Real[ArrayLike, 'm p'] | DataFrame, 

646 *, 

647 type: Literal['posterior', 'mean'] = 'posterior', # noqa: A002 

648 terms: Literal['y_hat', 'mean_forest', 'all'] 

649 | Sequence[Literal['y_hat', 'mean_forest', 'all']] = 'all', 

650 scale: Literal['linear', 'probability', 'class'] = 'linear', 

651 ) -> ( 

652 Shaped[Array, 'm num_samples'] 

653 | Shaped[Array, ' m'] 

654 | dict[str, Shaped[Array, 'm num_samples']] 

655 | dict[str, Shaped[Array, ' m']] 

656 ): 

657 """Predict at new covariates. 

658 

659 Parameters 

660 ---------- 

661 X 

662 New covariates with shape ``(m, p)``. 

663 type 

664 ``'posterior'`` returns one prediction per posterior sample, with 

665 shape ``(m, num_samples)``. ``'mean'`` averages the posterior 

666 samples, returning a vector of shape ``(m,)``. 

667 terms 

668 One of ``'y_hat'``, ``'mean_forest'``, ``'all'``, or a list. Since 

669 random effects and a variance forest are not supported, ``'y_hat'`` 

670 and ``'mean_forest'`` produce the same result. 

671 scale 

672 For probit binary regression: ``'linear'`` returns the eta values, 

673 ``'probability'`` returns ``Phi(eta)``, ``'class'`` returns 0 / 1. 

674 Only ``'linear'`` is valid for continuous outcomes. 

675 

676 Returns 

677 ------- 

678 Either a single jax array (for a single requested term) or a dict keyed by term name. 

679 

680 Raises 

681 ------ 

682 NotSampledError 

683 If `sample` has not been called yet. 

684 """ 

685 if not self.sampled: (empty)

686 msg = ( (empty)

687 "This BARTModel instance is not fitted yet. Call 'sample' before" 

688 ' using this model.' 

689 ) 

690 raise NotSampledError(msg) (empty)

691 terms_tuple = check_predict_args(type, scale, terms, self.probit_outcome_model) (empty)

692 

693 pred = self._predict_y_hat_internal(self._prepare_x(X).T) (empty)

694 

695 if self.probit_outcome_model and scale in ('probability', 'class'): (empty)

696 prob = ndtr(pred) (empty)

697 pred_out = jnp.where(prob < 0.5, 0, 1) if scale == 'class' else prob (empty)

698 else: 

699 pred_out = pred (empty)

700 

701 if type == 'mean': (empty)

702 pred_out = jnp.mean(pred_out, axis=1) (empty)

703 

704 wants_y_hat = ('y_hat' in terms_tuple) or ('all' in terms_tuple) (empty)

705 wants_mean_forest = ('mean_forest' in terms_tuple) or ('all' in terms_tuple) (empty)

706 single = sum([wants_y_hat, wants_mean_forest]) == 1 (empty)

707 if single: (empty)

708 return pred_out (empty)

709 result: dict[str, Shaped[Array, '...']] = {} (empty)

710 if wants_y_hat: 710 ↛ 712line 710 didn't jump to line 712 because the condition on line 710 was always true(empty)

711 result['y_hat'] = pred_out (empty)

712 if wants_mean_forest: 712 ↛ 714line 712 didn't jump to line 714 because the condition on line 712 was always true(empty)

713 result['mean_forest_predictions'] = pred_out (empty)

714 return result (empty)

715 

716 def _prepare_x(self, X: Real[ArrayLike, 'm p'] | DataFrame) -> Real[Array, 'm p']: (empty)

717 """Convert covariates to a 2-D jax array, replaying the fitted preprocessor if any.""" 

718 if self._preprocessor is None: (empty)

719 return check_X(X) (empty)

720 if make_preprocessor(X) is None: (empty)

721 msg = ( (empty)

722 'this model was fit on a DataFrame, so prediction covariates must' 

723 ' also be a pandas/polars DataFrame with the same columns; got a' 

724 ' non-DataFrame. Passing a raw array would bypass the fitted' 

725 ' preprocessing (e.g. one-hot encoding) and silently misalign the' 

726 ' features.' 

727 ) 

728 raise TypeError(msg) (empty)

729 return jnp.asarray(self._preprocessor.transform(X)) (empty)

730 

731 def _predict_y_hat_internal( (empty)

732 self, x: Real[ArrayLike, 'p m'] | Literal['train'] 

733 ) -> Float32[Array, 'm num_samples']: 

734 """Return predictions on the original outcome scale, layout ``(m, num_samples)``.""" 

735 latent = self._bart.predict(x, kind=PredictKind.latent_samples) (empty)

736 if self.probit_outcome_model: (empty)

737 # bartz integrates the binary offset into latent; result already on probit scale. 

738 return latent.T (empty)

739 if self.standardize: 739 ↛ 741line 739 didn't jump to line 741 because the condition on line 739 was always true(empty)

740 return (latent * self.y_std + self.y_bar).T (empty)

741 return latent.T 

742 

743 

744def standardize_y( (empty)

745 y_train: Real[ArrayLike, ' n'], is_probit: bool, standardize: bool 

746) -> tuple[Float32[Array, ''], Float32[Array, ''], Float32[Array, ' n']]: 

747 """Return ``(y_bar, y_std, y_for_bartz)`` matching stochtree's standardization.""" 

748 y = jnp.asarray(y_train, jnp.float32) (empty)

749 if is_probit: (empty)

750 return ndtri(y.mean()), jnp.float32(1.0), (y != 0).astype(jnp.float32) (empty)

751 if standardize: 751 ↛ 756line 751 didn't jump to line 756 because the condition on line 751 was always true(empty)

752 y_bar = y.mean() (empty)

753 y_std_val = y.std() (empty)

754 y_std = jnp.where(y_std_val > 0, y_std_val, 1.0) (empty)

755 return y_bar, y_std, (y - y_bar) / y_std (empty)

756 return jnp.float32(0.0), jnp.float32(1.0), y 

757 

758 

759def resolve_sigma2_leaf_init( (empty)

760 sigma2_leaf_init: FloatLike | None, 

761 num_trees: int, 

762 is_probit: bool, 

763 var_resid_train: FloatLike, 

764) -> FloatLike: 

765 """Default `sigma2_leaf_init` per stochtree: probit→2/num_trees, continuous→var(resid)/num_trees.""" 

766 if sigma2_leaf_init is not None: (empty)

767 return sigma2_leaf_init (empty)

768 if is_probit: (empty)

769 return 2.0 / num_trees (empty)

770 return var_resid_train / num_trees (empty)

771 

772 

773def resolve_leaf_variance_prior( (empty)

774 shape: FloatLike, 

775 scale: FloatLike | None, 

776 sigma2_leaf_init: FloatLike, 

777 *, 

778 num_trees: int, 

779 is_probit: bool, 

780 var_resid_train: FloatLike, 

781) -> Wishart: 

782 """Translate stochtree's leaf variance prior to a `Wishart` on the leaf precision. 

783 

784 Despite the documented ``IG(shape, scale)``, stochtree's sampler uses the 

785 posterior ``IG(shape / 2 + L / 2, scale / 2 + S / 2)``, i.e., the prior is 

786 ``IG(shape / 2, scale / 2)``, which is `Wishart` with ``nu = shape`` and 

787 ``rate = scale``. 

788 """ 

789 if scale is None: (empty)

790 scale = 1.0 / num_trees if is_probit else var_resid_train / num_trees (empty)

791 return Wishart(nu=shape, rate=scale, value=jnp.reciprocal(sigma2_leaf_init)) (empty)

792 

793 

794@jit (empty)

795def get_leaf_scale_samples(trace: MainTrace) -> Float32[Array, ' num_samples']: (empty)

796 """Return the leaf variance samples, chains concatenated.""" 

797 prec = trace.leaf_prior_cov_inv (empty)

798 if trace.has_chains: (empty)

799 prec = chain_to_axis(prec, chain_vmap_axes(trace).leaf_prior_cov_inv) (empty)

800 prec = lax.collapse(prec, 0, 2) (empty)

801 return jnp.reciprocal(prec) (empty)

802 

803 

804def resolve_variance_prior( (empty)

805 shape: FloatLike, 

806 scale: FloatLike, 

807 sigma2_init: FloatLike | None, 

808 var_resid_train: FloatLike, 

809) -> tuple[Float32[Array, ''], Float32[Array, ''], Float32[Array, ''], FloatLike]: 

810 """Translate stochtree's IG(shape, scale) prior to bartz's error variance prior. 

811 

812 The IG(shape, scale) prior on σ² is the scaled-inverse-χ² with 

813 ``sigma_df = 2*shape`` and prior harmonic mean ``square(sigma_scale) = 

814 scale/shape``; the chain starts at `sigma2_init` (default 

815 ``var(resid_train)``), decoupled from the prior. The mapping is branchless so 

816 `shape` / `scale` may be traced; the unrepresentable IG(0, scale>0) (positive 

817 rate, zero df) yields a NaN that surfaces downstream rather than an error. 

818 

819 Parameters 

820 ---------- 

821 shape 

822 Stochtree's ``sigma2_global_shape``. 

823 scale 

824 Stochtree's ``sigma2_global_scale``. 

825 sigma2_init 

826 Stochtree's ``sigma2_init``. If `None`, defaults to `var_resid_train`. 

827 var_resid_train 

828 Variance of the residual, the default chain start for σ². 

829 

830 Returns 

831 ------- 

832 sigma_df : Float32[Array, ''] 

833 Degrees of freedom of bartz's error variance prior. 

834 sigma_scale : Float32[Array, ''] 

835 Scale of bartz's prior (sqrt of the prior harmonic mean of the variance). 

836 sigma_init : Float32[Array, ''] 

837 Initial error standard deviation seeding the chain. 

838 sigma2_init_stored : FloatLike 

839 The chain starting value of σ², suitable for ``BARTModel.sigma2_init``. 

840 """ 

841 shape = jnp.asarray(shape, jnp.float32) (empty)

842 scale = jnp.asarray(scale, jnp.float32) (empty)

843 sigma2_start = sigma2_init if sigma2_init is not None else var_resid_train (empty)

844 sigma_init = jnp.sqrt(jnp.asarray(sigma2_start, jnp.float32)) (empty)

845 # IG(shape, scale) <=> scaled-inv-chi2(df=2*shape, harmonic mean=scale/shape). 

846 # The `scale > 0` guard keeps IG(0, 0) at harmonic mean 0 (avoiding 0/0) while 

847 # letting IG(0, scale>0) overflow to inf -> NaN rate, flagging it as invalid. 

848 harmonic_mean = jnp.where(scale > 0, scale / shape, 0.0) (empty)

849 return 2.0 * shape, jnp.sqrt(harmonic_mean), sigma_init, sigma2_start (empty)

850 

851 

852def check_variable_weights( (empty)

853 variable_weights: Float[ArrayLike, ' p'] | None, p: int 

854) -> Float32[Array, ' p'] | None: 

855 """Validate `variable_weights`, returning the jax array (or None).""" 

856 if variable_weights is None: (empty)

857 return None (empty)

858 arr = jnp.asarray(variable_weights, jnp.float32) (empty)

859 if arr.shape != (p,): 859 ↛ 860line 859 didn't jump to line 860 because the condition on line 859 was never true(empty)

860 msg = f'variable_weights must have shape (p,)=({p},), got {arr.shape}' 

861 raise ValueError(msg) 

862 return arr (empty)