Coverage for src/lsqfitgp/_Kernel/_crosskernel.py: 99%

1# lsqfitgp/_Kernel/_crosskernel.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 enum 6 ctx1feabcd

21import functools 6 ctx1feabcd

22import sys 6 ctx1feabcd

23import collections 6 ctx1feabcd

24import types 6 ctx1feabcd

25import abc 6 ctx1feabcd

26import warnings 6 ctx1feabcd

27 

28import numpy 6 ctx1feabcd

29from jax import numpy as jnp 6 ctx1feabcd

30 

31from .. import _array 6 ctx1feabcd

32from .. import _jaxext 6 ctx1feabcd

33from .. import _utils 6 ctx1feabcd

34 

35from . import _util 6 ctx1feabcd

36 

37@functools.lru_cache(maxsize=None) 6 ctx1feabcd

38def least_common_superclass(*classes): 6 ctx1feabcd

39 """ 

40 Find a "least" common superclass. The class is searched in all the MROs, 

41 but the comparison is done with `issubclass` to support virtual inheritance. 

42 """ 

43 mros = [c.__mro__ for c in classes] 6 ctx1feabcd

44 indices = [0] * len(mros) 6 ctx1feabcd

45 for i, mroi in enumerate(mros): 6 ctx1feabcd

46 for j, mroj in enumerate(mros): 6 ctx1feabcd

47 if i == j: 6 ctx1feabcd

48 continue 6 ctx1feabcd

49 while not issubclass(mroi[0], mroj[indices[j]]): 6 ctx1feabcd

50 indices[j] += 1 6 ctx1feabcd

51 idx = numpy.argmin(indices) 6 ctx1feabcd

52 return mros[idx][indices[idx]] 6 ctx1feabcd

53 

54class CrossKernel: 6 ctx1feabcd

55 r""" 

56  

57 Base class to represent kernels, i.e., covariance functions. 

58 

59 A kernel is a two-argument function that computes the covariance between 

60 two functions at some points according to some probability distribution: 

61 

62 .. math:: 

63 \mathrm{kernel}(x, y) = \mathrm{Cov}[f(x), g(y)]. 

64  

65 `CrossKernel` objects are callable, the signature is ``obj(x, y)``, and 

66 they can be summed and multiplied between them and with scalars. They 

67 are immutable; all operations return new objects. 

68 

69 Parameters 

70 ---------- 

71 core : callable 

72 A function with signature ``core(x, y)``, where ``x`` and ``y`` 

73 are two broadcastable numpy arrays, which computes the value of the 

74 kernel. 

75 scale, loc, derivable, maxdim, dim : 

76 If specified, these arguments are passed as arguments to the 

77 correspondingly named operators, in the order listed here. See `linop`. 

78 Briefly: the kernel selects only fields `dim` in the input, checks the 

79 dimensionality against `maxdim`, checks there are not too many 

80 derivatives taken on the arguments, then transforms as ``(x - loc) / 

81 scale``. If any argument is callable, it is passed `**kw` and must 

82 return the actual argument. If an argument is a tuple, it is interpreted 

83 as a pair of arguments. 

84 forcekron : bool, default False 

85 If True, apply ``.transf('forcekron')`` to the kernel, before the 

86 operations above. Available only for `Kernel`. 

87 batchbytes : number, optional 

88 If specified, apply ``.batch(batchbytes)`` to the kernel. 

89 dynkw : dict, optional 

90 Additional keyword arguments passed to `core` that can be modified 

91 by transformations. Deleted by transformations by default. 

92 **initkw : 

93 Additional keyword arguments passed to `core` that can be read but not 

94 changed by transformations. 

95 

96 Attributes 

97 ---------- 

98 initkw : dict 

99 The `initkw` argument. 

100 dynkw : dict 

101 The `dynkw` argument, or a modification of it if the object has been 

102 transformed. 

103 core : callable 

104 The `core` argument partially evaluated on `initkw`, or another 

105 function wrapping it if the object has been transformed. 

106  

107 Methods 

108 ------- 

109 batch 

110 linop 

111 algop 

112 transf 

113 register_transf 

114 register_linop 

115 register_corelinop 

116 register_xtransf 

117 register_algop 

118 register_ufuncalgop 

119 transf_help 

120 has_transf 

121 list_transf 

122 super_transf 

123 make_linop_family 

124 

125 See also 

126 -------- 

127 Kernel 

128 

129 Notes 

130 ----- 

131 The predefined class hierarchy and the class logic of the transformations 

132 assume that each kernel class corresponds to a subalgebra, i.e., addition 

133 and multiplication preserve the class. 

134  

135 """ 

