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

1# lsqfitgp/_Kernel/_crosskernel.py 

2# 

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

21import collections 3 ctx1cba

22import enum 3 ctx1cba

23import functools 3 ctx1cba

24import types 3 ctx1cba

25import warnings 3 ctx1cba

26 

27import numpy 3 ctx1cba

28from jax import numpy as jnp 3 ctx1cba

29 

30from lsqfitgp import _array, _jaxext, _utils 3 ctx1cba

31from lsqfitgp._Kernel import _util 3 ctx1cba

32 

33# subclasses defined in other modules, which set these names when imported 

34Kernel: type = None 3 ctx1cba

35Constant: type = None 3 ctx1cba

36CrossConstant: type = None 3 ctx1cba

37 

38 

39@functools.cache 3 ctx1cba

40def least_common_superclass(*classes): 3 ctx1cba

41 """ 

42 Find a "least" common superclass. 

43 

44 The class is searched in all the MROs, but the comparison is done with 

45 `issubclass` to support virtual inheritance. 

46 """ 

47 mros = [c.__mro__ for c in classes] 3 ctx1cba

48 indices = [0] * len(mros) 3 ctx1cba

49 for i, mroi in enumerate(mros): 3 ctx1cba

50 for j, mroj in enumerate(mros): 3 ctx1cba

51 if i == j: 3 ctx1cba

52 continue 3 ctx1cba

53 while not issubclass(mroi[0], mroj[indices[j]]): 3 ctx1cba

54 indices[j] += 1 3 ctx1cba

55 idx = numpy.argmin(indices) 3 ctx1cba

56 return mros[idx][indices[idx]] 3 ctx1cba

57 

58 

59class CrossKernel: 3 ctx1cba

60 r""" 

61 

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

63 

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

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

66 

67 .. math:: 

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

69 

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

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

72 are immutable; all operations return new objects. 

73 

74 Parameters 

75 ---------- 

76 core : callable 

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

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

79 kernel. 

80 scale, loc, derivable, maxdim, dim : 

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

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

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

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

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

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

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

88 as a pair of arguments. 

89 forcekron : bool, default False 

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

91 operations above. Available only for `Kernel`. 

92 batchbytes : number, optional 

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

94 dynkw : dict, optional 

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

96 by transformations. Deleted by transformations by default. 

97 **initkw : 

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

99 changed by transformations. 

100 

101 Attributes 

102 ---------- 

103 initkw : dict 

104 The `initkw` argument. 

105 dynkw : dict 

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

107 transformed. 

108 core : callable 

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

110 function wrapping it if the object has been transformed. 

111 

112 Methods 

113 ------- 

114 batch 

115 linop 

116 algop 

117 transf 

118 register_transf 

119 register_linop 

120 register_corelinop 

121 register_xtransf 

122 register_algop 

123 register_ufuncalgop 

124 transf_help 

125 has_transf 

126 list_transf 

127 super_transf 

128 make_linop_family 

129 

130 See Also 

131 -------- 

132 Kernel 

133 

134 Notes 

135 ----- 

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

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

138 and multiplication preserve the class. 

139 

140 """ 

141 

142 __slots__ = '_core', '_dynkw', '_initkw' 3 ctx1cba

143 # only __new__ and _clone shall access these attributes 

144 

145 @property 3 ctx1cba

146 def initkw(self): 3 ctx1cba

147 return types.MappingProxyType(self._initkw) 3 ctx1cba

148 

149 @property 3 ctx1cba

150 def dynkw(self): 3 ctx1cba

151 return types.MappingProxyType(self._dynkw) 3 ctx1cba

152 

153 @property 3 ctx1cba

154 def core(self): 3 ctx1cba

155 return self._core 3 ctx1cba

156 

157 def __new__( 3 ctx1cba

158 cls, 

159 core, 

160 *, 

161 scale=None, 

162 loc=None, 

163 derivable=None, 

164 maxdim=None, 

165 dim=None, 

166 forcekron=False, 

167 batchbytes=None, 

168 dynkw={}, # noqa: B006, read only 

169 **initkw, 

170 ): 

