Coverage for src / bartz / mcmcstep / _state.py: 95%

1# bartz/src/bartz/mcmcstep/_state.py 

2# 

3# Copyright (c) 2024-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"""Module defining the BART MCMC state and initialization.""" 

26 

27from collections.abc import Callable, Hashable, Sequence (empty)

28from dataclasses import fields, replace (empty)

29from enum import Enum (empty)

30from functools import partial, wraps (empty)

31from math import log2 (empty)

32from typing import Any, Literal, TypedDict, TypeVar (empty)

33 

34import numpy (empty)

35from equinox import Module, error_if, filter_jit (empty)

36from equinox import field as eqx_field (empty)

37from jax import ( (empty)

38 NamedSharding, 

39 device_put, 

40 eval_shape, 

41 jit, 

42 lax, 

43 make_mesh, 

44 random, 

45 tree, 

46 vmap, 

47) 

48from jax import numpy as jnp (empty)

49from jax.scipy.linalg import solve_triangular (empty)

50from jax.sharding import AxisType, Mesh, PartitionSpec (empty)

51from jaxtyping import Array, Bool, Float32, Int32, Integer, PyTree, Shaped, UInt (empty)

52 

53from bartz.grove import tree_depths (empty)

54from bartz.jaxext import get_default_device, is_key, minimal_unsigned_dtype (empty)

55 

56 

57class OutcomeType(Enum): (empty)

58 """Whether the regression outcome is continuous or binary (probit).""" 

59 

60 continuous = 'continuous' (empty)

61 binary = 'binary' (empty)

62 

63 

64def field(*, chains: bool = False, data: bool = False, **kwargs: Any): # noqa: ANN202 (empty)

65 """Extend `equinox.field` with two new parameters. 

66 

67 Parameters 

68 ---------- 

69 chains 

70 Whether the arrays in the field have an optional first axis that 

71 represents independent Markov chains. 

72 data 

73 Whether the last axis of the arrays in the field represent units of 

74 the data. 

75 **kwargs 

76 Other parameters passed to `equinox.field`. 

77 

78 Returns 

79 ------- 

80 A dataclass field descriptor with the special attributes in the metadata, unset if False. 

81 """ 

82 metadata = dict(kwargs.pop('metadata', {})) (empty)

83 assert 'chains' not in metadata (empty)

84 assert 'data' not in metadata (empty)

85 if chains: (empty)

86 metadata['chains'] = True (empty)

87 if data: (empty)

88 metadata['data'] = True (empty)

89 return eqx_field(metadata=metadata, **kwargs) (empty)

90 

91 

92def chain_vmap_axes(x: PyTree[Module | Any, 'T']) -> PyTree[int | None, 'T']: (empty)

93 """Determine vmapping axes for chains. 

94 

95 This function determines the argument to the `in_axes` or `out_axes` 

96 parameter of `jax.vmap` to vmap over all and only the chain axes found in the 

97 pytree `x`. 

98 

99 Parameters 

100 ---------- 

101 x 

102 A pytree. Subpytrees that are Module attributes marked with 

103 ``field(..., chains=True)`` are considered to have a leading chain axis. 

104 

105 Returns 

106 ------- 

107 A pytree with the same structure as `x` with 0 or None in the leaves. 

108 """ 

109 return _find_metadata(x, 'chains', 0, None) 2 ctx1ab

110 

111 

112def data_vmap_axes(x: PyTree[Module | Any, 'T']) -> PyTree[int | None, 'T']: (empty)

113 """Determine vmapping axes for data. 

114 

115 This is analogous to `chain_vmap_axes` but returns -1 for all fields 

116 marked with ``field(..., data=True)``. 

117 """ 

118 return _find_metadata(x, 'data', -1, None) 2 ctx1ab

119 

120 

121T = TypeVar('T') (empty)

122 

123 

124def _find_metadata( (empty)

125 x: PyTree[Any, ' S'], key: Hashable, if_true: T, if_false: T 

126) -> PyTree[T, ' S']: 

127 """Replace all subtrees of x marked with a metadata key.""" 

128 

129 def is_lazy_array(x: object) -> bool: 2 ctx1ab

130 return isinstance(x, _LazyArray) 2 ctx1ab

131 

132 def is_module(x: object) -> bool: 2 ctx1ab

133 return isinstance(x, Module) and not is_lazy_array(x) 2 ctx1ab

134 

135 if is_module(x): 2 ctx1ab

136 args = [] 2 ctx1ab

137 for f in fields(x): 2 ctx1ab

138 v = getattr(x, f.name) 2 ctx1ab

139 if f.metadata.get('static', False): 2 ctx1ab

140 args.append(v) 2 ctx1ab

141 elif f.metadata.get(key, False): 2 ctx1ab

142 subtree = tree.map(lambda _: if_true, v, is_leaf=is_lazy_array) 2 ctx1ab