136 

137 __slots__ = '_initkw', '_dynkw', '_core' 6 ctx1feabcd

138 # only __new__ and _clone shall access these attributes 

139 

140 @property 6 ctx1feabcd

141 def initkw(self): 6 ctx1feabcd

142 return types.MappingProxyType(self._initkw) 6 ctx1feabcd

143 

144 @property 6 ctx1feabcd

145 def dynkw(self): 6 ctx1feabcd

146 return types.MappingProxyType(self._dynkw) 6 ctx1feabcd

147 

148 @property 6 ctx1feabcd

149 def core(self): 6 ctx1feabcd

150 return self._core 6 ctx1feabcd

151 

152 def __new__(cls, core, *, 6 ctx1feabcd

153 scale=None, 

154 loc=None, 

155 derivable=None, 

156 maxdim=None, 

157 dim=None, 

158 forcekron=False, 

159 batchbytes=None, 

160 dynkw={}, 

161 **initkw, 

162 ): 

163 self = super().__new__(cls) 6 ctx1feabcd

164 

165 self._initkw = initkw 6 ctx1feabcd

166 self._dynkw = dict(dynkw) 6 ctx1feabcd

167 self._core = lambda x, y, **dynkw: core(x, y, **initkw, **dynkw) 6 ctx1feabcd

168 

169 if forcekron: 6 ctx1feabcd

170 self = self.transf('forcekron') 4 ctx1abcd

171 

172 linop_args = { 6 ctx1feabcd

173 'scale': scale, 

174 'loc': loc, 

175 'derivable': derivable, 

176 'maxdim': maxdim, 

177 'dim': dim, 

178 } 

179 for transfname, arg in linop_args.items(): 6 ctx1feabcd

180 if callable(arg): 6 ctx1feabcd

181 arg = arg(**initkw) 6 ctx1feabcd

182 if isinstance(arg, tuple): 6 ctx1feabcd

183 self = self.linop(transfname, *arg) 5 ctx1eabcd

184 else: 

185 self = self.linop(transfname, arg) 6 ctx1feabcd

186 

187 if batchbytes is not None: 6 ctx1feabcd

188 self = self.batch(batchbytes) 5 ctx1eabcd

189 

190 return self 6 ctx1feabcd

191 

192 def __call__(self, x, y): 6 ctx1feabcd

193 x = _array.asarray(x) 6 ctx1feabcd

194 y = _array.asarray(y) 6 ctx1feabcd

195 shape = _array.broadcast(x, y).shape 6 ctx1feabcd

196 result = self.core(x, y, **self.dynkw) 6 ctx1feabcd

197 assert isinstance(result, (numpy.ndarray, jnp.number, jnp.ndarray)) 6 ctx1feabcd

198 assert jnp.issubdtype(result.dtype, jnp.number), result.dtype 6 ctx1feabcd

199 assert result.shape == shape, (result.shape, shape) 6 ctx1feabcd

200 return result 6 ctx1feabcd

201 

202 def _clone(self, cls=None, *, initkw=None, dynkw=None, core=None): 6 ctx1feabcd

203 newself = object.__new__(self.__class__ if cls is None else cls) 6 ctx1feabcd

204 newself._initkw = self._initkw if initkw is None else dict(initkw) 6 ctx1feabcd

205 newself._dynkw = {} if dynkw is None else dict(dynkw) 6 ctx1feabcd

206 newself._core = self._core if core is None else core 6 ctx1feabcd

207 return newself 6 ctx1feabcd

208 

209 class _side(enum.Enum): 6 ctx1feabcd

210 LEFT = 0 6 ctx1feabcd

211 RIGHT = 1 6 ctx1feabcd

212 

213 @classmethod 6 ctx1feabcd

214 def _nary(cls, op, kernels, side): 6 ctx1feabcd

215 

216 if side is cls._side.LEFT: 6 ctx1feabcd

217 wrapper = lambda c, _, y, **kw: lambda x: c(x, y, **kw) 6 ctx1feabcd

218 arg = lambda x, _: x 6 ctx1feabcd

219 elif side is cls._side.RIGHT: 4 ctx1abcd