171 self = super().__new__(cls) 3 ctx1cba

172 

173 self._initkw = initkw 3 ctx1cba

174 self._dynkw = dict(dynkw) 3 ctx1cba

175 self._core = lambda x, y, **dynkw: core(x, y, **initkw, **dynkw) 3 ctx1cba

176 

177 if forcekron: 3 ctx1cba

178 self = self.transf('forcekron') 1 ctx1a

179 

180 linop_args = { 3 ctx1cba

181 'scale': scale, 

182 'loc': loc, 

183 'derivable': derivable, 

184 'maxdim': maxdim, 

185 'dim': dim, 

186 } 

187 for transfname, arg in linop_args.items(): 3 ctx1cba

188 if callable(arg): 3 ctx1cba

189 arg = arg(**initkw) 3 ctx1cba

190 if isinstance(arg, tuple): 3 ctx1cba

191 self = self.linop(transfname, *arg) 2 ctx1ba

192 else: 

193 self = self.linop(transfname, arg) 3 ctx1cba

194 

195 if batchbytes is not None: 3 ctx1cba

196 self = self.batch(batchbytes) 3 ctx1cba

197 

198 return self 3 ctx1cba

199 

200 def __call__(self, x, y): 3 ctx1cba

201 x = _array.asarray(x) 3 ctx1cba

202 y = _array.asarray(y) 3 ctx1cba

203 shape = _array.broadcast(x, y).shape 3 ctx1cba

204 result = self.core(x, y, **self.dynkw) 3 ctx1cba

205 assert isinstance(result, (numpy.ndarray, jnp.number, jnp.ndarray)) 3 ctx1cba

206 assert jnp.issubdtype(result.dtype, jnp.number), result.dtype 3 ctx1cba

207 assert result.shape == shape, (result.shape, shape) 3 ctx1cba

208 return result 3 ctx1cba

209 

210 def _clone(self, cls=None, *, initkw=None, dynkw=None, core=None): 3 ctx1cba

211 newself = object.__new__(self.__class__ if cls is None else cls) 3 ctx1cba

212 newself._initkw = self._initkw if initkw is None else dict(initkw) 3 ctx1cba

213 newself._dynkw = {} if dynkw is None else dict(dynkw) 3 ctx1cba

214 newself._core = self._core if core is None else core 3 ctx1cba

215 return newself 3 ctx1cba

216 

217 class _side(enum.Enum): 3 ctx1cba

218 LEFT = 0 3 ctx1cba

219 RIGHT = 1 3 ctx1cba

220 

221 @classmethod 3 ctx1cba

222 def _nary(cls, op, kernels, side): 3 ctx1cba

223 

224 if side is cls._side.LEFT: 3 ctx1cba

225 wrapper = lambda c, _, y, **kw: lambda x: c(x, y, **kw) 3 ctx1cba

226 arg = lambda x, _: x 3 ctx1cba

227 elif side is cls._side.RIGHT: 1 ctx1a

228 wrapper = lambda c, x, _, **kw: lambda y: c(x, y, **kw) 1 ctx1a

229 arg = lambda _, y: y 1 ctx1a

230 else: # pragma: no cover 

231 raise KeyError(side) 

232 

233 cores = [k.core for k in kernels] 3 ctx1cba

234 

235 def core(x, y, **kw): 3 ctx1cba

236 wrapped = [wrapper(c, x, y, **kw) for c in cores] 3 ctx1cba

237 transformed = op(*wrapped) 3 ctx1cba

238 return transformed(arg(x, y)) 3 ctx1cba

239 

240 return __class__(core) 3 ctx1cba

241 

242 def __add__(self, other): 3 ctx1cba

243 return self.algop('add', other) 3 ctx1cba

244 

245 __radd__ = __add__ 3 ctx1cba

