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

1# lsqfitgp/bayestree/_bcf.py 

2# 

3# Copyright (c) 2023, 2024, 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 6 ctx1fabcde

21import warnings 6 ctx1fabcde

22 

23import numpy 6 ctx1fabcde

24from scipy import stats 6 ctx1fabcde

25from jax import numpy as jnp 6 ctx1fabcde

26import jax 6 ctx1fabcde

27import gvar 6 ctx1fabcde

28 

29from .. import copula 6 ctx1fabcde

30from .. import _kernels 6 ctx1fabcde

31from .. import _fit 6 ctx1fabcde

32from .. import _array 6 ctx1fabcde

33from .. import _GP 6 ctx1fabcde

34from .. import _fastraniter 6 ctx1fabcde

35from .. import _jaxext 6 ctx1fabcde

36from .. import _gvarext 6 ctx1fabcde

37from .. import _utils 6 ctx1fabcde

38 

39# TODO add methods or options to do causal inference stuff, e.g., impute missing 

40# outcomes, or ate, att, cate, catt, sate, satt. Remember that the effect may 

41# also depend on aux. See bartCause, possibly copy its naming. 

42 

43def _recursive_cast(dtype, default, mapping): 6 ctx1fabcde

44 if dtype in mapping: 44 ↛ 45line 44 didn't jump to line 45 because the condition on line 44 was never true5 ctx1abcde

45 return mapping[dtype] 

46 elif dtype.names is not None: 5 ctx1abcde

