Coverage for src/lsqfitgp/_kernels/_bart.py: 99%

1# lsqfitgp/_kernels/_bart.py 

2# 

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

21import itertools 3 ctx1bca

22 

23from jax import jit, lax 3 ctx1bca

24from jax import numpy as jnp 3 ctx1bca

25from jax.scipy import special as jspecial 3 ctx1bca

26from numpy.lib import recfunctions 3 ctx1bca

27 

28from lsqfitgp import _array, _jaxext 3 ctx1bca

29from lsqfitgp._Kernel import kernel 3 ctx1bca

30 

31 

32@kernel(derivable=False, batchbytes=10e6) 3 ctx1bca

33def _BARTBase( 3 ctx1bca

34 x, 

35 y, 

36 alpha=0.95, 

37 beta=2, 

38 maxd=2, 

39 gamma=1, 

40 splits=None, 

41 pnt=None, 

42 intercept=True, 

43 weights=None, 

44 reset=None, 

45 indices=False, 

46): 

47 r""" 

48 BART kernel. 

49 

50 Good default parameters: ``maxd=4, reset=2`` if ``alpha`` and ``beta`` are 

51 kept fixed at the default values, ``maxd=10, reset=[2,4,6,8]`` otherwise. 

52 Derivatives are faster with forward autodiff. 

53 

54 Parameters 

55 ---------- 

56 x, y : arrays 

57 Input points. The array type can be structured, in which case every leaf 

58 field represents a dimension; or unstructured, which specifies a single 

59 dimension. 

60 alpha, beta : scalar 

61 The parameters of the branching probability. 

62 maxd : int 

63 The maximum depth of the trees. 

64 splits : pair of arrays 

65 The first is an int (p,) array containing the number of splitting 

66 points along each dimension, the second has shape (n, p) and contains 

67 the sorted splitting points in each column, filled with high values 

68 after the length. Use `BART.splits_from_coord` to produce them. 

69 gamma : scalar or str 

70 Interpolation coefficient in [0, 1] between a lower and a upper 

71 bound on the infinite maxd limit, or a string 'auto' indicating to 

72 use a formula which depends on alpha, beta, maxd and the number of 

73 covariates, empirically calibrated on maxd from 1 to 3. Default 1 

74 (upper bound). 

75 pnt : (maxd + 1,) array, optional 

76 Nontermination probabilities at depths 0...maxd. If specified, 

77 ``alpha``, ``beta`` and ``maxd`` are ignored. 

78 intercept : bool, default True 

79 The correlation is in [1 - alpha, 1] (or [1 - pnt[0], 1] when using 

80 pnt). If intercept=False, it is rescaled to [0, 1]. 

81 weights : (p,) array, optional 

82 Unnormalized selection probabilities for the covariate axes. If not 

83 specified, all axes have the same probability to be selected for 

84 splitting. 

85 reset : int or sequence of int, optional 

86 List of depths at which the recursion is reset, in the sense that the 

87 function value at a reset depth is evaluated on the initial inputs for 

88 all recursion paths, instead of the modified input handed down by the 

89 recursion. Default none. 

90 indices : bool, default False 

91 If False, the inputs `x`, `y` represent coordinate values. If True, 

92 they are taken to be already the indices of the points in the splitting 

93 grid, as can be obtained with `BART.indices_from_coord`. 

94 

95 Methods 

96 ------- 

97 splits_from_coord 

98 indices_from_coord 

99 correlation 

100 

101 Notes 

102 ----- 

103 This is the covariance function of the latent mean prior of BART (Bayesian 

104 Additive Regression Trees) [1]_ with an upper bound :math:`D` on the depth 

105 of the trees. This prior is the distribution of the function 

106 

107 .. math:: 

108 f(\mathbf x) = \lim_{m\to\infty} 

109 \sum_{j=1}^m g(\mathbf x; T_j, M_j), 

110 

111 where each :math:`g(\mathbf x; T_j, M_j)` is a decision tree evaluated at 

112 :math:`\mathbf x`, with structure :math:`T_j` and leaf values :math:`M_j`. 

113 The trees are i.i.d., with the following distribution for :math:`T_j`: for 

114 a node at depth :math:`d`, with :math:`d = 0` for the root, the probability 

115 of not being a leaf, conditional on its existence and its ancestors only, is 

116 

117 .. math:: 

118 P_d = \alpha (1+d)^{-\beta}, \quad 

119 \alpha \in [0, 1], \quad \beta \ge 0. 

120 

121 For a non-leaf node, conditional on existence and ancestors, the splitting 

122 variable has uniform distribution amongst the variables with any splitting 

123 points not used by ancestors, and the splitting point has uniform 

124 distribution amongst the available ones. The splitting points are fixed, 

125 tipically from the data. 

126 

127 The distribution of leaves :math:`M_j` is i.i.d. Normal with variance 

128 :math:`1/m`, such that :math:`f(x)` has variance 1. In the limit 

129 :math:`m\to\infty`, the distribution of :math:`f(x)` becomes a Gaussian 

130 process. 

131 

132 Since the trees are independent, the covariance function can be computed 

133 for a single tree. Consider two coordinates :math:`x` and :math:`y`, with 

134 :math:`x \le y`. Let :math:`n^-`, :math:`n^0` and :math:`n^+` be the 

135 number of splitting points respectively before :math:`x`, between 

136 :math:`x`, :math:`y` and after :math:`y`. Next, define :math:`\mathbf 

137 n^-`, :math:`\mathbf n^0` and :math:`\mathbf n^+` as the vectors of such 

138 quantities for each dimension, with a total of :math:`p` dimensions, and 

139 :math:`\mathbf n = \mathbf n^- + \mathbf n^0 + \mathbf n^+`. Then the 

140 covariance function can be written recursively as 

141 

142 .. math:: 

143 \newcommand{\nvecs}{\mathbf n^-, \mathbf n^0, \mathbf n^+} 

144 k(\mathbf x, \mathbf y) &= k_0(\nvecs), \\ 

145 k_D(\nvecs) &= 1 - (1 - \gamma) P_D, 

146 \quad \mathbf n^0 \ne \mathbf 0, \\ 

147 k_d(\mathbf 0, \mathbf 0, \mathbf 0) &= 1, \\ 

148 k_d(\nvecs) &= 1 - P_d \Bigg(1 - \frac1{W(\mathbf n)} 

149 \sum_{\substack{i=1 \\ n_i\ne 0}}^p 

150 \frac{w_i}{n_i} \Bigg( \\ 

151 &\qquad \sum_{k=0}^{n^-_i - 1} 

152 k_{d+1}(\mathbf n^-_{n^-_i=k}, \mathbf n^0, \mathbf n^+) 

153 + {} \\ 

154 &\qquad \sum_{k=0}^{n^+_i - 1} 

155 k_{d+1}(\mathbf n^-, \mathbf n^0, \mathbf n^+_{n^+_i=k}) 

156 \Bigg) 

157 \Bigg), \quad d < D, \\ 

158 W(\mathbf n) &= \sum_{\substack{i=1 \\ n_i\ne 0}}^p w_i. 

159 

160 The introduction of a maximum depth :math:`D` is necessary for 

161 computational feasibility. As :math:`D` increases, the result converges to 

162 the one without depth limit. For :math:`D \le 2` (the default value), the 

163 covariance is implemented in closed form and takes :math:`O(p)` to compute. 

164 For :math:`D > 2`, the computational complexity grows exponentially as 

165 :math:`O(p(\bar np)^{D-2})`, where :math:`\bar n` is the average number of splitting 

166 points along a dimension. 

167 

168 In the maximum allowed depth is 1, i.e., either :math:`D = 1` or 

169 :math:`\beta\to\infty`, the kernel assumes the simple form 

170 

171 .. math:: 

172 k(\mathbf x, \mathbf y) &= 1 - P_0 \left( 

173 1 - Q + \frac Q{W(\mathbf n)} 

174 \sum_{\substack{i=1 \\ n_i\ne 0}}^p w_i 

175 \frac{n^0_i}{n_i} \right), \\ 

176 Q &= \begin{cases} 

177 1 - (1 - \gamma) P_1 & \mathbf n^0 \ne \mathbf 0, \\ 

178 1 & \mathbf n^0 = \mathbf 0, 

179 \end{cases} 

180 

181 which is separable along dimensions, i.e., it has no interactions. 

182 

183 References 

184 ---------- 

185 .. [1] Hugh A. Chipman, Edward I. George, Robert E. McCulloch "BART: 

186 Bayesian additive regression trees," The Annals of Applied Statistics, 

187 Ann. Appl. Stat. 4(1), 266-298, (March 2010). 

188 """ # noqa: DOC001 

