Coverage for src/bartz/BART/

1# bartz/src/bartz/BART/_gbart.py

5# This file is part of bartz.

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 classes `mc_gbart` and `gbart` that mimic the R BART3 package."""

27from collections.abc import Mapping (empty)

28from functools import cached_property, partial (empty)

29from types import MappingProxyType (empty)

30from typing import Any, Literal (empty)

32import jax.numpy as jnp (empty)

33from equinox import Module, field (empty)

34from jax.scipy.special import ndtr (empty)

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

37from bartz._interface import ( (empty)

38 ArrayLike,

39 Bart,

40 DataFrame,

41 FloatLike,

42 PredictKind,

43 Series,

44 SparseConfig,

45 _process_predictor_input,

46 _process_response_input,

47)

48from bartz._jaxext.scipy.stats import invgamma (empty)

49from bartz.mcmcloop import BurninTrace, MainTrace (empty)

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

51from bartz.mcmcstep._state import State (empty)

52from bartz.prepcovars import GivenSplitsBinner, RangeEvenBinner, UniqueQuantileBinner (empty)

53from bartz.prepcovars._prepcovars import _sigma2_from_ols (empty)

56class mc_gbart(Module): (empty)

57 R"""

58 Nonparametric regression with Bayesian Additive Regression Trees (BART).

60 Regress `y_train` on `x_train` with a latent mean function represented as

61 a sum of decision trees [2]_. The inference is carried out by sampling the

62 posterior distribution of the tree ensemble with an MCMC.

64 Parameters

65 ----------

66 x_train

67 The training predictors.

68 y_train

69 The training responses.

70 x_test

71 The test predictors.

72 type

73 The type of regression. 'wbart' for continuous regression, 'pbart' for

74 binary regression with probit link.

75 sparse

76 Whether to activate variable selection on the predictors as done in

77 [1]_.

78 theta

79 a

80 b

81 rho

82 Hyperparameters of the sparsity prior used for variable selection.

84 The prior distribution on the choice of predictor for each decision rule

85 is

87 .. math::

88 (s_1, \ldots, s_p) \sim

89 \operatorname{Dirichlet}(\mathtt{theta}/p, \ldots, \mathtt{theta}/p).

91 If `theta` is not specified, it's a priori distributed according to

93 .. math::

94 \frac{\mathtt{theta}}{\mathtt{theta} + \mathtt{rho}} \sim

95 \operatorname{Beta}(\mathtt{a}, \mathtt{b}).

97 If not specified, `rho` is set to the number of predictors p. To tune

98 the prior, consider setting a lower `rho` to prefer more sparsity.

99 If setting `theta` directly, it should be in the ballpark of p or lower

100 as well.

101 augment

102 Whether to account exactly for the decision rules forbidden by the

103 ancestors of each node when updating the variable selection

104 probabilities, using data augmentation. Only relevant if ``sparse=True``.

105 Like the ``augment`` option of R BART3, but sampling the exact full

106 conditional rather than substituting expected counts.

107 varprob

108 The probability distribution over the `p` predictors for choosing a

109 predictor to split on in a decision node a priori. Must be > 0. It does

110 not need to be normalized to sum to 1. If not specified, use a uniform

111 distribution. If ``sparse=True``, this is used as initial value for the

112 MCMC.

113 xinfo

114 A matrix with the cutpoins to use to bin each predictor. If not

115 specified, it is generated automatically according to `usequants` and

116 `numcut`.

117

118 Each row shall contain a sorted list of cutpoints for a predictor. If

119 there are less cutpoints than the number of columns in the matrix,

120 fill the remaining cells with NaN.

121

122 `xinfo` shall be a matrix even if `x_train` is a dataframe.

123 usequants

124 Whether to use predictors quantiles instead of a uniform grid to bin

125 predictors. Ignored if `xinfo` is specified.

126 rm_const

127 How to treat predictors with no associated decision rules (i.e., there

128 are no available cutpoints for that predictor). If `True` (default),

129 they are ignored. If `False`, an error is raised if there are any.

