Coverage for src/lsqfitgp/bayestree/_bcf.py: 79%

1# lsqfitgp/bayestree/_bcf.py 

2# 

3# Copyright (c) 2023, 2024, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

7# lsqfitgp is free software: you can redistribute it and/or modify 

8# it under the terms of the GNU General Public License as published by 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# lsqfitgp is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>. 

19 

20import functools 3 ctx1cab

21import warnings 3 ctx1cab

22 

23import gvar 3 ctx1cab

24import numpy 3 ctx1cab

25from jax import jit 3 ctx1cab

26from jax import numpy as jnp 3 ctx1cab

27 

28from lsqfitgp import ( 3 ctx1cab

29 _GP, 

30 _array, 

31 _fastraniter, 

32 _fit, 

33 _gvarext, 

34 _jaxext, 

35 _kernels, 

36 _utils, 

37 copula, 

38) 

39 

40 

41def _recursive_cast(dtype, default, mapping): 3 ctx1cab

42 if dtype in mapping: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true2 ctx1ab

43 return mapping[dtype] 

44 elif dtype.names is not None: 2 ctx1ab

45 return numpy.dtype( 2 ctx1ab

46 [ 

47 (name, _recursive_cast(dtype[name], default, mapping)) 

48 for name in dtype.names 

49 ] 

50 ) 

51 elif dtype.subdtype is not None: 51 ↛ 53line 51 didn't jump to line 53 because the condition on line 51 was never true2 ctx1ab

52 # note: has names => does not have subdtype 

53 return numpy.dtype((_recursive_cast(dtype.base, default, mapping), dtype.shape)) 

54 elif default is None: 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true2 ctx1ab

55 return dtype 

56 else: 

57 return default 2 ctx1ab

58 

59 

60def cast(dtype, default, mapping={}): # noqa: B006, read only 3 ctx1cab

61 """ 

62 Recursively cast a numpy data type. 

63 

64 Parameters 

65 ---------- 

66 dtype : dtype 

67 The data type to cast. 

68 default : dtype or None 

69 The leaf fields of `dtype` are casted to `default`, which can be 

70 structured, unless they appear in `mapping`. If None, dtypes not in 

71 `mapping` are left unchanged. 

72 mapping : dict 

73 A dictionary from dtypes to dtypes, indicating specific casting rules. 

74 The dtypes can be structured, a match of a structured dtype takes 

75 precedence over matches in its leaves, and the converted dtype is not 

76 further searched for matches. 

77 

78 Returns 

79 ------- 

80 casted_dtype : dtype 

81 The casted version of `dtype`. May not have the same structure if 

82 `mapping` contains structured dtypes. 

83 """ 

84 mapping = {numpy.dtype(k): numpy.dtype(v) for k, v in mapping.items()} 2 ctx1ab

85 default = None if default is None else numpy.dtype(default) 2 ctx1ab

86 return _recursive_cast(numpy.dtype(dtype), default, mapping) 2 ctx1ab

87 

88 

89class bcf: 3 ctx1cab

90 def __init__( # noqa: C901, PLR0915 3 ctx1cab

91 self, 

92 *, 

93 y, 

94 z, 

95 x_mu, 

96 x_tau=None, 

97 pihat, 

98 include_pi='mu', 

99 weights=None, 

100 fitkw={}, # noqa: B006, read only 

101 kernelkw_mu={}, # noqa: B006, read only 

102 kernelkw_tau={}, # noqa: B006, read only 

103 marginalize_mean=True, 

104 gpaux=None, 

105 x_aux=None, 

106 otherhp={}, # noqa: B006, read only 

107 transf='standardize', 

108 ): 