246 

247 def __mul__(self, other): 3 ctx1cba

248 return self.algop('mul', other) 3 ctx1cba

249 

250 __rmul__ = __mul__ 3 ctx1cba

251 

252 def __pow__(self, other): 3 ctx1cba

253 return self.algop('pow', exponent=other) 1 ctx1a

254 

255 def __rpow__(self, other): 3 ctx1cba

256 return self.algop('rpow', base=other) 1 ctx1a

257 

258 def _swap(self): 3 ctx1cba

259 """Permute the arguments.""" 

260 core = self.core 3 ctx1cba

261 return self._clone(__class__, core=lambda x, y, **kw: core(y, x, **kw)) 3 ctx1cba

262 

263 def batch(self, maxnbytes): 3 ctx1cba

264 """ 

265 Return a batched version of the kernel. 

266 

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

268 usage. 

269 

270 Parameters 

271 ---------- 

272 maxnbytes : number 

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

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

275 not induced by the operations in the kernel. 

276 

277 Returns 

278 ------- 

279 batched_kernel : CrossKernel 

280 The same kernel but with batched computations. 

281 """ 

282 core = _jaxext.batchufunc(self.core, maxnbytes=maxnbytes) 3 ctx1cba

283 return self._clone(core=core) 3 ctx1cba

284 

285 @classmethod 3 ctx1cba

286 def _crossmro(cls): 3 ctx1cba

287 """MRO iterator excluding subclasses of Kernel.""" 

288 for c in cls.mro(): # pragma: no branch 3 ctx1cba

289 if not issubclass(c, Kernel): 3 ctx1cba

290 yield c 3 ctx1cba

291 if c is __class__: 3 ctx1cba

292 break 1 ctx1a

293 

294 _transf = {} # noqa: RUF012, registry shared by design, replaced in subclasses 3 ctx1cba

295 

296 _Transf = collections.namedtuple('_Transf', ['func', 'doc', 'kind']) 3 ctx1cba

297 

298 def __init_subclass__(cls, **kw): 3 ctx1cba

299 super().__init_subclass__(**kw) 3 ctx1cba

300 cls._transf = {} 3 ctx1cba

301 cls.__slots__ = () 3 ctx1cba

302 

303 @classmethod 3 ctx1cba

304 def _transfmro(cls): 3 ctx1cba

305 """Iterate over the superclasses with a _transf attribute.""" 

306 for c in cls.mro(): # pragma: no branch 3 ctx1cba

307 yield c 3 ctx1cba

308 if c is __class__: 3 ctx1cba

309 break 3 ctx1cba

310 

311 @classmethod 3 ctx1cba

312 def _settransf(cls, transfname, transf): 3 ctx1cba

313 if transfname in cls._transf: 3 ctx1cba

314 msg = f'transformation {transfname!r} already registered for {cls.__name__}' 1 ctx1a

315 raise KeyError(msg) 1 ctx1a

316 cls._transf[transfname] = cls._Transf(*transf) 3 ctx1cba

317 

318 @classmethod 3 ctx1cba

319 def _alltransf(cls): 3 ctx1cba

320 """List all accessible transfs as dict name -> (tcls, transf).""" 

321 transfs = {} 3 ctx1cba

322 for tcls in cls._transfmro(): 3 ctx1cba

323 for name, transf in tcls._transf.items(): 3 ctx1cba

324 transfs.setdefault(name, (tcls, transf)) 3 ctx1cba

325 return transfs 3 ctx1cba

326 

327 @classmethod 3 ctx1cba

328 def _gettransf(cls, transfname, transfmro=None): 3 ctx1cba

329 """ 

330 Find a transformation. 

331 

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

333 

334 Parameters 

335 ---------- 

336 transfname : hashable 

337 The transformation name. 

338 transfmro : iterable of types, optional 

339 The classes to search, by default the MRO up to `CrossKernel`. 

340 

341 Returns 

342 ------- 

343 cls : type 

344 The class where the transformation was found. 

345 transf, doc, kind : tuple 

346 The objects set by `register_transf`. 

347 

348 Raises 

349 ------ 

350 KeyError : 

351 The transformation was not found. 

352 """ 