143 args.append(subtree) 2 ctx1ab

144 else: 

145 args.append(_find_metadata(v, key, if_true, if_false)) 2 ctx1ab

146 return x.__class__(*args) 2 ctx1ab

147 

148 def get_axes(x: object) -> PyTree[T]: 2 ctx1ab

149 if is_module(x): 3 ctx1aeb

150 return _find_metadata(x, key, if_true, if_false) 2 ctx1ae

151 else: 

152 return tree.map(lambda _: if_false, x, is_leaf=is_lazy_array) 2 ctx1ab

153 

154 def is_leaf(x: object) -> bool: 2 ctx1ab

155 return isinstance(x, Module) # this catches _LazyArray as well 2 ctx1ab

156 

157 return tree.map(get_axes, x, is_leaf=is_leaf) 2 ctx1ab

158 

159 

160class Forest(Module): (empty)

161 """Represents the MCMC state of a sum of trees.""" 

162 

163 leaf_tree: ( (empty)

164 Float32[Array, '*chains num_trees 2**d'] 

165 | Float32[Array, '*chains num_trees k 2**d'] 

166 ) = field(chains=True) 

167 """The leaf values.""" (empty)

168 

169 var_tree: UInt[Array, '*chains num_trees 2**(d-1)'] = field(chains=True) (empty)

170 """The decision axes.""" (empty)

171 

172 split_tree: UInt[Array, '*chains num_trees 2**(d-1)'] = field(chains=True) (empty)

173 """The decision boundaries.""" (empty)

174 

175 affluence_tree: Bool[Array, '*chains num_trees 2**(d-1)'] = field(chains=True) (empty)

176 """Marks leaves that can be grown.""" (empty)

177 

178 max_split: UInt[Array, ' p'] 

179 """The maximum split index for each predictor.""" (empty)

180 

181 blocked_vars: UInt[Array, ' q'] | None 

182 """Indices of variables that are not used. This shall include at least (empty)

183 the `i` such that ``max_split[i] == 0``, otherwise behavior is 

184 undefined.""" 

185 

186 p_nonterminal: Float32[Array, ' 2**d'] 

187 """The prior probability of each node being nonterminal, conditional on (empty)

188 its ancestors. Includes the nodes at maximum depth which should be set 

189 to 0.""" 

190 

191 p_propose_grow: Float32[Array, ' 2**(d-1)'] 

192 """The unnormalized probability of picking a leaf for a grow proposal.""" (empty)

193 

194 leaf_indices: UInt[Array, '*chains num_trees n'] = field(chains=True, data=True) (empty)

195 """The index of the leaf each datapoints falls into, for each tree.""" (empty)

196 

197 min_points_per_decision_node: Int32[Array, ''] | None 

198 """The minimum number of data points in a decision node.""" (empty)

199 

200 min_points_per_leaf: Int32[Array, ''] | None 

201 """The minimum number of data points in a leaf node.""" (empty)

202 

203 log_trans_prior: Float32[Array, '*chains num_trees'] | None = field(chains=True) (empty)

204 """The log transition and prior Metropolis-Hastings ratio for the (empty)

205 proposed move on each tree.""" 

206 

207 log_likelihood: Float32[Array, '*chains num_trees'] | None = field(chains=True) (empty)

208 """The log likelihood ratio.""" (empty)

209 

210 grow_prop_count: Int32[Array, '*chains'] = field(chains=True) (empty)

211 """The number of grow proposals made during one full MCMC cycle.""" (empty)

212 

213 prune_prop_count: Int32[Array, '*chains'] = field(chains=True) (empty)

214 """The number of prune proposals made during one full MCMC cycle.""" (empty)

215 

216 grow_acc_count: Int32[Array, '*chains'] = field(chains=True) (empty)

217 """The number of grow moves accepted during one full MCMC cycle.""" (empty)

218 

219 prune_acc_count: Int32[Array, '*chains'] = field(chains=True) (empty)

220 """The number of prune moves accepted during one full MCMC cycle.""" (empty)

221 

222 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'] | None 

223 """The prior precision matrix of a leaf, conditional on the tree structure. (empty)

224 For the univariate case (k=1), this is a scalar (the inverse variance). 

225 The prior covariance of the sum of trees is 

226 ``num_trees * leaf_prior_cov_inv^-1``.""" 

227 

228 log_s: Float32[Array, '*chains p'] | None = field(chains=True) (empty)

229 """The logarithm of the prior probability for choosing a variable to split (empty)

230 along in a decision rule, conditional on the ancestors. Not normalized. 

231 If `None`, use a uniform distribution.""" 

232 

233 theta: Float32[Array, '*chains'] | None = field(chains=True) (empty)

234 """The concentration parameter for the Dirichlet prior on the variable (empty)

235 distribution `s`. Required only to update `log_s`.""" 