109 r""" 

110 Nonparametric Bayesian regression with a GP version of BCF. 

111 

112 BCF (Bayesian Causal Forests) is a regression method for observational 

113 causal inference studies introduced in [1]_ based on a pair of BART 

114 models. 

115 

116 This class evaluates a Gaussian process regression with a kernel which 

117 accurately approximates BCF in the infinite trees limit of each BART 

118 model. The hyperparameters are optimized to their marginal MAP. 

119 

120 The model is (loosely, see notes below) :math:`y = \mu(x) + z\tau(x)`, 

121 so :math:`\tau(x)` is the expected causal effect of :math:`z` on 

122 :math:`y` at location :math:`x`. 

123 

124 Parameters 

125 ---------- 

126 y : (n,) array, series or dataframe 

127 Outcome. 

128 z : (n,) array, series or dataframe 

129 Binary treatment status: 0 control group, 1 treatment group. 

130 x_mu : (n, p) array, series or dataframe 

131 Covariates for the :math:`\mu` model. 

132 x_tau : (n, q) array, series or dataframe, optional 

133 Covariates for the :math:`\tau` model. If not specified, use `x_mu`. 

134 pihat : (n,) array, series or dataframe 

135 Estimated propensity score, i.e., P(Z=1|X). 

136 include_pi : {'mu', 'tau', 'both'}, optional 

137 Whether to include the propensity score in the :math:`\mu` model, 

138 the :math:`\tau` model, or both. Default is ``'mu'``. 

139 weights : (n,) array, series or dataframe 

140 Weights used to rescale the error variance (as 1 / weight). 

141 fitkw : dict 

142 Additional arguments passed to `~lsqfitgp.empbayes_fit`, overrides 

143 the defaults. 

144 kernelkw_mu, kernelkw_tau : dict 

145 Additional arguments passed to `~lsqfitgp.BART` for each model, 

146 overrides the defaults. 

147 marginalize_mean : bool 

148 If True (default), marginalize the intercept of the model. 

149 gpaux : callable, optional 

150 If specified, this function is called with a pair ``(hp, gp)``, 

151 where ``hp`` is a dictionary of hyperparameters, and ``gp`` is a 

152 `~lsqfitgp.GP` object under construction, and is expected to return 

153 a modified ``gp`` with a new process named ``'aux'`` defined with 

154 `~lsqfitgp.GP.defproc` or similar. The process is added to the 

155 regression model. The input to the process is a structured array 

156 with fields: 

157 

158 'train' : bool 

159 Indicates whether the data is training set (the one passed on 

160 initialization) or test set (the one passed to `pred` or `gp`). 

161 'i' : int 

162 Index of the flattened array. 

163 'z' : int 

164 Treatment status. 

165 'mu', 'tau' : structured 

166 The values in `x_mu` and `x_tau`, converted to indices according 

167 to the BART grids. Where `pihat` has been added, there are two 

168 subfields: ``'x'`` which contains the covariates, and 

169 ``'pihat'``, the latter expressed in indices as well. 

170 'pihat' : float 

171 The `pihat` argument. Contrary to the subfield included under 

172 ``'mu'`` and/or ``'tau'``, this field contains the original 

173 values. 

174 'aux' : structured 

175 The values in `x_aux`, if specified. 

176 

177 x_aux : (n, k) array, series or dataframe, optional 

178 Additional covariates for the ``'aux'`` process. 

179 otherhp : dictionary of gvar 

180 A dictionary with the prior of arbitrary additional hyperpameters, 

181 intended to be used by ``gpaux`` or ``transf``. 

182 transf : (list of) str or pair of callable 

183 Data transformation. Either a string indicating a pre-defined 

184 transformation, or a pair ``(from_data, to_data)``, two functions 

185 with signatures ``from_data(hp, y) -> eta`` and ``to_data(hp, eta) 

186 -> y``, where ``eta`` is the value to which the model is fit, and 

187 ``hp`` is the dictionary of hyperparameters. The functions must be 

188 ufuncs and one the inverse of the other w.r.t. the second parameter. 

189 ``from_data`` must be derivable with `jax` w.r.t. ``y``. 

190 

191 If a list of such specifications is provided, the transformations 

192 are applied in order, with the first one being the outermost, i.e., 

193 the one applied first to the data. 

194 

195 If a transformation uses additional hyperparameters, either 

196 predefined automatically or passed by the user through `otherhp`, 

197 they are inferred with the rest of the hyperparameters. 

198 

199 The pre-defined transformations are: 

200 

201 'standardize' (default) 

202 eta = (y - mean(train_y)) / sdev(train_y) 

203 'yeojohnson' 

204 The Yeo-Johnson transformation [2]_ to reduce skewness. The 

205 :math:`\lambda` parameter is bounded in :math:`(0, 2)` 

206 for implementation convenience, this restriction may be lifted 

207 in future versions. 

208 

209 Notes 

210 ----- 

211 The regression model is: 

212 

213 .. math:: 

214 \eta_i = g(y_i; \ldots) &= m + {} \\ 

215 &\phantom{{}={}} + 

216 \lambda_\mu 

217 \mu(\mathbf x^\mu_i, \hat\pi_i?) + {} \\ 

218 &\phantom{{}={}} + 

219 \lambda_\tau 

220 \tau(\mathbf x^\tau_i, \hat\pi_i?) (z_i - z_0) + {} \\ 

221 &\phantom{{}={}} + 

222 \mathrm{aux}(i, z_i, \mathbf x^\mu_i, \mathbf x^\tau_i, 

223 \hat\pi_i, \mathbf x^\text{aux}_i) + {} \\ 

224 &\phantom{{}={}} + 

225 \varepsilon_i, \\ 

226 \varepsilon_i &\sim 

227 N(0, \sigma^2 / w_i), \\ 

228 m &\sim N(0, 1), \\ 

229 \log \sigma^2 &\sim N(\log\bar w, 4), \\ 

230 \lambda_\mu 

231 &\sim \mathrm{HalfCauchy}(2), \\ 

232 \lambda_\tau 

233 &\sim \mathrm{HalfNormal}(1.48), \\ 

234 \mu &\sim \mathrm{GP}(0, 

235 \mathrm{BART}(\alpha_\mu, \beta_\mu) ), \\ 

236 \tau &\sim \mathrm{GP}(0, 

237 \mathrm{BART}(\alpha_\tau, \beta_\tau) ), \\ 

238 \mathrm{aux} & \sim \mathrm{GP}(0, \text{<user defined>}), \\ 

239 \alpha_\mu, \alpha_\tau &\sim \mathrm{Beta}(2, 1), \\ 

240 \beta_\mu, \beta_\tau &\sim \mathrm{InvGamma}(1, 1), \\ 

241 z_0 &\sim U(0, 1), 

242 

243 To make the inference, :math:`(\mu, \tau, \boldsymbol\varepsilon, m, 

244 \mathrm{aux})` are marginalized analytically, and the marginal posterior 

245 mode of 

246 :math:`(\sigma, \lambda_*, \alpha_*, \beta_*, z_0, \ldots)` is found by 

247 numerical minimization, after transforming them to express their prior 

248 as a Gaussian copula. Their marginal posterior covariance matrix is 

249 estimated with an approximation of the hessian inverse. See 

250 `~lsqfitgp.empbayes_fit` and use the parameter ``fitkw`` to customize 

251 this procedure. 

252 

253 The tree splitting grid of the BART kernel is set using quantiles of the 

254 observed covariates. This corresponds to settings ``usequants=True``, 

255 ``numcut=inf`` in the R packages BayesTree and BART. Use the parameters 

256 `kernelkw_mu` and `kernelkw_tau` to customize the grids. 

257 

258 The difference between the regression model evaluated at :math:`Z=1` vs. 

259 :math:`Z=0` can be interpreted as the causal effect :math:`Z \rightarrow 

260 Y` if the unconfoundedness assumption is made: 

261 

262 .. math:: 

263 \{Y(Z=0), Y(Z=1)\} \perp\!\!\!\perp Z \mid X. 

264 

265 In practical terms, this holds when: 

266 

267 1) :math:`X` are pre-treatment variables, i.e., they represent 

268 quantities causally upstream of :math:`Z`. 

269 

270 2) :math:`X` are sufficient to adjust for all common causes of 

271 :math:`Z` and :math:`Y`, such that the only remaining difference 

272 is the causal effect and not just a correlation. 

273 

274 Here :math:`X` consists in `x_tau`, `x_mu` and `x_aux`. However these 

275 arrays may also be used to pass "technical" values used to set up the 

276 model, that do not satisfy the uncounfoundedness assumption, if you know 

277 what you are doing. 

278 

279 Attributes 

280 ---------- 

281 m : float or gvar 

282 The prior mean :math:`m`. 

283 sigma : gvar 

284 The error term standard deviation :math:`\sigma`. If there are 

285 weights, the sdev for each unit is obtained dividing ``sigma`` by 

286 sqrt(weight). 

287 alpha_mu, alpha_tau : gvar 

288 The numerator of the tree spawn probability :math:`\alpha_*` (named 

289 ``base`` in R bcf). 

290 beta_mu, beta_tau : gvar 

291 The depth exponent of the tree spawn probability :math:`\beta_*` 

292 (named ``power`` in R bcf). 

293 lambda_mu, lambda_tau : gvar 

294 The prior standard deviation :math:`\lambda_*`. 

295 z_0 : gvar 

296 The treatment coding parameter. 

297 fit : empbayes_fit 

298 The hyperparameters fit object. 

299 

300 Methods 

301 ------- 

302 gp : 

303 Create a GP object. 

304 data : 

305 Creates the dictionary to be passed to `GP.pred` to represent data. 

306 pred : 

307 Evaluate the regression function at given locations. 

308 from_data : 

309 Convert :math:`y` to :math:`\eta`. 

310 to_data : 

311 Convert :math:`\eta` to :math:`y`. 

312 

313 See Also 

314 -------- 

315 lsqfitgp.BART 

316 

317 References 

318 ---------- 

319 .. [1] P. Richard Hahn, Jared S. Murray, Carlos M. Carvalho "Bayesian 

320 Regression Tree Models for Causal Inference: Regularization, 

321 Confounding, and Heterogeneous Effects (with Discussion)," Bayesian 

322 Analysis 15(3), 965-1056, September 2020, 

323 https://doi.org/10.1214/19-BA1195 

324 .. [2] Yeo, In-Kwon; Johnson, Richard A. (2000). "A New Family of Power 

325 Transformations to Improve Normality or Symmetry". Biometrika. 87 

326 (4): 954–959. https://doi.org/10.1093/biomet/87.4.954 

327 """ # noqa: DOC001 

