Coverage for src/bartz/bcf/_loop.py: 99%

1# bartz/src/bartz/bcf/_loop.py 

2# 

3# Copyright (c) 2026, The Bartz Contributors 

4# 

5# This file is part of bartz. 

6# 

7# Permission is hereby granted, free of charge, to any person obtaining a copy 

8# of this software and associated documentation files (the "Software"), to deal 

9# in the Software without restriction, including without limitation the rights 

10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

11# copies of the Software, and to permit persons to whom the Software is 

12# furnished to do so, subject to the following conditions: 

13# 

14# The above copyright notice and this permission notice shall be included in all 

15# copies or substantial portions of the Software. 

16# 

17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

23# SOFTWARE. 

24 

25"""Implement the BCF MCMC step and traces, to be run with `run_mcmc`.""" 

26 

27from dataclasses import replace (empty)

28from typing import cast (empty)

29 

30import jax.numpy as jnp (empty)

31from equinox import tree_at (empty)

32from jax import lax, random, vmap (empty)

33from jaxtyping import Array, Float, Float32, Int32, Key, UInt (empty)

34 

35from bartz._jaxext import field, float32_matmuls, jit, sliced_map, split (empty)

36from bartz.bcf._state import BCFState (empty)

37from bartz.mcmcloop._trace import BurninTrace, MainTrace, Trace (empty)

38from bartz.mcmcstep._axes import CHAIN_AXIS (empty)

39from bartz.mcmcstep._state import ( (empty)

40 Forest, 

41 State, 

42 StepConfig, 

43 split_key_for_chains, 

44 vmap_chains, 

45) 

46from bartz.mcmcstep._step import step, step_leaf_prior_cov_inv, step_trees, sum_resid (empty)

47 

48 

49def recompute_prec_trees( (empty)

50 forest: Forest, prec_scale: Float[Array, ' n'], config: StepConfig 

51) -> Float32[Array, 'num_trees tree_size']: 

52 """ 

53 Rebuild `Forest.prec_tree` from scratch for a new `prec_scale`. 

54 

55 Like the incremental update in `step`, the result is valid at the leaves of 

56 the largest version of each tree and at the node of the last move. 

57 

58 Parameters 

59 ---------- 

60 forest 

61 The forest whose per-leaf precision cache is stale. 

62 prec_scale 

63 The new per-datapoint precision scale. 

64 config 

65 The MCMC configuration, for the reduction settings. 

66 

67 Returns 

68 ------- 

69 The per-leaf sums of `prec_scale`. 

70 """ 

71 _, tree_size = forest.leaf_tree.shape (empty)

72 

73 def one_tree( (empty)

74 args: tuple[UInt[Array, ' n'], Int32[Array, '']], 

75 ) -> Float32[Array, ' tree_size']: 

76 leaf_indices, move_node = args (empty)

77 # the sum over all bins has the shape of the residual reduction, so use 

78 # its settings rather than the ones tuned for the two-bin `prec_tree` update 

79 tree = sum_resid( (empty)

80 prec_scale, 

81 leaf_indices, 

82 tree_size, 

83 config.resid_reduction_config, 

84 config.data_sharded, 

85 ) 

86 children = 2 * move_node + jnp.arange(2) (empty)

87 return tree.at[move_node].set(tree[children].sum()) (empty)

88 

89 xs = (forest.leaf_indices, forest.move_node) (empty)

90 

91 def all_trees() -> Float32[Array, 'num_trees tree_size']: (empty)

92 return vmap(one_tree)(xs) (empty)

93 

94 if config.prec_count_num_trees is None: 94 ↛ 99line 94 didn't jump to line 99 because the condition on line 94 was always true(empty)

95 return all_trees() (empty)

96 

97 else: 

98 

99 def tree_batches( (empty)

100 batch_size: int = config.prec_count_num_trees, 

101 ) -> Float32[Array, 'num_trees tree_size']: 

102 return sliced_map(one_tree, xs, batch_size=batch_size) (empty)

103 

104 # like `compute_prec_trees`, batch the trees on cpu to bound the reduction 

105 # temporaries 

106 return lax.platform_dependent(cpu=tree_batches, cuda=all_trees) (empty)

107 

108 

109def bcf_step_mu(key: Key[Array, ''], state: BCFState) -> BCFState: (empty)

110 """Update the prognostic forest and its leaf prior precision.""" 

111 # `step` rebuilds the state with `replace`, so it preserves the subclass. 

112 # WORKAROUND(python<3.12): type `step` as generic over the state subclass 

113 # (PEP 695) instead of casting here, since a TypeVar renders badly in the 

114 # html documentation. 

115 return cast(BCFState, step(key, state)) (empty)

116 

117 

