Coverage for src/bartz/bcf/_bcf.py: 98%

1# bartz/src/bartz/bcf/_bcf.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"""Bayesian Causal Forests (BCF) interface.""" 

26 

27import dataclasses (empty)

28import json (empty)

29from operator import attrgetter (empty)

30from pathlib import Path (empty)

31from typing import Any, Literal, cast (empty)

32 

33import equinox as eqx (empty)

34import jax (empty)

35import jax.numpy as jnp (empty)

36import numpy as np (empty)

37from jax import device_put, lax, random, tree (empty)

38from jax.scipy import special (empty)

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

40 

41from bartz._interface import ( (empty)

42 ArrayLike, 

43 DataFrame, 

44 FloatLike, 

45 Series, 

46 _flatten_chain_sample, 

47 _guarded_response_variance, 

48 _process_error_variance_settings, 

49 _process_leaf_variance_settings, 

50 _process_offset_settings, 

51 _process_predictor_input, 

52 _process_response_input, 

53 predict_latent, 

54) 

55from bartz._jaxext import split (empty)

56from bartz.bcf._loop import BCFBurninTrace, BCFMainTrace, bcf_step (empty)

57from bartz.bcf._state import init_bcf (empty)

58from bartz.mcmcloop import MainTrace, run_mcmc (empty)

59from bartz.mcmcloop._trace import Trace (empty)

60from bartz.mcmcstep import OutcomeType, Wishart (empty)

61from bartz.mcmcstep._axes import chain_vmap_axes, trace_sample_axes (empty)

62from bartz.mcmcstep._state import make_p_nonterminal (empty)

63from bartz.prepcovars import RangeEvenBinner, UniqueQuantileBinner (empty)

64 

65 

