Coverage for src/lsqfitgp/_GP/_elements.py: 100%

1# lsqfitgp/_GP/_elements.py 

2# 

3# Copyright (c) 2020, 2022, 2023, 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 abc 6 ctx1feabcd

21import functools 6 ctx1feabcd

22import warnings 6 ctx1feabcd

23import math 6 ctx1feabcd

24 

25import gvar 6 ctx1feabcd

26import numpy 6 ctx1feabcd

27from scipy import sparse 6 ctx1feabcd

28import jax 6 ctx1feabcd

29from jax import numpy as jnp 6 ctx1feabcd

30 

31from .. import _Deriv 6 ctx1feabcd

32from .. import _array 6 ctx1feabcd

33from .. import _jaxext 6 ctx1feabcd

34from .. import _gvarext 6 ctx1feabcd

35from .. import _linalg 6 ctx1feabcd

36 

37from . import _base 6 ctx1feabcd

38 

39class GPElements(_base.GPBase): 6 ctx1feabcd

40 

41 def __init__(self, *, checkpos, checksym, posepsfac, halfmatrix): 6 ctx1feabcd

42 self._elements = dict() # key -> _Element 6 ctx1feabcd

43 self._covblocks = dict() # (key, key) -> matrix (2d flattened) 6 ctx1feabcd

44 self._priordict = {} # key -> gvar array (shaped) 6 ctx1feabcd

45 self._checkpositive = bool(checkpos) 6 ctx1feabcd

46 self._posepsfac = float(posepsfac) 6 ctx1feabcd

47 self._checksym = bool(checksym) 6 ctx1feabcd

48 self._halfmatrix = bool(halfmatrix) 6 ctx1feabcd

49 self._dtype = None 6 ctx1feabcd

50 assert not (halfmatrix and checksym) 6 ctx1feabcd

51 

52 def _clone(self): 6 ctx1feabcd

53 newself = super()._clone() 6 ctx1feabcd

54 newself._elements = self._elements.copy() 6 ctx1feabcd

55 newself._covblocks = self._covblocks.copy() 6 ctx1feabcd

56 newself._priordict = self._priordict.copy() 6 ctx1feabcd

57 newself._checkpositive = self._checkpositive 6 ctx1feabcd

58 newself._posepsfac = self._posepsfac 6 ctx1feabcd

59 newself._checksym = self._checksym 6 ctx1feabcd

60 newself._halfmatrix = self._halfmatrix 6 ctx1feabcd

61 newself._dtype = self._dtype 6 ctx1feabcd

62 return newself 6 ctx1feabcd

63 

64 @staticmethod 6 ctx1feabcd

65 def _concatenate(alist): 6 ctx1feabcd

66 """ 

67 Decides to use numpy.concatenate or jnp.concatenate depending on the 

68 input to support gvars. 

69 """ 

70 if any(a.dtype == object for a in alist): 6 ctx1feabcd

71 return numpy.concatenate(alist) 6 ctx1feabcd

72 else: 

73 return jnp.concatenate(alist) 6 ctx1feabcd

74 

75 @staticmethod 6 ctx1feabcd

76 def _triu_indices_and_back(n): 6 ctx1feabcd

77 """ 

78 Return indices to get the upper triangular part of a matrix, and indices 

79 to convert a flat array of upper triangular elements to a symmetric 

80 matrix. 

81 """ 

82 ix, iy = jnp.triu_indices(n) 4 ctx1abcd

83 q = jnp.empty((n, n), ix.dtype) 4 ctx1abcd

84 a = jnp.arange(ix.size) 4 ctx1abcd

85 q = q.at[ix, iy].set(a) 4 ctx1abcd

86 q = q.at[iy, ix].set(a) 4 ctx1abcd

87 return ix, iy, q 4 ctx1abcd

88 

89 class _Element(abc.ABC): 6 ctx1feabcd

90 """ 

91 Abstract class for an object holding information associated to a key in 

92 a GP object. 

93 """ 

94 

95 @property 6 ctx1feabcd

96 @abc.abstractmethod 6 ctx1feabcd

97 def shape(self): # pragma: no cover 6 ctx1feabcd

98 """Output shape""" 

99 pass 

100 

101 @property 6 ctx1feabcd

102 def size(self): 6 ctx1feabcd

103 return math.prod(self.shape) 6 ctx1feabcd

104 

105 class _Points(_Element): 6 ctx1feabcd

106 """Points where the process is evaluated""" 