118def bcf_step_tau_0(key: Key[Array, ''], state: BCFState) -> BCFState: (empty)

119 """Update the treatment effect intercept.""" 

120 # `resid` is stored scaled: ``resid_unit * resid = data residual``, whereas 

121 # `tau_0`, `tau_X`, `b`, `error_cov_inv` and `tau_0_prior_cov_inv` are 

122 # on the data scale (matching `_bcf.predict`). Convert `resid` in and out of data 

123 # units so the scalar Gibbs updates here and in `bcf_step_b` are 

124 # unit-consistent for any `resid_unit` (no-op when it is 1). 

125 

126 if state.tau_0_prior_cov_inv is None: (empty)

127 return state (empty)

128 

129 else: 

130 # get coding basis, possibly adaptive so not just 0 and 1 

131 b_z = state.b[state.trt.astype(int)] (empty)

132 

133 # partial residual removing current tau_0 effect, on the data scale 

134 partial_resid = state.resid * state.resid_unit + state.tau_0 * b_z (empty)

135 

136 # determine full conditional of tau_0 

137 prec = ( (empty)

138 jnp.sum(jnp.square(b_z)) * state.error_cov_inv.value 

139 + state.tau_0_prior_cov_inv 

140 ) 

141 mean = jnp.sum(b_z * partial_resid) * state.error_cov_inv.value / prec (empty)

142 

143 # sample tau_0 from full conditional 

144 tau_0_new = mean + random.normal(key) * lax.rsqrt(prec) (empty)

145 

146 # update state to reflect new tau_0 

147 return replace( (empty)

148 state, 

149 tau_0=tau_0_new, 

150 resid=state.resid - b_z * (tau_0_new - state.tau_0) / state.resid_unit, 

151 ) 

152 

153 

154def bcf_step_tau(key: Key[Array, ''], state: BCFState) -> BCFState: (empty)

155 """Update the treatment effect forest and its leaf prior precision.""" 

156 keys = split(key, 2) (empty)

157 

158 # get coding basis, possibly adaptive so not just 0 and 1 

159 b_z = state.b[state.trt.astype(int)] (empty)

160 

161 # Target for tau is (Y - mu - b_z * tau_0) / b_z. 

162 # Its residual is target - tau = (Y - mu - b_z*tau_0 - b_z*tau) / b_z 

163 b_z_zero = jnp.abs(b_z) < 1e-10 (empty)

164 b_z_safe = jnp.where(b_z_zero, 1.0, b_z) (empty)

165 mu_resid = state.resid (empty)

166 initial_resid_tau = jnp.where(b_z_zero, 0.0, mu_resid / b_z_safe) (empty)

167 

168 # Swap the tau forest into the forest slot and run only the tree and leaf 

169 # prior precision steps on it; the mu forest rides along in `forest_tau` 

170 # and is swapped back at the end. The other sub-steps of `step` (latent 

171 # outcome, error precision, sparsity, step counter) belong to the mu phase 

172 # alone. 

173 mu_prec_scale = state.prec_scale (empty)

174 state = replace( (empty)

175 state, 

176 forest=state.forest_tau, 

177 forest_tau=state.forest, 

178 resid=initial_resid_tau, 

179 prec_scale=jnp.square(b_z), 

180 ) 

181 

182 state = cast(BCFState, step_trees(keys.pop(), state)) (empty)

183 state = cast(BCFState, step_leaf_prior_cov_inv(keys.pop(), state)) (empty)

184 

185 # Update tau_X! (the residual difference is scaled, bring it to data units) 

186 if state.tau_X is not None: (empty)

187 state = replace( (empty)

188 state, 

189 tau_X=state.tau_X + (initial_resid_tau - state.resid) * state.resid_unit, 

190 ) 

191 

192 # Swap the forests back and restore the mu-side fields 

193 return replace( (empty)

194 state, 

195 forest=state.forest_tau, 

196 forest_tau=state.forest, 

197 resid=jnp.where(b_z_zero, mu_resid, state.resid * b_z_safe), 

198 prec_scale=mu_prec_scale, 

199 ) 

200 

201 

202def bcf_step_b(key: Key[Array, ''], state: BCFState) -> BCFState: (empty)

203 """Update the adaptive coding weights.""" 

204 if state.b_prior_cov_inv is None: (empty)

205 return state (empty)

206 

207 else: 

208 assert state.tau_X is not None (empty)

209 

210 # get coding basis 

211 b_z = state.b[state.trt.astype(int)] (empty)

212 

213 # partial residual removing current b effect, on the data scale (see 

214 # `bcf_step_tau_0` about units) 

215 tau_full = state.tau_0 + state.tau_X (empty)

216 partial_resid = state.resid * state.resid_unit + tau_full * b_z (empty)

217 