189 splits = BART._check_splits(splits, indices) 3 ctx1bca

190 if not x.dtype.names: 3 ctx1bca

191 x = x[..., None] 2 ctx1ba

192 if not y.dtype.names: 3 ctx1bca

193 y = y[..., None] 2 ctx1ba

194 if indices: 3 ctx1bca

195 ix = BART._check_x(x) 2 ctx1ca

196 iy = BART._check_x(y) 2 ctx1ca

197 else: 

198 ix = BART._indices_from_coord(x, splits) 2 ctx1ba

199 iy = BART._indices_from_coord(y, splits) 2 ctx1ba

200 return BART.correlation( 3 ctx1bca

201 splits[0], 

202 ix, 

203 iy, 

204 pnt=pnt, 

205 alpha=alpha, 

206 beta=beta, 

207 gamma=gamma, 

208 maxd=maxd, 

209 intercept=intercept, 

210 weights=weights, 

211 reset=reset, 

212 altinput=True, 

213 ) 

214 

215 

216class BART(_BARTBase): 3 ctx1bca

217 __doc__ = _BARTBase.__doc__ 3 ctx1bca

218 

219 @classmethod 3 ctx1bca

220 def splits_from_coord(cls, x): 3 ctx1bca

221 """ 

222 Generate splitting points from data. 

223 

224 Parameters 

225 ---------- 

226 x : array of numbers 

227 The data. Can be passed in two formats: 1) a structured array where 

228 each leaf field represents a dimension, 2) a normal array where the 

229 last axis runs over dimensions. In the structured case, each 

230 index in any shaped field is a different dimension. 

231 

232 Returns 

233 ------- 

234 length : int (p,) array 

235 The number of splitting points along each of ``p`` dimensions. 

236 splits : (n, p) array 

237 Each column contains the sorted splitting points along a dimension. 

238 The splitting points are the midpoints between consecutive values 

239 appearing in `x` for that dimension. Column ``splits[:, i]`` 

240 contains splitting points only up to ``length[i]``, while afterward 

241 it is filled with a very large value. 

242 """ 