107 

108 def __init__(self, x, deriv, proc): 6 ctx1feabcd

109 assert isinstance(x, (numpy.ndarray, jnp.ndarray, _array.StructuredArray)) 6 ctx1feabcd

110 assert isinstance(deriv, _Deriv.Deriv) 6 ctx1feabcd

111 self.x = x 6 ctx1feabcd

112 self.deriv = deriv 6 ctx1feabcd

113 self.proc = proc 6 ctx1feabcd

114 

115 @property 6 ctx1feabcd

116 def shape(self): 6 ctx1feabcd

117 return self.x.shape 6 ctx1feabcd

118 

119 class _LinTransf(_Element): 6 ctx1feabcd

120 """Linear transformation of other _Element objects""" 

121 

122 shape = None 6 ctx1feabcd

123 

124 def __init__(self, transf, keys, shape): 6 ctx1feabcd

125 self.transf = transf 6 ctx1feabcd

126 self.keys = keys 6 ctx1feabcd

127 self.shape = shape 6 ctx1feabcd

128 

129 def matrices(self, gp): 6 ctx1feabcd

130 """ 

131 Matrix coefficients of the transformation (with flattened inputs 

132 and output) 

133 """ 

134 elems = [gp._elements[key] for key in self.keys] 4 ctx1abcd

135 matrices = [] 4 ctx1abcd

136 transf = jax.vmap(self.transf, 0, 0) 4 ctx1abcd

137 for i, elem in enumerate(elems): 4 ctx1abcd

138 inputs = [ 4 ctx1abcd

139 jnp.eye(elem.size).reshape((elem.size,) + elem.shape) 

140 if j == i else 

141 jnp.zeros((elem.size,) + ej.shape) 

142 for j, ej in enumerate(elems) 

143 ] 

144 output = transf(*inputs).reshape(elem.size, self.size).T 4 ctx1abcd

145 matrices.append(output) 4 ctx1abcd

146 return matrices 4 ctx1abcd

147 

148 class _Cov(_Element): 6 ctx1feabcd

149 """User-provided covariance matrix block(s)""" 

150 

151 shape = None 6 ctx1feabcd

152 

153 def __init__(self, blocks, shape): 6 ctx1feabcd

154 """ blocks = dict (key, key) -> matrix """ 

155 self.blocks = blocks 5 ctx1eabcd

156 self.shape = shape 5 ctx1eabcd

157 

158 @_base.newself 6 ctx1feabcd

159 def addx(self, x, key=None, *, deriv=0, proc=_base.GPBase.DefaultProcess): 6 ctx1feabcd

160 """ 

161  

162 Add points where the Gaussian process is evaluated. 

163  

164 The GP object keeps the various x arrays in a dictionary. If ``x`` is an 

165 array, you have to specify its dictionary key with the ``key`` parameter. 

166 Otherwise, you can directly pass a dictionary for ``x``. 

167  

168 To specify that on the given ``x`` a derivative of the process instead of 

169 the process itself should be evaluated, use the parameter ``deriv``. 

170  

171 `addx` may or may not copy the input arrays. 

172  

173 Parameters 

174 ---------- 

175 x : array or dictionary of arrays 

176 The points to be added. 

177 key : hashable 

178 If ``x`` is an array, the dictionary key under which ``x`` is added. 

179 Can not be specified if ``x`` is a dictionary. 

180 deriv : Deriv-like 

181 Derivative specification. A `Deriv` object or something that 

182 can be converted to `Deriv`. 

183 proc : hashable 

184 The process to be evaluated on the points. If not specified, use 

185 the default process. 

186  

187 """ 

188 

189 # TODO after I implement block solving, add per-key covariance matrix 

190 # flags. 

191 

192 # TODO add `copy` parameter, default False, to copy the input arrays 

193 # if they are numpy arrays. 

194 

195 # this interface does not allow adding a single dictionary as x element 

196 # unless it's wrapped as a 0d numpy array, but this is for the best 

197 

198 deriv = _Deriv.Deriv(deriv) 6 ctx1feabcd

199 

200 if proc not in self._procs: 6 ctx1feabcd

201 raise KeyError(f'process named {proc!r} not found') 4 ctx1abcd

202 

203 if hasattr(x, 'keys'): 6 ctx1feabcd

204 if key is not None: 4 ctx1abcd

205 raise ValueError('can not specify key if x is a dictionary') 4 ctx1abcd

206 if None in x: 4 ctx1abcd