220 wrapper = lambda c, x, _, **kw: lambda y: c(x, y, **kw) 4 ctx1abcd

221 arg = lambda _, y: y 4 ctx1abcd

222 else: # pragma: no cover 

223 raise KeyError(side) 

224 

225 cores = [k.core for k in kernels] 6 ctx1feabcd

226 def core(x, y, **kw): 6 ctx1feabcd

227 wrapped = [wrapper(c, x, y, **kw) for c in cores] 6 ctx1feabcd

228 transformed = op(*wrapped) 6 ctx1feabcd

229 return transformed(arg(x, y)) 6 ctx1feabcd

230 

231 return __class__(core) 6 ctx1feabcd

232 

233 # TODO instead of wrapping the cores just by closing over an argument, 

234 # define a `PartialKernel` object that has linop defined, with one arg 

235 # fixed to None, and __call__ that closes over an argument. This way I 

236 # could apply ops directly with deflintransf. 

237 

238 # TODO move in the binary op methods the logic that decides wether to 

239 # return NotImplemented, while the algops will raise an exception if the 

240 # argument is not to their liking 

241 

242 def __add__(self, other): 6 ctx1feabcd

243 return self.algop('add', other) 6 ctx1feabcd

244 

245 __radd__ = __add__ 6 ctx1feabcd

246 

247 def __mul__(self, other): 6 ctx1feabcd

248 return self.algop('mul', other) 6 ctx1feabcd

249 

250 __rmul__ = __mul__ 6 ctx1feabcd

251 

252 def __pow__(self, other): 6 ctx1feabcd

253 return self.algop('pow', exponent=other) 4 ctx1abcd

254 

255 def __rpow__(self, other): 6 ctx1feabcd

256 return self.algop('rpow', base=other) 4 ctx1abcd

257 

258 def _swap(self): 6 ctx1feabcd

259 """ permute the arguments """ 

260 core = self.core 6 ctx1feabcd

261 return self._clone( 6 ctx1feabcd

262 __class__, 

263 core=lambda x, y, **kw: core(y, x, **kw), 

264 ) 

265 

266 # TODO make _swap a transf inherited by CrossIsotropicKernel => messes 

267 # up with Kernel subclasses. I want to make Kernel subclasses identity 

268 # on swap, so it can't be an inherited transf, because Kernel appears 

269 # as the second base. 

270 

271 def batch(self, maxnbytes): 6 ctx1feabcd

272 """ 

273 Return a batched version of the kernel. 

274 

275 The batched kernel processes its inputs in chunks to try to limit memory 

276 usage. 

277 

278 Parameters 

279 ---------- 

280 maxnbytes : number 

281 The maximum number of input bytes per chunk, counted after 

282 broadcasting the input shapes. Actual broadcasting may not occur if 

283 not induced by the operations in the kernel. 

284 

285 Returns 

286 ------- 

287 batched_kernel : CrossKernel 

288 The same kernel but with batched computations. 

289 """ 

290 core = _jaxext.batchufunc(self.core, maxnbytes=maxnbytes) 5 ctx1eabcd

291 return self._clone(core=core) 5 ctx1eabcd

292 

293 @classmethod 6 ctx1feabcd

294 def _crossmro(cls): 6 ctx1feabcd

295 """ MRO iterator excluding subclasses of Kernel """ 

296 for c in cls.mro(): # pragma: no branch 6 ctx1feabcd

297 if not issubclass(c, Kernel): 6 ctx1feabcd

298 yield c 6 ctx1feabcd

299 if c is __class__: 6 ctx1feabcd

300 break 4 ctx1abcd

301 

302 _transf = {} 6 ctx1feabcd

303 

304 _Transf = collections.namedtuple('_Transf', ['func', 'doc', 'kind']) 6 ctx1feabcd

305 

306 def __init_subclass__(cls, **kw): 6 ctx1feabcd

307 super().__init_subclass__(**kw) 6 ctx1feabcd

308 cls._transf = {} 6 ctx1feabcd

309 cls.__slots__ = () 6 ctx1feabcd

310 

311 @classmethod 6 ctx1feabcd

312 def _transfmro(cls): 6 ctx1feabcd

313 """ Iterator of superclasses with a _transf attribute """ 

314 for c in cls.mro(): # pragma: no branch 6 ctx1feabcd

315 yield c 6 ctx1feabcd