243 x = cls._check_x(x) 3 ctx1bca

244 return cls._splits_from_coord(x) 3 ctx1bca

245 

246 @staticmethod 3 ctx1bca

247 @jit 3 ctx1bca

248 def _splits_from_coord(x): 3 ctx1bca

249 """ 

250 Jitted implementation of splits_from_coord. 

251 

252 Applying jit avoids the recompilation in lax.scan each time the method 

253 is called, and splits_from_coord can not be jitted directly because x 

254 could be a numpy structured array. 

255 """ 

256 x = x.reshape(-1, x.shape[-1]) if x.size else x.reshape(1, x.shape[-1]) 3 ctx1bca

257 if jnp.issubdtype(x.dtype, jnp.inexact): 3 ctx1bca

258 info = jnp.finfo 3 ctx1bca

259 else: 

260 info = jnp.iinfo 1 ctx1a

261 fill = info(x.dtype).max 3 ctx1bca

262 

263 def loop(_, xi): 3 ctx1bca

264 u = jnp.unique(xi, size=xi.size, fill_value=fill) 3 ctx1bca

265 m = jnp.where(u[1:] < fill, (u[1:] + u[:-1]) / 2, fill) 3 ctx1bca

266 l = jnp.searchsorted(m, fill) 3 ctx1bca

267 return _, (l, m) 3 ctx1bca

268 

269 _, (length, midpoints) = lax.scan(loop, None, x.T) 3 ctx1bca

270 return length, midpoints.T 3 ctx1bca

271 

272 @classmethod 3 ctx1bca

273 def indices_from_coord(cls, x, splits): 3 ctx1bca

274 """ 

275 Convert coordinates to indices w.r.t. splitting points. 

276 

277 Parameters 

278 ---------- 

279 x : array of numbers 

280 The coordinates. Can be passed in two formats: 1) a structured 

281 array where each leaf field represents a dimension, 2) a normal 

282 array where the last axis runs over dimensions. In the structured 

283 case, each index in any shaped field is a different dimension. 

284 splits : pair of arrays 

285 The first is an int (p,) array containing the number of splitting 

286 points along each dimension, the second has shape (n, p) and 

287 contains the sorted splitting points in each column, filled with 

288 high values after the length. 

289 

290 Returns 

291 ------- 

292 ix : int array 

293 An array with the same shape as ``x``, unless ``x`` is a structured 

294 array, in which case the last axis of ``ix`` is the flattened version 

295 of the structured type. ``ix`` contains indices mapping ``x`` to 

296 positions between splitting points along each coordinate, with the 

297 following convention: index 0 means before the first split, index 

298 i > 0 means between split i - 1 and split i. 

299 """ 

300 splits = cls._check_splits(splits, False) 2 ctx1ca

301 return cls._indices_from_coord(x, splits) 2 ctx1ca

302 

303 @classmethod 3 ctx1bca

304 def _indices_from_coord(cls, x, checked_splits): 3 ctx1bca

305 x = cls._check_x(x) 3 ctx1bca