236 

237 a: Float32[Array, ''] | None 

238 """Parameter of the prior on `theta`. Required only to sample `theta`. (empty)

239 See `step_theta`.""" 

240 

241 b: Float32[Array, ''] | None 

242 """Parameter of the prior on `theta`. Required only to sample `theta`. (empty)

243 See `step_theta`.""" 

244 

245 rho: Float32[Array, ''] | None 

246 """Parameter of the prior on `theta`. Required only to sample `theta`. (empty)

247 See `step_theta`.""" 

248 

249 def num_chains(self) -> int | None: (empty)

250 """Return the number of chains, or `None` if not multichain.""" 

251 # maybe this should be replaced by chain_shape() -> () | (int,) 

252 if self.var_tree.ndim == 2: 4 ctx1areo

253 return None 2 ctx1ro

254 else: 

255 return self.var_tree.shape[0] 2 ctx1ae

256 

257 

258class StepConfig(Module): (empty)

259 """Options for the MCMC step.""" 

260 

261 steps_done: Int32[Array, ''] 

262 """The number of MCMC steps completed so far.""" (empty)

263 

264 sparse_on_at: Int32[Array, ''] | None 

265 """After how many steps to turn on variable selection.""" (empty)

266 

267 resid_num_batches: int | None = field(static=True) (empty)

268 """The number of batches for computing the sum of residuals. If (empty)

269 `None`, they are computed with no batching.""" 

270 

271 count_num_batches: int | None = field(static=True) (empty)

272 """The number of batches for computing counts. If (empty)

273 `None`, they are computed with no batching.""" 

274 

275 prec_num_batches: int | None = field(static=True) (empty)

276 """The number of batches for computing precision scales. If (empty)

277 `None`, they are computed with no batching.""" 

278 

279 prec_count_num_trees: int | None = field(static=True) (empty)

280 """Batch size for processing trees to compute count and prec trees.""" (empty)

281 

282 mesh: Mesh | None = field(static=True) (empty)

283 """The mesh used to shard data and computation across multiple devices.""" (empty)

284 

285 

286class State(Module): (empty)

287 """Represents the MCMC state of BART.""" 

288 

289 X: UInt[Array, 'p n'] = field(data=True) (empty)

290 """The predictors.""" (empty)

291 

292 binary_y: None | Bool[Array, ' n'] | Bool[Array, 'k n'] = field(data=True) (empty)

293 """The response as booleans for binary regression, `None` for continuous. (empty)

294 In the mixed binary-continuous case, only the binary outcome components 

295 are stored, with shape ``(kb, n)``.""" 

296 

297 z: None | Float32[Array, '*chains n'] | Float32[Array, '*chains k n'] = field( (empty)

298 chains=True, data=True 

299 ) 

300 """The latent variable for binary regression. `None` in continuous (empty)

301 regression. In the mixed binary-continuous case, only the binary outcome 

302 components are stored, with shape ``(*chains, kb, n)``.""" 

303 

304 binary_indices: None | Int32[Array, ' kb'] 

305 """The indices of binary outcome components in the full list of outcome (empty)

306 components. `None` when there are no binary components. Filled in by 

307 `init` and used by `step_z` to update only the binary rows of `resid`.""" 

308 

309 offset: Float32[Array, ''] | Float32[Array, ' k'] 

310 """Constant shift added to the sum of trees.""" (empty)

311 

312 resid: Float32[Array, '*chains n'] | Float32[Array, '*chains k n'] = field( (empty)

313 chains=True, data=True 

314 ) 

315 """The residuals (`y` or `z` minus sum of trees).""" (empty)

316 

317 error_cov_inv: Float32[Array, '*chains'] | Float32[Array, '*chains k k'] = field( (empty)

318 chains=True 

319 ) 

320 """The inverse error covariance (scalar for univariate, matrix for multivariate). (empty)

321 Identity in binary regression.""" 

322 

323 prec_scale: Float32[Array, ' n'] | None = field(data=True) (empty)

324 """The scale on the error precision, i.e., ``1 / error_scale ** 2``. (empty)

325 `None` in binary regression.""" 

326 

327 error_cov_df: Float32[Array, ''] | None 

328 """The df parameter of the inverse Wishart prior on the noise (empty)

329 covariance. For the univariate case, the relationship to the inverse 

330 gamma prior parameters is ``alpha = df / 2``. 

331 `None` in binary regression.""" 

332 

333 error_cov_scale: Float32[Array, ''] | Float32[Array, 'k k'] | None 

334 """The scale parameter of the inverse Wishart prior on the noise (empty)

335 covariance. For the univariate case, the relationship to the inverse 

336 gamma prior parameters is ``beta = scale / 2``. 

337 `None` in binary regression.""" 

338 

339 forest: Forest 

340 """The sum of trees model.""" (empty)