207 raise ValueError('None key in x not allowed') 4 ctx1abcd

208 else: 

209 if key is None: 6 ctx1feabcd

210 raise ValueError('x is not dictionary but key is None') 4 ctx1abcd

211 x = {key: x} 6 ctx1feabcd

212 

213 for key in x: 6 ctx1feabcd

214 if key in self._elements: 6 ctx1feabcd

215 raise KeyError('key {!r} already in GP'.format(key)) 4 ctx1abcd

216 

217 gx = x[key] 6 ctx1feabcd

218 

219 # Convert to JAX array, numpy array or StructuredArray. 

220 # convert eagerly to jax to avoid problems with tracing. 

221 gx = _array._asarray_jaxifpossible(gx) 6 ctx1feabcd

222 

223 # Check dtype is compatible with previous arrays. 

224 # TODO since we never concatenate arrays we could allow a less 

225 # strict compatibility. In principle we could allow really anything 

226 # as long as the kernel eats it, but this probably would let bugs 

227 # through without being really ever useful. What would make sense 

228 # is checking the dtype structure matches recursively and check 

229 # concrete dtypes of fields can be casted. 

230 # TODO result_type is too lax. Examples: str, float -> str, 

231 # object, float -> object. I should use something like the 

232 # ordering function in updowncast.py. 

233 if self._dtype is not None: 6 ctx1feabcd

234 try: 6 ctx1feabcd

235 self._dtype = numpy.result_type(self._dtype, gx.dtype) 6 ctx1feabcd

236 # do not use jnp.result_type, it does not support 

237 # structured types 

238 except TypeError: 4 ctx1abcd

239 msg = 'x[{!r}].dtype = {!r} not compatible with {!r}' 4 ctx1abcd

240 msg = msg.format(key, gx.dtype, self._dtype) 4 ctx1abcd

241 raise TypeError(msg) 4 ctx1abcd

242 else: 

243 self._dtype = gx.dtype 6 ctx1feabcd

244 

245 # Check that the derivative specifications are compatible with the 

246 # array data type. 

247 if gx.dtype.names is None: 6 ctx1feabcd

248 if not deriv.implicit: 6 ctx1feabcd

249 raise ValueError('x has no fields but derivative has') 4 ctx1abcd

250 else: 

251 for dim in deriv: 6 ctx1feabcd

252 if dim not in gx.dtype.names: 6 ctx1feabcd

253 raise ValueError(f'deriv field {dim!r} not in x') 4 ctx1abcd

254 

255 self._elements[key] = self._Points(gx, deriv, proc) 6 ctx1feabcd

256 

257 def _get_x_dtype(self): 6 ctx1feabcd

258 """ Get the data type of x points """ 

259 return self._dtype 6 ctx1feabcd

260 

261 def addtransf(self, tensors, key, *, axes=1): 6 ctx1feabcd

262 """ 

263  

264 Apply a linear transformation to already specified process points. The 

265 result of the transformation is represented by a new key. 

266  

267 Parameters 

268 ---------- 

269 tensors : dict 

270 Dictionary mapping keys of the GP to arrays/scalars. Each array is 

271 matrix-multiplied with the process array represented by its key, 

272 while scalars are just multiplied. Finally, the keys are summed 

273 over. 

274 key : hashable 

275 A new key under which the transformation is placed. 

276 axes : int 

277 Number of axes to be summed over for matrix multiplication, 

278 referring to trailing axes for tensors in ` tensors``, and to 

279 heading axes for process points. Default 1. 

280  

281 Returns 

282 ------- 

283 gp : GP 

284 A new GP object with the applied modifications. 

285  

286 Notes 

287 ----- 

288 The multiplication between the tensors and the process is done with 

289 np.tensordot with, by default, 1-axis contraction. For >2d arrays this 

290 is different from numpy's matrix multiplication, which would act on the 

291 second-to-last dimension of the second array. 

292  

293 """ 

294 # Note: it may seem nice that when an array has less axes than `axes`, 

295 # the summation would be restricted only on the existing axes. However 

296 # this brings about the ambiguous case where only one of the factors has 

297 # not enough axes. How many axes do you sum over on the other? 

298 

299 # Check axes. 

300 assert isinstance(axes, int) and axes >= 0, axes 5 ctx1eabcd

301 

302 # Check key. 

303 if key is None: 5 ctx1eabcd

304 raise ValueError('key can not be None') 4 ctx1abcd