306 if x.shape[-1] != checked_splits[0].size: 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true3 ctx1bca

307 msg = ( 

308 f'splitting grid is for {checked_splits[0].size} ' 

309 f'dimensions, found {x.shape[-1]}' 

310 ) 

311 raise ValueError(msg) 

312 return cls._searchsorted_vectorized(checked_splits[1], x) 3 ctx1bca

313 

314 @classmethod 3 ctx1bca

315 def correlation( # noqa: C901, PLR0915 3 ctx1bca

316 cls, 

317 splitsbefore_or_totalsplits, 

318 splitsbetween_or_index1, 

319 splitsafter_or_index2, 

320 *, 

321 alpha=0.95, 

322 beta=2, 

323 gamma=1, 

324 maxd=2, 

325 debug=False, 

326 pnt=None, 

327 intercept=True, 

328 weights=None, 

329 reset=None, 

330 altinput=False, 

331 ): 

332 """ 

333 Compute the BART prior correlation between two points. 

334 

335 Apart from arguments ``maxd``, ``debug`` and ``reset``, this method is fully 

336 vectorized. 

337 

338 Parameters 

339 ---------- 

340 splitsbefore_or_totalsplits : int (p,) array 

341 The number of splitting points less than the two points, separately 

342 along each coordinate, or the total number of splits if ``altinput``. 

343 splitsbetween_or_index1 : int (p,) array 

344 The number of splitting points between the two points, separately 

345 along each coordinate, or the index in the splitting bins of the 

346 first point if ``altinput``, where 0 means to the left of the leftmost 

347 splitting point. 

348 splitsafter_or_index2 : int (p,) array 

349 The number of splitting points greater than the two points, 

350 separately along each coordinate, or the index in the splitting bins 

351 of the second point if ``altinput``. 

352 alpha : scalar 

353 See `BART`. 

354 beta : scalar 

355 See `BART`. 

356 gamma : scalar or str 

357 See `BART`. 

358 maxd : int 

359 See `BART`. 

360 debug : bool 

361 If True, disable shortcuts in the tree recursion. Default False. 

362 pnt : (maxd + 1,) array, optional 

363 See `BART`. 

364 intercept : bool, default True 

365 See `BART`. 

366 weights : (p,) array, optional 

367 See `BART`. 

368 reset : int or sequence of int, optional 

369 See `BART`. 

370 altinput : bool 

371 If True, take as input the indices in the splitting bins of the 

372 points instead of the counts of splitting points separating them, 

373 and use a different implementation optimized for that case. Default 

374 False. The `BART` kernel uses ``altinput=True``. 

375 

376 Returns 

377 ------- 

378 corr : scalar 

379 The prior correlation. 

380 """ 

381 # check splitting indices are integers 

382 splitsbefore_or_totalsplits = jnp.asarray(splitsbefore_or_totalsplits) 3 ctx1bca

383 splitsbetween_or_index1 = jnp.asarray(splitsbetween_or_index1) 3 ctx1bca

384 splitsafter_or_index2 = jnp.asarray(splitsafter_or_index2) 3 ctx1bca

385 assert jnp.issubdtype(splitsbefore_or_totalsplits.dtype, jnp.integer) 3 ctx1bca

386 assert jnp.issubdtype(splitsbetween_or_index1.dtype, jnp.integer) 3 ctx1bca

387 assert jnp.issubdtype(splitsafter_or_index2.dtype, jnp.integer) 3 ctx1bca

388 

389 # check splitting indices 

390 with _jaxext.skipifabstract(): 3 ctx1bca

391 assert jnp.all(splitsbefore_or_totalsplits >= 0), ( 3 ctx1bca

392 'splitting counts must be nonnegative' 

393 ) 

394 if altinput: 3 ctx1bca

395 assert jnp.all( 3 ctx1bca

396 (splitsbetween_or_index1 >= 0) 

397 & (splitsbetween_or_index1 <= splitsbefore_or_totalsplits) 

398 ), 'splitting index must be in [0, n]' 

399 assert jnp.all( 3 ctx1bca

400 (splitsafter_or_index2 >= 0) 

401 & (splitsafter_or_index2 <= splitsbefore_or_totalsplits) 

402 ), 'splitting index must be in [0, n]' 

403 else: 

404 assert jnp.all(splitsbetween_or_index1 >= 0), ( 1 ctx1a

405 'splitting counts must be nonnegative' 

406 ) 

407 assert jnp.all(splitsafter_or_index2 >= 0), ( 1 ctx1a

408 'splitting counts must be nonnegative' 

409 ) 

410 

411 # get splitting probabilities 

412 if pnt is None: 3 ctx1bca