130 sigest

131 An estimate of the residual standard deviation on `y_train`, used to set

132 `lambda_`. If not specified, it is estimated by linear regression (with

133 intercept, and without taking into account `w`). Ignored if `lambda_` is

134 specified.

135 sigdf

136 The degrees of freedom of the scaled inverse-chisquared prior on the

137 noise variance.

138 sigquant

139 The quantile of the prior on the noise variance that shall match

140 `sigest` to set the scale of the prior. Ignored if `lambda_` is specified.

141 k

142 The inverse scale of the prior standard deviation on the latent mean

143 function, relative to half the observed range of `y_train`. If `y_train`

144 has less than two elements, `k` is ignored and the scale is set to 1.

145 power

146 base

147 Parameters of the prior on tree node generation. The probability that a

148 node at depth `d` (0-based) is non-terminal is ``base / (1 + d) **

149 power``.

150 lambda_

151 The prior harmonic mean of the error variance. (The harmonic mean of x

152 is 1/mean(1/x).) If not specified, it is set based on `sigest` and

153 `sigquant`.

154 tau_num

155 The numerator in the expression that determines the prior standard

156 deviation of leaves. If not specified, default to ``(max(y_train) -

157 min(y_train)) / 2`` (or 1 if `y_train` has less than two elements) for

158 continuous regression, and 3 for binary regression.

159 offset

160 The prior mean of the latent mean function. If not specified, it is set

161 to the mean of `y_train` for continuous regression, and to

162 ``Phi^-1(mean(y_train))`` for binary regression. If `y_train` is empty,

163 `offset` is set to 0. With binary regression, if `y_train` is all

164 `False` or `True`, it is set to ``Phi^-1(1/(n+1))`` or

165 ``Phi^-1(n/(n+1))``, respectively.

166 w

167 Coefficients that rescale the error standard deviation on each

168 datapoint. Not specifying `w` is equivalent to setting it to 1 for all

169 datapoints. Note: `w` is ignored in the automatic determination of

170 `sigest`, so either the weights should be O(1), or `sigest` should be

171 specified by the user. Not supported with binary regression

172 (``type='pbart'``).

173 ntree

174 The number of trees used to represent the latent mean function. By

175 default 200 for continuous regression and 50 for binary regression.

176 numcut

177 If `usequants` is `False`: the exact number of cutpoints used to bin the

178 predictors, ranging between the minimum and maximum observed values

179 (excluded).

180

181 If `usequants` is `True`: the maximum number of cutpoints to use for

182 binning the predictors. Each predictor is binned such that its

183 distribution in `x_train` is approximately uniform across bins. The

184 number of bins is at most the number of unique values appearing in

185 `x_train`, or ``numcut + 1``.

186

187 Before running the algorithm, the predictors are compressed to the

188 smallest integer type that fits the bin indices, so `numcut` is best set

189 to the maximum value of an unsigned integer type, like 255.

190

191 Ignored if `xinfo` is specified.

192 ndpost

193 The number of MCMC samples to save, after burn-in. `ndpost` is the

194 total number of samples across all chains. `ndpost` is rounded up to the

195 first multiple of `mc_cores`.

196 nskip

197 The number of initial MCMC samples to discard as burn-in. This number

198 of samples is discarded from each chain.

199 keepevery

200 The thinning factor for the MCMC samples, after burn-in. By default, 1

201 for continuous regression and 10 for binary regression.

202 printevery

203 The number of iterations (including thinned-away ones) between each log

204 line. Set to `None` to disable logging. ^C interrupts the MCMC only

205 every `printevery` iterations, so with logging disabled it's impossible

206 to kill the MCMC conveniently.

207 mc_cores

208 The number of independent MCMC chains.

209 seed

210 The seed for the random number generator.

211 bart_kwargs

212 Additional arguments passed to `bartz.Bart`.

213

214 Raises

215 ------

216 ValueError

217 If `w` is set with binary regression (``type='pbart'``).

218

219 Notes

220 -----