47 return numpy.dtype([ 5 ctx1abcde

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

49 for name in dtype.names 

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 true5 ctx1abcde

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 true5 ctx1abcde

55 return dtype 

56 else: 

57 return default 5 ctx1abcde

58 

59def cast(dtype, default, mapping={}): 6 ctx1fabcde

60 """ 

61 Recursively cast a numpy data type. 

62 

63 Parameters 

64 ---------- 

65 dtype : dtype 

66 The data type to cast. 

67 default : dtype or None 

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

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

70 `mapping` are left unchanged. 

71 mapping : dict 

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

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

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

75 further searched for matches. 

76 

77 Returns 

78 ------- 

79 casted_dtype : dtype 

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

81 `mapping` contains structured dtypes. 

82 """ 

83 mapping = {numpy.dtype(k): numpy.dtype(v) for k, v in mapping.items()} 5 ctx1abcde

84 default = None if default is None else numpy.dtype(default) 5 ctx1abcde

85 return _recursive_cast(numpy.dtype(dtype), default, mapping) 5 ctx1abcde

86 

87 # TODO 

88 # - move this to generic utils 

89 # - make unit tests 

90 

91class bcf: 6 ctx1fabcde

92 

93 def __init__(self, *, 6 ctx1fabcde

94 y, 

95 z, 

96 x_mu, 

97 x_tau=None, 

98 pihat, 

99 include_pi='mu', 

100 weights=None, 

101 fitkw={}, 

102 kernelkw_mu={}, 

103 kernelkw_tau={}, 

104 marginalize_mean=True, 

105 gpaux=None, 

106 x_aux=None, 

107 otherhp={}, 

108 transf='standardize', 

109 ): 

110 r""" 

111 Nonparametric Bayesian regression with a GP version of BCF. 

112 

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

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

115 models. 

116 

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

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

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

120 

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

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

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

124 

125 Parameters 

126 ---------- 

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

128 Outcome. 

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

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

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

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

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

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

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

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

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

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

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

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

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

142 fitkw : dict 

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

144 the defaults. 

145 kernelkw_mu, kernelkw_tau : dict 

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

147 overrides the defaults. 

148 marginalize_mean : bool 

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

150 gpaux : callable, optional 

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

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

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

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

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

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

157 with fields: 

158 

159 'train' : bool 

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

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

162 'i' : int 

163 Index of the flattened array. 

164 'z' : int 

165 Treatment status. 

166 'mu', 'tau' : structured 

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

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

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

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

171 'pihat' : float 

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

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

174 values. 

175 'aux' : structured 

176 The values in `x_aux`, if specified. 

177  

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

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

180 otherhp : dictionary of gvar 

181 A dictionary with the prior of arbitrary additional hyperpameters, 

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

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

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

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

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

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

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

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

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

191 

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

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

194 the one applied first to the data. 

195 

196 If a transformation uses additional hyperparameters, either 

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

198 they are inferred with the rest of the hyperparameters. 

199 

200 The pre-defined transformations are: 

201 

202 'standardize' (default) 

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

204 'yeojohnson' 

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

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

207 for implementation convenience, this restriction may be lifted 

208 in future versions. 

209  

210 Notes 

211 ----- 

212 The regression model is: 

213 

214 .. math:: 

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

216 &\phantom{{}={}} + 

217 \lambda_\mu 

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

219 &\phantom{{}={}} + 

220 \lambda_\tau 

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

222 &\phantom{{}={}} + 

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

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

225 &\phantom{{}={}} + 

226 \varepsilon_i, \\ 

227 \varepsilon_i &\sim 

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

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

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

231 \lambda_\mu 

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

233 \lambda_\tau 

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

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

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

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

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

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

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

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

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

243 

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

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

246 mode of 

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

248 numerical minimization, after transforming them to express their prior 

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

250 estimated with an approximation of the hessian inverse. See 

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

252 this procedure. 

253 

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

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

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

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

258 

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

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

261 Y` if the unconfoundedness assumption is made: 

262 

263 .. math:: 

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

265 

266 In practical terms, this holds when: 

267 

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

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

270 

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

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

273 is the causal effect and not just a correlation. 

274 

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

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

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

278 what you are doing. 

279 

280 Attributes 

281 ---------- 

282 m : float or gvar 

283 The prior mean :math:`m`. 

284 sigma : gvar 

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

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

287 sqrt(weight). 

288 alpha_mu, alpha_tau : gvar 

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

290 ``base`` in R bcf). 

291 beta_mu, beta_tau : gvar 

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

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

294 lambda_mu, lambda_tau : gvar 

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

296 z_0 : gvar 

297 The treatment coding parameter. 

298 fit : empbayes_fit 

299 The hyperparameters fit object. 

300 

301 Methods 

302 ------- 

303 gp : 

304 Create a GP object. 

305 data : 

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

307 pred : 

308 Evaluate the regression function at given locations. 

309 from_data : 

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

311 to_data : 

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

313 

314 See also 

315 -------- 

316 lsqfitgp.BART 

317  

318 References 

319 ---------- 

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

321 Regression Tree Models for Causal Inference: Regularization, 

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

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

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

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

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

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

328 """ 

329 

330 # convert covariates to StructuredArray 

331 x_mu = self._to_structured(x_mu) 5 ctx1abcde

332 if x_tau is not None: 5 ctx1abcde

333 x_tau = self._to_structured(x_tau) 4 ctx1bcde

334 assert x_tau.shape == x_mu.shape 4 ctx1bcde

335 if x_aux is not None: 5 ctx1abcde

336 x_aux = self._to_structured(x_aux) 4 ctx1bcde

337 assert x_aux.shape == x_mu.shape 4 ctx1bcde

338 

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

340 y = self._to_vector(y) 5 ctx1abcde

341 z = self._to_vector(z) 5 ctx1abcde

342 pihat = self._to_vector(pihat) 5 ctx1abcde

343 assert y.shape == z.shape == pihat.shape == x_mu.shape 5 ctx1abcde

344 if weights is not None: 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true5 ctx1abcde

345 weights = self._to_vector(weights) 

346 assert weights.shape == x_mu.shape 

347 

348 # check include_pi 

349 if include_pi not in ('mu', 'tau', 'both'): 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true5 ctx1abcde

350 raise KeyError(f'invalid value include_pi={include_pi!r}') 

351 self._include_pi = include_pi 5 ctx1abcde

352 

353 # add pihat to covariates 

354 x_mu, x_tau = self._append_pihat(x_mu, x_tau, pihat) 5 ctx1abcde

355 

356 # grid and indices 

357 splits_mu = _kernels.BART.splits_from_coord(x_mu) 5 ctx1abcde

358 i_mu = self._toindices(x_mu, splits_mu) 5 ctx1abcde

359 if x_tau is None: 5 ctx1abcde

360 splits_tau = splits_mu 5 ctx1abcde

361 i_tau = None 5 ctx1abcde

362 else: 

363 splits_tau = _kernels.BART.splits_from_coord(x_tau) 4 ctx1bcde

364 i_tau = self._toindices(x_tau, splits_tau) 4 ctx1bcde

365 

366 # get functions for data transformation 

367 from_data, to_data, transfloss, transfhp = self._get_transf( 5 ctx1abcde

368 transf=transf, weights=weights, y=y) 

369 

370 # scale of error variance 

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

372 

373 # prior on hyperparams 

374 hyperprior = copula.makedict({ 5 ctx1abcde

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

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

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

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

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

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

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

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

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

384 }) 

385 

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

387 if marginalize_mean: 5 ctx1abcde

388 hyperprior.pop('m') 5 ctx1abcde

389 

390 # add data transformation and user hyperparameters 

391 def update_hyperparams(new, newname, raises): 5 ctx1abcde

392 new = gvar.BufferDict(new) 5 ctx1abcde

393 for key in new.all_keys(): 5 ctx1abcde

394 if hyperprior.has_dictkey(key): 394 ↛ 395line 394 didn't jump to line 395 because the condition on line 394 was never true5 ctx1abcde

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

396 if raises: 

397 raise ValueError(message) 

398 else: 

399 warnings.warn(message) 

400 hyperprior.update(new) 5 ctx1abcde

401 update_hyperparams(transfhp, 'data transformation', True) 5 ctx1abcde

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

403 update_hyperparams(otherhp, 'user', False) 5 ctx1abcde

404 

405 # GP factory 

406 def gpfactory(hp, *, z, i_mu, i_tau, pihat, x_aux, weights, 5 ctx1abcde

407 splits_mu, splits_tau, **_): 

408 

409 # TODO maybe I should pass kernelkw_* as arguments, but they may not 

410 # be jittable. I need jitkw in empbayes_fit for that. 

411 

412 kw_overridable = dict( 5 ctx1abcde

413 maxd=10, 

414 reset=[2, 4, 6, 8], 

415 intercept=False, 

416 ) 

417 kw_not_overridable = dict(indices=True) 5 ctx1abcde

418 

419 gp = _GP.GP(checkpos=False, checksym=False, solver='chol') 5 ctx1abcde

420 

421 for name, kernelkw in dict(mu=kernelkw_mu, tau=kernelkw_tau).items(): 5 ctx1abcde

422 kw = dict( 5 ctx1abcde

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

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

425 dim=name, 

426 splits=eval(f'splits_{name}'), 

427 **kw_overridable, 

428 ) 

429 kw.update(kernelkw) 5 ctx1abcde

430 kernel = _kernels.BART(**kw, **kw_not_overridable) 5 ctx1abcde

431 kernel *= hp[f'lambda_{name}'] ** 2 5 ctx1abcde

432 

433 gp = gp.defproc(name, kernel) 5 ctx1abcde

434 

435 if 'm' in hp: 5 ctx1abcde

436 kernel_mean = 0 * _kernels.Constant() 4 ctx1bcde

437 else: 

438 kernel_mean = _kernels.Constant() 5 ctx1abcde

439 gp = gp.defproc('m', kernel_mean) 5 ctx1abcde

440 

441 if gpaux is None: 5 ctx1abcde

442 gp = gp.defproc('aux', 0 * _kernels.Constant()) 5 ctx1abcde

443 else: 

444 gp = gpaux(hp, gp) 4 ctx1bcde

445 

446 gp = gp.deflintransf( 5 ctx1abcde

447 gp.DefaultProcess, 

448 lambda m, mu, tau, aux: lambda x: 

449 m(x) + mu(x) + tau(x) * (x['z'] - hp['z_0']) + aux(x), 

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

451 ) 

452 

453 x = self._join_points(True, z, i_mu, i_tau, pihat, x_aux) 5 ctx1abcde

454 gp = gp.addx(x, 'trainmean') 5 ctx1abcde

455 errcov = self._error_cov(hp, weights, x) 5 ctx1abcde

456 return (gp 5 ctx1abcde

457 .addcov(errcov, 'trainnoise') 

458 .addtransf({'trainmean': 1, 'trainnoise': 1}, 'train') 

459 ) 

460 

461 # data factory 

462 def data(hp, *, y, **_): 5 ctx1abcde

463 return {'train': from_data(hp, y) - hp.get('m', 0)} 5 ctx1abcde

464 

465 # fit hyperparameters 

466 gpkw = dict( 5 ctx1abcde

467 y=y, 

468 z=z, 

469 i_mu=i_mu, 

470 i_tau=i_tau, 

471 pihat=pihat, 

472 x_aux=x_aux, 

473 weights=weights, 

474 splits_mu=splits_mu, 

475 splits_tau=splits_tau, 

476 ) 

477 options = dict( 5 ctx1abcde

478 verbosity=3, 

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

480 mlkw=dict(epsrel=0), 

481 forward=True, 

482 gpfactorykw=gpkw, 

483 additional_loss=transfloss, 

484 ) 

485 options.update(fitkw) 5 ctx1abcde

486 fit = _fit.empbayes_fit(hyperprior, gpfactory, data, **options) 5 ctx1abcde

487 

488 # extract hyperparameters from minimization result 

489 self.m = fit.p.get('m', 0) 5 ctx1abcde

490 self.sigma = gvar.sqrt(fit.p['sigma^2']) 5 ctx1abcde

491 self.lambda_mu = fit.p['lambda_mu'] 5 ctx1abcde

492 self.lambda_tau = fit.p['lambda_tau'] 5 ctx1abcde

493 self.alpha_mu = fit.p['alpha_mu'] 5 ctx1abcde

494 self.alpha_tau = fit.p['alpha_tau'] 5 ctx1abcde

495 self.beta_mu = fit.p['beta_mu'] 5 ctx1abcde

496 self.beta_tau = fit.p['beta_tau'] 5 ctx1abcde

497 self.z_0 = fit.p['z_0'] 5 ctx1abcde

498 

499 # save other attributes 

500 self.fit = fit 5 ctx1abcde

501 self._from_data = from_data 5 ctx1abcde

502 self._to_data = to_data 5 ctx1abcde

503 

504 def _append_pihat(self, x_mu, x_tau, pihat): 6 ctx1fabcde

505 ip = self._include_pi 5 ctx1abcde

506 if ip == 'mu' or ip == 'both': 5 ctx1abcde

507 x_mu = _array.StructuredArray.from_dict(dict( 5 ctx1abcde

508 x=x_mu, 

509 pihat=pihat, 

510 )) 

511 if x_tau is not None and (ip == 'tau' or ip == 'both'): 5 ctx1abcde

512 x_tau = _array.StructuredArray.from_dict(dict( 4 ctx1bcde

513 x=x_tau, 

514 pihat=pihat, 

515 )) 

516 return x_mu, x_tau 5 ctx1abcde

517 

518 @staticmethod 6 ctx1fabcde

519 def _join_points(train, z, i_mu, i_tau, pihat, x_aux): 6 ctx1fabcde

520 """ join covariates into a single StructuredArray """ 

521 columns = dict( 5 ctx1abcde

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

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

524 z=z, 

525 mu=i_mu, 

526 tau=i_mu if i_tau is None else i_tau, 

527 pihat=pihat, 

528 ) 

529 if x_aux is not None: 5 ctx1abcde

530 columns.update(aux=x_aux) 4 ctx1bcde

531 return _array.StructuredArray.from_dict(columns) 5 ctx1abcde

532 

533 @staticmethod 6 ctx1fabcde

534 def _error_cov(hp, weights, x): 6 ctx1fabcde

535 """ fill error covariance matrix """ 

536 if weights is None: 536 ↛ 539line 536 didn't jump to line 539 because the condition on line 536 was always true5 ctx1abcde

537 error_var = jnp.broadcast_to(hp['sigma^2'], len(x)) 5 ctx1abcde

538 else: 

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

540 return jnp.diag(error_var) 5 ctx1abcde

541 

542 def _gethp(self, hp, rng): 6 ctx1fabcde

543 if not isinstance(hp, str): 543 ↛ 544line 543 didn't jump to line 544 because the condition on line 543 was never true5 ctx1abcde

544 return hp 

545 elif hp == 'map': 5 ctx1abcde

546 return self.fit.pmean 5 ctx1abcde

547 elif hp == 'sample': 547 ↛ 550line 547 didn't jump to line 550 because the condition on line 547 was always true5 ctx1abcde

548 return _fastraniter.sample(self.fit.pmean, self.fit.pcov, rng=rng) 5 ctx1abcde

549 else: 

550 raise KeyError(hp) 

551 

552 def gp(self, *, hp='map', z=None, x_mu=None, x_tau=None, pihat=None, 6 ctx1fabcde

553 x_aux=None, weights=None, rng=None): 

554 """ 

555 Create a Gaussian process with the fitted hyperparameters. 

556 

557 Parameters 

558 ---------- 

559 hp : str or dict 

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

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

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

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

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

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

566 initialization) must be specified. 

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

568 Control model covariates at test points. 

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

570 Moderating model covariates at test points. 

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

572 Estimated propensity score at test points. 

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

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

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

576 Weights for the error variance on the test points. 

577 rng : numpy.random.Generator, optional 

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

579 

580 Returns 

581 ------- 

582 gp : GP 

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

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

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

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

587 

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

589 """ 

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): 6 ctx1fabcde

595 """ 

596 Internal function to create the GP object. This function must work 

597 both if the arguments are user-provided and need to be checked and 

598 converted to standard format, or if they are traced jax values. 

599 """ 

600 

601 # create GP object 

602 gp = self.fit.gpfactory(hp, **gpfactorykw) 5 ctx1abcde

603 

604 # add test points 

605 if z is not None: 5 ctx1abcde

606 

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 1 ctx1a

646 .addcov(errcov, 'testnoise') 

647 .addtransf({'testmean': 1, 'testnoise': 1}, 'test') 

648 ) 