413 assert maxd == int(maxd), maxd 3 ctx1bca

414 assert maxd >= 0, maxd 3 ctx1bca

415 alpha = jnp.asarray(alpha) 3 ctx1bca

416 beta = jnp.asarray(beta) 3 ctx1bca

417 with _jaxext.skipifabstract(): 3 ctx1bca

418 assert jnp.all((alpha >= 0) & (alpha <= 1)), 'alpha must be in [0, 1]' 3 ctx1bca

419 assert jnp.all(beta >= 0), 'beta must be in [0, inf)' 3 ctx1bca

420 d = jnp.arange(maxd + 1) 3 ctx1bca

421 alpha = alpha[..., None] 3 ctx1bca

422 beta = beta[..., None] 3 ctx1bca

423 pnt = alpha / (1 + d) ** beta 3 ctx1bca

424 else: 

425 pnt = jnp.asarray(pnt) 1 ctx1a

426 

427 # get covariate weights 

428 if weights is None: 3 ctx1bca

429 weights = jnp.ones(splitsbefore_or_totalsplits.shape[-1], pnt.dtype) 3 ctx1bca

430 else: 

431 weights = jnp.asarray(weights) 1 ctx1a

432 

433 # get interpolation coefficients 

434 if isinstance(gamma, str): 3 ctx1bca

435 if gamma == 'auto': 1 ctx1a

436 assert reset is None 1 ctx1a

437 assert 1 <= pnt.shape[-1] - 1 <= 3 1 ctx1a

438 p = weights.shape[-1] 1 ctx1a

439 gamma = cls._gamma(p, pnt) 1 ctx1a

440 else: 

441 raise KeyError(gamma) 1 ctx1a

442 else: 

443 gamma = jnp.asarray(gamma) 3 ctx1bca

444 

445 # check values are in range 

446 with _jaxext.skipifabstract(): 3 ctx1bca

447 assert jnp.all((gamma >= 0) & (gamma <= 1)), 'gamma must be in [0, 1]' 3 ctx1bca

448 assert jnp.all((pnt >= 0) & (pnt <= 1)), 'pnt must be in [0, 1]' 3 ctx1bca

449 assert jnp.all(weights >= 0), 'weights must be in [0, inf)' 3 ctx1bca

450 

451 # set first splitting probability to 1 to remove flat baseline (keep 

452 # last!) 

453 if not intercept: 3 ctx1bca

454 pnt = pnt.at[..., 0].set(1) 2 ctx1ca

455 

456 # expand and check recursion reset depths 

457 if reset is None: 3 ctx1bca

458 reset = [] 2 ctx1ba

459 if not hasattr(reset, '__len__'): 3 ctx1bca

460 reset = [reset] 1 ctx1a

461 reset = [0, *reset, pnt.shape[-1] - 1] 3 ctx1bca

462 for i, j in itertools.pairwise(reset): 3 ctx1bca

463 assert int(j) == j, (i, j) 3 ctx1bca

464 assert i <= j, (i, j) 3 ctx1bca

465 

466 # convert reset depths list to brackets with repetition 

467 brackets_norep = list(itertools.pairwise(reset)) 3 ctx1bca

468 brackets = [brackets_norep[0] + (1,)] 3 ctx1bca

469 for t, b in brackets_norep[1:]: 3 ctx1bca

470 lt, lb, lr = brackets[-1] 2 ctx1ca

471 if altinput and not debug and lr * (b - t) == lb - lt and b - t <= 2: 2 ctx1ca

472 brackets[-1] = lt, b, lr + 1 2 ctx1ca

473 else: 

474 brackets.append((t, b, 1)) 1 ctx1a

475 

476 # call recursive function for each recursion slice 

477 corr = gamma 3 ctx1bca

478 for t, b, repeat in reversed(brackets): 3 ctx1bca

479 probs = pnt[..., t : b + 1] 3 ctx1bca

480 if t > 0: 3 ctx1bca

481 probs = probs.at[..., 0].set(1) 1 ctx1a

482 if repeat > 1: 3 ctx1bca

483 head = probs[..., 0:1] 2 ctx1ca

484 one = jnp.ones_like(head) 2 ctx1ca

485 probs = jnp.concatenate( 2 ctx1ca

486 list( 

487 itertools.chain.from_iterable( 

488 reversed( 

489 [ 

490 [head if i == 0 else one, p] 

491 for i, p in enumerate( 

492 jnp.split(probs[..., 1:], repeat, axis=-1) 

493 ) 

494 ] 

495 ) 

496 ) 

497 ), 

498 axis=-1, 

499 ) 

500 else: 

501 repeat = None 2 ctx1ba