221 This interface imitates the function ``mc_gbart`` from the R package `BART3

222 <https://github.com/rsparapa/bnptools>`_, but with these differences:

223

224 - If ``usequants=False``, R BART3 switches to quantiles anyway if there are

225 less predictor values than the required number of bins, while bartz

226 always follows the specification.

227 - Some functionality is missing.

228 - The error variance parameter is called `lambda_` instead of `lambda`,

229 since the latter is a reserved word in Python.

230 - There are some additional attributes, and some missing.

231 - The trees have a maximum depth of 6.

232 - `rm_const` refers to predictors without decision rules instead of

233 predictors that are constant in `x_train`.

234 - If `rm_const=True` and some variables are dropped, the predictors

235 matrix/dataframe passed to `predict` should still include them.

236

237 References

238 ----------

239 .. [1] Linero, Antonio R. (2018). "Bayesian Regression Trees for

240 High-Dimensional Prediction and Variable Selection". In: Journal of the

241 American Statistical Association 113.522, pp. 626-636.

242 .. [2] Hugh A. Chipman, Edward I. George, Robert E. McCulloch "BART:

243 Bayesian additive regression trees," The Annals of Applied Statistics,

244 Ann. Appl. Stat. 4(1), 266-298, (March 2010).

245 """

246

247 _bart: Bart

248 _x_train_fmt: Any = field(static=True, default=None) (empty)

249 _yhat_test: Float32[Array, 'ndpost m'] | None = None (empty)

250

251 sigest: Float32[Array, ''] | None = None (empty)

252 """The estimated standard deviation of the error used to set `lambda_`.""" (empty)

253

254 def __init__( (empty)

255 self,

256 x_train: Real[ArrayLike, 'n p'] | DataFrame,

257 y_train: Float32[ArrayLike, ' n'] | Series,

258 *,

259 x_test: Real[ArrayLike, 'm p'] | DataFrame | None = None,

260 type: Literal['wbart', 'pbart'] = 'wbart', # noqa: A002

261 sparse: bool = False,

262 theta: FloatLike | None = None,

263 a: FloatLike = 0.5,

264 b: FloatLike = 1.0,

265 rho: FloatLike | None = None,

266 augment: bool = False,

267 varprob: Float[ArrayLike, ' p'] | None = None,

268 xinfo: Float[ArrayLike, 'p ncut'] | None = None,

269 usequants: bool = False,

270 rm_const: bool = True,

271 sigest: FloatLike | None = None,

272 sigdf: FloatLike = 3.0,

273 sigquant: FloatLike = 0.9,

274 k: FloatLike = 2.0,

275 power: FloatLike = 2.0,

276 base: FloatLike = 0.95,

277 lambda_: FloatLike | None = None,

278 tau_num: FloatLike | None = None,

279 offset: FloatLike | None = None,

280 w: Float[ArrayLike, ' n'] | Series | None = None,

281 ntree: int | None = None,

282 numcut: int = 100,

283 ndpost: int = 1000,

284 nskip: int = 100,

285 keepevery: int | None = None,

286 printevery: int | None = 100,

287 mc_cores: int = 2,

288 seed: int | Key[Array, ''] = 0,

289 bart_kwargs: Mapping = MappingProxyType({}),

290 ) -> None:

291 # BART3 does not support heteroskedastic probit

292 if type == 'pbart' and w is not None: (empty)

293 msg = ( (empty)

294 "w is not supported with binary regression (type='pbart');"

295 ' BART3 has no heteroskedastic probit.'

296 )

297 raise ValueError(msg) (empty)

298

299 # set defaults that depend on type of regression

300 if keepevery is None: (empty)

301 keepevery = 10 if type == 'pbart' else 1 (empty)

302 if ntree is None: (empty)

303 ntree = 50 if type == 'pbart' else 200 (empty)

304

305 # pre-process the data to numeric arrays once, so the OLS estimate of

306 # `sigest` and `Bart` share a single copy of the (memory-heavy) X matrix.

307 # `Bart` records the format as plain arrays, so `predict` re-implements

308 # the input-format consistency check against the original format here.

309 x_train, self._x_train_fmt = _process_bart3_predictor_input(x_train) (empty)

310 y_train = _process_response_input(y_train) (empty)

311

312 # map the BART3 error-variance settings to Bart's sigma prior, estimating

313 # `sigest` by linear regression on x_train when needed

314 sigma_kw, self.sigest = _resolve_sigma_prior( (empty)

315 x_train,

316 y_train,

317 type=type,

318 sigest=sigest,

319 sigdf=sigdf,

320 sigquant=sigquant,

321 lambda_=lambda_,

322 )

323

324 # convert to per-chain n_save for Bart

325 num_chains = None if mc_cores == 1 else mc_cores (empty)

326 actual_num_chains = num_chains or 1 (empty)

327 n_save = ndpost // actual_num_chains + bool(ndpost % actual_num_chains) (empty)

328

329 # translate xinfo/usequants/numcut to a binner factory

330 if xinfo is not None: (empty)

331 binner = partial(GivenSplitsBinner, xinfo=jnp.asarray(xinfo)) (empty)

332 elif usequants: (empty)

333 binner = partial( (empty)

334 UniqueQuantileBinner, max_bins=numcut + 1, max_subsample=None

335 )

336 else:

337 binner = partial(RangeEvenBinner, max_bins=numcut + 1) (empty)

338

339 # set most calling arguments for Bart

340 kwargs: dict = dict( (empty)

341 x_train=x_train,

342 y_train=y_train,

343 outcome_type=dict(wbart='continuous', pbart='binary')[type],

344 sparse=SparseConfig(

345 enabled=sparse, theta=theta, a=a, b=b, rho=rho, augment=augment

346 ),

347 varprob=varprob,

348 binner=binner,

349 rm_const=rm_const,

350 **sigma_kw,

351 k=k,

352 power=power,

353 base=base,

354 tau_num=tau_num,

355 offset=offset,

356 error_scale=w,

357 num_trees=ntree,

358 n_save=n_save,

359 n_burn=nskip,

360 n_skip=keepevery,

361 printevery=printevery,

362 seed=seed,

363 maxdepth=6,

364 num_chains=num_chains,

365 precompute_predict_train=True,

366 )

367

368 # default min_points_per_leaf to 5 (unless set by the user) to match

369 # BART3's hard-coded nl>=5 && nr>=5 birth check.

370 # min_points_per_decision_node keeps the Bart default of 10

371 # (= 2 * min_points_per_leaf): it makes the proposal efficient by not

372 # trying to grow leaves too small to split, without changing the target

373 # posterior, which thus matches BART3.

374 if 'min_points_per_leaf' not in bart_kwargs.get('init_kw', {}): (empty)

375 bart_kwargs = dict( (empty)

376 bart_kwargs,

377 init_kw=dict(bart_kwargs.get('init_kw', {}), min_points_per_leaf=5),

378 )

379

380 # add user arguments

381 kwargs.update(bart_kwargs) (empty)

382

383 # invoke Bart

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

385

386 # predict at test points

387 if x_test is not None: (empty)

388 self._yhat_test = self.predict(x_test) (empty)

389

390 # Public attributes from Bart

391

392 @property (empty)

393 def ndpost(self) -> int: (empty)

394 """The number of MCMC samples saved, after burn-in."""

395 return self._bart.ndpost (empty)

396

397 @property (empty)

398 def offset(self) -> Float32[Array, '']: (empty)

399 """The prior mean of the latent mean function."""

400 return self._bart.offset (empty)

401

402 # Private attributes from Bart

403

404 @property (empty)

405 def _main_trace(self) -> MainTrace: (empty)

406 return self._bart._main_trace # noqa: SLF001 (empty)

407

408 @property (empty)

409 def _burnin_trace(self) -> BurninTrace: (empty)

410 return self._bart._burnin_trace # noqa: SLF001 (empty)

411

412 @property (empty)

413 def _mcmc_state(self) -> State: (empty)

414 return self._bart._mcmc_state # noqa: SLF001 (empty)

415

416 @property (empty)

417 def _splits(self) -> Real[Array, 'p max_num_splits']: (empty)

418 return self._bart._binner._splits # noqa: SLF001 (empty)

419

420 # Properties

421

422 @property (empty)

423 def yhat_test(self) -> Float32[Array, 'ndpost m'] | None: (empty)

424 """The conditional posterior mean at `x_test` for each MCMC iteration."""

425 return self._yhat_test (empty)

426

427 @cached_property (empty)

428 def accept( (empty)

429 self,

430 ) -> (

431 Float32[Array, ' nskip_plus_ndpost']

432 | Float32[Array, 'nskip_plus_ndpost_per_core mc_cores']

433 ):

434 """The fraction of trees with an accepted move, including burn-in samples.