649 

650 return gp 5 ctx1abcde

651 

652 def _check_coherent_covariates(self, z, x_mu, x_tau, pihat, x_aux): 6 ctx1fabcde

653 if z is None: 5 ctx1abcde

654 assert x_mu is None 4 ctx1bcde

655 assert x_tau is None 4 ctx1bcde

656 assert pihat is None 4 ctx1bcde

657 assert x_aux is None 4 ctx1bcde

658 else: 

659 assert x_mu is not None 1 ctx1a

660 assert pihat is not None 1 ctx1a

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

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

663 assert train_tau is None 1 ctx1a

664 else: 

665 assert train_tau is not None 

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

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

668 assert train_aux is None 1 ctx1a

669 else: 

670 assert train_aux is not None 

671 

672 def data(self, *, hp='map', rng=None): 6 ctx1fabcde

673 """ 

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

675 

676 Parameters 

677 ---------- 

678 hp : str or dict 

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

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

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

682 rng : numpy.random.Generator, optional 

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

684 

685 Returns 

686 ------- 

687 data : dict 

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

689 `GP.pred` method. 

690 """ 

691 

692 hp = self._gethp(hp, rng) 

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

694 

695 def pred(self, *, hp='map', error=False, z=None, x_mu=None, x_tau=None, 6 ctx1fabcde

696 pihat=None, x_aux=None, weights=None, transformed=True, samples=None, 

697 gvars=False, rng=None): 