316 if c is __class__: 6 ctx1feabcd

317 break 6 ctx1feabcd

318 

319 @classmethod 6 ctx1feabcd

320 def _settransf(cls, transfname, transf): 6 ctx1feabcd

321 if transfname in cls._transf: 6 ctx1feabcd

322 raise KeyError(f'transformation {transfname!r} already registered ' 4 ctx1abcd

323 f'for {cls.__name__}') 

324 cls._transf[transfname] = cls._Transf(*transf) 6 ctx1feabcd

325 

326 @classmethod 6 ctx1feabcd

327 def _alltransf(cls): 6 ctx1feabcd

328 """ list all accessible transfs as dict name -> (tcls, transf) """ 

329 transfs = {} 6 ctx1feabcd

330 for tcls in cls._transfmro(): 6 ctx1feabcd

331 for name, transf in tcls._transf.items(): 6 ctx1feabcd

332 transfs.setdefault(name, (tcls, transf)) 6 ctx1feabcd

333 return transfs 6 ctx1feabcd

334 

335 @classmethod 6 ctx1feabcd

336 def _gettransf(cls, transfname, transfmro=None): 6 ctx1feabcd

337 """ 

338 Find a transformation. 

339 

340 The transformation is searched following the MRO up to `CrossKernel`. 

341 

342 Parameters 

343 ---------- 

344 transfname : hashable 

345 The transformation name. 

346 

347 Returns 

348 ------- 

349 cls : type 

350 The class where the transformation was found. 

351 transf, doc, kind : tuple 

352 The objects set by `register_transf`. 

353 

354 Raises 

355 ------ 

356 KeyError : 

357 The transformation was not found. 

358 """ 

359 if transfmro is None: 6 ctx1feabcd

360 transfmro = cls._transfmro() 6 ctx1feabcd

361 for c in transfmro: 6 ctx1feabcd

362 try: 6 ctx1feabcd

363 return c, c._transf[transfname] 6 ctx1feabcd

364 except KeyError: 6 ctx1feabcd

365 pass 6 ctx1feabcd

366 raise KeyError(transfname) 4 ctx1abcd

367 

368 @classmethod 6 ctx1feabcd

369 def inherit_transf(cls, transfname, *, intermediates=False): 6 ctx1feabcd

370 """ 

371  

372 Inherit a transformation from a superclass. 

373 

374 Parameters 

375 ---------- 

376 transfname : hashable 

377 The name of the transformation. 

378 intermediates : bool, default False 

379 If True, make all superclasses up to the one definining the 

380 transformation inherit it too. 

381 

382 Raises 

383 ------ 

384 KeyError : 

385 The transformation was not found in any superclass, or the 

386 transformation is already registered on any of the target classes. 

387 

388 See also 

389 -------- 

390 transf 

391 

392 """ 

393 tcls, transf = cls._gettransf(transfname) 6 ctx1feabcd

394 cls._settransf(transfname, transf) 6 ctx1feabcd

395 if intermediates: 6 ctx1feabcd

396 for c in cls.mro()[1:]: # pragma: no branch 6 ctx1feabcd

397 if c is tcls: 6 ctx1feabcd

398 break 6 ctx1feabcd

399 c._settransf(transfname, transf) 6 ctx1feabcd

400 

401 @classmethod 6 ctx1feabcd

402 def inherit_all_algops(cls, intermediates=False): 6 ctx1feabcd

403 """ 

404 

405 Inherit all algebraic operations from superclasses. 

406 

407 This makes sense if the class represents a subalgebra, i.e., it 

408 should be preserved by addition and multiplication. 

409 

410 Parameters 

411 ---------- 

412 intermediates : bool, default False 

413 If True, make all superclasses up to the one definining the 

414 transformation inherit it too. 

415 

416 Raises 

417 ------ 

418 KeyError : 

419 An algebraic operation is already registered for one of the target 

420 classes. 

421 

422 See also 

423 -------- 

424 transf 

425 

426 """ 

427 mro = cls._transfmro() 6 ctx1feabcd

428 next(mro) 6 ctx1feabcd

429 for name, (_, transf) in next(mro)._alltransf().items(): 6 ctx1feabcd

430 if transf.kind is cls._algopmarker: 6 ctx1feabcd

431 cls.inherit_transf(name, intermediates=intermediates) 6 ctx1feabcd

432 

433 @classmethod 6 ctx1feabcd