305 if key in self._elements: 5 ctx1eabcd

306 raise KeyError(f'key {key!r} already in GP') 4 ctx1abcd

307 

308 # Check keys. 

309 for k in tensors: 5 ctx1eabcd

310 if k not in self._elements: 5 ctx1eabcd

311 raise KeyError(k) 4 ctx1abcd

312 

313 # Check tensors and convert them to jax arrays. 

314 if len(tensors) == 0: 5 ctx1eabcd

315 raise ValueError('empty tensors, undetermined output shape') 4 ctx1abcd

316 tens = {} 5 ctx1eabcd

317 for k, t in tensors.items(): 5 ctx1eabcd

318 t = jnp.asarray(t) 5 ctx1eabcd

319 # no need to check dtype since jax supports only numerical arrays 

320 with _jaxext.skipifabstract(): 5 ctx1eabcd

321 if self._checkfinite and not jnp.all(jnp.isfinite(t)): 5 ctx1eabcd

322 raise ValueError(f'tensors[{k!r}] contains infs/nans') 4 ctx1abcd

323 rshape = self._elements[k].shape 5 ctx1eabcd

324 if t.shape and t.shape[t.ndim - axes:] != rshape[:axes]: 5 ctx1eabcd

325 raise ValueError(f'tensors[{k!r}].shape = {t.shape!r} can not be multiplied with shape {rshape!r} with {axes}-axes contraction') 4 ctx1abcd

326 tens[k] = t 5 ctx1eabcd

327 

328 # Check shapes broadcast correctly. 

329 arrays = tens.values() 5 ctx1eabcd

330 elements = (self._elements[k] for k in tens) 5 ctx1eabcd

331 shapes = ( 5 ctx1eabcd

332 t.shape[:t.ndim - axes] + e.shape[axes:] if t.shape else e.shape 

333 for t, e in zip(arrays, elements) 

334 ) 

335 try: 5 ctx1eabcd

336 shape = jnp.broadcast_shapes(*shapes) 5 ctx1eabcd

337 except ValueError: 4 ctx1abcd

338 msg = 'can not broadcast tensors with shapes [' 4 ctx1abcd

339 msg += ', '.join(repr(t.shape) for t in arrays) 4 ctx1abcd

340 msg += '] contracted with arrays with shapes [' 4 ctx1abcd

341 msg += ', '.join(repr(e.shape) for e in elements) + ']' 4 ctx1abcd

342 raise ValueError(msg) 4 ctx1abcd

343 

344 # Define linear transformation. 

345 def equiv_lintransf(*args): 5 ctx1eabcd

346 assert len(args) == len(tens) 5 ctx1eabcd

347 out = None 5 ctx1eabcd

348 for a, (k, t) in zip(args, tens.items()): 5 ctx1eabcd

349 if t.shape: 5 ctx1eabcd

350 b = jnp.tensordot(t, a, axes) 5 ctx1eabcd

351 else: 

352 b = t * a 5 ctx1eabcd

353 if out is None: 5 ctx1eabcd

354 out = b 5 ctx1eabcd

355 else: 

356 out = out + b 5 ctx1eabcd

357 return out 5 ctx1eabcd

358 keys = list(tens.keys()) 5 ctx1eabcd

359 return self.addlintransf(equiv_lintransf, keys, key, checklin=False) 5 ctx1eabcd

360 

361 @_base.newself 6 ctx1feabcd

362 def addlintransf(self, transf, keys, key, *, checklin=None): 6 ctx1feabcd

363 """ 

364  

365 Define a finite linear transformation of the evaluated process. 

366  

367 Parameters 

368 ---------- 

369 transf : callable 

370 A function with signature ``f(array1, array2, ...) -> array`` which 

371 computes the linear transformation. The function must be 

372 jax-traceable, i.e., use jax.numpy instead of numpy. 

373 keys : sequence 

374 Keys of parts of the process to be passed as inputs to the 

375 transformation. 

376 key : hashable 

377 The key of the newly defined points. 

378 checklin : bool 

379 If True (default), check that the given function is linear in its 

380 inputs. The default can be overridden at initialization of the GP 

381 object. Note that an affine function (x -> a + bx) is not linear. 

382  

383 Raises 

384 ------ 

385 RuntimeError : 

386 The transformation seems not to be linear. To disable the linearity 

387 check, initialize the GP with ``checklin=False``. 

388  

389 """ 

390 

391 # TODO elementwise operations can be applied more efficiently to 