353 if transfmro is None: 3 ctx1cba

354 transfmro = cls._transfmro() 3 ctx1cba

355 for c in transfmro: 3 ctx1cba

356 if transfname in c._transf: 3 ctx1cba

357 return c, c._transf[transfname] 3 ctx1cba

358 raise KeyError(transfname) 1 ctx1a

359 

360 @classmethod 3 ctx1cba

361 def inherit_transf(cls, transfname, *, intermediates=False): 3 ctx1cba

362 """ 

363 Inherit a transformation from a superclass. 

364 

365 Parameters 

366 ---------- 

367 transfname : hashable 

368 The name of the transformation. 

369 intermediates : bool, default False 

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

371 transformation inherit it too. 

372 

373 Raises 

374 ------ 

375 KeyError : 

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

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

378 

379 See Also 

380 -------- 

381 transf 

382 """ 

383 tcls, transf = cls._gettransf(transfname) 3 ctx1cba

384 cls._settransf(transfname, transf) 3 ctx1cba

385 if intermediates: 3 ctx1cba

386 for c in cls.mro()[1:]: # pragma: no branch 3 ctx1cba

387 if c is tcls: 3 ctx1cba

388 break 3 ctx1cba

389 c._settransf(transfname, transf) 3 ctx1cba

390 

391 @classmethod 3 ctx1cba

392 def inherit_all_algops(cls, intermediates=False): 3 ctx1cba

393 """ 

394 Inherit all algebraic operations from superclasses. 

395 

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

397 should be preserved by addition and multiplication. 

398 

399 Parameters 

400 ---------- 

401 intermediates : bool, default False 

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

403 transformation inherit it too. 

404 

405 Raises 

406 ------ 

407 KeyError : 

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

409 classes. 

410 

411 See Also 

412 -------- 

413 transf 

414 """ 

415 mro = cls._transfmro() 3 ctx1cba

416 next(mro) 3 ctx1cba

417 for name, (_, transf) in next(mro)._alltransf().items(): 3 ctx1cba

418 if transf.kind is cls._algopmarker: 3 ctx1cba

419 cls.inherit_transf(name, intermediates=intermediates) 3 ctx1cba

420 

421 @classmethod 3 ctx1cba

422 def list_transf(cls, superclasses=True): 3 ctx1cba

423 """ 

424 List all the available transformations. 

425 

426 Parameters 

427 ---------- 

428 superclasses : bool, default True 

429 Include transformations defined in superclasses. 

430 

431 Returns 

432 ------- 

433 transfs : dict of Transf 

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

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

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

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

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

439 """ 

440 if superclasses: 2 ctx1ca

441 source = cls._alltransf().items 1 ctx1a

442 else: 

443 

444 def source(): 2 ctx1ca

445 for name, transf in cls._transf.items(): 2 ctx1ca

446 yield name, (cls, transf) 2 ctx1ca

447 

448 return { 2 ctx1ca

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

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

451 } 

452 

453 Transf = collections.namedtuple('Transf', ['tcls', 'kind', 'func', 'doc']) 3 ctx1cba

454 

455 @classmethod 3 ctx1cba

456 def has_transf(cls, transfname): 3 ctx1cba

457 """ 

458 Check if a transformation is registered. 

459 

460 Parameters 

461 ---------- 

462 transfname : hashable 

463 The transformation name. 

464 

465 Returns 

466 ------- 

467 has_transf : bool 

468 Whether the transformation is registered. 

469 

470 See Also 

471 -------- 

472 transf 

473 """ 

474 try: 1 ctx1a

475 cls._gettransf(transfname) 1 ctx1a

476 except KeyError as exc: 1 ctx1a