434 def list_transf(cls, superclasses=True): 6 ctx1feabcd

435 """ 

436 List all the available transformations. 

437 

438 Parameters 

439 ---------- 

440 superclasses : bool, default True 

441 Include transformations defined in superclasses. 

442 

443 Returns 

444 ------- 

445 transfs: dict of Transf 

446 The dictionary keys are the transformation names, the values are 

447 named tuples ``(tcls, kind, impl, doc)`` where ``tcls`` is the class 

448 defining the transformation, ``kind`` is the kind of transformation, 

449 ``impl`` is the implementation with signature ``impl(tcls, self, 

450 *args, **kw)``, ``doc`` is the docstring. 

451 """ 

452 if superclasses: 4 ctx1abcd

453 source = cls._alltransf().items 4 ctx1abcd

454 else: 

455 def source(): 4 ctx1abcd

456 for name, transf in cls._transf.items(): 4 ctx1abcd

457 yield name, (cls, transf) 4 ctx1abcd

458 return { 4 ctx1abcd

459 name: cls.Transf(tcls, transf.kind, transf.func, transf.doc) 

460 for name, (tcls, transf) in source() 

461 } 

462 

463 Transf = collections.namedtuple('Transf', ['tcls', 'kind', 'func', 'doc']) 6 ctx1feabcd

464 

465 @classmethod 6 ctx1feabcd

466 def has_transf(cls, transfname): 6 ctx1feabcd

467 """ 

468 Check if a transformation is registered. 

469 

470 Parameters 

471 ---------- 

472 transfname : hashable 

473 The transformation name. 

474 

475 Returns 

476 ------- 

477 has_transf : bool 

478 Whether the transformation is registered. 

479 

480 See also 

481 -------- 

482 transf 

483 """ 

484 try: 4 ctx1abcd

485 cls._gettransf(transfname) 4 ctx1abcd

486 except KeyError as exc: 4 ctx1abcd

487 if exc.args == (transfname,): 4 ctx1abcd

488 return False 4 ctx1abcd

489 else: # pragma: no cover 

490 raise 

491 else: 

492 return True 4 ctx1abcd

493 

494 @classmethod 6 ctx1feabcd

495 def transf_help(cls, transfname): 6 ctx1feabcd

496 """ 

497  

498 Return the documentation of a transformation. 

499 

500 Parameters 

501 ---------- 

502 transfname : hashable 

503 The name of the transformation. 

504 

505 Returns 

506 ------- 

507 doc : str 

508 The documentation of the transformation. 

509 

510 See also 

511 -------- 

512 transf 

513 

514 """ 

515 _, transf = cls._gettransf(transfname) 4 ctx1abcd

516 return transf.doc 4 ctx1abcd

517 

518 def transf(self, transfname, *args, **kw): 6 ctx1feabcd

519 """ 

520 

521 Return a transformed kernel. 

522 

523 Parameters 

524 ---------- 

525 transfname : hashable 

526 A name identifying the transformation. 

527 *args, **kw : 

528 Arguments to the transformation. 

529 

530 Returns 

531 ------- 

532 newkernel : object 

533 The output of the transformation. 

534 

535 Raises 

536 ------ 

537 KeyError 

538 The transformation is not defined in this class or any superclass. 

539 

540 See also 

541 -------- 

542 linop, algop, transf_help, has_transf, list_transf, super_transf, register_transf, register_linop, register_corelinop, register_xtransf, register_algop, register_ufuncalgop 

543 

544 """ 

545 tcls, transf = self._gettransf(transfname) 4 ctx1abcd

546 return transf.func(tcls, self, *args, **kw) 4 ctx1abcd

547 

548 @classmethod 6 ctx1feabcd

549 def super_transf(cls, transfname, self, *args, **kw): 6 ctx1feabcd

550 """ 

551  

552 Transform the kernel using a superclass transformation. 

553 

554 This is equivalent to `transf` but is invoked on a class and the 

555 object is passed explicitly. The definition of the transformation is 

556 searched starting from the next class in the MRO of `self`. 

557 

558 Parameters 

559 ---------- 

560 transfname, *args, **kw : 

561 See `transf`. 

562 self : CrossKernel 

563 The object to transform. 

564 

565 Returns 

566 ------- 

567 newkernel : object 

568 The output of the transformation. 

569 

570 """ 

571 mro = list(self._transfmro()) 5 ctx1eabcd

