Coverage for src/lsqfitgp/bayestree/_bart.py: 83%
139 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-04 11:43 +0000
1# lsqfitgp/bayestree/_bart.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/>.
20import functools 1cab
22import gvar 1cab
23import numpy 1cab
24from jax import jit 1cab
25from jax import numpy as jnp 1cab
27from lsqfitgp import _GP, _array, _fastraniter, _fit, _kernels, copula 1cab
30class bart: 1cab
31 def __init__( 1cab
32 self,
33 x_train,
34 y_train,
35 *,
36 weights=None,
37 fitkw={}, # noqa: B006, read only
38 kernelkw={}, # noqa: B006, read only
39 marginalize_mean=True,
40 ):
41 r"""
42 Nonparametric Bayesian regression with a GP version of BART.
44 Evaluate a Gaussian process regression with a kernel which accurately
45 approximates the infinite trees limit of BART. The hyperparameters are
46 optimized to their marginal MAP.
48 Parameters
49 ----------
50 x_train : (n, p) array or dataframe
51 Observed covariates.
52 y_train : (n,) array
53 Observed outcomes.
54 weights : (n,) array
55 Weights used to rescale the error variance (as 1 / weight).
56 fitkw : dict
57 Additional arguments passed to `~lsqfitgp.empbayes_fit`, overrides
58 the defaults.
59 kernelkw : dict
60 Additional arguments passed to `~lsqfitgp.BART`, overrides the
61 defaults.
62 marginalize_mean : bool
63 If True (default), marginalize the intercept of the model.
65 Attributes
66 ----------
67 mean : gvar
68 The prior mean :math:`\mu`.
69 sigma : float or gvar
70 The error term standard deviation :math:`\sigma`. If there are
71 weights, the sdev for each unit is obtained dividing ``sigma`` by
72 sqrt(weight).
73 alpha : gvar
74 The numerator of the tree spawn probability :math:`\alpha` (named
75 ``base`` in BayesTree and BART).
76 beta : gvar
77 The depth exponent of the tree spawn probability :math:`\beta`
78 (named ``power`` in BayesTree and BART).
79 meansdev : gvar
80 The prior standard deviation :math:`\lambda` of the latent
81 regression function.
82 fit : empbayes_fit
83 The hyperparameters fit object.
85 Methods
86 -------
87 gp :
88 Create a GP object.
89 data :
90 Creates the dictionary to be passed to `GP.pred` to represent
91 ``y_train``.
92 pred :
93 Evaluate the regression function at given locations.
95 See Also
96 --------
97 lsqfitgp.BART
99 Notes
100 -----
101 The regression model is:
103 .. math::
104 y_i &= \mu + \lambda f(\mathbf x_i) + \varepsilon_i, \\
105 \varepsilon_i &\overset{\mathrm{i.i.d.}}{\sim}
106 N(0, \sigma^2 / w_i), \\
107 \mu &\sim N(
108 (\max(\mathbf y) + \min(\mathbf y)) / 2,
109 (\max(\mathbf y) - \min(\mathbf y))^2 / 4
110 ), \\
111 \log \sigma^2 &\sim N(
112 \log(\overline{w(y - \bar y)^2}),
113 4
114 ), \\
115 \log \lambda &\sim N(
116 \log ((\max(\mathbf y) - \min(\mathbf y)) / 4),
117 4
118 ), \\
119 f &\sim \mathrm{GP}(
120 0,
121 \mathrm{BART}(\alpha,\beta)
122 ), \\
123 \alpha &\sim \mathrm{B}(2, 1), \\
124 \beta &\sim \mathrm{IG}(1, 1).
126 To make the inference, :math:`(f, \boldsymbol\varepsilon, \mu)` are
127 marginalized analytically, and the marginal posterior mode of
128 :math:`(\sigma, \lambda, \alpha, \beta)` is found by numerical
129 minimization, after transforming them to express their prior as a
130 Gaussian copula. Their marginal posterior covariance matrix is estimated
131 with an approximation of the hessian inverse. See
132 `~lsqfitgp.empbayes_fit` and use the parameter ``fitkw`` to customize
133 this procedure.
135 The tree splitting grid of the BART kernel is set using quantiles of the
136 observed covariates. This corresponds to settings ``usequants=True``,
137 ``numcut=inf`` in the R packages BayesTree and BART. Use the
138 ``kernelkw`` parameter to customize the grid.
139 """ # noqa: DOC001
140 # convert covariates to StructuredArray
141 x_train = self._to_structured(x_train) 1ab
143 # convert outcomes to 1d array
144 if hasattr(y_train, 'to_numpy'): 1ab
145 y_train = y_train.to_numpy() 1a
146 y_train = y_train.squeeze() # for dataframes 1a
147 y_train = jnp.asarray(y_train) 1ab
148 assert y_train.shape == x_train.shape 1ab
150 # check weights
151 self._no_weights = weights is None 1ab
152 if self._no_weights: 152 ↛ 154line 152 didn't jump to line 154 because the condition on line 152 was always true1ab
153 weights = jnp.ones_like(y_train) 1ab
154 assert weights.shape == y_train.shape 1ab
156 # prior mean and variance
157 ymin = jnp.min(y_train) 1ab
158 ymax = jnp.max(y_train) 1ab
159 mu_mu = (ymax + ymin) / 2 1ab
160 k_sigma_mu = (ymax - ymin) / 2 1ab
162 # splitting points and indices
163 splits = _kernels.BART.splits_from_coord(x_train) 1ab
164 i_train = self._toindices(x_train, splits) 1ab
166 # prior on hyperparams
167 sigma2_priormean = numpy.mean((y_train - y_train.mean()) ** 2 * weights) 1ab
168 hyperprior = copula.makedict( 1ab
169 {
170 'alpha': copula.beta(2, 1), # base of tree gen prob
171 'beta': copula.invgamma(1, 1), # exponent of tree gen prob
172 'log(k)': gvar.gvar(numpy.log(2), 2), # denominator of prior sdev
173 'log(sigma2)': gvar.gvar(numpy.log(sigma2_priormean), 2),
174 # i.i.d. error variance, scaled with weights
175 'mean': gvar.gvar(mu_mu, k_sigma_mu), # mean of the GP
176 }
177 )
178 if marginalize_mean: 178 ↛ 182line 178 didn't jump to line 182 because the condition on line 178 was always true1ab
179 hyperprior.pop('mean') 1ab
181 # GP factory
182 def makegp(hp, *, i_train, weights, splits, **_): 1ab
183 kw = dict(alpha=hp['alpha'], beta=hp['beta'], maxd=10, reset=[2, 4, 6, 8]) 1ab
184 kw.update(kernelkw) 1ab
185 kernel = _kernels.BART(splits=splits, indices=True, **kw) 1ab
186 kernel *= (k_sigma_mu / hp['k']) ** 2 1ab
188 gp = ( 1ab
189 _GP.GP(kernel, checkpos=False, checksym=False, solver='chol')
190 .addx(i_train, 'trainmean')
191 .addcov(jnp.diag(hp['sigma2'] / weights), 'trainnoise')
192 )
193 pieces = {'trainmean': 1, 'trainnoise': 1} 1ab
194 if 'mean' not in hp: 194 ↛ 197line 194 didn't jump to line 197 because the condition on line 194 was always true1ab
195 gp = gp.addcov(k_sigma_mu**2, 'mean') 1ab
196 pieces.update({'mean': 1}) 1ab
197 return gp.addtransf(pieces, 'train') 1ab
199 # data factory
200 def info(hp, *, mu_mu, **_): 1ab
201 return {'train': y_train - hp.get('mean', mu_mu)} 1ab
203 # fit hyperparameters
204 gpkw = dict(i_train=i_train, weights=weights, splits=splits, mu_mu=mu_mu) 1ab
205 options = dict( 1ab
206 verbosity=3,
207 raises=False,
208 minkw=dict(method='l-bfgs-b', options=dict(maxls=4, maxiter=100)),
209 mlkw=dict(epsrel=0),
210 forward=True,
211 gpfactorykw=gpkw,
212 )
213 options.update(fitkw) 1ab
214 fit = _fit.empbayes_fit(hyperprior, makegp, info, **options) 1ab
216 # extract hyperparameters from minimization result
217 self.sigma = gvar.sqrt(fit.p['sigma2']) 1ab
218 self.alpha = fit.p['alpha'] 1ab
219 self.beta = fit.p['beta'] 1ab
220 self.meansdev = k_sigma_mu / fit.p['k'] 1ab
221 self.mean = fit.p.get('mean', mu_mu) 1ab
223 # set public attributes
224 self.fit = fit 1ab
226 # set private attributes
227 self._ystd = y_train.std() 1ab
229 def _gethp(self, hp, rng): 1cab
230 if not isinstance(hp, str): 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true1ab
231 return hp
232 elif hp == 'map': 232 ↛ 234line 232 didn't jump to line 234 because the condition on line 232 was always true1ab
233 return self.fit.pmean 1ab
234 elif hp == 'sample':
235 return _fastraniter.sample(self.fit.pmean, self.fit.pcov, rng=rng)
236 else:
237 raise KeyError(hp)
239 def gp(self, *, hp='map', x_test=None, weights=None, rng=None): 1cab
240 """
241 Create a Gaussian process with the fitted hyperparameters.
243 Parameters
244 ----------
245 hp : str or dict
246 The hyperparameters to use. If ``'map'``, use the marginal maximum a
247 posteriori. If ``'sample'``, sample hyperparameters from the
248 posterior. If a dict, use the given hyperparameters.
249 x_test : array or dataframe, optional
250 Additional covariates for "test points".
251 weights : array, optional
252 Weights for the error variance on the test points.
253 rng : numpy.random.Generator, optional
254 Random number generator, used if ``hp == 'sample'``.
256 Returns
257 -------
258 gp : GP
259 A centered Gaussian process object. To add the mean, use the
260 ``mean`` attribute of the `bart` object. The keys of the GP are
261 'Xmean', 'Xnoise', and 'X', where the "X" stands either for 'train'
262 or 'test', and X = Xmean + Xnoise.
263 """
264 hp = self._gethp(hp, rng)
265 return self._gp(hp, x_test, weights, self.fit.gpfactorykw)
267 def _gp(self, hp, x_test, weights, gpfactorykw): 1cab
269 # create GP object
270 gp = self.fit.gpfactory(hp, **gpfactorykw) 1ab
272 # add test points
273 if x_test is not None: 1ab
274 # convert covariates to indices
275 x_test = self._to_structured(x_test) 1a
276 i_test = self._toindices(x_test, gpfactorykw['splits']) 1a
277 assert i_test.dtype == gpfactorykw['i_train'].dtype 1a
279 # check weights
280 if weights is not None: 280 ↛ 281line 280 didn't jump to line 281 because the condition on line 280 was never true1a
281 weights = jnp.asarray(weights)
282 assert weights.shape == i_test.shape
283 else:
284 weights = jnp.ones(i_test.shape) 1a
286 # add test points
287 gp = gp.addx(i_test, 'testmean').addcov( 1a
288 jnp.diag(hp['sigma2'] / weights), 'testnoise'
289 )
290 pieces = {'testmean': 1, 'testnoise': 1} 1a
291 if 'mean' not in hp: 291 ↛ 293line 291 didn't jump to line 293 because the condition on line 291 was always true1a
292 pieces.update({'mean': 1}) 1a
293 gp = gp.addtransf(pieces, 'test') 1a
295 return gp 1ab
297 def data(self, *, hp='map', rng=None): 1cab
298 """
299 Get the data to be passed to `GP.pred` on a GP object returned by `gp`.
301 Parameters
302 ----------
303 hp : str or dict
304 The hyperparameters to use. If ``'map'``, use the marginal maximum a
305 posteriori. If ``'sample'``, sample hyperparameters from the
306 posterior. If a dict, use the given hyperparameters.
307 rng : numpy.random.Generator, optional
308 Random number generator, used if ``hp == 'sample'``.
310 Returns
311 -------
312 data : dict
313 A dictionary representing ``y_train`` in the format required by the
314 `GP.pred` method.
315 """
316 hp = self._gethp(hp, rng)
317 return self.fit.data(hp, **self.fit.gpfactorykw)
319 def pred( 1cab
320 self,
321 *,
322 hp='map',
323 error=False,
324 format='matrices', # noqa: A002, public parameter
325 x_test=None,
326 weights=None,
327 rng=None,
328 ):
329 """
330 Predict the outcome at given locations.
332 Parameters
333 ----------
334 hp : str or dict
335 The hyperparameters to use. If ``'map'``, use the marginal maximum a
336 posteriori. If ``'sample'``, sample hyperparameters from the
337 posterior. If a dict, use the given hyperparameters.
338 error : bool
339 If ``False`` (default), make a prediction for the latent mean. If
340 ``True``, add the error term.
341 format : {'matrices', 'gvar'}
342 If 'matrices' (default), return the mean and covariance matrix
343 separately. If 'gvar', return an array of gvars.
344 x_test : array or dataframe, optional
345 Covariates for the locations where the prediction is computed. If
346 not specified, predict at the data covariates.
347 weights : array, optional
348 Weights for the error variance on the test points.
349 rng : numpy.random.Generator, optional
350 Random number generator, used if ``hp == 'sample'``.
352 Returns
353 -------
354 mean, cov : arrays
355 If ``format`` is 'matrices' (default), the mean and covariance
356 matrix of the Normal posterior distribution over the regression
357 function at the specified locations.
358 out : array of `GVar`
359 If ``format`` is 'gvar', the same distribution represented as an
360 array of `GVar` objects.
361 """
362 hp = self._gethp(hp, rng) 1ab
363 if x_test is not None: 1ab
364 x_test = self._to_structured(x_test) 1a
365 mean, cov = self._pred(hp, x_test, weights, self.fit.gpfactorykw, bool(error)) 1ab
367 if format == 'gvar': 367 ↛ 368line 367 didn't jump to line 368 because the condition on line 367 was never true1ab
368 return gvar.gvar(mean, cov, fast=True)
369 elif format == 'matrices': 369 ↛ 372line 369 didn't jump to line 372 because the condition on line 369 was always true1ab
370 return mean, cov 1ab
371 else:
372 raise KeyError(format)
374 @functools.cached_property 1cab
375 def _pred(self): 1cab
377 @functools.partial(jit, static_argnums=(4,)) 1ab
378 def _pred(hp, x_test, weights, gpfactorykw, error): 1ab
379 gp = self._gp(hp, x_test, weights, gpfactorykw) 1ab
380 data = self.fit.data(hp, **gpfactorykw) 1ab
381 if x_test is None: 1ab
382 label = 'train' 1b
383 else:
384 label = 'test' 1a
385 if not error: 1ab
386 label += 'mean' 1b
387 outmean, outcov = gp.predfromdata(data, label, raw=True) 1ab
388 return outmean + hp.get('mean', gpfactorykw['mu_mu']), outcov 1ab
390 return _pred 1ab
392 @classmethod 1cab
393 def _to_structured(cls, x): 1cab
395 # convert to StructuredArray
396 if hasattr(x, 'columns'): 1ab
397 x = _array.StructuredArray.from_dataframe(x) 1a
398 elif x.dtype.names is None: 1ab
399 x = _array.unstructured_to_structured(x) 1b
400 else:
401 x = _array.StructuredArray(x) 1a
403 # check
404 assert x.ndim == 1 1ab
406 def check_numerical(path, dtype): 1ab
407 if not numpy.issubdtype(dtype, numpy.number): 407 ↛ 408line 407 didn't jump to line 408 because the condition on line 407 was never true1ab
408 msg = f'covariate `{path}` is not numerical'
409 raise TypeError(msg)
411 cls._walk_dtype(x.dtype, check_numerical) 1ab
413 return x 1ab
415 @classmethod 1cab
416 def _walk_dtype(cls, dtype, task, path=None): 1cab
417 if dtype.names is None: 1ab
418 task(path, dtype) 1ab
419 else:
420 for name in dtype.names: 1ab
421 subpath = name if path is None else path + ':' + name 1ab
422 cls._walk_dtype(dtype[name], task, subpath) 1ab
424 @staticmethod 1cab
425 def _toindices(x, splits): 1cab
426 ix = _kernels.BART.indices_from_coord(x, splits) 1ab
427 return _array.unstructured_to_structured(ix, names=x.dtype.names) 1ab
429 def __repr__(self): 1cab
430 out = f"""BART fit: 1a
431alpha = {self.alpha} (0 -> intercept only, 1 -> any)
432beta = {self.beta} (0 -> any, ∞ -> no interactions)
433mean = {self.mean}
434latent sdev = {self.meansdev} (large -> conservative extrapolation)
435data total sdev = {self._ystd:.3g}"""
437 if self._no_weights: 437 ↛ 441line 437 didn't jump to line 441 because the condition on line 437 was always true1a
438 out += f""" 1a
439error sdev = {self.sigma}"""
440 else:
441 weights = numpy.array(self.fit.gpfactorykw['weights'])
442 avgsigma = numpy.sqrt(numpy.mean(self.sigma**2 / weights))
443 out += f"""
444error sdev (avg weighted) = {avgsigma}
445error sdev (unweighted) = {self.sigma}"""
447 return out 1a