66def _process_bcf_predictor_input( (empty)

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

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

69 """ 

70 Process predictors (one predictor per column) to bartz layout (p, n). 

71 

72 Parameters 

73 ---------- 

74 x 

75 The predictor data. 

76 

77 Returns 

78 ------- 

79 tuple[Shaped[Array, "p n"], Any] 

80 A tuple containing the predictors transposed to (p, n) shape and their 

81 original format metadata. 

82 """ 

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

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

85 return _process_predictor_input(x) (empty)

86 

87 

88def _fold_chains(trace: Trace, path: str) -> Float32[Array, 'num_samples ...']: (empty)

89 """Fold the chain axis of a trace field into its sample axis, like `predict_latent`.""" 

90 get = attrgetter(path) (empty)

91 return _flatten_chain_sample( (empty)

92 get(trace), get(chain_vmap_axes(trace)), get(trace_sample_axes(trace)) 

93 ) 

94 

95 

96def _serialize_binner(binner: Any, max_split: Any = None) -> dict[str, Any]: # noqa: ANN401 (empty)

97 """ 

98 Serialize binner state to a dictionary of NumPy arrays. 

99 

100 Parameters 

101 ---------- 

102 binner 

103 The binner object to serialize. 

104 max_split 

105 Optional pre-extracted max_split array. 

106 

107 Returns 

108 ------- 

109 dict[str, Any] 

110 A dictionary containing the serialized binner attributes. 

111 

112 Raises 

113 ------ 

114 RuntimeError 

115 If max_split was donated to JAX and not provided explicitly. 

116 """ 

117 binner_dict = {} (empty)

118 binner_dict['class'] = binner.__class__.__name__ (empty)

119 

120 if max_split is None: # pragma: no cover (empty)

121 try: 

122 max_split = np.asarray(binner.max_split) 

123 except RuntimeError as exc: 

124 msg = 'max_split array was donated to JAX. Pass max_split explicitly.' 

125 raise RuntimeError(msg) from exc 

126 binner_dict['max_split'] = np.asarray(max_split) (empty)

127 

128 if hasattr(binner, '_splits'): (empty)

129 binner_dict['_splits'] = np.asarray(binner._splits) # noqa: SLF001 (empty)

130 

131 if hasattr(binner, '_low'): # pragma: no cover (empty)

132 binner_dict['_low'] = np.asarray(binner._low) # noqa: SLF001 

133 binner_dict['_high'] = np.asarray(binner._high) # noqa: SLF001 

134 binner_dict['_max_bins'] = np.asarray(binner._max_bins) # noqa: SLF001 

135 

136 return binner_dict (empty)

137 

138 

139def _deserialize_binner(data: Any) -> Any: # noqa: ANN401 (empty)

140 """ 

141 Reconstruct binner instance from serialized NPZ data dictionary. 

142 

143 Parameters 

144 ---------- 

145 data 

146 A mapping containing the serialized binner fields from NPZ archive. 

147 

148 Returns 

149 ------- 

150 binner : Any 

151 A reconstituted binner instance. 

152 """ 

153 binner_cls_name = str(data.get('binner.class', 'UniqueQuantileBinner')) (empty)

154 if binner_cls_name == 'UniqueQuantileBinner': (empty)

155 binner = object.__new__(UniqueQuantileBinner) (empty)

156 object.__setattr__(binner, '_splits', jnp.asarray(data['binner._splits'])) (empty)

157 object.__setattr__(binner, 'max_split', jnp.asarray(data['binner.max_split'])) (empty)

158 else: # pragma: no cover 

159 binner = object.__new__(RangeEvenBinner) 

160 object.__setattr__(binner, '_low', jnp.asarray(data['binner._low'])) 

161 object.__setattr__(binner, '_high', jnp.asarray(data['binner._high'])) 

162 object.__setattr__(binner, '_max_bins', int(data['binner._max_bins'])) 

163 object.__setattr__(binner, 'max_split', jnp.asarray(data['binner.max_split'])) 

164 return binner (empty)

165 

166 

167def make_leaf_prior_cov_inv( (empty)

168 value: FloatLike, sample: bool, shape: FloatLike, scale: FloatLike 

169) -> Wishart: 

170 """Build the leaf precision prior from the inverse-gamma prior on the variance.""" 

171 if sample: (empty)

172 return Wishart(nu=2 * shape, rate=2 * scale, value=value) (empty)

173 else: 

174 return Wishart(nu=None, rate=None, value=value) (empty)

175 

176 

177class bcf(eqx.Module): (empty)

178 R""" 

179 Bayesian Causal Forests (BCF). 

180 

181 Regress `y_train` on `x_train` and `z_train` (treatment) with two latent 

182 mean functions represented as sums of decision trees: 

183 Y = mu(X, pihat) + tau(X, pihat) * Z + error 

184 

185 For continuous outcomes, the hyperparameters with units are on the scale 

186 of the standardized response, see `standardize`, while the outputs are on 

187 the scale of `y_train`. 

188 

189 Parameters 

190 ---------- 

191 x_train 

192 The training predictors (confounders/modifiers). 

193 y_train 

194 The training responses. 

195 z_train 

196 The binary treatment assignment (0 or 1). 

197 pihat_train 

198 The estimated propensity scores. If provided, appended to `x_train`. 

199 x_test 

200 The test predictors. If provided, predictions at these points are 

201 computed and stored in `mu_test` and `tau_test`. 

202 z_test 

203 The test treatment assignment (0 or 1). If provided together with 

204 `x_test`, the outcome prediction is stored in `yhat_test`. 

205 pihat_test 

206 The test propensity scores. Must be passed together with `pihat_train`. 

207 include_pihat_in_mu 

208 Whether to include propensity scores in the prognostic forest. 

209 include_pihat_in_tau 

210 Whether to include propensity scores in the treatment effect forest. 

211 num_trees_mu 

212 The number of trees used for the prognostic forest `mu`. 

213 num_trees_tau 

214 The number of trees used for the treatment effect forest `tau`. 

215 ndpost 

216 The number of MCMC samples to save, after burn-in, per chain. The 

217 posterior samples of all chains are concatenated, for a total of 

218 ``num_chains * ndpost``. 

219 nskip 

220 The number of initial MCMC samples to discard as burn-in, per chain. 

221 num_chains 

222 The number of independent MCMC chains. `None` for a single chain 

223 without an explicit chain axis, which is equivalent to 1. 

224 k_mu 

225 Prior parameter k for prognostic forest. 

226 k_tau 

227 Prior parameter k for treatment forest. 

228 sigma_df 

229 Prior degrees of freedom for error variance. 

230 sigma_scale 

231 Prior scale of the error standard deviation. 

232 sigma_init 

233 Initial value of the error standard deviation. 

234 leaf_prior_cov_inv_mu 

235 Custom leaf prior precision for the prognostic forest. 

236 leaf_prior_cov_inv_tau 

237 Custom leaf prior precision for the treatment effect forest. 

238 min_points_per_leaf_mu 

239 Minimum data points per leaf for prognostic forest. 

240 min_points_per_leaf_tau 

241 Minimum data points per leaf for treatment forest. 

242 tau_0_prior_var 

243 Prior variance for global intercept tau_0. 

244 sample_intercept 

245 Whether to sample a global treatment intercept `tau_0`. 

246 adaptive_coding 

247 Whether to use adaptive coding for the treatment effect. 

248 sample_sigma2_leaf_mu 

249 Whether to sample the leaf parameter variance for the prognostic forest. 

250 sigma2_leaf_shape_mu 

251 The shape parameter for the Inverse-Gamma prior on the prognostic forest leaf variance. 

252 sigma2_leaf_scale_mu 

253 The scale parameter for the Inverse-Gamma prior on the prognostic forest leaf variance. 

254 sample_sigma2_leaf_tau 

255 Whether to sample the leaf parameter variance for the treatment effect forest. 

256 sigma2_leaf_shape_tau 

257 The shape parameter for the Inverse-Gamma prior on the treatment effect forest leaf variance. 

258 sigma2_leaf_scale_tau 

259 The scale parameter for the Inverse-Gamma prior on the treatment effect forest leaf variance. 

260 standardize 

261 Whether to standardize `y_train` internally, ignored for binary 

262 outcomes. If `False`, the hyperparameters with units (`sigma_scale`, 

263 `sigma_init`, `leaf_prior_cov_inv_*`, `tau_0_prior_var`, 

264 `sigma2_leaf_scale_*`) are on the scale of `y_train` instead of the 

265 standardized one. 

266 outcome_type 

267 Either 'continuous' or 'binary' (probit link). 

268 delta_max 

269 Maximum plausible treatment effect on the probability scale for binary probit models. 

270 seed 

271 The seed for the random number generator. 

272 

273 Raises 

274 ------ 

275 ValueError 

276 If binary outcome is specified but `y_train` contains values other 

277 than 0 or 1, or if the format of `x_test` does not match `x_train` 

278 format, or if `z_test` or `pihat_test` is passed without `x_test`, or 

279 if only one of `pihat_train` and `pihat_test` is passed, or if 

280 `z_test` or `pihat_test` does not match the length of `x_test`. 

281 """ 

282 

283 _mcmc_state: Any 

284 _binner: Any 

285 _main_trace: Any 

286 _burnin_trace: Any 

287 _tau_0_trace: Any 

288 _b_trace: Any 

289 _leaf_prior_cov_inv_mu_trace: Any 

290 _leaf_prior_cov_inv_tau_trace: Any 

291 _x_train_fmt: Any = eqx.field(static=True, default=None) (empty)

292 _standardize: bool = eqx.field(static=True, default=False) (empty)

293 _y_mean: Float32[ArrayLike, ''] | float = eqx.field(default=0.0) (empty)

294 _y_std: Float32[ArrayLike, ''] | float = eqx.field(default=1.0) (empty)

295 _outcome_type: str = eqx.field(static=True, default='continuous') (empty)

296 _offset: Float32[ArrayLike, ''] | float = eqx.field(default=0.0) (empty)

297 _mu_test: Float32[Array, 'num_samples m'] | None = eqx.field(default=None) (empty)

298 _tau_test: Float32[Array, 'num_samples m'] | None = eqx.field(default=None) (empty)

299 _yhat_test: Float32[Array, 'num_samples m'] | None = eqx.field(default=None) (empty)

300 

301 def __init__( # noqa: C901, PLR0915 (empty)

302 self, 

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

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

305 z_train: Float32[ArrayLike, ' n'] | Series, 

306 *, 

307 pihat_train: Float32[ArrayLike, ' n'] | Series | None = None, 

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

309 z_test: Float32[ArrayLike, ' m'] | Series | None = None, 

310 pihat_test: Float32[ArrayLike, ' m'] | Series | None = None, 

311 include_pihat_in_mu: bool = True, 

312 include_pihat_in_tau: bool = False, 

313 num_trees_mu: int = 250, 

314 num_trees_tau: int = 100, 

315 ndpost: int = 1000, 

316 nskip: int = 100, 

317 num_chains: int | None = None, 

318 k_mu: float = 2.0, 

319 k_tau: float = 10.0, 

320 sigma_df: float = 3.0, 

321 sigma_scale: float | Literal['auto'] = 'auto', 

322 sigma_init: float | Literal['auto'] = 'auto', 

323 leaf_prior_cov_inv_mu: FloatLike | None = None, 

324 leaf_prior_cov_inv_tau: FloatLike | None = None, 

325 min_points_per_leaf_mu: int = 5, 

326 min_points_per_leaf_tau: int = 5, 

327 tau_0_prior_var: float | None = None, 

328 sample_intercept: bool = True, 

329 adaptive_coding: bool = False, 

330 sample_sigma2_leaf_mu: bool = True, 

331 sigma2_leaf_shape_mu: float = 3.0, 

332 sigma2_leaf_scale_mu: FloatLike | None = None, 

333 sample_sigma2_leaf_tau: bool = False, 

334 sigma2_leaf_shape_tau: float = 3.0, 

335 sigma2_leaf_scale_tau: FloatLike | None = None, 

336 standardize: bool = True, 

337 outcome_type: Literal['continuous', 'binary'] = 'continuous', 

338 delta_max: float = 0.9, 

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

340 ) -> None: 

341 

342 # 1. Pre-process the data (convert to arrays and transpose X to (p, n)) 

343 x_train, self._x_train_fmt = _process_bcf_predictor_input(x_train) (empty)

344 y_train = _process_response_input(y_train) (empty)

345 z_train = _process_response_input(z_train) (empty)

346 z_train = eqx.error_if( (empty)

347 z_train, 

348 jnp.any((z_train != 0) & (z_train != 1)), 

349 'Values in `z_train` must be 0 or 1.', 

350 ).astype(bool) 

351 

352 self._outcome_type = outcome_type (empty)

353 

354 if outcome_type == 'binary': (empty)

355 if not jnp.all((y_train == 0) | (y_train == 1)): (empty)

356 msg = 'Values in `y_train` must be strictly 0 or 1 for binary outcomes.' (empty)

357 raise ValueError(msg) (empty)

358 standardize = False (empty)

359 

360 if standardize: (empty)

361 y_mean = jnp.mean(y_train) (empty)

362 y_std = jnp.std(y_train) (empty)

363 y_std = jnp.where(y_std == 0, 1.0, y_std) (empty)

364 y_train_internal = (y_train - y_mean) / y_std (empty)

365 else: 

366 y_mean = jnp.float32(0.0) (empty)

367 y_std = jnp.float32(1.0) (empty)

368 y_train_internal = y_train (empty)

369 

370 self._standardize = standardize (empty)

371 self._y_mean = y_mean (empty)

372 self._y_std = y_std (empty)

373 

374 if pihat_train is not None: (empty)

375 pihat_train = _process_response_input(pihat_train) (empty)

376 

377 if x_test is None: (empty)

378 if z_test is not None or pihat_test is not None: (empty)

379 msg = '`z_test` and `pihat_test` require `x_test`.' (empty)

380 raise ValueError(msg) (empty)

381 else: 

382 x_test_preprocessed, x_test_fmt = _process_bcf_predictor_input(x_test) (empty)

383 _, m = x_test_preprocessed.shape (empty)

384 del x_test_preprocessed (empty)

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

386 msg = ( (empty)

387 f'Format of x_test {x_test_fmt} does not match x_train' 

388 f' {self._x_train_fmt}' 

389 ) 

390 raise ValueError(msg) (empty)

391 if (pihat_train is None) != (pihat_test is None): (empty)

392 msg = '`pihat_train` and `pihat_test` must be passed together.' (empty)

393 raise ValueError(msg) (empty)

394 if z_test is not None: (empty)

395 z_test = _process_response_input(z_test) (empty)

396 (len_z,) = z_test.shape (empty)

397 if len_z != m: (empty)

398 msg = f'`z_test` has length {len_z}, but `x_test` has {m} rows.' (empty)

399 raise ValueError(msg) (empty)

400 z_test = eqx.error_if( (empty)

401 z_test, 

402 jnp.any((z_test != 0) & (z_test != 1)), 

403 'Values in `z_test` must be 0 or 1.', 

404 ).astype(bool) 

405 if pihat_test is not None: 405 ↛ 416line 405 didn't jump to line 416 because the condition on line 405 was always true(empty)

406 pihat_test = _process_response_input(pihat_test) (empty)

407 (len_pihat,) = pihat_test.shape (empty)

408 if len_pihat != m: (empty)

409 msg = ( (empty)

410 f'`pihat_test` has length {len_pihat}, but `x_test` has' 

411 f' {m} rows.' 

412 ) 

413 raise ValueError(msg) (empty)

414 

415 # 2. Append pihat to X to create unified predictor matrix 

416 x_train_unified = x_train (empty)

417 pihat_index = None (empty)

418 

419 if pihat_train is not None: 419 ↛ 427line 419 didn't jump to line 427 because the condition on line 419 was always true(empty)

420 # x_train is (p, n), pihat_train is (n,). Add a channel dim to pihat to 

421 # make it (1, n) 

422 pihat_row = pihat_train[jnp.newaxis, :] (empty)

423 x_train_unified = jnp.concatenate([x_train_unified, pihat_row], axis=0) (empty)

424 pihat_index = x_train_unified.shape[0] - 1 (empty)

425 

426 # 3. Resolve priors for both mu and tau forests 

427 binary_mask = ( (empty)

428 jnp.ones((), dtype=bool) 

429 if outcome_type == 'binary' 

430 else jnp.zeros((), dtype=bool) 

431 ) 

432 

433 offset_val = _process_offset_settings(y_train_internal, binary_mask, None, None) (empty)

434 self._offset = offset_val (empty)

435 

436 if leaf_prior_cov_inv_mu is None: (empty)

437 if outcome_type == 'binary': (empty)

438 leaf_prior_cov_inv_mu = jnp.array(num_trees_mu, jnp.float32) (empty)

439 else: 

440 leaf_prior_cov_inv_mu = _process_leaf_variance_settings( (empty)

441 y_train_internal, 

442 binary_mask, 

443 missing=None, 

444 k=jnp.array(k_mu), 

445 num_trees=num_trees_mu, 

446 tau_num=None, 

447 sigma_mu_df=None, 

448 ).value 

449 if leaf_prior_cov_inv_tau is None: (empty)

450 if outcome_type == 'binary': (empty)

451 p_val = 0.6827 (empty)

452 q_quantile = special.ndtri((p_val + 1) / 2.0) (empty)

453 phi_0 = 1.0 / jnp.sqrt(2 * jnp.pi) (empty)

454 sigma2_tau = ((delta_max / (q_quantile * phi_0)) ** 2) / num_trees_tau (empty)

455 leaf_prior_cov_inv_tau = jnp.reciprocal(sigma2_tau) (empty)

456 else: 

457 leaf_prior_cov_inv_tau = _process_leaf_variance_settings( (empty)

458 y_train_internal, 

459 binary_mask, 

460 missing=None, 

461 k=jnp.array(k_tau), 

462 num_trees=num_trees_tau, 

463 tau_num=None, 

464 sigma_mu_df=None, 

465 ).value 

466 

467 error_cov_inv = _process_error_variance_settings( (empty)

468 y_train_internal, 

469 binary_mask, 

470 None, 

471 outcome_type=OutcomeType(outcome_type), 

472 sigma_df=sigma_df, 

473 sigma_scale=sigma_scale, 

474 sigma_init=sigma_init, 

475 error_scale=None, 

476 ) 

477 

478 p_nonterminal_mu = make_p_nonterminal(d=10, alpha=0.95, beta=2.0) (empty)

479 p_nonterminal_tau = make_p_nonterminal(d=5, alpha=0.25, beta=3.0) (empty)

480 

481 if outcome_type == 'binary': (empty)

482 var_y = 1.0 (empty)

483 else: 

484 var_y = _guarded_response_variance(y_train_internal, None, None) (empty)

485 if sigma2_leaf_scale_mu is None: (empty)

486 sigma2_leaf_scale_mu = var_y / num_trees_mu (empty)

487 if sigma2_leaf_scale_tau is None: (empty)

488 sigma2_leaf_scale_tau = 0.5 * var_y / num_trees_tau (empty)

489 

490 # 3.5 Bin the unified data 

491 rng = random.key(seed) if not isinstance(seed, jax.Array) else seed (empty)

492 keys = split(rng) (empty)

493 

494 binner = UniqueQuantileBinner(x_train_unified, key=keys.pop()) (empty)

495 x_train_binned = binner.bin(x_train_unified) (empty)

496 # copies because `init_bcf` may donate them 

497 max_split_mu = jnp.copy(binner.max_split) (empty)

498 max_split_tau = jnp.copy(binner.max_split) (empty)

499 

500 if pihat_index is not None: 500 ↛ 508line 500 didn't jump to line 508 because the condition on line 500 was always true(empty)

501 if not include_pihat_in_mu: (empty)

502 max_split_mu = max_split_mu.at[pihat_index].set(0) (empty)

503 if not include_pihat_in_tau: 503 ↛ 508line 503 didn't jump to line 508 because the condition on line 503 was always true(empty)

504 # Block splits on propensity score for tau 

505 max_split_tau = max_split_tau.at[pihat_index].set(0) (empty)

506 

507 # 4. Initialize BCFState 

508 initial_state = init_bcf( (empty)

509 X_unified=x_train_binned, 

510 trt=z_train, 

511 y=y_train_internal, 

512 outcome_type=outcome_type, 

513 offset=offset_val, 

514 max_split_mu=max_split_mu, 

515 max_split_tau=max_split_tau, 

516 num_trees_mu=num_trees_mu, 

517 num_trees_tau=num_trees_tau, 

518 p_nonterminal_mu=p_nonterminal_mu, 

519 p_nonterminal_tau=p_nonterminal_tau, 

520 leaf_prior_cov_inv_mu=make_leaf_prior_cov_inv( 

521 leaf_prior_cov_inv_mu, 

522 sample_sigma2_leaf_mu, 

523 sigma2_leaf_shape_mu, 

524 sigma2_leaf_scale_mu, 

525 ), 

526 leaf_prior_cov_inv_tau=make_leaf_prior_cov_inv( 

527 leaf_prior_cov_inv_tau, 

528 sample_sigma2_leaf_tau, 

529 sigma2_leaf_shape_tau, 

530 sigma2_leaf_scale_tau, 

531 ), 

532 min_points_per_leaf_mu=min_points_per_leaf_mu, 

533 min_points_per_leaf_tau=min_points_per_leaf_tau, 

534 # ignore all predictors without splits, like `Bart(..., rm_const=True)` 

535 filter_splitless_vars_mu=jnp.sum(max_split_mu == 0).item(), 

536 filter_splitless_vars_tau=jnp.sum(max_split_tau == 0).item(), 

537 tau_0_prior_var=tau_0_prior_var, 

538 sample_intercept=sample_intercept, 

539 adaptive_coding=adaptive_coding, 

540 error_cov_inv=error_cov_inv, 

541 num_chains=num_chains, 

542 ) 

543 

544 # 5. Run the MCMC loop 

545 # WORKAROUND(python<3.12): once `run_mcmc` is generic over the state 

546 # subclass (PEP 695), `bcf_step` will type-check as its `step` and the 

547 # traces will come out typed, dropping the ignore and the casts. 

548 final_state, burnin_trace, main_trace = run_mcmc( (empty)

549 keys.pop(), 

550 initial_state, 

551 ndpost, 

552 n_burn=nskip, 

553 step=bcf_step, # ty: ignore[invalid-argument-type] 

554 burnin_trace_type=BCFBurninTrace, 

555 main_trace_type=BCFMainTrace, 

556 ) 

557 burnin_trace = cast(BCFBurninTrace, burnin_trace) (empty)

558 main_trace = cast(BCFMainTrace, main_trace) (empty)

559 self._mcmc_state = final_state (empty)

560 self._binner = binner (empty)

561 # fold the chains like `predict_latent` does to align the samples 

562 self._tau_0_trace = _fold_chains(main_trace, 'tau_0') (empty)

563 self._b_trace = _fold_chains(main_trace, 'b') (empty)

564 self._leaf_prior_cov_inv_mu_trace = _fold_chains( (empty)

565 main_trace, 'mu.leaf_prior_cov_inv' 

566 ) 

567 self._leaf_prior_cov_inv_tau_trace = _fold_chains( (empty)

568 main_trace, 'tau.leaf_prior_cov_inv' 

569 ) 

570 self._main_trace = {'mu': main_trace.mu, 'tau': main_trace.tau} (empty)

571 self._burnin_trace = {'mu': burnin_trace.mu, 'tau': burnin_trace.tau} (empty)

572 

573 # 6. Predict at the test points, now that the traces are available 

574 if x_test is not None: (empty)

575 test_pred = self.predict(x_test, pihat_test=pihat_test) (empty)

576 self._mu_test = test_pred['mu'] (empty)

577 self._tau_test = test_pred['tau'] (empty)

578 if z_test is not None: 578 ↛ exitline 578 didn't return from function '__init__' because the condition on line 578 was always true(empty)

579 self._yhat_test = self._mu_test + z_test * self._tau_test (empty)

580 

581 @classmethod (empty)

582 def _from_saved_state( (empty)

583 cls, 

584 *, 

585 binner: Any, # noqa: ANN401 

586 tau_0_trace: Any, # noqa: ANN401 

587 b_trace: Any, # noqa: ANN401 

588 main_trace: Any, # noqa: ANN401 

589 burnin_trace: Any = None, # noqa: ANN401 

590 mcmc_state: Any = None, # noqa: ANN401 

591 leaf_prior_cov_inv_mu_trace: Any = None, # noqa: ANN401 

592 leaf_prior_cov_inv_tau_trace: Any = None, # noqa: ANN401 

593 x_train_fmt: Any = None, # noqa: ANN401 

594 standardize: bool = False, 

595 y_mean: Float32[ArrayLike, ''] | float = 0.0, 

596 y_std: Float32[ArrayLike, ''] | float = 1.0, 

597 outcome_type: str = 'continuous', 

598 offset: Float32[ArrayLike, ''] | float = 0.0, 

599 mu_test: Float32[Array, 'num_samples m'] | None = None, 

600 tau_test: Float32[Array, 'num_samples m'] | None = None, 

601 yhat_test: Float32[Array, 'num_samples m'] | None = None, 

602 ) -> 'bcf': 

603 """ 

604 Private factory constructor to initialize bcf instance from restored state. 

605 

606 Parameters 

607 ---------- 

608 binner 

609 The binner instance for continuous predictor transforms. 

610 tau_0_trace 

611 Posterior trace of the tau_0 intercept. 

612 b_trace 

613 Posterior trace of the adaptive coding weights, columns b0 and b1. 

614 main_trace 

615 Posterior traces for mu and tau forests. 

616 burnin_trace 

617 Optional burn-in trace data. 

618 mcmc_state 

619 Optional final MCMC state. 

620 leaf_prior_cov_inv_mu_trace 

621 Optional prior variance trace for mu forest. 

622 leaf_prior_cov_inv_tau_trace 

623 Optional prior variance trace for tau forest. 

624 x_train_fmt 

625 Formatting metadata of training predictors. 

626 standardize 

627 Whether predictions should be unscaled back to original outcome units. 

628 y_mean 

629 Original training outcome mean for unscaling. 

630 y_std 

631 Original training outcome standard deviation for unscaling. 

632 outcome_type 

633 The regression target type ('continuous' or 'binary'). 

634 offset 

635 Probit latent scale offset (0.0 for continuous). 

636 mu_test 

637 Optional control mean at the test points. 

638 tau_test 

639 Optional treatment effect at the test points. 

640 yhat_test 

641 Optional outcome at the test points. 

642 

643 Returns 

644 ------- 

645 bcf 

646 A reconstituted `bcf` model ready for prediction. 

647 """ 

648 model = object.__new__(cls) (empty)

649 object.__setattr__(model, '_mcmc_state', mcmc_state) (empty)

650 object.__setattr__(model, '_binner', binner) (empty)

651 object.__setattr__(model, '_main_trace', main_trace) (empty)

652 object.__setattr__(model, '_burnin_trace', burnin_trace) (empty)

653 object.__setattr__(model, '_tau_0_trace', tau_0_trace) (empty)

654 object.__setattr__(model, '_b_trace', b_trace) (empty)

655 object.__setattr__( (empty)

656 model, '_leaf_prior_cov_inv_mu_trace', leaf_prior_cov_inv_mu_trace 

657 ) 

658 object.__setattr__( (empty)

659 model, '_leaf_prior_cov_inv_tau_trace', leaf_prior_cov_inv_tau_trace 

660 ) 

661 object.__setattr__(model, '_x_train_fmt', x_train_fmt) (empty)

662 object.__setattr__(model, '_standardize', standardize) (empty)

663 object.__setattr__(model, '_y_mean', jnp.asarray(y_mean)) (empty)

664 object.__setattr__(model, '_y_std', jnp.asarray(y_std)) (empty)

665 object.__setattr__(model, '_outcome_type', outcome_type) (empty)

666 object.__setattr__(model, '_offset', jnp.asarray(offset)) (empty)

667 object.__setattr__(model, '_mu_test', mu_test) (empty)

668 object.__setattr__(model, '_tau_test', tau_test) (empty)

669 object.__setattr__(model, '_yhat_test', yhat_test) (empty)

670 return model (empty)

671 

672 def save_npz(self, path: str | Path) -> None: (empty)

673 """ 

674 Save the loaded BCF traces to an NPZ archive. 

675 

676 Parameters 

677 ---------- 

678 path 

679 The file path to save the NPZ archive. 

680 """ 

681 state = {} (empty)

682 

683 # Explicit schema versioning 

684 state['schema_version'] = np.array(1) (empty)

685 

686 # Serialize binner attributes 

687 try: (empty)

688 max_split = np.asarray(self._binner.max_split) (empty)

689 except RuntimeError: # pragma: no cover 

690 # Array was donated to JAX during fit(), recover from state 

691 max_split = ( 

692 np.asarray(self._mcmc_state.forest_tau.max_split) 

693 if self._mcmc_state is not None 

694 else None 

695 ) 

696 

697 for k, v in _serialize_binner(self._binner, max_split=max_split).items(): (empty)

698 state[f'binner.{k}'] = v (empty)

699 

700 # Save scalar traces 

701 state['tau_0_trace'] = np.asarray(self._tau_0_trace) (empty)

702 state['b0_trace'] = np.asarray(self._b_trace[:, 0]) (empty)

703 state['b1_trace'] = np.asarray(self._b_trace[:, 1]) (empty)

704 

705 # Save standardization metadata 

706 state['standardize'] = np.array(self._standardize) (empty)

707 state['_y_mean'] = np.asarray(self._y_mean) (empty)

708 state['_y_std'] = np.asarray(self._y_std) (empty)

709 

710 # Save binary / probit metadata 

711 state['_outcome_type'] = np.array(self._outcome_type) (empty)

712 state['_offset'] = np.asarray(self._offset) (empty)

713 

714 # Save main_trace (mu and tau forests) based on dataclass fields 

715 for forest_name, trace in self._main_trace.items(): (empty)

716 state[f'main_trace.{forest_name}.class'] = trace.__class__.__name__ (empty)

717 for field_name in trace.__dataclass_fields__: (empty)

718 # Skip mesh because we cannot easily serialize JAX mesh obj 

719 if field_name == 'mesh': (empty)

720 continue (empty)

721 val = getattr(trace, field_name) (empty)

722 if val is not None: (empty)

723 state[f'main_trace.{forest_name}.{field_name}'] = np.asarray(val) (empty)

724 

725 # Save format string if any 

726 if self._x_train_fmt is not None: 726 ↛ 730line 726 didn't jump to line 730 because the condition on line 726 was always true(empty)

727 state['x_train_fmt'] = json.dumps(self._x_train_fmt) (empty)

728 

729 # Save test predictions, present only if `x_test` was passed 

730 for key in ('_mu_test', '_tau_test', '_yhat_test'): (empty)

731 val = getattr(self, key) (empty)

732 if val is not None: (empty)

733 state[key] = np.asarray(val) (empty)

734 

735 np.savez_compressed(path, allow_pickle=True, **state) (empty)

736 

737 @classmethod (empty)

738 def load_npz(cls, path: str | Path) -> 'bcf': (empty)

739 """ 

740 Load BCF traces from an NPZ archive, bypassing __init__ MCMC. 

741 

742 Parameters 

743 ---------- 

744 path 

745 The file path to the saved NPZ archive. 

746 

747 Returns 

748 ------- 

749 bcf 

750 A reconstituted `bcf` object ready for prediction. 

751 

752 Raises 

753 ------ 

754 ValueError 

755 If the schema version in the archive is unsupported. 

756 """ 

757 with np.load(path, allow_pickle=False) as data: (empty)

758 # Inspect schema version 

759 schema_version = int(data.get('schema_version', 1)) (empty)

760 if schema_version != 1: (empty)

761 msg = f'Unsupported schema version: {schema_version}' (empty)

762 raise ValueError(msg) (empty)

763 

764 # Reconstruct binner 

765 binner = _deserialize_binner(data) (empty)

766 

767 # Scalar traces 

768 tau_0_trace = jnp.asarray(data['tau_0_trace']) (empty)

769 b_trace = jnp.stack([data['b0_trace'], data['b1_trace']], axis=1) (empty)

770 

771 # Standardization metadata 

772 standardize = bool(data.get('standardize', False)) (empty)

773 y_mean = data.get('_y_mean', 0.0) (empty)

774 y_std = data.get('_y_std', 1.0) (empty)

775 outcome_type = str(data.get('_outcome_type', 'continuous')) (empty)

776 offset = data.get('_offset', 0.0) (empty)

777 

778 # Reconstruct main_trace respecting dataclass field defaults 

779 main_trace = {} (empty)

780 for forest_name in ['mu', 'tau']: (empty)

781 trace = object.__new__(MainTrace) (empty)

782 for field_name, field_def in MainTrace.__dataclass_fields__.items(): (empty)

783 key = f'main_trace.{forest_name}.{field_name}' (empty)

784 if key in data: (empty)

785 val = data[key] (empty)

786 if field_name == 'has_chains': (empty)

787 val = bool(val) (empty)

788 else: 

789 val = jnp.asarray(val) (empty)

790 object.__setattr__(trace, field_name, val) (empty)

791 else: 

792 # Respect dataclass default or default_factory if defined 

793 if field_def.default is not dataclasses.MISSING: (empty)

794 default_val = field_def.default # pragma: no cover 

795 elif field_def.default_factory is not dataclasses.MISSING: (empty)

796 default_val = ( 

797 field_def.default_factory() 

798 ) # pragma: no cover 

799 else: 

800 default_val = None (empty)

801 object.__setattr__(trace, field_name, default_val) (empty)

802 main_trace[forest_name] = trace (empty)

803 

804 fmt_str = str(data.get('x_train_fmt', 'None')) (empty)

805 x_train_fmt = None if fmt_str == 'None' else json.loads(fmt_str) (empty)

806 

807 mu_test, tau_test, yhat_test = ( (empty)

808 jnp.asarray(data[key]) if key in data else None 

809 for key in ('_mu_test', '_tau_test', '_yhat_test')