572 idx = mro.index(cls) 5 ctx1eabcd

573 tcls, transf = self._gettransf(transfname, mro[idx + 1:]) 5 ctx1eabcd

574 return transf.func(tcls, self, *args, **kw) 5 ctx1eabcd

575 

576 def linop(self, transfname, *args, **kw): 6 ctx1feabcd

577 r""" 

578 

579 Transform kernels to represent the application of a linear operator. 

580 

581 .. math:: 

582 \text{kernel}_1(x, y) &= \mathrm{Cov}[f_1(x), g_1(y)], \\ 

583 \text{kernel}_2(x, y) &= \mathrm{Cov}[f_2(x), g_2(y)], \\ 

584 &\ldots \\ 

585 \text{newkernel}(x, y) &= 

586 \mathrm{Cov}[T_f(f_1, f_2, \ldots)(x), T_g(g_1, g_2, \ldots)(y)] 

587 

588 Parameters 

589 ---------- 

590 transfname : hashable 

591 The name of the transformation. 

592 *args : 

593 A sequence of `CrossKernel` instances, representing the operands, 

594 followed by one or two non-kernel arguments, indicating how to act 

595 on each side of the kernels. If both arguments represent the 

596 identity, this is a no-op. If there is only one argument, it is 

597 intended that the two arguments are equal. `None` always represents 

598 the identity. 

599 

600 Returns 

601 ------- 

602 newkernel : CrossKernel 

603 The transformed kernel. 

604 

605 Raises 

606 ------ 

607 ValueError : 

608 The transformation exists but was not defined by `register_linop`. 

609 

610 See also 

611 -------- 

612 transf 

613 

614 Notes 

615 ----- 

616 The linear operator is defined on the function the kernel represents the 

617 distribution of, not in full generality on the kernel itself. When 

618 multiple kernels are involved, their distributions may be considered 

619 independent or not, depending on the specific operation. 

620 

621 If the result is a subclass of the class defining the transformation, 

622 the result is casted to the latter. Then, if the result and all the 

623 operands are instances of `Kernel`, but the two operator arguments 

624 differ, the result is casted to its first non-`Kernel` superclass. 

625  

626 """ 

627 tcls, transf = self._gettransf(transfname) 6 ctx1feabcd

628 if transf.kind is not self._linopmarker: 6 ctx1feabcd

629 raise ValueError(f'the transformation {transfname!r} was not ' 4 ctx1abcd

630 f'defined with register_linop and so can not be invoked ' 

631 f'by linop') 

632 return transf.func(tcls, self, *args) 6 ctx1feabcd

633 

634 def algop(self, transfname, *operands, **kw): 6 ctx1feabcd

635 r""" 

636 

637 Return a nonnegative algebraic transformation of the input kernels. 

638 

639 .. math:: 

640 \mathrm{newkernel}(x, y) &= 

641 f(\mathrm{kernel}_1(x, y), \mathrm{kernel}_2(x, y), \ldots), \\ 

642 f(z_1, z_2, \ldots) &= \sum_{k_1,k_2,\ldots=0}^\infty 

643 a_{k_1 k_2 \ldots} z_1^{k_1} z_2^{k_2} \ldots, 

644 \quad a_* \ge 0. 

645 

646 Parameters 

647 ---------- 

648 transfname : hashable 

649 A name identifying the transformation. 

650 *operands : CrossKernel, scalar 

651 Arguments to the transformation in addition to self. 

652 **kw : 

653 Additional arguments to the transformation, not considered as 

654 operands. 

655 

656 Returns 

657 ------- 

658 newkernel : CrossKernel or NotImplemented 

659 The transformed kernel, or NotImplemented if the operation is 

660 not supported. 

661 

662 See also 

663 -------- 

664 transf 

665 

666 Notes 

667 ----- 

668 The class of `newkernel` is the least common superclass of: the 

669 "natural" output of the operation, the class defining the 

670 transformation, and the classes of the operands. 

671 

672 For class determination, scalars in the input count as `Constant` 

673 if nonnegative or traced by jax, else `CrossConstant`. 

674 

675 """ 

676 tcls, transf = self._gettransf(transfname) 6 ctx1feabcd

677 if transf.kind is not self._algopmarker: 6 ctx1feabcd

678 raise ValueError(f'the transformation {transfname!r} was not ' 4 ctx1abcd

679 f'defined with register_algop and so can not be invoked ' 

680 f'by algop') 