435

436 Unlike BART3, the iterations thinned away by `keepevery` are not

437 recorded.

438 """

439 # `Bart.accept` is (mc_cores, samples) or (samples,); the public layout

440 # is (samples, mc_cores), like `sigma`

441 return self._bart.accept.T (empty)

442

443 @cached_property (empty)

444 def prob_test(self) -> Float32[Array, 'ndpost m'] | None: (empty)

445 """The posterior probability of y being True at `x_test` for each MCMC iteration."""

446 if self._yhat_test is None or self._mcmc_state.z is None: (empty)

447 return None (empty)

448 return ndtr(self._yhat_test) (empty)

449

450 @cached_property (empty)

451 def prob_test_mean(self) -> Float32[Array, ' m'] | None: (empty)

452 """The marginal posterior probability of y being True at `x_test`."""

453 if self.prob_test is None: (empty)

454 return None (empty)

455 return self.prob_test.mean(axis=0) (empty)

456

457 @cached_property (empty)

458 def prob_train(self) -> Float32[Array, 'ndpost n'] | None: (empty)

459 """The posterior probability of y being True at `x_train` for each MCMC iteration."""

460 if self._mcmc_state.z is not None: (empty)

461 return ndtr(self.yhat_train) (empty)

462 else:

463 return None (empty)

464

465 @cached_property (empty)

466 def prob_train_mean(self) -> Float32[Array, ' n'] | None: (empty)

467 """The marginal posterior probability of y being True at `x_train`."""

468 if self.prob_train is None: (empty)

469 return None (empty)

470 else:

471 return self.prob_train.mean(axis=0) (empty)

472

473 @cached_property (empty)

474 def sigma( (empty)

475 self,

476 ) -> (

477 Float32[Array, ' nskip_plus_ndpost']

478 | Float32[Array, 'nskip_plus_ndpost_per_core mc_cores']

479 | None

480 ):

481 """The standard deviation of the error, including burn-in samples."""

482 if self._mcmc_state.z is not None: (empty)

483 return None (empty)

484 assert self._burnin_trace.error_cov_inv.ndim <= 2 # chains and samples (empty)

485 tc = chain_vmap_axes(self._main_trace).error_cov_inv (empty)

486

487 def arrange(arr: Shaped[Array, '...']) -> Shaped[Array, '...']: (empty)

488 # Public output is (nskip+ndpost, mc_cores) = (samples, chains).

489 return chain_to_axis(arr, tc, target=-1) (empty)

490

491 return jnp.sqrt( (empty)

492 jnp.reciprocal(

493 jnp.concatenate(

494 [

495 arrange(self._burnin_trace.error_cov_inv),

496 arrange(self._main_trace.error_cov_inv),

497 ],

498 axis=0,

499 )

500 )

501 )

502

503 @cached_property (empty)

504 def sigma_(self) -> Float32[Array, 'ndpost'] | None: (empty)

505 """The standard deviation of the error, only over the post-burnin samples and flattened."""

506 if self._mcmc_state.z is not None: (empty)

507 return None (empty)

508 assert self._main_trace.error_cov_inv.ndim <= 2 # chains and samples (empty)

509 arr = chain_to_axis( (empty)

510 self._main_trace.error_cov_inv,

511 chain_vmap_axes(self._main_trace).error_cov_inv,

512 )

513 return jnp.sqrt(jnp.reciprocal(arr)).reshape(-1) (empty)

514

515 @cached_property (empty)

516 def sigma_mean(self) -> Float32[Array, ''] | None: (empty)

517 """The mean of `sigma`, only over the post-burnin samples."""

518 if self.sigma_ is None: (empty)

519 return None (empty)

520 return self.sigma_.mean() (empty)

521

522 @cached_property (empty)

523 def varcount(self) -> Int32[Array, 'ndpost p']: (empty)

524 """Histogram of predictor usage for decision rules in the trees."""

525 return self._bart.varcount (empty)

526

527 @cached_property (empty)

528 def varcount_mean(self) -> Float32[Array, ' p']: (empty)

529 """Average of `varcount` across MCMC iterations."""

530 return self._bart.varcount_mean (empty)

531

532 @cached_property (empty)

533 def varprob(self) -> Float32[Array, 'ndpost p']: (empty)

534 """Posterior samples of the probability of choosing each predictor for a decision rule."""

535 return self._bart.varprob (empty)

536

537 @cached_property (empty)

538 def varprob_mean(self) -> Float32[Array, ' p']: (empty)

539 """The marginal posterior probability of each predictor being chosen for a decision rule."""

540 return self._bart.varprob_mean (empty)

541

542 @cached_property (empty)

543 def yhat_test_mean(self) -> Float32[Array, ' m'] | None: (empty)

544 """The marginal posterior mean at `x_test`.