698 r""" 

699 Predict the transformed outcome at given locations. 

700 

701 Parameters 

702 ---------- 

703 hp : str or dict 

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

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

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

707 error : bool, default False 

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

709 add the error term. 

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

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

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

713 initialization) must be specified. 

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

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

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

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

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

719 Estimated propensity score at test points. 

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

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

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

723 Weights for the error variance on the test points. 

724 transformed : bool, default True 

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

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

727 samples : int, optional 

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

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

730 posterior. 

731 gvars : bool, default False 

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

733 as an array of `GVar` variables. 

734 rng : numpy.random.Generator, optional 

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

736 is not `None`. 

737 

738 Returns 

739 ------- 

740 If ``samples`` is `None` and ``gvars`` is `False` (default): 

741 

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

743 The mean and covariance matrix of the Normal posterior distribution 

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

745 locations. 

746 

747 If ``samples`` is `None` and ``gvars`` is `True`: 

748 

749 out : (m,) array of gvars 

750 The same distribution represented as an array of `~gvar.GVar` 

751 objects. 

752 

753 If ``samples`` is an integer: 

754 

755 sample : (samples, m) array 

756 Posterior samples over either the regression function, :math:`\eta`, 

757 or :math:`y`. 

758 """ 

759 

760 # check consistency of output choice 