681 return transf.func(tcls, self, *operands, **kw) 6 ctx1feabcd

682 

683 @classmethod 6 ctx1feabcd

684 def register_transf(cls, func, transfname=None, doc=None, kind=None): 6 ctx1feabcd

685 """ 

686  

687 Register a transformation for use with `transf`. 

688 

689 The transformation will be accessible to subclasses. 

690 

691 Parameters 

692 ---------- 

693 func : callable 

694 A function ``func(tcls, self, *args, **kw) -> object`` implementing 

695 the transformation, where ``tcls`` is the class that defines the 

696 transformation. 

697 transfname : hashable, optional. 

698 The `transfname` parameter to `transf` this transformation will 

699 be accessible under. If not specified, use the name of `func`. 

700 doc : str, optional 

701 The documentation of the transformation for `transf_help`. If not 

702 specified, use the docstring of `func`. 

703 kind : object, optional 

704 An arbitrary marker. 

705 

706 Returns 

707 ------- 

708 func : callable 

709 The `func` argument as is. 

710 

711 Raises 

712 ------ 

713 KeyError : 

714 The name is already in use for another transformation in the same 

715 class. 

716 

717 See also 

718 -------- 

719 transf 

720 

721 """ 

722 if transfname is None: 6 ctx1feabcd

723 transfname = func.__name__ 6 ctx1feabcd

724 if doc is None: 6 ctx1feabcd

725 doc = func.__doc__ 6 ctx1feabcd

726 cls._settransf(transfname, (func, doc, kind)) 6 ctx1feabcd

727 return func 6 ctx1feabcd

728 

729 # TODO forbid to override with a different kind? 

730 

731 @classmethod 6 ctx1feabcd

732 def register_linop(cls, op, transfname=None, doc=None, argparser=None): 6 ctx1feabcd

733 """ 

734  

735 Register a transformation for use with `linop`. 

736 

737 Parameters 

738 ---------- 

739 op : callable 

740 A function ``op(tcls, self, arg1, arg2, *operands) -> CrossKernel`` 

741 that returns the new kernel, where ``arg1`` and ``arg2`` represent 

742 the operators acting on each side of the kernels, and ``operands`` 

743 are the other kernels beyond ``self``. 

744 transfname, doc : optional 

745 See `register_transf`. 

746 argparser : callable, optional 

747 A function applied to ``arg1`` and ``arg2``. Not called if the 

748 argument is `None`. It should map the identity to `None`. 

749 

750 Returns 

751 ------- 

752 op : callable 

753 The `op` argument as is. 

754 

755 Notes 

756 ----- 

757 The function `op` is called only if ``arg1`` or ``arg2`` is not `None` 

758 after potential conversion with `argparser`. 

759 

760 See also 

761 -------- 

762 transf 

763 

764 """ 

765 

766 if transfname is None: 6 ctx1feabcd

767 transfname = op.__name__ # for result type error message 6 ctx1feabcd

768 

769 @functools.wraps(op) 6 ctx1feabcd

770 def func(tcls, self, *allargs): 6 ctx1feabcd

771 

772 # split the arguments in kernels and non-kernels 

773 for pos, arg in enumerate(allargs): 6 ctx1feabcd

774 if not isinstance(arg, __class__): 6 ctx1feabcd

775 break 6 ctx1feabcd

776 else: 

777 pos = len(allargs) 4 ctx1abcd

778 operands = allargs[:pos] 6 ctx1feabcd

779 args = allargs[pos:] 6 ctx1feabcd

780 

781 # check the arguments from the first non-kernel onwards are 1 or 2 

782 if len(args) not in (1, 2): 6 ctx1feabcd

783 raise ValueError(f'incorrect number of non-kernel tail ' 4 ctx1abcd

784 f'arguments {len(args)}, expected 1 or 2') 

785 

786 # wrap argument parser to enforce preserving None 

787 if argparser: 6 ctx1feabcd

788 conv = lambda x: None if x is None else argparser(x) 6 ctx1feabcd

789 else: 

790 conv = lambda x: x 5 ctx1eabcd

791 

792 # determine if the two arguments count as "identical" or not 

793 if len(args) == 1: 6 ctx1feabcd

794 arg = conv(*args) 6 ctx1feabcd

795 different = False 6 ctx1feabcd