392 # primary gvars (tipical case), so the method could use an option 

393 # `elementwise`. What is the reliable way to check it is indeed 

394 # elementwise with a single random vector? Zero items of the tangent 

395 # at random with p=0.5 and check they stay zero? (And of course check 

396 # the shape is preserved.) 

397 

398 # Check key. 

399 if key is None: 6 ctx1feabcd

400 raise ValueError('key can not be None') 4 ctx1abcd

401 if key in self._elements: 6 ctx1feabcd

402 raise KeyError(f'key {key!r} already in GP') 4 ctx1abcd

403 

404 # Check keys. 

405 for k in keys: 6 ctx1feabcd

406 if k not in self._elements: 6 ctx1feabcd

407 raise KeyError(k) 4 ctx1abcd

408 

409 # Determine shape. 

410 class ArrayMockup: 6 ctx1feabcd

411 def __init__(self, elem): 6 ctx1feabcd

412 self.shape = elem.shape 6 ctx1feabcd

413 self.dtype = float 6 ctx1feabcd

414 inp = [ArrayMockup(self._elements[k]) for k in keys] 6 ctx1feabcd

415 out = jax.eval_shape(transf, *inp) 6 ctx1feabcd

416 shape = out.shape 6 ctx1feabcd

417 

418 # Check that the transformation is linear. 

419 if checklin is None: 6 ctx1feabcd

420 checklin = self._checklin 6 ctx1feabcd

421 if checklin: 6 ctx1feabcd

422 shapes = [self._elements[k].shape for k in keys] 6 ctx1feabcd

423 self._checklinear(transf, shapes) 6 ctx1feabcd

424 

425 self._elements[key] = self._LinTransf(transf, keys, shape) 6 ctx1feabcd

426 

427 @_base.newself 6 ctx1feabcd

428 def addcov(self, covblocks, key=None, *, decomps=None): 6 ctx1feabcd

429 """ 

430  

431 Add user-defined prior covariance matrix blocks. 

432  

433 Covariance matrices defined with `addcov` represent arbitrary 

434 finite-dimensional zero-mean Gaussian variables, assumed independent 

435 from all other variables in the GP object. 

436  

437 Parameters 

438 ---------- 

439 covblocks : array or dictionary of arrays 

440 If an array: a covariance matrix (or tensor) to be added under key 

441 ``key``. If a dictionary: a mapping from pairs of keys to the 

442 corresponding covariance matrix blocks. A missing off-diagonal 

443 block in the dictionary is interpreted as a matrix of zeros, 

444 unless the corresponding transposed block is specified. 

445 key : hashable 

446 If ``covblocks`` is an array, the dictionary key under which 

447 ``covblocks`` is added. Can not be specified if ``covblocks`` is a 

448 dictionary. 

449 decomps : Decomposition or dict of Decompositions 

450 Pre-computed decompositions of (not necessarily all) diagonal 

451 blocks, as produced by `decompose`. The keys are single 

452 GP keys and not pairs like in ``covblocks``. 

453  

454 Raises 

455 ------ 

456 KeyError : 

457 A key is already used in the GP. 

458 ValueError : 

459 ``covblocks`` and/or ``key`` and ``decomps`` are malformed or 

460 inconsistent. 

461 TypeError : 

462 Wrong type of ``covblocks`` or ``decomps``. 

463  

464 """ 

465 

466 # TODO maybe allow passing only the lower/upper triangular part for 

467 # the diagonal blocks, like I meta-allow for out of diagonal blocks? 

468 

469 # TODO with multiple blocks and a single decomp, the decomp could be 

470 # interpreted as the decomposition of the whole block matrix. 

471 

472 # Check type of `covblocks` and standardize it to dictionary. 

473 if hasattr(covblocks, 'keys'): 5 ctx1eabcd

474 if key is not None: 4 ctx1abcd

475 raise ValueError('can not specify key if covblocks is a dictionary') 4 ctx1abcd

476 if None in covblocks: 4 ctx1abcd

477 raise ValueError('None key in covblocks not allowed') 4 ctx1abcd

478 if decomps is not None and not hasattr(decomps, 'keys'): 4 ctx1abcd

479 raise TypeError('covblocks is dictionary but decomps is not') 4 ctx1abcd

480 else: 

481 if key is None: 5 ctx1eabcd

482 raise ValueError('covblocks is not dictionary but key is None') 4 ctx1abcd

483 covblocks = {(key, key): covblocks} 5 ctx1eabcd