341 

342 config: StepConfig 

343 """Metadata and configurations for the MCMC step.""" (empty)

344 

345 

346def _init_shape_shifting_parameters( (empty)

347 y: Float32[Array, ' n'] | Float32[Array, 'k n'], 

348 outcome_type: OutcomeType | list[OutcomeType], 

349 offset: Float32[Array, ''] | Float32[Array, ' k'], 

350 error_scale: Float32[Any, ' n'] | None, 

351 error_cov_df: float | Float32[Any, ''] | None, 

352 error_cov_scale: float | Float32[Any, ''] | Float32[Any, 'k k'] | None, 

353 leaf_prior_cov_inv: Float32[Array, ''] | Float32[Array, 'k k'], 

354) -> tuple[ 

355 bool, 

356 tuple[()] | tuple[int], 

357 None | Float32[Array, ''], 

358 None | Float32[Array, ''], 

359 None | Float32[Array, ''], 

360 None | Int32[Array, ' kb'], 

361]: 

362 """ 

363 Check and initialize parameters that change array type/shape based on outcome kind. 

364 

365 Parameters 

366 ---------- 

367 y 

368 The response variable (used only for shape checks). 

369 outcome_type 

370 Whether the regression is continuous or binary. Can be a list of 

371 `OutcomeType` for per-component specification in the multivariate case. 

372 offset 

373 The offset to add to the predictions. 

374 error_scale 

375 Per-observation error scale (univariate only). 

376 error_cov_df 

377 The error covariance degrees of freedom. 

378 error_cov_scale 

379 The error covariance scale. 

380 leaf_prior_cov_inv 

381 The inverse of the leaf prior covariance. 

382 

383 Returns 

384 ------- 

385 is_binary 

386 Whether all outcomes are binary. 

387 kshape 

388 The outcome shape, empty for univariate, (k,) for multivariate. 

389 error_cov_inv 

390 The initialized error covariance inverse. 

391 error_cov_df 

392 The error covariance degrees of freedom (as array). 

393 error_cov_scale 

394 The error covariance scale (as array). 

395 binary_indices 

396 The indices of binary outcome components, or `None` if there are none. 

397 """ 

398 kshape = offset.shape 2 ctx1ab

399 

400 # determine per-component outcome kinds 

401 if isinstance(outcome_type, list): 5 ctx1afgbh

402 assert kshape, 'per-component outcome_type requires multivariate y' 3 ctx1fgh

403 (k,) = kshape 3 ctx1fgh

404 assert len(outcome_type) == k 4 ctx1fguh

405 binary_mask = [t is OutcomeType.binary for t in outcome_type] 4 ctx1fguh

406 is_binary = all(binary_mask) 3 ctx1fgh

407 is_mixed = any(binary_mask) and not is_binary 4 ctx1fguh

408 else: 

409 is_binary = outcome_type is OutcomeType.binary 2 ctx1ab

410 is_mixed = False 2 ctx1ab

411 

412 if is_mixed: 6 ctx1afgbuh

413 binary_indices = jnp.array([i for i, b in enumerate(binary_mask) if b]) 3 ctx1fgh

414 else: 

415 binary_indices = None 2 ctx1ab

416 

417 # All-binary 

418 if is_binary: 7 ctx1adfgbhk

419 assert error_scale is None 2 ctx1dk

420 assert error_cov_df is None 2 ctx1dk

421 assert error_cov_scale is None 2 ctx1dk

422 if kshape: 3 ctx1dAk

423 error_cov_inv = jnp.eye(kshape[0]) 1 ctx1A

424 else: 

425 error_cov_inv = jnp.array(1.0) 2 ctx1dk

426 

427 # Mixed binary-continuous (multivariate, diagonal error covariance) 

428 elif is_mixed: 5 ctx1afgbh

429 assert error_scale is None, ( 4 ctx1fgBh

430 'error_scale is not supported for mixed binary-continuous' 

431 ) 

432 error_cov_df = jnp.asarray(error_cov_df) 4 ctx1fgBh

433 error_cov_scale = jnp.asarray(error_cov_scale) 3 ctx1fgh

434 assert error_cov_scale.shape == 2 * kshape 3 ctx1fgh

435 

436 # enforce diagonal error_cov_scale 

437 diag = jnp.diag(jnp.diag(error_cov_scale)) 3 ctx1fgh

438 error_cov_scale = error_if( 3 ctx1fgh

439 error_cov_scale, 

440 jnp.any(error_cov_scale != diag), 

441 'error_cov_scale must be diagonal for mixed binary-continuous', 

442 ) 

443 

444 # initialize diagonal error_cov_inv: use inv-gamma mode for continuous 

445 # components, 1.0 for binary components 

446 scale_diag = jnp.diag(error_cov_scale) 3 ctx1fgh