502 corr = cls._correlation_vectorized( 3 ctx1bca

503 splitsbefore_or_totalsplits, 

504 splitsbetween_or_index1, 

505 splitsafter_or_index2, 

506 probs, 

507 corr, 

508 weights, 

509 debug, 

510 altinput, 

511 repeat, 

512 ) 

513 return corr 3 ctx1bca

514 

515 @staticmethod 3 ctx1bca

516 def _gamma(p, pnt): 3 ctx1bca

517 # gamma(alpha, beta, maxd) = 

518 # = (gamma_0 - gamma_d maxd) (1 - alpha^s 2^(-t beta)) = 

519 # = (gamma_0 - gamma_d maxd) (1 - P0^s-t P1^t) 

520 

521 gamma_0 = 0.611 + 0.021 * jnp.exp(-1.3 * (p - 1)) 1 ctx1a

522 gamma_d = -0.0034 + 0.084 * jnp.exp(-2.02 * (p - 1)) 1 ctx1a

523 s = 2.03 - 0.69 * jnp.exp(-0.72 * (p - 1)) 1 ctx1a

524 t = 4.01 - 1.49 * jnp.exp(-0.77 * (p - 1)) 1 ctx1a

525 

526 maxd = pnt.shape[-1] - 1 1 ctx1a

527 floor = jnp.clip(gamma_0 - gamma_d * maxd, 0, 1) 1 ctx1a

528 

529 P0 = pnt[..., 0] 1 ctx1a

530 P1 = jnp.minimum(P0, pnt[..., 1]) 1 ctx1a

531 corner = jnp.where(P0, 1 - P0 ** (s - t) * P1**t, 1) 1 ctx1a

532 

533 return floor * corner 1 ctx1a

534 

535 @staticmethod 3 ctx1bca

536 def _check_x(x): 3 ctx1bca

537 x = _array.asarray(x) 3 ctx1bca

538 if x.dtype.names: 3 ctx1bca

539 x = recfunctions.structured_to_unstructured(x) 2 ctx1ca

540 return x 3 ctx1bca

541 

542 @staticmethod 3 ctx1bca

543 def _check_splits(splits, indices): 3 ctx1bca

544 l, s = splits 3 ctx1bca

545 l = jnp.asarray(l) 3 ctx1bca

546 assert l.ndim == 1 3 ctx1bca

547 if not indices: 3 ctx1bca

548 s = jnp.asarray(s) 3 ctx1bca

549 assert 1 <= s.ndim <= 2 3 ctx1bca

550 if s.ndim == 1: 3 ctx1bca

551 s = s[:, None] 2 ctx1ba

552 assert l.size == s.shape[1] 3 ctx1bca

553 with _jaxext.skipifabstract(): 3 ctx1bca

554 assert jnp.all((l >= 0) & (l <= s.shape[0])), 'length out of bounds' 3 ctx1bca

555 if not indices: 3 ctx1bca

556 assert jnp.all(jnp.sort(s, axis=0) == s), 'unsorted splitting points' 3 ctx1bca

557 return l, s 3 ctx1bca

558 

559 @staticmethod 3 ctx1bca

560 @functools.partial(jit, static_argnames=('side',)) 3 ctx1bca

561 def _searchsorted_vectorized(A, V, **kw): 3 ctx1bca

562 """ 

563 Apply `jnp.searchsorted` separately on each column. 

564 

565 A : (n, p) 

566 V : (..., p) 

567 out : (..., p). 

568 """ 

569 

570 def loop(_, av): 3 ctx1bca

571 return _, jnp.searchsorted(*av, **kw) 3 ctx1bca

572 

573 _, out = lax.scan(loop, None, (A.T, V.T)) 3 ctx1bca

574 return out.T 3 ctx1bca

575 

576 @classmethod 3 ctx1bca

577 @functools.partial(jit, static_argnums=(0, 7)) 3 ctx1bca

578 def _correlation_old(cls, nminus, n0, nplus, pnt, gamma, w, debug): # noqa: PLR0915 3 ctx1bca

579 """Old version, kept around for cross-checking.""" 

580 assert nminus.shape == n0.shape == nplus.shape == w.shape 1 ctx1a

581 assert nminus.ndim == 1 1 ctx1a

582 assert nminus.size >= 0 1 ctx1a

583 assert pnt.ndim == 1 1 ctx1a

584 assert pnt.size > 0 1 ctx1a

585 

586 # optimization to avoid looping over ignored axes 

587 nminus = jnp.where(w, nminus, 0) 1 ctx1a

588 n0 = jnp.where(w, n0, 0) 1 ctx1a