477 if exc.args == (transfname,): 1 ctx1a

478 return False 1 ctx1a

479 else: # pragma: no cover 

480 raise 

481 else: 

482 return True 1 ctx1a

483 

484 @classmethod 3 ctx1cba

485 def transf_help(cls, transfname): 3 ctx1cba

486 """ 

487 Return the documentation of a transformation. 

488 

489 Parameters 

490 ---------- 

491 transfname : hashable 

492 The name of the transformation. 

493 

494 Returns 

495 ------- 

496 doc : str 

497 The documentation of the transformation. 

498 

499 See Also 

500 -------- 

501 transf 

502 """ 

503 _, transf = cls._gettransf(transfname) 1 ctx1a

504 return transf.doc 1 ctx1a

505 

506 def transf(self, transfname, *args, **kw): 3 ctx1cba

507 """ 

508 Return a transformed kernel. 

509 

510 Parameters 

511 ---------- 

512 transfname : hashable 

513 A name identifying the transformation. 

514 *args : 

515 Positional arguments to the transformation. 

516 **kw : 

517 Keyword arguments to the transformation. 

518 

519 Returns 

520 ------- 

521 newkernel : object 

522 The output of the transformation. 

523 

524 Raises 

525 ------ 

526 KeyError 

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

528 

529 See Also 

530 -------- 

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

532 """ 

533 tcls, transf = self._gettransf(transfname) 1 ctx1a

534 return transf.func(tcls, self, *args, **kw) 1 ctx1a

535 

536 @classmethod 3 ctx1cba

537 def super_transf(cls, transfname, self, *args, **kw): 3 ctx1cba

538 """ 

539 Transform the kernel using a superclass transformation. 

540 

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

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

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

544 

545 Parameters 

546 ---------- 

547 transfname : hashable 

548 See `transf`. 

549 self : CrossKernel 

550 The object to transform. 

551 *args : 

552 See `transf`. 

553 **kw : 

554 See `transf`. 

555 

556 Returns 

557 ------- 

558 newkernel : object 

559 The output of the transformation. 

560 """ 

561 mro = list(self._transfmro()) 2 ctx1ba

562 idx = mro.index(cls) 2 ctx1ba

563 tcls, transf = self._gettransf(transfname, mro[idx + 1 :]) 2 ctx1ba

564 return transf.func(tcls, self, *args, **kw) 2 ctx1ba

565 

566 def linop(self, transfname, *args, **kw): 3 ctx1cba

567 r""" 

568 

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

570 

571 .. math:: 

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

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

574 &\ldots \\ 

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

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

577 

578 Parameters 

579 ---------- 

580 transfname : hashable 

581 The name of the transformation. 

582 *args : 

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

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

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

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

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

588 the identity. 

589 **kw : 

590 Ignored. 

591 

592 Returns 

593 ------- 

594 newkernel : CrossKernel 

595 The transformed kernel. 

596 

597 Raises 

598 ------ 

599 ValueError : 

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

601 

602 See Also 

603 -------- 

604 transf 

605 

606 Notes 

607 ----- 

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

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

610 multiple kernels are involved, their distributions may be considered 

611 independent or not, depending on the specific operation. 

612 

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

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

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

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

617 

618 """ 

619 tcls, transf = self._gettransf(transfname) 3 ctx1cba

620 if transf.kind is not self._linopmarker: 3 ctx1cba

621 msg = ( 1 ctx1a

622 f'the transformation {transfname!r} was not ' 

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

624 f'by linop' 

625 ) 

626 raise ValueError(msg) 1 ctx1a

627 return transf.func(tcls, self, *args) 3 ctx1cba

628 

629 def algop(self, transfname, *operands, **kw): 3 ctx1cba