447 inv_diag = jnp.where( 3 ctx1fgh

448 jnp.array(binary_mask), 

449 1.0, 

450 error_cov_df / jnp.where(scale_diag, scale_diag, 1.0), 

451 ) 

452 error_cov_inv = jnp.diag(inv_diag) 3 ctx1fgh

453 

454 # All-continuous 

455 else: 

456 error_cov_df = jnp.asarray(error_cov_df) 2 ctx1ab

457 error_cov_scale = jnp.asarray(error_cov_scale) 2 ctx1ab

458 assert error_cov_scale.shape == 2 * kshape 2 ctx1ab

459 

460 # Multivariate vs univariate 

461 if kshape: 4 ctx1almb

462 error_cov_inv = error_cov_df * _inv_via_chol_with_gersh(error_cov_scale) 2 ctx1lm

463 else: 

464 # inverse gamma prior: alpha = df / 2, beta = scale / 2 

465 error_cov_inv = error_cov_df / error_cov_scale 2 ctx1ab

466 

467 assert y.shape[:-1] == kshape 2 ctx1ab

468 assert leaf_prior_cov_inv.shape == 2 * kshape 2 ctx1ab

469 

470 return ( 2 ctx1ab

471 is_binary, 

472 kshape, 

473 error_cov_inv, 

474 error_cov_df, 

475 error_cov_scale, 

476 binary_indices, 

477 ) 

478 

479 

480def _check_splitless_vars( (empty)

481 filter_splitless_vars: int, 

482 max_split: UInt[Array, ' p'], 

483 offset: Float32[Array, ''] | Float32[Array, ' k'], 

484) -> Float32[Array, ''] | Float32[Array, ' k']: 

485 """Check there aren't too many deactivated predictors.""" 

486 msg = ( 2 ctx1ab

487 f'there are more than {filter_splitless_vars=} predictors with no splits, ' 

488 'please increase `filter_splitless_vars` or investigate the missing splits' 

489 ) 

490 return error_if(offset, jnp.sum(max_split == 0) > filter_splitless_vars, msg) 2 ctx1ab

491 

492 

493def _parse_outcome_type( (empty)

494 outcome_type: 'OutcomeType | str | Sequence[OutcomeType | str]', 

495) -> 'OutcomeType | list[OutcomeType]': 

496 """Normalize outcome_type to enum (or list of enums).""" 

497 if isinstance(outcome_type, Sequence) and not isinstance(outcome_type, str): 7 ctx1afpgqbh

498 return [OutcomeType(t) for t in outcome_type] 3 ctx1fgh

499 else: 

500 return OutcomeType(outcome_type) 4 ctx1apqb

501 

502 

503def _parse_p_nonterminal( (empty)

504 p_nonterminal: Float32[Any, ' d_minus_1'], 

505) -> Float32[Array, ' d_minus_1+1']: 

506 """Check it's in (0, 1) and pad with a 0 at the end.""" 

507 p_nonterminal = jnp.asarray(p_nonterminal) 2 ctx1ab

508 ok = (p_nonterminal > 0) & (p_nonterminal < 1) 2 ctx1ab

509 p_nonterminal = error_if(p_nonterminal, ~ok, 'p_nonterminal must be in (0, 1)') 2 ctx1ab

510 return jnp.pad(p_nonterminal, (0, 1)) 2 ctx1ab

511 

512 

513def make_p_nonterminal( (empty)

514 d: int, 

515 alpha: float | Float32[Array, ''] = 0.95, 

516 beta: float | Float32[Array, ''] = 2.0, 

517) -> Float32[Array, ' {d}-1']: 

518 """Prepare the `p_nonterminal` argument to `init`. 

519 

520 It is calculated according to the formula: 

521 

522 P_nt(depth) = alpha / (1 + depth)^beta, with depth 0-based 

523 

524 Parameters 

525 ---------- 

526 d 

527 The maximum depth of the trees (d=1 means tree with only root node) 

528 alpha 

529 The a priori probability of the root node having children, conditional 

530 on it being possible 

531 beta 

532 The exponent of the power decay of the probability of having children 

533 with depth. 

534 

535 Returns 

536 ------- 

537 An array of probabilities, one per tree level but the last. 

538 """ 

539 assert d >= 1 2 ctx1ab

540 depth = jnp.arange(d - 1) 2 ctx1ab

541 return alpha / (1 + depth).astype(float) ** beta 2 ctx1ab

542 

543 

544class _LazyArray(Module): (empty)

545 """Like `functools.partial` but specialized to array-creating functions like `jax.numpy.zeros`.""" 

546 

547 array_creator: Callable 

548 shape: tuple[int, ...] 

549 args: tuple 

550 

551 def __init__( (empty)

552 self, array_creator: Callable, shape: tuple[int, ...], *args: Any 

553 ) -> None: 

554 self.array_creator = array_creator 2 ctx1ab