484 if decomps is not None: 5 ctx1eabcd

485 decomps = {key: decomps} 4 ctx1abcd

486 

487 if decomps is None: 5 ctx1eabcd

488 decomps = {} 5 ctx1eabcd

489 

490 # Convert blocks to jax arrays and determine shapes from diagonal 

491 # blocks. 

492 shapes = {} 5 ctx1eabcd

493 preblocks = {} 5 ctx1eabcd

494 for keys, block in covblocks.items(): 5 ctx1eabcd

495 # TODO maybe check that keys is a 2-tuple 

496 for key in keys: 5 ctx1eabcd

497 if key in self._elements: 5 ctx1eabcd

498 raise KeyError(f'key {key!r} already in GP') 4 ctx1abcd

499 xkey, ykey = keys 5 ctx1eabcd

500 if block is None: 5 ctx1eabcd

501 raise TypeError(f'block {keys!r} is None') 4 ctx1abcd

502 # because jnp.asarray(None) interprets None as nan 

503 # (see jax issue #14506) 

504 block = jnp.asarray(block) 5 ctx1eabcd

505 

506 if xkey == ykey: 5 ctx1eabcd

507 

508 if block.ndim % 2 == 1: 5 ctx1eabcd

509 raise ValueError(f'diagonal block {key!r} has odd number of axes') 4 ctx1abcd

510 

511 half = block.ndim // 2 5 ctx1eabcd

512 head = block.shape[:half] 5 ctx1eabcd

513 tail = block.shape[half:] 5 ctx1eabcd

514 if head != tail: 5 ctx1eabcd

515 raise ValueError(f'shape {block.shape!r} of diagonal block {key!r} is not symmetric') 4 ctx1abcd

516 shapes[xkey] = head 5 ctx1eabcd

517 

518 with _jaxext.skipifabstract(): 5 ctx1eabcd

519 if self._checksym and not jnp.allclose(block, block.T): 5 ctx1eabcd

520 raise ValueError(f'diagonal block {key!r} is not symmetric') 4 ctx1abcd

521 

522 preblocks[keys] = block 5 ctx1eabcd

523 

524 # Check decomps is consistent with covblocks. 

525 for key, dec in decomps.items(): 5 ctx1eabcd

526 if key not in shapes: 4 ctx1abcd

527 raise KeyError(f'key {key!r} in decomps not found in diagonal blocks') 4 ctx1abcd

528 if not isinstance(dec, _linalg.Decomposition): 4 ctx1abcd

529 raise TypeError(f'decomps[{key!r}] = {dec!r} is not a decomposition') 4 ctx1abcd

530 n = math.prod(shapes[key]) 4 ctx1abcd

531 if dec.n != n: 4 ctx1abcd

532 raise ValueError(f'decomposition matrix size {dec.n} != diagonal block size {n} for key {key!r}') 4 ctx1abcd

533 

534 # Reshape blocks to square matrices and check that the shapes of out of 

535 # diagonal blocks match those of diagonal ones. 

536 blocks = {} 5 ctx1eabcd

537 for keys, block in preblocks.items(): 5 ctx1eabcd

538 with _jaxext.skipifabstract(): 5 ctx1eabcd

539 if self._checkfinite and not jnp.all(jnp.isfinite(block)): 5 ctx1eabcd

540 raise ValueError(f'block {keys!r} not finite') 4 ctx1abcd

541 xkey, ykey = keys 5 ctx1eabcd

542 if xkey == ykey: 5 ctx1eabcd

543 size = math.prod(shapes[xkey]) 5 ctx1eabcd

544 blocks[keys] = block.reshape((size, size)) 5 ctx1eabcd

545 else: 

546 for key in keys: 4 ctx1abcd

547 if key not in shapes: 4 ctx1abcd

548 raise KeyError(f'key {key!r} from off-diagonal block {keys!r} not found in diagonal blocks') 4 ctx1abcd

549 eshape = shapes[xkey] + shapes[ykey] 4 ctx1abcd

550 if block.shape != eshape: 4 ctx1abcd

551 raise ValueError(f'shape {block.shape!r} of block {keys!r} is not {eshape!r} as expected from diagonal blocks') 4 ctx1abcd

552 xsize = math.prod(shapes[xkey]) 4 ctx1abcd

553 ysize = math.prod(shapes[ykey]) 4 ctx1abcd

554 block = block.reshape((xsize, ysize)) 4 ctx1abcd