761 if samples is None: 761 ↛ 762line 761 didn't jump to line 762 because the condition on line 761 was never true5 ctx1abcde

762 if not transformed: 

763 raise ValueError('Posterior is required in analytical form ' 

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

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

766 'the transformation model space -> data space ' 

767 'is arbitrary. Either sample the posterior or ' 

768 'get the result in model space.') 

769 else: 

770 if not transformed and not error: 770 ↛ 771line 770 didn't jump to line 771 because the condition on line 770 was never true5 ctx1abcde

771 raise ValueError('Posterior is required in data space ' 

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

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

774 'transformation model space -> data space ' 

775 'applies after adding the error.') 

776 assert not gvars, 'can not represent posterior samples as gvars' 5 ctx1abcde

777 

778 # TODO allow exceptions to these rules when there are no transformations 

779 # or the only transformation is 'standardize'. 

780 

781 # get hyperparameters 

782 hp = self._gethp(hp, rng) 5 ctx1abcde

783 

784 # check presence of covariates is coherent 

785 self._check_coherent_covariates(z, x_mu, x_tau, pihat, x_aux) 5 ctx1abcde

786 

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

788 # compiled implementation 

789 if z is not None: 5 ctx1abcde

790 z = self._to_vector(z) 1 ctx1a

791 pihat = self._to_vector(pihat) 1 ctx1a