328 # convert covariates to StructuredArray 

329 x_mu = self._to_structured(x_mu) 2 ctx1ab

330 if x_tau is not None: 2 ctx1ab

331 x_tau = self._to_structured(x_tau) 1 ctx1b

332 assert x_tau.shape == x_mu.shape 1 ctx1b

333 if x_aux is not None: 2 ctx1ab

334 x_aux = self._to_structured(x_aux) 1 ctx1b

335 assert x_aux.shape == x_mu.shape 1 ctx1b

336 

337 # convert outcomes, treatment, propensity score, weights to 1d arrays 

338 y = self._to_vector(y) 2 ctx1ab

339 z = self._to_vector(z) 2 ctx1ab

340 pihat = self._to_vector(pihat) 2 ctx1ab

341 assert y.shape == z.shape == pihat.shape == x_mu.shape 2 ctx1ab

342 if weights is not None: 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true2 ctx1ab

343 weights = self._to_vector(weights) 

344 assert weights.shape == x_mu.shape 

345 

346 # check include_pi 

347 if include_pi not in ('mu', 'tau', 'both'): 347 ↛ 348line 347 didn't jump to line 348 because the condition on line 347 was never true2 ctx1ab

348 msg = f'invalid value include_pi={include_pi!r}' 

349 raise KeyError(msg) 