589 nplus = jnp.where(w, nplus, 0) 1 ctx1a

590 

591 float_type = _jaxext.float_type(pnt, gamma, w) 1 ctx1a

592 

593 if nminus.size == 0: 1 ctx1a

594 return jnp.array(1, float_type) 1 ctx1a

595 

596 anyn0 = jnp.any(jnp.logical_and(n0, w)) 1 ctx1a

597 

598 if pnt.size == 1: 1 ctx1a

599 return jnp.where(anyn0, 1 - (1 - gamma) * pnt[0], 1) 1 ctx1a

600 

601 nout = nminus + nplus 1 ctx1a

602 n = nout + n0 1 ctx1a

603 Wn = jnp.sum(jnp.where(n, w, 0)) # <-- @ 1 ctx1a

604 

605 if pnt.size == 2 and not debug: 1 ctx1a

606 Q = 1 - (1 - gamma) * pnt[1] 1 ctx1a

607 sump = Q * jnp.sum(jnp.where(n, w * nout / n, 0)) # <-- @ 1 ctx1a

608 return jnp.where(anyn0, 1 - pnt[0] * (1 - sump / Wn), 1) 1 ctx1a

609 

610 if pnt.size == 3 and not debug: 1 ctx1a

611 Q = 1 - (1 - gamma) * pnt[2] 1 ctx1a

612 s = w * nout / n 1 ctx1a

613 S = jnp.sum(jnp.where(n, s, 0)) # <-- @ 1 ctx1a

614 t = w * n0 / n 1 ctx1a

615 psin = jspecial.digamma(n.astype(float_type)) 1 ctx1a

616 

617 def terms(nminus, nplus): 1 ctx1a

618 nminus0 = nminus + n0 1 ctx1a

619 Wnmod = Wn - jnp.where(nminus0, 0, w) 1 ctx1a

620 frac = jnp.where(nminus0, w * nminus / nminus0, 0) 1 ctx1a

621 terms1 = (S - s + frac) / Wnmod 1 ctx1a

622 psi1nminus0 = jspecial.digamma((1 + nminus0).astype(float_type)) 1 ctx1a

623 terms2 = ((nplus - 1) * (S + t) - w * n0 * (psin - psi1nminus0)) / Wn 1 ctx1a

624 return jnp.where(nplus, terms1 + terms2, 0) 1 ctx1a

625 

626 tplus = terms(nminus, nplus) 1 ctx1a

627 tminus = terms(nplus, nminus) 1 ctx1a

628 tall = jnp.where(n, w * (tplus + tminus) / n, 0) 1 ctx1a

629 sump = (1 - pnt[1]) * S + pnt[1] * Q * jnp.sum(tall) # <-- @ 1 ctx1a

630 return jnp.where(anyn0, 1 - pnt[0] * (1 - sump / Wn), 1) 1 ctx1a

631 

632 p = len(nminus) 1 ctx1a

633 

634 val = (0.0, nminus, n0, nplus) 1 ctx1a

635 

636 def loop(i, val): 1 ctx1a

637 sump, nminus, n0, nplus = val 1 ctx1a

638 

639 nminusi = nminus[i] 1 ctx1a

640 n0i = n0[i] 1 ctx1a

641 nplusi = nplus[i] 1 ctx1a

642 ni = nminusi + n0i + nplusi 1 ctx1a

643 

644 val = (0.0, nminus, n0, nplus, i, nminusi) 1 ctx1a

645 

646 def loop(k, val): 1 ctx1a

647 sumn, nminus, n0, nplus, i, nminusi = val 1 ctx1a

648 

649 # here I use the fact that .at[].set won't set the value if the 

650 # index is out of bounds 

651 nminus = nminus.at[jnp.where(k < nminusi, i, i + p)].set(k) 1 ctx1a

652 nplus = nplus.at[jnp.where(k >= nminusi, i, i + p)].set(k - nminusi) 1 ctx1a

653 

654 sumn += cls._correlation_old( 1 ctx1a

655 nminus, n0, nplus, pnt[1:], gamma, w, debug 

656 ) 

657 

658 nminus = nminus.at[i].set(nminusi) 1 ctx1a

659 nplus = nplus.at[i].set(nplusi) 1 ctx1a

660 

661 return sumn, nminus, n0, nplus, i, nminusi 1 ctx1a

662 

663 # if ni == 0 I skip recursion by passing 0 as iteration end 

664 end = jnp.where(ni, nminusi + nplusi, 0) 1 ctx1a

665 start = jnp.zeros_like(end) 1 ctx1a

666 sumn, nminus, n0, nplus, _, _ = lax.fori_loop(start, end, loop, val) 1 ctx1a