555 self.shape = shape 2 ctx1ab

556 self.args = args 2 ctx1ab

557 

558 def __call__(self, **kwargs: Any) -> T: (empty)

559 return self.array_creator(self.shape, *self.args, **kwargs) 2 ctx1ab

560 

561 @property (empty)

562 def ndim(self) -> int: (empty)

563 return len(self.shape) 3 ctx1jdi

564 

565 

566def init( (empty)

567 *, 

568 X: UInt[Any, 'p n'], 

569 y: Float32[Any, ' n'] | Float32[Any, ' k n'], 

570 outcome_type: OutcomeType | str | Sequence[OutcomeType | str] = 'continuous', 

571 offset: float | Float32[Any, ''] | Float32[Any, ' k'], 

572 max_split: UInt[Any, ' p'], 

573 num_trees: int, 

574 p_nonterminal: Float32[Any, ' d_minus_1'], 

575 leaf_prior_cov_inv: float | Float32[Any, ''] | Float32[Array, 'k k'], 

576 error_cov_df: float | Float32[Any, ''] | None = None, 

577 error_cov_scale: float | Float32[Any, ''] | Float32[Array, 'k k'] | None = None, 

578 error_scale: Float32[Any, ' n'] | None = None, 

579 min_points_per_decision_node: int | Integer[Any, ''] | None = None, 

580 resid_num_batches: int | None | Literal['auto'] = 'auto', 

581 count_num_batches: int | None | Literal['auto'] = 'auto', 

582 prec_num_batches: int | None | Literal['auto'] = 'auto', 

583 prec_count_num_trees: int | None | Literal['auto'] = 'auto', 

584 save_ratios: bool = False, 

585 filter_splitless_vars: int = 0, 

586 min_points_per_leaf: int | Integer[Any, ''] | None = None, 

587 log_s: Float32[Any, ' p'] | None = None, 

588 theta: float | Float32[Any, ''] | None = None, 

589 a: float | Float32[Any, ''] | None = None, 

590 b: float | Float32[Any, ''] | None = None, 

591 rho: float | Float32[Any, ''] | None = None, 

592 sparse_on_at: int | Integer[Any, ''] | None = None, 

593 num_chains: int | None = None, 

594 mesh: Mesh | dict[str, int] | None = None, 

595 target_platform: Literal['cpu', 'gpu'] | None = None, 

596) -> State: 