545

546 Not defined with binary regression because it's error-prone, typically

547 the right thing to consider would be `prob_test_mean`.

548 """

549 if self._yhat_test is None or self._mcmc_state.z is not None: (empty)

550 return None (empty)

551 return self._yhat_test.mean(axis=0) (empty)

552

553 @cached_property (empty)

554 def yhat_train(self) -> Float32[Array, 'ndpost n']: (empty)

555 """The conditional posterior mean at `x_train` for each MCMC iteration."""

556 return self._bart.predict('train', kind=PredictKind.latent_samples) (empty)

557

558 @cached_property (empty)

559 def yhat_train_mean(self) -> Float32[Array, ' n'] | None: (empty)

560 """The marginal posterior mean at `x_train`.

561

562 Not defined with binary regression because it's error-prone, typically

563 the right thing to consider would be `prob_train_mean`.

564 """

565 if self._mcmc_state.z is not None: (empty)

566 return None (empty)

567 else:

568 return self.yhat_train.mean(axis=0) (empty)

569

570 # Public methods from Bart

571

572 def predict( (empty)

573 self, x_test: Real[ArrayLike, 'm p'] | DataFrame

574 ) -> Float32[Array, 'ndpost m']:

575 """

576 Evaluate the sum-of-trees at `x_test` for each MCMC iteration.