630 r""" 

631 

632 Return a nonnegative algebraic transformation of the input kernels. 

633 

634 .. math:: 

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

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

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

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

639 \quad a_* \ge 0. 

640 

641 Parameters 

642 ---------- 

643 transfname : hashable 

644 A name identifying the transformation. 

645 *operands : CrossKernel, scalar 

646 Arguments to the transformation in addition to self. 

647 **kw : 

648 Additional arguments to the transformation, not considered as 

649 operands. 

650 

651 Returns 

652 ------- 

653 newkernel : CrossKernel or NotImplemented 

654 The transformed kernel, or NotImplemented if the operation is 

655 not supported. 

656 

657 See Also 

658 -------- 

659 transf 

660 

661 Notes 

662 ----- 

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

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

665 transformation, and the classes of the operands. 

666 

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

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

669 

670 """ 

671 tcls, transf = self._gettransf(transfname) 3 ctx1cba

672 if transf.kind is not self._algopmarker: 3 ctx1cba

673 msg = ( 1 ctx1a

674 f'the transformation {transfname!r} was not ' 

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

676 f'by algop' 

677 ) 

678 raise ValueError(msg) 1 ctx1a

679 return transf.func(tcls, self, *operands, **kw) 3 ctx1cba

680 

681 @classmethod 3 ctx1cba

682 def register_transf(cls, func, transfname=None, doc=None, kind=None): 3 ctx1cba

683 """ 

684 Register a transformation for use with `transf`. 

685 

686 The transformation will be accessible to subclasses. 

687 

688 Parameters 

689 ---------- 

690 func : callable 

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

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

693 transformation. 

694 transfname : hashable, optional. 

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

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

697 doc : str, optional 

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

699 specified, use the docstring of `func`. 

700 kind : object, optional 

701 An arbitrary marker. 

702 

703 Returns 

704 ------- 

705 func : callable 

706 The `func` argument as is. 

707 

708 Raises 

709 ------ 

710 KeyError : 

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

712 class. 

713 

714 See Also 

715 -------- 

716 transf 

717 """ 

718 if transfname is None: 3 ctx1cba

719 transfname = func.__name__ 3 ctx1cba

720 if doc is None: 3 ctx1cba

721 doc = func.__doc__ 3 ctx1cba

722 cls._settransf(transfname, (func, doc, kind)) 3 ctx1cba

723 return func 3 ctx1cba

724 

725 @classmethod 3 ctx1cba

726 def register_linop(cls, op, transfname=None, doc=None, argparser=None): # noqa: C901 3 ctx1cba

727 """ 

728 Register a transformation for use with `linop`. 

729 

730 Parameters 

731 ---------- 

732 op : callable 

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

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

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

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

737 transfname : hashable, optional 

738 See `register_transf`. 

739 doc : str, optional 

740 See `register_transf`. 

741 argparser : callable, optional 

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

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

744 

745 Returns 

746 ------- 

747 op : callable 

748 The `op` argument as is. 

749 

750 See Also 

751 -------- 

752 transf 

753 

754 Notes 

755 ----- 

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

757 after potential conversion with `argparser`. 

758 """ 

759 if transfname is None: 3 ctx1cba

760 transfname = op.__name__ # for result type error message 3 ctx1cba

761 

762 @functools.wraps(op) 3 ctx1cba

763 def func(tcls, self, *allargs): # noqa: C901 3 ctx1cba

764 

765 # split the arguments in kernels and non-kernels 

766 for pos, arg in enumerate(allargs): # noqa: B007, `pos` used after the loop 3 ctx1cba

767 if not isinstance(arg, __class__): # ty: ignore[unresolved-reference] 3 ctx1cba

768 break 3 ctx1cba

769 else: 

770 pos = len(allargs) 1 ctx1a

771 operands = allargs[:pos] 3 ctx1cba

772 args = allargs[pos:] 3 ctx1cba

773 

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

775 if len(args) not in (1, 2): 3 ctx1cba

776 msg = ( 1 ctx1a

777 f'incorrect number of non-kernel tail ' 

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

779 ) 

780 raise ValueError(msg) 1 ctx1a

781 

782 # wrap argument parser to enforce preserving None 