350 self._include_pi = include_pi 2 ctx1ab

351 

352 # add pihat to covariates 

353 x_mu, x_tau = self._append_pihat(x_mu, x_tau, pihat) 2 ctx1ab

354 

355 # grid and indices 

356 splits_mu = _kernels.BART.splits_from_coord(x_mu) 2 ctx1ab

357 i_mu = self._toindices(x_mu, splits_mu) 2 ctx1ab

358 if x_tau is None: 2 ctx1ab

359 splits_tau = splits_mu 2 ctx1ab

360 i_tau = None 2 ctx1ab

361 else: 

362 splits_tau = _kernels.BART.splits_from_coord(x_tau) 1 ctx1b

363 i_tau = self._toindices(x_tau, splits_tau) 1 ctx1b

364 

365 # get functions for data transformation 

366 from_data, to_data, transfloss, transfhp = self._get_transf( 2 ctx1ab

367 transf=transf, weights=weights, y=y 

368 ) 

369 

370 # scale of error variance 

371 logsigma2_loc = 0 if weights is None else numpy.log(jnp.mean(weights)) 2 ctx1ab

372 

373 # prior on hyperparams 

374 hyperprior = copula.makedict( 2 ctx1ab

375 { 

376 'm': gvar.gvar(0, 1), 

377 'sigma^2': copula.lognorm(logsigma2_loc, 2), 

378 'lambda_mu': copula.halfcauchy(2), 

379 'lambda_tau': copula.halfnorm(1.48), 

380 'alpha_mu': copula.beta(2, 1), 

381 'alpha_tau': copula.beta(2, 1), 

382 'beta_mu': copula.invgamma(1, 1), 

383 'beta_tau': copula.invgamma(1, 1), 

384 'z_0': copula.uniform(0, 1), 

385 } 

386 ) 

387 

388 # remove explicit mean parameter if it's baked into the Gaussian process 

389 if marginalize_mean: 2 ctx1ab

390 hyperprior.pop('m') 2 ctx1ab

391 

392 # add data transformation and user hyperparameters 

393 def update_hyperparams(new, newname, raises): 2 ctx1ab

394 new = gvar.BufferDict(new) 2 ctx1ab

395 for key in new.all_keys(): 2 ctx1ab

396 if hyperprior.has_dictkey(key): 396 ↛ 397line 396 didn't jump to line 397 because the condition on line 396 was never true2 ctx1ab

397 message = f'{newname} hyperparameter {key!r} overrides existing one' 

398 if raises: 

399 raise ValueError(message) 

400 else: 

401 warnings.warn(message) 

402 hyperprior.update(new) 2 ctx1ab

403 

404 update_hyperparams(transfhp, 'data transformation', True) 2 ctx1ab

405 # the hypers handed by _get_transf are not allowed to override 

406 update_hyperparams(otherhp, 'user', False) 2 ctx1ab

407 

408 # GP factory 

409 def gpfactory( 2 ctx1ab

410 hp, *, z, i_mu, i_tau, pihat, x_aux, weights, splits_mu, splits_tau, **_ 

411 ): 

412 

413 kw_overridable = dict(maxd=10, reset=[2, 4, 6, 8], intercept=False) 2 ctx1ab

414 kw_not_overridable = dict(indices=True) 2 ctx1ab

415 

416 gp = _GP.GP(checkpos=False, checksym=False, solver='chol') 2 ctx1ab

417 

418 splits = dict(mu=splits_mu, tau=splits_tau) 2 ctx1ab

419 for name, kernelkw in dict(mu=kernelkw_mu, tau=kernelkw_tau).items(): 2 ctx1ab

420 kw = dict( 2 ctx1ab

421 alpha=hp[f'alpha_{name}'], 

422 beta=hp[f'beta_{name}'], 

423 dim=name, 

424 splits=splits[name], 

425 **kw_overridable, 

426 ) 

427 kw.update(kernelkw) 2 ctx1ab

428 kernel = _kernels.BART(**kw, **kw_not_overridable) 2 ctx1ab

429 kernel *= hp[f'lambda_{name}'] ** 2 2 ctx1ab

430 

431 gp = gp.defproc(name, kernel) 2 ctx1ab

432 

433 if 'm' in hp: 2 ctx1ab

434 kernel_mean = 0 * _kernels.Constant() 1 ctx1b

435 else: 

436 kernel_mean = _kernels.Constant() 2 ctx1ab

437 gp = gp.defproc('m', kernel_mean) 2 ctx1ab

438 

439 if gpaux is None: 2 ctx1ab

440 gp = gp.defproc('aux', 0 * _kernels.Constant()) 2 ctx1ab

441 else: 

442 gp = gpaux(hp, gp) 1 ctx1b