796 arg1 = arg2 = arg 6 ctx1feabcd

797 else: 

798 arg1, arg2 = args 6 ctx1feabcd

799 different = arg1 is not arg2 6 ctx1feabcd

800 arg1 = conv(arg1) 6 ctx1feabcd

801 arg2 = conv(arg2) 6 ctx1feabcd

802 different &= arg1 is not arg2 6 ctx1feabcd

803 # they must be not identical both before and after to handle 

804 # these cases: 

805 # - if the user passes identical arguments, but argparser 

806 # makes copies, it must still count as identical 

807 # - if the user passes arguments which are not identical, 

808 # but argparser sends them to the same object, then they 

809 # surely represent the same transf 

810 

811 # handle no-op case 

812 if arg1 is None and arg2 is None: 6 ctx1feabcd

813 return self 6 ctx1feabcd

814 

815 # invoke implementation 

816 result = op(tcls, self, arg1, arg2, *operands) 6 ctx1feabcd

817 

818 # check result is a kernel 

819 if not isinstance(result, __class__): 6 ctx1feabcd

820 raise TypeError(f'linop {transfname!r} returned ' 4 ctx1abcd

821 f'object of type {result.__class__.__name__}, expected ' 

822 f'subclass of {__class__.__name__}') 

823 

824 # modify class of the result 

825 rcls = result.__class__ 6 ctx1feabcd

826 if issubclass(rcls, tcls): 6 ctx1feabcd

827 rcls = tcls 6 ctx1feabcd

828 all_operands_kernel = all(isinstance(o, Kernel) for o in operands) 6 ctx1feabcd

829 if isinstance(self, Kernel) and all_operands_kernel and different: 6 ctx1feabcd

830 rcls = next(rcls._crossmro()) 6 ctx1feabcd

831 if rcls is not result.__class__: 6 ctx1feabcd

832 result = result._clone(rcls) 6 ctx1feabcd

833 

834 return result 6 ctx1feabcd

835 

836 cls.register_transf(func, transfname, doc, cls._linopmarker) 6 ctx1feabcd

837 return op 6 ctx1feabcd

838 

839 class _LinOpMarker(str): pass 6 ctx1feabcd

840 _linopmarker = _LinOpMarker('linop') 6 ctx1feabcd

841 

842 @classmethod 6 ctx1feabcd

843 def register_corelinop(cls, corefunc, transfname=None, doc=None, argparser=None): 6 ctx1feabcd

844 """ 

845 

846 Register a linear operator with a function that acts only on the core. 

847 

848 Parameters 

849 ---------- 

850 corefunc : callable 

851 A function ``corefunc(core, arg1, arg2, *cores) -> newcore``, where 

852 ``core`` is the function that implements the kernel passed at 

853 initialization, and ``cores`` for other operands. 

854 transfname, doc, argparser : 

855 See `register_linop`. 

856 

857 Returns 

858 ------- 

859 corefunc : callable 

860 The `corefunc` argument as is. 

861 

862 See also 

863 -------- 

864 transf 

865 

866 """ 

867 @functools.wraps(corefunc) 6 ctx1feabcd

868 def op(_, self, arg1, arg2, *operands): 6 ctx1feabcd

869 cores = (o.core for o in operands) 6 ctx1feabcd

870 core = corefunc(self.core, arg1, arg2, *cores) 6 ctx1feabcd

871 return self._clone(core=core) 6 ctx1feabcd

872 cls.register_linop(op, transfname, doc, argparser) 6 ctx1feabcd

873 return corefunc 6 ctx1feabcd

874 

875 @classmethod 6 ctx1feabcd

876 def register_xtransf(cls, xfunc, transfname=None, doc=None): 6 ctx1feabcd

877 """ 

878 

879 Register a linear operator that acts only on the input. 

880 

881 Parameters 

882 ---------- 

883 xfunc : callable 

884 A function ``xfunc(arg) -> (lambda x: newx)`` that takes in a 

885 `linop` argument and produces a function to transform the input. Not 

886 called if ``arg`` is `None`. To indicate the identity, return 

887 `None`. 

888 transfname, doc : 

889 See `register_linop`. `argparser` is not provided because its 

890 functionality can be included in `xfunc`. 

891 

892 Returns 

893 ------- 

894 xfunc : callable 

895 The `xfunc` argument as is. 

896 

897 See also 

898 -------- 

899 transf 

900