577

578 Parameters

579 ----------

580 x_test

581 The test predictors.

582

583 Returns

584 -------

585 Posterior samples of the latent function value at `x_test`. In the continuous case, this is the conditional mean.

586

587 Raises

588 ------

589 ValueError

590 If `x_test` has a different format than `x_train`.

591 """

592 # pre-process and check the format matches x_train; Bart only sees plain

593 # arrays, so this consistency check is re-implemented here

594 x_test, x_test_fmt = _process_bart3_predictor_input(x_test) (empty)

595 if x_test_fmt != self._x_train_fmt: (empty)

596 msg = ( (empty)

597 f'Input format mismatch: {x_test_fmt=} '

598 f'!= x_train_fmt={self._x_train_fmt!r}'

599 )

600 raise ValueError(msg) (empty)

601 return self._bart.predict(x_test, kind=PredictKind.latent_samples) (empty)

602

603

604class gbart(mc_gbart): (empty)

605 """Subclass of `mc_gbart` that forces `mc_cores=1`."""

606

607 def __init__(self, *args: Any, **kwargs: Any) -> None: (empty)

608 if 'mc_cores' in kwargs: 608 ↛ 611line 608 didn't jump to line 611 because the condition on line 608 was always true(empty)

609 msg = "gbart.__init__() got an unexpected keyword argument 'mc_cores'" (empty)

610 raise TypeError(msg) (empty)

611 kwargs.update(mc_cores=1)

612 super().__init__(*args, **kwargs)

613

614

615def _process_bart3_predictor_input( (empty)

616 x: Real[ArrayLike, 'n p'] | DataFrame,

617) -> tuple[Shaped[Array, 'p n'], Any]:

618 """Process BART3-style predictors (one predictor per column) to bartz layout.