555 blocks[keys] = block 4 ctx1abcd

556 revkeys = keys[::-1] 4 ctx1abcd

557 blockT = preblocks.get(revkeys) 4 ctx1abcd

558 if blockT is None: 4 ctx1abcd

559 blocks[revkeys] = block.T 4 ctx1abcd

560 

561 # Check symmetry of out of diagonal blocks. 

562 if self._checksym: 5 ctx1eabcd

563 with _jaxext.skipifabstract(): 4 ctx1abcd

564 for keys, block in blocks.items(): 4 ctx1abcd

565 xkey, ykey = keys 4 ctx1abcd

566 if xkey != ykey: 4 ctx1abcd

567 blockT = blocks[ykey, xkey] 4 ctx1abcd

568 if not jnp.allclose(block.T, blockT): 4 ctx1abcd

569 raise ValueError(f'block {keys!r} is not the transpose of block {revkeys!r}') 4 ctx1abcd

570 

571 # Create _Cov objects. 

572 for key, shape in shapes.items(): 5 ctx1eabcd

573 self._elements[key] = self._Cov(blocks, shape) 5 ctx1eabcd

574 decomp = decomps.get(key) 5 ctx1eabcd

575 if decomp is not None: 5 ctx1eabcd

576 self._decompcache[key,] = decomp 4 ctx1abcd

577 

578 def _makecovblock_points(self, xkey, ykey): 6 ctx1feabcd

579 x = self._elements[xkey] 6 ctx1feabcd

580 y = self._elements[ykey] 6 ctx1feabcd

581 

582 assert isinstance(x, self._Points) 6 ctx1feabcd

583 assert isinstance(y, self._Points) 6 ctx1feabcd

584 

585 kernel = self._crosskernel(x.proc, y.proc) 6 ctx1feabcd

586 if kernel is self._zerokernel: 6 ctx1feabcd

587 # TODO handle zero cov block efficiently 

588 return jnp.zeros((x.size, y.size)) 6 ctx1feabcd

589 

590 kernel = kernel.linop('diff', x.deriv, y.deriv) 6 ctx1feabcd

591 

592 if x is y and not self._checksym and self._halfmatrix: 6 ctx1feabcd

593 ix, iy, back = self._triu_indices_and_back(x.size) 4 ctx1abcd

594 flat = x.x.reshape(-1) 4 ctx1abcd

595 ax = flat[ix] 4 ctx1abcd

596 ay = flat[iy] 4 ctx1abcd

597 halfcov = kernel(ax, ay) 4 ctx1abcd

598 cov = halfcov[back] 4 ctx1abcd

599 # TODO to avoid inefficiencies like in BART, maybe _Kernel should 

600 # have a method outer(x) that by default simply does self(x[None, 

601 # :], x[:, None]) but can be overwritten. This halfmatrix impl could 

602 # be moved there with an option outer(x, *, half=False). To carry 

603 # over custom implementations of outer, there should be a callable 

604 # attribute _outer, optionally set at initialization, that is 

605 # transformed by kernel operations. 

606 else: 

607 ax = x.x.reshape(-1)[:, None] 6 ctx1feabcd

608 ay = y.x.reshape(-1)[None, :] 6 ctx1feabcd

609 cov = kernel(ax, ay) 6 ctx1feabcd

610 

611 return cov 6 ctx1feabcd

612 

613 def _makecovblock_lintransf_any(self, xkey, ykey): 6 ctx1feabcd

614 x = self._elements[xkey] 6 ctx1feabcd

615 y = self._elements[ykey] 6 ctx1feabcd

616 assert isinstance(x, self._LinTransf) 6 ctx1feabcd

617 

618 # Gather covariance matrices to be transformed. 

619 covs = [] 6 ctx1feabcd

620 for k in x.keys: 6 ctx1feabcd

621 elem = self._elements[k] 6 ctx1feabcd

622 cov = self._covblock(k, ykey) 6 ctx1feabcd

623 assert cov.shape == (elem.size, y.size) 6 ctx1feabcd

624 cov = cov.reshape(elem.shape + (y.size,)) 6 ctx1feabcd

625 covs.append(cov) 6 ctx1feabcd

626 

627 # Apply transformation. 

628 t = jax.vmap(x.transf, -1, -1) 6 ctx1feabcd

629 cov = t(*covs) 6 ctx1feabcd

630 assert cov.shape == x.shape + (y.size,) 6 ctx1feabcd