443 

444 gp = gp.deflintransf( 2 ctx1ab

445 gp.DefaultProcess, 

446 lambda m, mu, tau, aux: ( 

447 lambda x: m(x) + mu(x) + tau(x) * (x['z'] - hp['z_0']) + aux(x) 

448 ), 

449 ['m', 'mu', 'tau', 'aux'], 

450 ) 

451 

452 x = self._join_points(True, z, i_mu, i_tau, pihat, x_aux) 2 ctx1ab

453 gp = gp.addx(x, 'trainmean') 2 ctx1ab

454 errcov = self._error_cov(hp, weights, x) 2 ctx1ab

455 return gp.addcov(errcov, 'trainnoise').addtransf( 2 ctx1ab

456 {'trainmean': 1, 'trainnoise': 1}, 'train' 

457 ) 

458 

459 # data factory 

460 def data(hp, *, y, **_): 2 ctx1ab

461 return {'train': from_data(hp, y) - hp.get('m', 0)} 2 ctx1ab

462 

463 # fit hyperparameters 

464 options = dict( 2 ctx1ab

465 verbosity=3, 

466 minkw=dict(method='l-bfgs-b', options=dict(maxls=4, maxiter=100)), 

467 mlkw=dict(epsrel=0), 

468 forward=True, 

469 gpfactorykw=dict( 

470 y=y, 

471 z=z, 

472 i_mu=i_mu, 

473 i_tau=i_tau, 

474 pihat=pihat, 

475 x_aux=x_aux, 

476 weights=weights, 

477 splits_mu=splits_mu, 

478 splits_tau=splits_tau, 

479 ), 

480 additional_loss=transfloss, 

481 ) 

482 options.update(fitkw) 2 ctx1ab

483 fit = _fit.empbayes_fit(hyperprior, gpfactory, data, **options) 2 ctx1ab

484 

485 # extract hyperparameters from minimization result 

486 self.m = fit.p.get('m', 0) 2 ctx1ab

487 self.sigma = gvar.sqrt(fit.p['sigma^2']) 2 ctx1ab

488 self.lambda_mu = fit.p['lambda_mu'] 2 ctx1ab

489 self.lambda_tau = fit.p['lambda_tau'] 2 ctx1ab

490 self.alpha_mu = fit.p['alpha_mu'] 2 ctx1ab

491 self.alpha_tau = fit.p['alpha_tau'] 2 ctx1ab

492 self.beta_mu = fit.p['beta_mu'] 2 ctx1ab

493 self.beta_tau = fit.p['beta_tau'] 2 ctx1ab

494 self.z_0 = fit.p['z_0'] 2 ctx1ab

495 

496 # save other attributes 

497 self.fit = fit 2 ctx1ab

498 self._from_data = from_data 2 ctx1ab

499 self._to_data = to_data 2 ctx1ab

500 

501 def _append_pihat(self, x_mu, x_tau, pihat): 3 ctx1cab

502 ip = self._include_pi 2 ctx1ab

503 if ip in {'mu', 'both'}: 2 ctx1ab

504 x_mu = _array.StructuredArray.from_dict(dict(x=x_mu, pihat=pihat)) 2 ctx1ab

505 if x_tau is not None and ip in {'tau', 'both'}: 2 ctx1ab

506 x_tau = _array.StructuredArray.from_dict(dict(x=x_tau, pihat=pihat)) 1 ctx1b

507 return x_mu, x_tau 2 ctx1ab

508 

509 @staticmethod 3 ctx1cab

510 def _join_points(train, z, i_mu, i_tau, pihat, x_aux): 3 ctx1cab

511 """Join covariates into a single StructuredArray.""" 

512 columns = dict( 2 ctx1ab

513 train=jnp.broadcast_to(bool(train), z.shape), 

514 i=jnp.arange(z.size).reshape(z.shape), 

515 z=z, 

516 mu=i_mu, 

517 tau=i_mu if i_tau is None else i_tau, 

518 pihat=pihat, 

519 ) 

520 if x_aux is not None: 2 ctx1ab

521 columns.update(aux=x_aux) 1 ctx1b

522 return _array.StructuredArray.from_dict(columns) 2 ctx1ab

523 

524 @staticmethod 3 ctx1cab

525 def _error_cov(hp, weights, x): 3 ctx1cab

526 """Fill error covariance matrix.""" 

527 if weights is None: 527 ↛ 530line 527 didn't jump to line 530 because the condition on line 527 was always true2 ctx1ab

528 error_var = jnp.broadcast_to(hp['sigma^2'], len(x)) 2 ctx1ab

529 else: 

530 error_var = hp['sigma^2'] / weights 

531 return jnp.diag(error_var) 2 ctx1ab

532 

533 def _gethp(self, hp, rng): 3 ctx1cab

534 if not isinstance(hp, str): 534 ↛ 535line 534 didn't jump to line 535 because the condition on line 534 was never true2 ctx1ab

535 return hp 

536 elif hp == 'map': 2 ctx1ab

537 return self.fit.pmean 2 ctx1ab