619

620 Unlike `bartz.Bart`, BART3 lays out predictor matrices with one predictor

621 per column, so plain arrays are transposed to bartz's (p, n) layout.

622 Dataframes already use one column per predictor, so they are left untouched.

623 """

624 if not isinstance(x, DataFrame): (empty)

625 x = jnp.asarray(x).T (empty)

626 return _process_predictor_input(x) (empty)

627

628

629def _resolve_sigma_prior( (empty)

630 x_train: Shaped[Array, 'p n'],

631 y_train: Float32[Array, ' n'],

632 *,

633 type: Literal['wbart', 'pbart'], # noqa: A002

634 sigest: FloatLike | None,

635 sigdf: FloatLike,

636 sigquant: FloatLike,

637 lambda_: FloatLike | None,

638) -> tuple[dict, Float32[Array, ''] | None]:

639 """Map the BART3 error-variance settings to Bart's sigma prior.

640

641 Returns (sigma_kwargs, sigest) where sigest is the error standard deviation

642 estimate, or None for binary regression or when `lambda_` is given.

643 """

644 if type == 'pbart': (empty)

645 if sigest is not None or lambda_ is not None: (empty)

646 msg = 'Do not set `sigest` or `lambda_` for binary regression, they are ignored' (empty)

647 raise ValueError(msg) (empty)

648 return {}, None (empty)

649

650 if lambda_ is None: (empty)

651 if sigest is None: (empty)

652 sigest2 = _sigest2_ols(x_train, y_train) (empty)

653 else:

654 sigest2 = jnp.square(jnp.asarray(sigest, jnp.float32)) (empty)

655 sigest_out = jnp.sqrt(sigest2) (empty)

656 # lambda_ such that the sigquant quantile of the prior matches sigest²

657 invchi2 = invgamma.ppf(sigquant, sigdf / 2) / 2 (empty)

658 lambda_ = sigest2 / (invchi2 * sigdf) (empty)

659 else:

660 if sigest is not None: (empty)

661 msg = "Do not set `sigest` if `lambda_` is specified, it's ignored" (empty)

662 raise ValueError(msg) (empty)

663 lambda_ = jnp.asarray(lambda_, jnp.float32) (empty)

664 sigest_out = None (empty)

665

666 # Bart's prior reduces to scaled-inv-χ²(sigma_df, sigma_scale²) on the error

667 # variance, matching BART3's scaled-inv-χ²(sigdf, lambda_); sigma_init keeps

668 # the initial precision at the prior mean nu/rate = 1 / lambda_

669 sigma_scale = jnp.sqrt(lambda_) (empty)

670 sigma_kw = dict(sigma_df=sigdf, sigma_scale=sigma_scale, sigma_init=sigma_scale) (empty)

671 return sigma_kw, sigest_out (empty)

672

673

674def _sigest2_ols( (empty)

675 x_train: Shaped[Array, 'p n'], y_train: Float32[Array, ' n']

676) -> Float32[Array, '']:

677 """Estimate the error variance by OLS with intercept."""

678 p, n = x_train.shape (empty)

679 if n <= p: (empty)

680 msg = ( (empty)

681 f'cannot estimate `sigest` by OLS with {n} datapoints and {p} '

682 'predictors (it requires more datapoints than predictors); '

683 'specify `sigest` or `lambda_` explicitly'

684 )

685 raise ValueError(msg) (empty)

686 return _sigma2_from_ols(x_train, y_train) (empty)

Coverage for src/bartz/BART/_gbart.py: 99%

180 statements