631 return cov.reshape((x.size, y.size)) # don't leave out the ()! 6 ctx1feabcd

632 # the () probably was an obscure autograd bug, I don't think it will 

633 # be a problem again with jax 

634 

635 def _makecovblock(self, xkey, ykey): 6 ctx1feabcd

636 x = self._elements[xkey] 6 ctx1feabcd

637 y = self._elements[ykey] 6 ctx1feabcd

638 if isinstance(x, self._Points) and isinstance(y, self._Points): 6 ctx1feabcd

639 cov = self._makecovblock_points(xkey, ykey) 6 ctx1feabcd

640 elif isinstance(x, self._LinTransf): 6 ctx1feabcd

641 cov = self._makecovblock_lintransf_any(xkey, ykey) 6 ctx1feabcd

642 elif isinstance(y, self._LinTransf): 6 ctx1feabcd

643 cov = self._makecovblock_lintransf_any(ykey, xkey) 6 ctx1feabcd

644 cov = cov.T 6 ctx1feabcd

645 elif isinstance(x, self._Cov) and isinstance(y, self._Cov) and x.blocks is y.blocks and (xkey, ykey) in x.blocks: 5 ctx1eabcd

646 cov = x.blocks[xkey, ykey] 5 ctx1eabcd

647 else: 

648 # TODO handle zero cov block efficiently 

649 cov = jnp.zeros((x.size, y.size)) 5 ctx1eabcd

650 

651 with _jaxext.skipifabstract(): 6 ctx1feabcd

652 if self._checkfinite and not jnp.all(jnp.isfinite(cov)): 6 ctx1feabcd

653 raise RuntimeError(f'covariance block {(xkey, ykey)!r} is not finite') 4 ctx1abcd

654 if self._checksym and xkey == ykey and not jnp.allclose(cov, cov.T): 6 ctx1feabcd

655 raise RuntimeError(f'covariance block {(xkey, ykey)!r} is not symmetric') 4 ctx1abcd

656 

657 return cov 6 ctx1feabcd

658 

659 def _covblock(self, row, col): 6 ctx1feabcd

660 

661 if (row, col) not in self._covblocks: 6 ctx1feabcd

662 block = self._makecovblock(row, col) 6 ctx1feabcd

663 if row != col: 6 ctx1feabcd

664 if self._checksym: 6 ctx1feabcd

665 with _jaxext.skipifabstract(): 6 ctx1feabcd

666 blockT = self._makecovblock(col, row) 6 ctx1feabcd

667 if not jnp.allclose(block.T, blockT): 6 ctx1feabcd

668 msg = 'covariance block {!r} is not symmetric' 4 ctx1abcd

669 raise RuntimeError(msg.format((row, col))) 4 ctx1abcd

670 self._covblocks[col, row] = block.T 6 ctx1feabcd

671 self._covblocks[row, col] = block 6 ctx1feabcd

672 

673 return self._covblocks[row, col] 6 ctx1feabcd

674 

675 def _assemblecovblocks(self, rowkeys, colkeys=None): 6 ctx1feabcd

676 if colkeys is None: 6 ctx1feabcd

677 colkeys = rowkeys 6 ctx1feabcd

678 blocks = [ 6 ctx1feabcd

679 [self._covblock(row, col) for col in colkeys] 

680 for row in rowkeys 

681 ] 

682 return jnp.block(blocks) 6 ctx1feabcd

683 

684 def _checkpos(self, cov): 6 ctx1feabcd

685 with _jaxext.skipifabstract(): 6 ctx1feabcd

686 # eigv = jnp.linalg.eigvalsh(cov) 

687 # mineigv, maxeigv = jnp.min(eigv), jnp.max(eigv) 

688 with warnings.catch_warnings(): 6 ctx1feabcd

689 warnings.filterwarnings('ignore', r'Exited at iteration .+? with accuracies') 6 ctx1feabcd

690 warnings.filterwarnings('ignore', r'Exited postprocessing with accuracies') 6 ctx1feabcd

691 X = numpy.random.randn(len(cov), 1) 6 ctx1feabcd

692 A = numpy.asarray(cov) 6 ctx1feabcd

693 (mineigv,), _ = sparse.linalg.lobpcg(A, X, largest=False) 6 ctx1feabcd

694 (maxeigv,), _ = sparse.linalg.lobpcg(A, X, largest=True) 6 ctx1feabcd

695 assert mineigv <= maxeigv 6 ctx1feabcd