538 elif hp == 'sample': 538 ↛ 541line 538 didn't jump to line 541 because the condition on line 538 was always true2 ctx1ab

539 return _fastraniter.sample(self.fit.pmean, self.fit.pcov, rng=rng) 2 ctx1ab

540 else: 

541 raise KeyError(hp) 

542 

543 def gp( 3 ctx1cab

544 self, 

545 *, 

546 hp='map', 

547 z=None, 

548 x_mu=None, 

549 x_tau=None, 

550 pihat=None, 

551 x_aux=None, 

552 weights=None, 

553 rng=None, 

554 ): 

555 """ 

556 Create a Gaussian process with the fitted hyperparameters. 

557 

558 Parameters 

559 ---------- 

560 hp : str or dict 

561 The hyperparameters to use. If ``'map'`` (default), use the marginal 

562 maximum a posteriori. If ``'sample'``, sample hyperparameters from 

563 the posterior. If a dict, use the given hyperparameters. 

564 z : (m,) array, series or dataframe, optional 

565 Treatment status at test points. If specified, also `x_mu`, `pihat`, 

566 `x_tau` and `x_aux` (the latter two if and only also specified at 

567 initialization) must be specified. 

568 x_mu : (m, p) array, series or dataframe, optional 

569 Control model covariates at test points. 

570 x_tau : (m, q) array, series or dataframe, optional 

571 Moderating model covariates at test points. 

572 pihat : (m,) array, series or dataframe, optional 

573 Estimated propensity score at test points. 

574 x_aux : (m, k) array, series or dataframe, optional 

575 Additional covariates for the ``'aux'`` process. 

576 weights : (m,) array, series or dataframe, optional 

577 Weights for the error variance on the test points. 

578 rng : numpy.random.Generator, optional 

579 Random number generator, used if ``hp == 'sample'``. 

580 

581 Returns 

582 ------- 

583 gp : GP 

584 A centered Gaussian process object. To add the mean, use the `m` 

585 attribute of the `bcf` object. The keys of the GP are ``'@mean'``, 

586 ``'@noise'``, and ``'@'``, where the "@" stands either for 'train' 

587 or 'test', and @ = @mean + @noise. 

588 

589 This Gaussian process is defined on the transformed data ``eta``. 

590 """ 

591 hp = self._gethp(hp, rng) 

592 return self._gp(hp, z, x_mu, x_tau, pihat, x_aux, weights, self.fit.gpfactorykw) 

593 

594 def _gp(self, hp, z, x_mu, x_tau, pihat, x_aux, weights, gpfactorykw): 3 ctx1cab

595 """ 

596 Create the GP object. 

597 

598 This function must work both if the arguments are user-provided and 

599 need to be checked and converted to standard format, or if they are 

600 traced jax values. 

601 """ 

602 # create GP object 

603 gp = self.fit.gpfactory(hp, **gpfactorykw) 2 ctx1ab

604 

605 # add test points 

606 if z is not None: 2 ctx1ab

607 # check presence/absence of arguments is coherent 

608 self._check_coherent_covariates(z, x_mu, x_tau, pihat, x_aux) 1 ctx1a

609 

610 # check treatment and propensity score 

611 z = self._to_vector(z) 1 ctx1a

612 pihat = self._to_vector(pihat) 1 ctx1a

613 assert pihat.shape == z.shape 1 ctx1a

614 

615 # check weights 

616 if weights is not None: 616 ↛ 617line 616 didn't jump to line 617 because the condition on line 616 was never true1 ctx1a

617 weights = self._to_vector(weights) 

618 assert weights.shape == z.shape 

619 

620 # add propensity score to covariates 

621 x_mu = self._to_structured(x_mu) 1 ctx1a

622 assert x_mu.shape == z.shape 1 ctx1a

623 if x_tau is not None: 623 ↛ 624line 623 didn't jump to line 624 because the condition on line 623 was never true1 ctx1a

624 x_tau = self._to_structured(x_tau) 

625 assert x_tau.shape == z.shape 

626 x_mu, x_tau = self._append_pihat(x_mu, x_tau, pihat) 1 ctx1a

627 

628 # convert covariates to indices 

629 i_mu = self._toindices(x_mu, gpfactorykw['splits_mu']) 1 ctx1a

630 assert i_mu.dtype == gpfactorykw['i_mu'].dtype 1 ctx1a

631 if x_tau is not None: 631 ↛ 632line 631 didn't jump to line 632 because the condition on line 631 was never true1 ctx1a

632 i_tau = self._toindices(x_tau, gpfactorykw['splits_tau']) 

633 assert i_tau.dtype == gpfactorykw['i_tau'].dtype 

634 else: 

635 i_tau = None 1 ctx1a

636 

637 # check auxiliary points 

638 if x_aux is not None: 638 ↛ 639line 638 didn't jump to line 639 because the condition on line 638 was never true1 ctx1a

639 x_aux = self._to_structured(x_aux) 

640 

641 # add test points 

642 x = self._join_points(False, z, i_mu, i_tau, pihat, x_aux) 1 ctx1a

643 gp = gp.addx(x, 'testmean') 1 ctx1a

644 errcov = self._error_cov(hp, weights, x) 1 ctx1a