667 

668 sump += jnp.where(ni, w[i] * sumn / ni, 0) 1 ctx1a

669 

670 return sump, nminus, n0, nplus 1 ctx1a

671 

672 # skip summation if all(n0 == 0) 

673 end = jnp.where(anyn0, p, 0) 1 ctx1a

674 sump, _, _, _ = lax.fori_loop(0, end, loop, val) 1 ctx1a

675 

676 return jnp.where(anyn0, 1 - pnt[0] * (1 - sump / Wn), 1) 1 ctx1a

677 

678 @staticmethod 3 ctx1bca

679 def _scan_but_first(f, init, xs): 3 ctx1bca

680 """ 

681 `lax.scan`, but execute separately the first cycle. 

682 

683 The point is that I use it when the first cycle works on smaller arrays 

684 due to broadcasting. 

685 """ 

686 assert isinstance(xs, jnp.ndarray) 3 ctx1bca

687 assert len(xs) > 0 3 ctx1bca

688 init, out = f(init, xs[0]) 3 ctx1bca

689 assert out is None 3 ctx1bca

690 if len(xs) == 1: 3 ctx1bca

691 return init, out 2 ctx1ba

692 elif len(xs) == 2: 2 ctx1ca

693 return f(init, xs[1]) 1 ctx1a

694 else: 

695 return lax.scan(f, init, xs[1:]) 2 ctx1ca

696 

697 @classmethod 3 ctx1bca

698 @functools.partial(jit, static_argnums=(0, 7, 8)) 3 ctx1bca

699 def _correlation(cls, n, ix, iy, pnt, gamma, w, debug, repeat): # noqa: C901, PLR0915 3 ctx1bca

700 # this implementation is optimized assuming that the shapes are as 

701 # follows: 

702 # n (p,) 

703 # ix (n, 1, p) 

704 # iy (1, n, p) 

705 # pnt (d,) 

706 # gamma () or (n, n) 

707 # w (p,) 

708 

709 assert n.ndim == 1 3 ctx1bca

710 assert n.shape == ix.shape == iy.shape == w.shape 3 ctx1bca

711 assert pnt.ndim == 1 3 ctx1bca

712 assert pnt.size > 0 3 ctx1bca

713 assert gamma.ndim == 0 3 ctx1bca

714 

715 # check the strict conditions under which `repeat` is implemented 

716 if repeat is not None: 3 ctx1bca

717 assert not debug 2 ctx1ca

718 assert repeat > 0 2 ctx1ca

719 assert pnt.size % repeat == 0 2 ctx1ca

720 assert pnt.size // repeat <= 3 2 ctx1ca

721 else: 

722 repeat = 1 2 ctx1ba

723 

724 # infer float type from float arguments 

725 flt = _jaxext.float_type(pnt, gamma, w) 3 ctx1bca

726 

727 # no covariates, always return 1 

728 if n.size == 0: 728 ↛ 729line 728 didn't jump to line 729 because the condition on line 728 was never true3 ctx1bca

729 return jnp.array(1, flt) 

730 

731 # pre-cast all floats to the common type, to avoid unwanted float32 

732 # calculations in mixed float-integer operations 

733 pnt = pnt.astype(flt) 3 ctx1bca

734 gamma = gamma.astype(flt) 3 ctx1bca

735 w = w.astype(flt) 3 ctx1bca

736 

737 # ignore zero-weight axes 

738 n = jnp.where(w, n, 0) 3 ctx1bca

739 ix = jnp.where(w, ix, 0) 3 ctx1bca

740 iy = jnp.where(w, iy, 0) 3 ctx1bca

741 

742 # check if the points coincide 

743 seed = jnp.uint64(16132933535611723338) 3 ctx1bca

744 hx = _jaxext.fasthash64(ix, seed) 3 ctx1bca

745 hy = _jaxext.fasthash64(iy, seed) 3 ctx1bca

746 anyn0 = hx != hy 3 ctx1bca

747 # no hash collision checking, it would be branchless because of vmap, 

748 # the probability of collision building a nxn matrix with n=10000 is 

749 # -expm1(10000**2 * log1p(-1/2**64)) = 5e-12. 

750 

751 # base case of the recursion, no dependence on points apart from the 

752 # case when they are equal 

753 if pnt.size // repeat == 1: 3 ctx1bca

754 

755 def loop(carry, pnt): 1 ctx1a

756 anyn0, gamma = carry 1 ctx1a

757 gamma = jnp.where(anyn0, 1 - (1 - gamma) * pnt[0], 1) 1 ctx1a

758 return (anyn0, gamma), None 1 ctx1a