792 x_mu = self._to_structured(x_mu) 1 ctx1a

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

794 x_tau = self._to_structured(x_tau) 

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

796 x_aux = self._to_structured(x_aux) 

797 if weights is not None: 797 ↛ 798line 797 didn't jump to line 798 because the condition on line 797 was never true5 ctx1abcde

798 weights = self._to_vector(weights) 

799 

800 # GP regression 

801 mean, cov = self._pred(hp, z, x_mu, x_tau, pihat, x_aux, weights, self.fit.gpfactorykw, bool(error)) 5 ctx1abcde

802 

803 # return Normal posterior moments 

804 if samples is None: 804 ↛ 805line 804 didn't jump to line 805 because the condition on line 804 was never true5 ctx1abcde

805 if gvars: 

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

807 else: 

808 return mean, cov 

809 

810 # sample from posterior 

811 sample = jnp.stack(list(_fastraniter.raniter(mean, cov, n=samples, rng=rng))) 5 ctx1abcde

812 # TODO when I add vectorized sampling, use it here 

813 if not transformed: 813 ↛ 815line 813 didn't jump to line 815 because the condition on line 813 was always true5 ctx1abcde

814 sample = self._to_data(hp, sample) 5 ctx1abcde

815 return sample 5 ctx1abcde

816 

817 # TODO the default should be something in data space, so with samples. 

818 # If I handle the analyitical posterior through standardize, I could 

819 # also make it without samples by default. Although I guess for 

820 # whatever calculations samples are more convenient (just do the 

821 # calculation on the samples.) 

822 

823 @functools.cached_property 6 ctx1fabcde

824 def _pred(self): 6 ctx1fabcde

825 

826 @functools.partial(jax.jit, static_argnums=(8,)) 5 ctx1abcde

827 def _pred(hp, z, x_mu, x_tau, pihat, x_aux, weights, gpfactorykw, error): 5 ctx1abcde

828 gp = self._gp(hp, z, x_mu, x_tau, pihat, x_aux, weights, gpfactorykw) 5 ctx1abcde

829 data = self.fit.data(hp, **gpfactorykw) 5 ctx1abcde

830 if z is None: 5 ctx1abcde

831 label = 'train' 4 ctx1bcde

832 else: 

833 label = 'test' 1 ctx1a

834 if not error: 834 ↛ 835line 834 didn't jump to line 835 because the condition on line 834 was never true5 ctx1abcde

835 label += 'mean' 

836 outmean, outcov = gp.predfromdata(data, label, raw=True) 5 ctx1abcde

837 return outmean + hp.get('m', 0), outcov 5 ctx1abcde

838 

839 # TODO make everything pure and jit this per class instead of per 

840 # instance 

841 

842 return _pred 5 ctx1abcde

843 

844 def from_data(self, y, *, hp='map', rng=None): 6 ctx1fabcde

845 """ 

846 Transforms outcomes :math:`y` to the regression variable :math:`\\eta`. 

847 

848 Parameters 

849 ---------- 

850 y : (n,) array 

851 Outcomes. 

852 hp : str or dict 

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

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

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

856 rng : numpy.random.Generator, optional 

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

858