645 gp = gp.addcov(errcov, 'testnoise').addtransf( 1 ctx1a

646 {'testmean': 1, 'testnoise': 1}, 'test' 

647 ) 

648 

649 return gp 2 ctx1ab

650 

651 def _check_coherent_covariates(self, z, x_mu, x_tau, pihat, x_aux): 3 ctx1cab

652 if z is None: 2 ctx1ab

653 assert x_mu is None 1 ctx1b

654 assert x_tau is None 1 ctx1b

655 assert pihat is None 1 ctx1b

656 assert x_aux is None 1 ctx1b

657 else: 

658 assert x_mu is not None 1 ctx1a

659 assert pihat is not None 1 ctx1a

660 train_tau = self.fit.gpfactorykw['i_tau'] 1 ctx1a

661 if x_tau is None: 661 ↛ 664line 661 didn't jump to line 664 because the condition on line 661 was always true1 ctx1a

662 assert train_tau is None 1 ctx1a

663 else: 

664 assert train_tau is not None 

665 train_aux = self.fit.gpfactorykw['x_aux'] 1 ctx1a

666 if x_aux is None: 666 ↛ 669line 666 didn't jump to line 669 because the condition on line 666 was always true1 ctx1a

667 assert train_aux is None 1 ctx1a

668 else: 

669 assert train_aux is not None 

670 

671 def data(self, *, hp='map', rng=None): 3 ctx1cab

672 """ 

673 Get the data to be passed to `GP.pred` on a GP object returned by `gp`. 

674 

675 Parameters 

676 ---------- 

677 hp : str or dict 

678 The hyperparameters to use. If ``'map'`` (default), use the marginal 

679 maximum a posteriori. If ``'sample'``, sample hyperparameters from 

680 the posterior. If a dict, use the given hyperparameters. 

681 rng : numpy.random.Generator, optional 

682 Random number generator, used if ``hp == 'sample'``. 

683 

684 Returns 

685 ------- 

686 data : dict 

687 A dictionary representing ``eta`` in the format required by the 

688 `GP.pred` method. 

689 """ 

690 hp = self._gethp(hp, rng) 

691 return self.fit.data(hp, **self.fit.gpfactorykw) 

692 

693 def pred( # noqa: C901 3 ctx1cab

694 self, 

695 *, 

696 hp='map', 

697 error=False, 

698 z=None, 

699 x_mu=None, 

700 x_tau=None, 

701 pihat=None, 

702 x_aux=None, 

703 weights=None, 

704 transformed=True, 

705 samples=None, 

706 gvars=False, 

707 rng=None, 

708 ): 

709 r""" 

710 Predict the transformed outcome at given locations. 

711 

712 Parameters 

713 ---------- 

714 hp : str or dict 

715 The hyperparameters to use. If ``'map'`` (default), use the marginal 

716 maximum a posteriori. If ``'sample'``, sample hyperparameters from 

717 the posterior. If a dict, use the given hyperparameters. 

718 error : bool, default False 

719 If ``False``, make a prediction for the latent mean. If ``True``, 

720 add the error term. 

721 z : (m,) array, series or dataframe, optional 

722 Treatment status at test points. If specified, also `x_mu`, `pihat`, 

723 `x_tau` and `x_aux` (the latter two if and only also specified at 

724 initialization) must be specified. 

725 x_mu : (m, p) array, series or dataframe, optional 

726 :math:`\mu` model covariates at test points. 

727 x_tau : (m, q) array, series or dataframe, optional 

728 :math:`\tau` model covariates at test points. 

729 pihat : (m,) array, series or dataframe, optional 

730 Estimated propensity score at test points. 

731 x_aux : (m, k) array, series or dataframe, optional 

732 Additional covariates for the ``'aux'`` process at test points. 

733 weights : (m,) array, series or dataframe, optional 

734 Weights for the error variance on the test points. 

735 transformed : bool, default True 

736 If ``True``, return the prediction on the transformed outcome 

737 :math:`\eta`, else the observable outcome :math:`y`. 

738 samples : int, optional 

739 If specified, indicates the number of samples to take from the 

740 posterior. If not, return the mean and covariance matrix of the 

741 posterior. 

742 gvars : bool, default False 

743 If ``True``, return the mean and covariance matrix of the posterior 

744 as an array of `GVar` variables. 

745 rng : numpy.random.Generator, optional 

746 Random number generator, used if ``hp == 'sample'`` or ``samples`` 

747 is not `None`. 

748 

749 Returns 

750 ------- 

751 mean, cov : (m,) and (m, m) arrays 

752 If ``samples`` is `None` and ``gvars`` is `False` (default), the 

753 mean and covariance matrix of the Normal posterior distribution 

754 over the regression function or :math:`\eta` at the specified 

755 locations. 

756 out : (m,) array of gvars 

757 If ``samples`` is `None` and ``gvars`` is `True`, the same 

758 distribution represented as an array of `~gvar.GVar` objects. 

759 sample : (samples, m) array 

760 If ``samples`` is an integer, posterior samples over either the 

761 regression function, :math:`\eta`, or :math:`y`. 

762 """ 