218 # determine full conditional of b, one Gibbs update per treatment group 

219 groups = jnp.stack([~state.trt, state.trt]) (empty)

220 prec = ( (empty)

221 jnp.sum(jnp.square(tau_full) * groups, axis=1) * state.error_cov_inv.value 

222 + state.b_prior_cov_inv 

223 ) 

224 mean = ( (empty)

225 jnp.sum(tau_full * partial_resid * groups, axis=1) 

226 * state.error_cov_inv.value 

227 / prec 

228 ) 

229 

230 # sample b from full conditional 

231 b_new = mean + random.normal(key, (2,)) * lax.rsqrt(prec) (empty)

232 b_z_new = b_new[state.trt.astype(int)] (empty)

233 

234 # update state to reflect new b 

235 state = replace( (empty)

236 state, 

237 b=b_new, 

238 resid=state.resid - tau_full * (b_z_new - b_z) / state.resid_unit, 

239 ) 

240 

241 # the tau precision scale b_z^2 changed on every datapoint, so the 

242 # incrementally maintained per-leaf cache of the tau forest is stale 

243 # everywhere 

244 return tree_at( (empty)

245 lambda s: s.forest_tau.prec_tree, 

246 state, 

247 recompute_prec_trees(state.forest_tau, jnp.square(b_z_new), state.config), 

248 ) 

249 

250 

251@jit(donate_argnums=(1,)) (empty)

252@split_key_for_chains (empty)

253@vmap_chains (empty)

254@float32_matmuls (empty)

255def bcf_step(key: Key[Array, ''], state: BCFState) -> BCFState: (empty)

256 """ 

257 Do one BCF MCMC step. 

258 

259 Parameters 

260 ---------- 

261 key 

262 A jax random key. 

263 state 

264 A BCF mcmc state, as created by `init_bcf`. 

265 

266 Returns 

267 ------- 

268 The new BCF mcmc state. 

269 

270 Notes 

271 ----- 

272 The memory of the input state is re-used for the output state, so the input 

273 state can not be used any more after calling `bcf_step`. All this applies 

274 outside of `jax.jit`. 

275 """ 

276 keys = split(key, 4) (empty)

277 state = bcf_step_mu(keys.pop(), state) (empty)

278 state = bcf_step_tau_0(keys.pop(), state) (empty)

279 state = bcf_step_tau(keys.pop(), state) (empty)

280 return bcf_step_b(keys.pop(), state) (empty)

281 

282 

283def _tau_view(state: BCFState) -> BCFState: (empty)

284 """Return the state with the tau forest in the mu forest slot.""" 

285 return replace(state, forest=state.forest_tau, forest_tau=state.forest) (empty)

286 

287 

288class BCFBurninTrace(Trace): (empty)

289 """Burn-in trace of the BCF MCMC, the per-forest diagnostics and the scalar parameters.""" 

290 

291 mu: BurninTrace 

292 """The trace of the prognostic forest.""" (empty)

293 

294 tau: BurninTrace 

295 """The trace of the treatment forest.""" (empty)

296 

297 tau_0: Float32[Array, '*chains_and_samples'] = field(chains=CHAIN_AXIS, samples=0) (empty)

298 """The treatment effect intercept.""" (empty)

299 

300 b: Float32[Array, '*chains_and_samples 2'] = field(chains=CHAIN_AXIS, samples=0) (empty)

301 """The adaptive coding weights for untreated and treated units.""" (empty)

302 

303 @classmethod (empty)

304 def from_state(cls, state: State) -> 'BCFBurninTrace': (empty)

305 """Create a single-item burn-in trace from a BCF state.""" 

306 assert isinstance(state, BCFState) (empty)

307 return cls( (empty)

308 mu=BurninTrace.from_state(state), 

309 tau=BurninTrace.from_state(_tau_view(state)), 

310 tau_0=state.tau_0, 

311 b=state.b, 

312 ) 

313 

314 

315class BCFMainTrace(BCFBurninTrace): (empty)

316 """Main trace of the BCF MCMC, with the trees of both forests.""" 

317 

318 mu: MainTrace 

319 """The trace of the prognostic forest.""" (empty)

320 

321 tau: MainTrace 

322 """The trace of the treatment forest.""" (empty)

323 

324 @classmethod (empty)

325 def from_state(cls, state: State) -> 'BCFMainTrace': (empty)

326 """Create a single-item main trace from a BCF state.""" 

327 assert isinstance(state, BCFState) (empty)

328 kw: dict = dict( (empty)

329 vars(BCFBurninTrace.from_state(state)), 

330 mu=MainTrace.from_state(state), 

331 tau=MainTrace.from_state(_tau_view(state)), 

332 ) 

333 return cls(**kw) (empty)