783 if argparser: 3 ctx1cba

784 conv = lambda x: None if x is None else argparser(x) 3 ctx1cba

785 else: 

786 conv = lambda x: x 3 ctx1cba

787 

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

789 if len(args) == 1: 3 ctx1cba

790 arg = conv(*args) 3 ctx1cba

791 different = False 3 ctx1cba

792 arg1 = arg2 = arg 3 ctx1cba

793 else: 

794 arg1, arg2 = args 3 ctx1cba

795 different = arg1 is not arg2 3 ctx1cba

796 arg1 = conv(arg1) 3 ctx1cba

797 arg2 = conv(arg2) 3 ctx1cba

798 different &= arg1 is not arg2 3 ctx1cba

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

800 # these cases: 

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

802 # makes copies, it must still count as identical 

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

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

805 # surely represent the same transf 

806 

807 # handle no-op case 

808 if arg1 is None and arg2 is None: 3 ctx1cba

809 return self 3 ctx1cba

810 

811 # invoke implementation 

812 result = op(tcls, self, arg1, arg2, *operands) 3 ctx1cba

813 

814 # check result is a kernel 

815 if not isinstance(result, __class__): # ty: ignore[unresolved-reference] 3 ctx1cba

816 msg = ( 1 ctx1a

817 f'linop {transfname!r} returned ' 

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

819 f'subclass of {__class__.__name__}' # ty: ignore[unresolved-reference] 

820 ) 

821 raise TypeError(msg) 1 ctx1a

822 

823 # modify class of the result 

824 rcls = result.__class__ 3 ctx1cba

825 if issubclass(rcls, tcls): 3 ctx1cba

826 rcls = tcls 3 ctx1cba

827 all_operands_kernel = all(isinstance(o, Kernel) for o in operands) 3 ctx1cba

828 if isinstance(self, Kernel) and all_operands_kernel and different: 3 ctx1cba

829 rcls = next(rcls._crossmro()) 3 ctx1cba

830 if rcls is not result.__class__: 3 ctx1cba

831 result = result._clone(rcls) 3 ctx1cba

832 

833 return result 3 ctx1cba

834 

835 cls.register_transf(func, transfname, doc, cls._linopmarker) 3 ctx1cba

836 return op 3 ctx1cba

837 

838 class _LinOpMarker(str): 3 ctx1cba

839 pass 3 ctx1cba

840 

841 _linopmarker = _LinOpMarker('linop') 3 ctx1cba

842 

843 @classmethod 3 ctx1cba

844 def register_corelinop(cls, corefunc, transfname=None, doc=None, argparser=None): 3 ctx1cba

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 : hashable, optional 

855 See `register_linop`. 

856 doc : str, optional 

857 See `register_linop`. 

858 argparser : callable, optional 

859 See `register_linop`. 

860 

861 Returns 

862 ------- 

863 corefunc : callable 

864 The `corefunc` argument as is. 

865 

866 See Also 

867 -------- 

868 transf 

869 """ 

870 

871 @functools.wraps(corefunc) 3 ctx1cba

872 def op(_, self, arg1, arg2, *operands): 3 ctx1cba

873 cores = (o.core for o in operands) 3 ctx1cba

874 core = corefunc(self.core, arg1, arg2, *cores) 3 ctx1cba

875 return self._clone(core=core) 3 ctx1cba

876 

877 cls.register_linop(op, transfname, doc, argparser) 3 ctx1cba

878 return corefunc 3 ctx1cba

879 

880 @classmethod 3 ctx1cba

881 def register_xtransf(cls, xfunc, transfname=None, doc=None): 3 ctx1cba

882 """ 

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

884 

885 Parameters 

886 ---------- 

887 xfunc : callable 

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

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

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

891 `None`. There is no `argparser` parameter (see `register_linop`) 

892 because its functionality can be included in `xfunc`. 

893 transfname : hashable, optional 

894 See `register_linop`.