597 """ 

598 Make a BART posterior sampling MCMC initial state. 

599 

600 Parameters 

601 ---------- 

602 X 

603 The predictors. Note this is trasposed compared to the usual convention. 

604 y 

605 The response. If two-dimensional, the outcome is multivariate with the 

606 first axis indicating the component. For binary data, non-zero means 1, 

607 zero means 0. 

608 outcome_type 

609 Whether the regression is continuous or binary (probit). Can also be a 

610 sequence of `OutcomeType` values, one per outcome component, for mixed 

611 binary-continuous multivariate regression. 

612 offset 

613 Constant shift added to the sum of trees. 0 if not specified. 

614 max_split 

615 The maximum split index for each variable. All split ranges start at 1. 

616 num_trees 

617 The number of trees in the forest. 

618 p_nonterminal 

619 The probability of a nonterminal node at each depth. The maximum depth 

620 of trees is fixed by the length of this array. Use `make_p_nonterminal` 

621 to set it with the conventional formula. 

622 leaf_prior_cov_inv 

623 The prior precision matrix of a leaf, conditional on the tree structure. 

624 For the univariate case (k=1), this is a scalar (the inverse variance). 

625 The prior covariance of the sum of trees is 

626 ``num_trees * leaf_prior_cov_inv^-1``. The prior mean of leaves is 

627 always zero. 

628 error_cov_df 

629 error_cov_scale 

630 The df and scale parameters of the inverse Wishart prior on the error 

631 covariance. For the univariate case, the relationship to the inverse 

632 gamma prior parameters is ``alpha = df / 2``, ``beta = scale / 2``. 

633 Leave unspecified for binary regression. 

634 error_scale 

635 Each error is scaled by the corresponding factor in `error_scale`, so 

636 the error variance for ``y[i]`` is ``sigma2 * error_scale[i] ** 2``. 

637 Not supported for binary regression. If not specified, defaults to 1 for 

638 all points, but potentially skipping calculations. 

639 min_points_per_decision_node 

640 The minimum number of data points in a decision node. 0 if not 

641 specified. 

642 resid_num_batches 

643 count_num_batches 

644 prec_num_batches 

645 The number of batches, along datapoints, for summing the residuals, 

646 counting the number of datapoints in each leaf, and computing the 

647 likelihood precision in each leaf, respectively. `None` for no batching. 

648 If 'auto', it's chosen automatically based on the target platform; see 

649 the description of `target_platform` below for how it is determined. 

650 prec_count_num_trees 

651 The number of trees to process at a time when counting datapoints or 

652 computing the likelihood precision. If `None`, do all trees at once, 

653 which may use too much memory. If 'auto' (default), it's chosen 

654 automatically. 

655 save_ratios 

656 Whether to save the Metropolis-Hastings ratios. 

657 filter_splitless_vars 

658 The maximum number of variables without splits that can be ignored. If 

659 there are more, `init` raises an exception. 

660 min_points_per_leaf 

661 The minimum number of datapoints in a leaf node. 0 if not specified. 

662 Unlike `min_points_per_decision_node`, this constraint is not taken into 

663 account in the Metropolis-Hastings ratio because it would be expensive 

664 to compute. Grow moves that would violate this constraint are vetoed. 

665 This parameter is independent of `min_points_per_decision_node` and 

666 there is no check that they are coherent. It makes sense to set 

667 ``min_points_per_decision_node >= 2 * min_points_per_leaf``. 

668 log_s 

669 The logarithm of the prior probability for choosing a variable to split 

670 along in a decision rule, conditional on the ancestors. Not normalized. 

671 If not specified, use a uniform distribution. If not specified and 

672 `theta` or `rho`, `a`, `b` are, it's initialized automatically. 

673 theta 

674 The concentration parameter for the Dirichlet prior on `s`. Required 

675 only to update `log_s`. If not specified, and `rho`, `a`, `b` are 

676 specified, it's initialized automatically. 

677 a 

678 b 

679 rho 

680 Parameters of the prior on `theta`. Required only to sample `theta`. 

681 sparse_on_at 

682 After how many MCMC steps to turn on variable selection. 

683 num_chains 

684 The number of independent MCMC chains to represent in the state. Single 

685 chain with scalar values if not specified. 

686 mesh 

687 A jax mesh used to shard data and computation across multiple devices. 

688 If it has a 'chains' axis, that axis is used to shard the chains. If it 

689 has a 'data' axis, that axis is used to shard the datapoints. 

690 

691 As a shorthand, if a dictionary mapping axis names to axis size is 

692 passed, the corresponding mesh is created, e.g., ``dict(chains=4, 

693 data=2)`` will let jax pick 8 devices to split chains (which must be a 

694 multiple of 4) across 4 pairs of devices, where in each pair the data is 

695 split in two. 

696 

697 Note: if a mesh is passed, the arrays are always sharded according to 

698 it. In particular even if the mesh has no 'chains' or 'data' axis, the 

699 arrays will be replicated on all devices in the mesh. 

700 target_platform 

701 Platform ('cpu' or 'gpu') used to determine the number of batches 

702 automatically. If `mesh` is specified, the platform is inferred from the 

703 devices in the mesh. Otherwise, if `y` is a concrete array (i.e., `init` 

704 is not invoked in a `jax.jit` context), the platform is set to the 

705 platform of `y`. Otherwise, use `target_platform`. 

706 

707 To avoid confusion, in all cases where the `target_platform` argument 

708 would be ignored, `init` raises an exception if `target_platform` is 

709 set. 

710 

711 Returns 

712 ------- 

713 An initialized BART MCMC state. 

714 

715 Raises 

716 ------ 

717 ValueError 

718 If arguments unused in binary regression are set. 

719 

720 Notes 

721 ----- 

722 In decision nodes, the values in ``X[i, :]`` are compared to a cutpoint out 

723 of the range ``[1, 2, ..., max_split[i]]``. A point belongs to the left 

724 child iff ``X[i, j] < cutpoint``. Thus it makes sense for ``X[i, :]`` to be 

725 integers in the range ``[0, 1, ..., max_split[i]]``. 

726 

727 In general the arrays passed to this function as arguments may be donated, 

728 invalidating them. Create copies before passing them to `init` if this 

729 happens and you need them again. 

730 """ 

731 # convert to array all array-like arguments that are used in other 

732 # configurations but don't need further processing themselves 

733 X = jnp.asarray(X) 2 ctx1ab

734 y = jnp.asarray(y) 2 ctx1ab

735 assert y.dtype == jnp.float32 2 ctx1ab

736 offset = jnp.asarray(offset) 2 ctx1ab

737 leaf_prior_cov_inv = jnp.asarray(leaf_prior_cov_inv) 2 ctx1ab

738 max_split = jnp.asarray(max_split) 2 ctx1ab

739 

740 # normalize outcome_type to enum (or list of enums) 

741 outcome_type = _parse_outcome_type(outcome_type) 2 ctx1ab

742 

743 # check p_nonterminal and pad it with a 0 at the end (still not final shape) 

744 p_nonterminal = _parse_p_nonterminal(p_nonterminal) 2 ctx1ab

745 

746 # process arguments that change depending on outcome type 