763 # check consistency of output choice 

764 if samples is None: 764 ↛ 765line 764 didn't jump to line 765 because the condition on line 764 was never true2 ctx1ab

765 if not transformed: 

766 msg = ( 

767 'Posterior is required in analytical form ' 

768 '(samples=None) and in data space ' 

769 '(transformed=False), this is not possible as ' 

770 'the transformation model space -> data space ' 

771 'is arbitrary. Either sample the posterior or ' 

772 'get the result in model space.' 

773 ) 

774 raise ValueError(msg) 

775 else: 

776 if not transformed and not error: 776 ↛ 777line 776 didn't jump to line 777 because the condition on line 776 was never true2 ctx1ab

777 msg = ( 

778 'Posterior is required in data space ' 

779 '(transformed=False) and without error term ' 

780 '(error=False), this is not possible as the ' 

781 'transformation model space -> data space ' 

782 'applies after adding the error.' 

783 ) 

784 raise ValueError(msg) 

785 assert not gvars, 'can not represent posterior samples as gvars' 2 ctx1ab

786 

787 # get hyperparameters 

788 hp = self._gethp(hp, rng) 2 ctx1ab

789 

790 # check presence of covariates is coherent 

791 self._check_coherent_covariates(z, x_mu, x_tau, pihat, x_aux) 2 ctx1ab

792 

793 # convert all inputs to arrays compatible with jax to pass them to the 

794 # compiled implementation 

795 if z is not None: 2 ctx1ab

796 z = self._to_vector(z) 1 ctx1a

797 pihat = self._to_vector(pihat) 1 ctx1a

798 x_mu = self._to_structured(x_mu) 1 ctx1a

799 if x_tau is not None: 799 ↛ 800line 799 didn't jump to line 800 because the condition on line 799 was never true1 ctx1a

800 x_tau = self._to_structured(x_tau) 

801 if x_aux is not None: 801 ↛ 802line 801 didn't jump to line 802 because the condition on line 801 was never true1 ctx1a

802 x_aux = self._to_structured(x_aux) 

803 if weights is not None: 803 ↛ 804line 803 didn't jump to line 804 because the condition on line 803 was never true2 ctx1ab

804 weights = self._to_vector(weights) 

805 

806 # GP regression 

807 mean, cov = self._pred( 2 ctx1ab

808 hp, z, x_mu, x_tau, pihat, x_aux, weights, self.fit.gpfactorykw, bool(error) 

809 ) 

810 

811 # return Normal posterior moments 

812 if samples is None: 812 ↛ 813line 812 didn't jump to line 813 because the condition on line 812 was never true2 ctx1ab

813 if gvars: 

814 return gvar.gvar(mean, cov, fast=True) 

815 else: 

816 return mean, cov 

817 

818 # sample from posterior 

819 sample = jnp.stack(list(_fastraniter.raniter(mean, cov, n=samples, rng=rng))) 2 ctx1ab

820 if not transformed: 820 ↛ 822line 820 didn't jump to line 822 because the condition on line 820 was always true2 ctx1ab

821 sample = self._to_data(hp, sample) 2 ctx1ab

822 return sample 2 ctx1ab

823 

824 @functools.cached_property 3 ctx1cab

825 def _pred(self): 3 ctx1cab

826 

827 @functools.partial(jit, static_argnums=(8,)) 2 ctx1ab

828 def _pred(hp, z, x_mu, x_tau, pihat, x_aux, weights, gpfactorykw, error): 2 ctx1ab

829 gp = self._gp(hp, z, x_mu, x_tau, pihat, x_aux, weights, gpfactorykw) 2 ctx1ab

830 data = self.fit.data(hp, **gpfactorykw) 2 ctx1ab

831 if z is None: 2 ctx1ab

832 label = 'train' 1 ctx1b

833 else: 

834 label = 'test' 1 ctx1a

835 if not error: 835 ↛ 836line 835 didn't jump to line 836 because the condition on line 835 was never true2 ctx1ab

836 label += 'mean' 

837 outmean, outcov = gp.predfromdata(data, label, raw=True) 2 ctx1ab

838 return outmean + hp.get('m', 0), outcov 2 ctx1ab

839 

840 return _pred 2 ctx1ab

841 

842 def from_data(self, y, *, hp='map', rng=None): 3 ctx1cab

843 r""" 

844 Transform outcomes :math:`y` to the regression variable :math:`\eta`. 

845 

846 Parameters 

847 ---------- 

848 y : (n,) array 

849 Outcomes. 

850 hp : str or dict 

851 The hyperparameters to use. If ``'map'`` (default), use the marginal 

852 maximum a posteriori. If ``'sample'``, sample hyperparameters from 

853 the posterior. If a dict, use the given hyperparameters. 

854 rng : numpy.random.Generator, optional 

855 Random number generator, used if ``hp == 'sample'``. 

856 

857 Returns 

858 ------- 

859 eta : (n,) array 

860 Transformed outcomes. 

861 """ 

862 hp = self._gethp(hp, rng) 1 ctx1b

863 return self._from_data(hp, y) 1 ctx1b