747 is_binary, kshape, error_cov_inv, error_cov_df, error_cov_scale, binary_indices = ( 2 ctx1ab

748 _init_shape_shifting_parameters( 

749 y, 

750 outcome_type, 

751 offset, 

752 error_scale, 

753 error_cov_df, 

754 error_cov_scale, 

755 leaf_prior_cov_inv, 

756 ) 

757 ) 

758 

759 # extract array sizes from arguments 

760 (max_depth,) = p_nonterminal.shape 2 ctx1ab

761 p, n = X.shape 2 ctx1ab

762 

763 # check and initialize sparsity parameters 

764 if not _all_none_or_not_none(rho, a, b): 764 ↛ 765line 764 didn't jump to line 765 because the condition on line 764 was never true2 ctx1ab

765 msg = 'rho, a, b are not either all `None` or all set' 

766 raise ValueError(msg) 

767 if theta is None and rho is not None: 7 ctx1satCpqb

768 theta = rho 3 ctx1apq

769 if log_s is None and theta is not None: 10 ctx1sadtDECpqb

770 log_s = jnp.zeros(max_split.size) 3 ctx1apq

771 if not _all_none_or_not_none(theta, sparse_on_at): 771 ↛ 772line 771 didn't jump to line 772 because the condition on line 771 was never true6 ctx1sadDEb

772 msg = 'sparsity params (either theta or rho,a,b) and sparse_on_at must be either all None or all set' 

773 raise ValueError(msg) 

774 

775 # process multichain settings 

776 chain_shape = () if num_chains is None else (num_chains,) 4 ctx1areb

777 resid_shape = chain_shape + y.shape 4 ctx1areb

778 add_chains = partial(_add_chains, chain_shape=chain_shape) 2 ctx1ab

779 

780 # determine batch sizes for reductions 

781 mesh = _parse_mesh(num_chains, mesh) 2 ctx1ab

782 target_platform = _parse_target_platform( 2 ctx1ab

783 y, mesh, target_platform, resid_num_batches, count_num_batches, prec_num_batches 

784 ) 

785 red_cfg = _parse_reduction_configs( 2 ctx1ab

786 resid_num_batches, 

787 count_num_batches, 

788 prec_num_batches, 

789 prec_count_num_trees, 

790 y, 

791 num_trees, 

792 mesh, 

793 target_platform, 

794 ) 

795 

796 # check there aren't too many deactivated predictors 

797 offset = _check_splitless_vars(filter_splitless_vars, max_split, offset) 2 ctx1ab

798 

799 # determine shapes for trees 

800 tree_shape = (*chain_shape, num_trees) 2 ctx1ab

801 tree_size = 2**max_depth 2 ctx1ab

802 

803 # initialize all remaining stuff and put it in an unsharded state 

804 state = State( 5 ctx1adpqb

805 X=X, 

806 binary_y=y, # temporary to be sharded together with everything else 

807 z=( 

808 _LazyArray(jnp.full, resid_shape, offset[..., None]) 

809 if is_binary 

810 else _LazyArray( 

811 jnp.full, 

812 (*chain_shape, binary_indices.size, n), 

813 offset[binary_indices, None], 

814 ) 

815 if binary_indices is not None 

816 else None 

817 ), 

818 binary_indices=binary_indices, 

819 offset=offset, 

820 resid=( 

821 _LazyArray(jnp.zeros, resid_shape) 

822 if is_binary 

823 else None # resid is created later after y and offset are sharded 

824 ), 

825 error_cov_inv=add_chains(error_cov_inv), 

826 prec_scale=error_scale, # temporarily set to error_scale, fix after sharding 

827 error_cov_df=error_cov_df, 

828 error_cov_scale=error_cov_scale, 

829 forest=Forest( 

830 leaf_tree=_LazyArray( 

831 jnp.zeros, (*tree_shape, *kshape, tree_size), jnp.float32 

832 ), 

833 var_tree=_LazyArray( 

834 jnp.zeros, (*tree_shape, tree_size // 2), minimal_unsigned_dtype(p - 1) 

835 ), 

836 split_tree=_LazyArray( 

837 jnp.zeros, (*tree_shape, tree_size // 2), max_split.dtype 

838 ), 

839 affluence_tree=_LazyArray( 

840 _initial_affluence_tree, 

841 (*tree_shape, tree_size // 2), 

842 n, 

843 min_points_per_decision_node, 

844 ), 

845 blocked_vars=_get_blocked_vars(filter_splitless_vars, max_split), 

846 max_split=max_split, 

847 grow_prop_count=_LazyArray(jnp.zeros, chain_shape, int), 

848 grow_acc_count=_LazyArray(jnp.zeros, chain_shape, int), 

849 prune_prop_count=_LazyArray(jnp.zeros, chain_shape, int), 

850 prune_acc_count=_LazyArray(jnp.zeros, chain_shape, int),