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

1# lsqfitgp/_Kernel/_ops.py 

2# 

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

20"""Register linops on CrossKernel and AffineSpan.""" 

21 

22import functools 3 ctx1cba

23import numbers 3 ctx1cba

24import sys 3 ctx1cba

25 

26from jax import numpy as jnp 3 ctx1cba

27 

28from lsqfitgp import _array, _Deriv, _jaxext 3 ctx1cba

29from lsqfitgp._Kernel import _util 3 ctx1cba

30from lsqfitgp._Kernel._crosskernel import AffineSpan, CrossKernel 3 ctx1cba

31 

32 

33def rescale_argparser(fun): 3 ctx1cba

34 if not callable(fun): 2 ctx1ba

35 msg = "argument to 'rescale' must be a function" 1 ctx1a

36 raise ValueError(msg) # noqa: TRY004, changing the type would break callers 1 ctx1a

37 return fun 2 ctx1ba

38 

39 

40@functools.partial(CrossKernel.register_corelinop, argparser=rescale_argparser) 3 ctx1cba

41def rescale(core, xfun, yfun): 3 ctx1cba

42 r""" 

43 

44 Rescale the output of the function. 

45 

46 .. math:: 

47 T(f)(x) = \mathrm{fun}(x) f(x) 

48 

49 Parameters 

50 ---------- 

51 xfun, yfun : callable or None 

52 Functions from the type of the arguments of the kernel to scalar. 

53 

54 """ # noqa: DOC101, DOC103, DOC201 

55 if not xfun: 2 ctx1ba

56 return lambda x, y, **kw: core(x, y, **kw) * yfun(y) 1 ctx1a

57 elif not yfun: 2 ctx1ba

58 return lambda x, y, **kw: xfun(x) * core(x, y, **kw) 2 ctx1ba

59 else: 

60 return lambda x, y, **kw: xfun(x) * core(x, y, **kw) * yfun(y) 2 ctx1ba

61 

62 

63@CrossKernel.register_xtransf 3 ctx1cba

64def derivable(derivable): 3 ctx1cba

65 """ 

66 Specify the degree of derivability of the function. 

67 

68 Parameters 

69 ---------- 

70 xderivable, yderivable : int or None 

71 Degree of derivability of the function. None means unknown. 

72 

73 Notes 

74 ----- 

75 The derivability check is hardcoded into the kernel core and it is not 

76 possible to remove it afterwards by applying ``'derivable'`` again with a 

77 higher limit. 

78 """ # noqa: DOC103, DOC201 

79 if isinstance(derivable, bool): 3 ctx1cba

80 derivable = sys.maxsize if derivable else 0 3 ctx1cba

81 elif not isinstance(derivable, numbers.Integral) or derivable < 0: 3 ctx1cba

82 msg = f'derivability degree {derivable!r} not valid' 1 ctx1a

83 raise ValueError(msg) 1 ctx1a

84 

85 def error_func(current, n): 3 ctx1cba

86 msg = ( 1 ctx1a

87 f'Took {current} derivatives > limit {n} on argument ' 

88 'of a kernel. This error may be spurious if there are ' 

89 'derivatives on values that define the input to the kernel, for ' 

90 'example if a hyperparameter enters the calculation of x. To ' 

91 'suppress the error, initialize the kernel with derivable=True.' 

92 ) 

93 raise ValueError(msg) 1 ctx1a

94 

95 def xtransf(x): 3 ctx1cba

96 if hasattr(x, 'dtype'): 96 ↛ 106line 96 didn't jump to line 106 because the condition on line 96 was always true3 ctx1cba

97 # this branch handles limit_derivatives not accepting non-jax types 

98 # because of being based on jax.custom_jvp; this restriction on 

99 # custom_jvp appeared in jax 0.4.17 

100 if x.dtype.names is not None: 3 ctx1cba

101 x = _array.StructuredArray(x) # structured arrays are not 3 ctx1cba

102 # compatible with jax but common in lsqfitgp, so I wrap them 

103 elif not _jaxext.is_jax_type(x.dtype): 3 ctx1cba

104 return x # since anyway there would be an error if a derivative 1 ctx1a

105 # tried to pass through a non-jax type 

106 return _jaxext.limit_derivatives(x, n=derivable, error_func=error_func) 3 ctx1cba

107 

108 return xtransf 3 ctx1cba

109 

110 

111def _asfloat(x): 3 ctx1cba

112 return x.astype(_jaxext.float_type(x)) 3 ctx1cba

113 

114 

115def diff_argparser(deriv): 3 ctx1cba

116 deriv = _Deriv.Deriv(deriv) 3 ctx1cba

117 return deriv or None 3 ctx1cba

118 

119 

120@functools.partial(CrossKernel.register_corelinop, argparser=diff_argparser) 3 ctx1cba

121def diff(core, xderiv, yderiv): # noqa: C901 3 ctx1cba

122 r""" 

123 

124 Derive the function. 

125 

126 .. math:: 

127 T(f)(x) = \frac{\partial^n f}{\partial x^n} (x) 

128 

129 Parameters 

130 ---------- 

131 xderiv, yderiv : Deriv_like 

132 A `Deriv` or something that can be converted to a `Deriv`. 

133 

134 Raises 

135 ------ 

136 RuntimeError 

137 The derivative orders are greater than the `derivative` attribute. 

138 

139 """ # noqa: DOC101, DOC103, DOC201 

140 # reparse derivatives because they could be None 

141 xderiv = _Deriv.Deriv(xderiv) 3 ctx1cba

142 yderiv = _Deriv.Deriv(yderiv) 3 ctx1cba

143 

144 # wrapper of kernel with derivable arguments unpacked 

145 def f(x, y, *args, **kw): 3 ctx1cba

146 i = -1 3 ctx1cba

147 if not xderiv.implicit: 3 ctx1cba

148 for i, dim in enumerate(xderiv): 3 ctx1cba

149 x = x.at[dim].set(args[i]) 3 ctx1cba

150 if not yderiv.implicit: 3 ctx1cba

151 for j, dim in enumerate(yderiv): 3 ctx1cba

152 y = y.at[dim].set(args[1 + i + j]) 3 ctx1cba

153 return core(x, y, **kw) 3 ctx1cba

154 

155 # last x index in case iteration on x does not run but does on y 

156 i = -1 3 ctx1cba

157 

158 # derive w.r.t. first argument 

159 if xderiv.implicit: 3 ctx1cba

160 for _ in range(xderiv.order): 3 ctx1cba

161 f = _jaxext.elementwise_grad(f, 0) 3 ctx1cba

162 else: 

163 for i, dim in enumerate(xderiv): 3 ctx1cba

164 for _ in range(xderiv[dim]): 3 ctx1cba

165 f = _jaxext.elementwise_grad(f, 2 + i) 3 ctx1cba

166 

167 # derive w.r.t. second argument 

168 if yderiv.implicit: 3 ctx1cba

169 for _ in range(yderiv.order): 3 ctx1cba

170 f = _jaxext.elementwise_grad(f, 1) 3 ctx1cba

171 else: 

172 for j, dim in enumerate(yderiv): 3 ctx1cba

173 for _ in range(yderiv[dim]): 3 ctx1cba

174 f = _jaxext.elementwise_grad(f, 2 + 1 + i + j) 3 ctx1cba

175 

176 # check derivatives are ok for actual input arrays, wrap structured arrays 

177 def process_arg(x, deriv, pos): 3 ctx1cba

178 if x.dtype.names is not None: 3 ctx1cba

179 for dim in deriv: 3 ctx1cba

180 if dim not in x.dtype.names: 3 ctx1cba

181 msg = f'derivative along missing field {dim!r} on {pos} argument' 1 ctx1a

182 raise ValueError(msg) 1 ctx1a

183 if not jnp.issubdtype(x.dtype[dim], jnp.number): 3 ctx1cba

184 msg = ( 1 ctx1a

185 f'derivative along non-numeric field {dim!r} on {pos} argument' 

186 ) 

187 raise TypeError(msg) 1 ctx1a

188 return _array.StructuredArray(x) 3 ctx1cba

189 elif not deriv.implicit: 3 ctx1cba

190 msg = ( 1 ctx1a

191 'derivative on named fields with non-structured ' 

192 f'array on {pos} argument' 

193 ) 

194 raise ValueError(msg) 1 ctx1a

195 elif not jnp.issubdtype(x.dtype, jnp.number): 3 ctx1cba

196 msg = f'derivative along non-numeric array on {pos} argument' 1 ctx1a

197 raise TypeError(msg) 1 ctx1a

198 return x 3 ctx1cba

199 

200 def newcore(x, y, **kw): 3 ctx1cba

201 x = process_arg(x, xderiv, 'left') 3 ctx1cba

202 y = process_arg(y, yderiv, 'right') 3 ctx1cba

203 

204 args = [] 3 ctx1cba

205 

206 if not xderiv.implicit: 3 ctx1cba

207 args.extend(_asfloat(x[dim]) for dim in xderiv) 3 ctx1cba

208 elif xderiv: 3 ctx1cba

209 x = _asfloat(x) 3 ctx1cba

210 

211 if not yderiv.implicit: 3 ctx1cba

212 args.extend(_asfloat(y[dim]) for dim in yderiv) 3 ctx1cba

213 elif yderiv: 3 ctx1cba

214 y = _asfloat(y) 3 ctx1cba

215 

216 return f(x, y, *args, **kw) 3 ctx1cba

217 

218 return newcore 3 ctx1cba

219 

220 

221@CrossKernel.register_xtransf 3 ctx1cba

222def xtransf(fun): 3 ctx1cba

223 r""" 

224 

225 Transform the inputs of the function. 

226 

227 .. math:: 

228 T(f)(x) = f(\mathrm{fun}(x)) 

229 

230 Parameters 

231 ---------- 

232 xfun, yfun : callable or None 

233 Functions mapping a new kind of input to the kind of input accepted by 

234 the kernel. 

235 

236 """ # noqa: DOC103, DOC201 

237 if not callable(fun): 1 ctx1a

238 msg = "argument to 'xtransf' must be a function" 1 ctx1a

239 raise ValueError(msg) # noqa: TRY004, changing the type would break callers 1 ctx1a

240 return fun 1 ctx1a

241 

242 

243@CrossKernel.register_xtransf 3 ctx1cba

244def dim(dim): 3 ctx1cba

245 """ 

246 Restrict the function to a field of a structured input. 

247 

248 That is:: 

249 

250 T(f)(x) = f(x[dim]) 

251 

252 If the array is not structured, an exception is raised. If the field for 

253 name `dim` has a nontrivial shape, the array passed to the kernel is still 

254 structured but has only field `dim`. 

255 

256 Parameters 

257 ---------- 

258 xdim, ydim : None, str, list of str 

259 Field names or lists of field names. 

260 """ # noqa: DOC103, DOC201 

261 if not isinstance(dim, (str, list)): 3 ctx1cba

262 msg = f'dim must be a (list of) string, found {dim!r}' 1 ctx1a

263 raise TypeError(msg) 1 ctx1a

264 

265 def fun(x): 3 ctx1cba

266 if x.dtype.names is None: 3 ctx1cba

267 msg = f'cannot get dim={dim!r} from non-structured input' 1 ctx1a

268 raise ValueError(msg) 1 ctx1a

269 elif x.dtype[dim].shape: 3 ctx1cba

270 return x[[dim]] 1 ctx1a

271 else: 

272 return x[dim] 3 ctx1cba

273 

274 return fun 3 ctx1cba

275 

276 

277@CrossKernel.register_xtransf 3 ctx1cba

278def maxdim(maxdim): 3 ctx1cba

279 """ 

280 Restrict the process to a maximum input dimensionality. 

281 

282 Parameters 

283 ---------- 

284 xmaxdim, ymaxdim : None, int 

285 Maximum dimensionality of the input. 

286 

287 Notes 

288 ----- 

289 Once applied a restriction, the check is hardcoded into the kernel core and 

290 it is not possible to remove it by applying again `maxdim` with a larger 

291 limit. 

292 """ # noqa: DOC103, DOC201 

293 if not isinstance(maxdim, numbers.Integral) or maxdim < 0: 3 ctx1cba

294 msg = f'maximum dimensionality {maxdim!r} not valid' 1 ctx1a

295 raise ValueError(msg) 1 ctx1a

296 

297 def fun(x): 3 ctx1cba

298 nd = _array._nd(x.dtype) 3 ctx1cba

299 with _jaxext.skipifabstract(): 3 ctx1cba

300 if nd > maxdim: 3 ctx1cba

301 msg = f'kernel applied to input with {nd} fields > maxdim={maxdim}' 1 ctx1a

302 raise ValueError(msg) 1 ctx1a

303 return x 3 ctx1cba

304 

305 return fun 3 ctx1cba

306 

307 

308@CrossKernel.register_xtransf 3 ctx1cba

309def loc(loc): 3 ctx1cba

310 r""" 

311 Translate the process inputs. 

312 

313 .. math:: 

314 T(f)(x) = f(x - \mathrm{loc}) 

315 

316 Parameters 

317 ---------- 

318 xloc, yloc: None, number 

319 Translations. 

320 

321 """ # noqa: DOC103, DOC201 

322 with _jaxext.skipifabstract(): 3 ctx1cba

323 assert -jnp.inf < loc < jnp.inf, loc 3 ctx1cba

324 return lambda x: _util.ufunc_recurse_dtype(lambda x: x - loc, x) 3 ctx1cba

325 

326 

327@CrossKernel.register_xtransf 3 ctx1cba

328def scale(scale): 3 ctx1cba

329 r""" 

330 Rescale the process inputs. 

331 

332 .. math:: 

333 T(f)(x) = f(x / \mathrm{scale}) 

334 

335 Parameters 

336 ---------- 

337 xscale, yscale: None, number 

338 Rescaling factors. 

339 

340 """ # noqa: DOC103, DOC201 

341 with _jaxext.skipifabstract(): 3 ctx1cba

342 assert 0 < scale < jnp.inf, scale 3 ctx1cba

343 return lambda x: _util.ufunc_recurse_dtype(lambda x: x / scale, x) 3 ctx1cba

344 

345 

346def normalize_argparser(do): 3 ctx1cba

347 return do or None 2 ctx1ca

348 

349 

350@functools.partial(CrossKernel.register_corelinop, argparser=normalize_argparser) 3 ctx1cba

351def normalize(core, dox, doy): 3 ctx1cba

352 r""" 

353 Rescale the process to unit variance. 

354 

355 .. math:: 

356 T(f)(x) &= f(x) / \sqrt{\mathrm{Std}[f(x)]} \\ 

357 &= f(x) / \sqrt{\mathrm{kernel}(x, x)} 

358 

359 Parameters 

360 ---------- 

361 dox, doy : bool 

362 Whether to rescale. 

363 """ # noqa: DOC101, DOC103, DOC201 

364 if dox and doy: 2 ctx1ca

365 return lambda x, y, **kw: ( 2 ctx1ca

366 core(x, y, **kw) / jnp.sqrt(core(x, x, **kw) * core(y, y, **kw)) 

367 ) 

368 elif dox: 1 ctx1a

369 return lambda x, y, **kw: core(x, y, **kw) / jnp.sqrt(core(x, x, **kw)) 1 ctx1a

370 else: 

371 return lambda x, y, **kw: core(x, y, **kw) / jnp.sqrt(core(y, y, **kw)) 1 ctx1a

372 

373 

374@CrossKernel.register_corelinop 3 ctx1cba

375def cond(core, cond1, cond2, other): 3 ctx1cba

376 r""" 

377 

378 Switch between two independent processes based on a condition. 

379 

380 .. math:: 

381 T(f, g)(x) = \begin{cases} 

382 f(x) & \text{if $\mathrm{cond}(x)$,} \\ 

383 g(x) & \text{otherwise.} 

384 \end{cases} 

385 

386 Parameters 

387 ---------- 

388 cond1, cond2 : callable 

389 Function that is applied on an array of points and must return 

390 a boolean array with the same shape. 

391 other : CrossKernel 

392 Kernel of the process used where the condition is false. 

393 

394 """ # noqa: DOC101, DOC103, DOC201 

395 

396 def newcore(x, y, **kw): 2 ctx1ba

397 xcond = cond1(x) 2 ctx1ba

398 ycond = cond2(y) 2 ctx1ba

399 r = jnp.where(xcond & ycond, core(x, y, **kw), other(x, y, **kw)) 2 ctx1ba

400 return jnp.where((xcond ^ ycond).astype('u8'), 0, r) 2 ctx1ba

401 # the .astype('u8') fixes a weird bug with some old dependencies 

402 

403 return newcore 2 ctx1ba

404 

405 # kernel0.linop('choose', kernel1, kernel2, ..., lambda x: x['index']) 

406 

407 

408AffineSpan.inherit_transf('maxdim') 3 ctx1cba

409AffineSpan.inherit_transf('derivable') 3 ctx1cba

410 

411 

412@functools.partial(AffineSpan.register_linop, transfname='loc') 3 ctx1cba

413def affine_loc(tcls, self, xloc, yloc): 3 ctx1cba

414 dynkw = dict(self.dynkw) 1 ctx1a

415 newself = tcls.super_transf('loc', self, xloc, yloc) 1 ctx1a

416 dynkw['lloc'] = dynkw['lloc'] + xloc * dynkw['lscale'] 1 ctx1a

417 dynkw['rloc'] = dynkw['rloc'] + yloc * dynkw['rscale'] 1 ctx1a

418 return newself._clone(self.__class__, dynkw=dynkw) 1 ctx1a

419 

420 

421@functools.partial(AffineSpan.register_linop, transfname='scale') 3 ctx1cba

422def affine_scale(tcls, self, xscale, yscale): 3 ctx1cba

423 dynkw = dict(self.dynkw) 2 ctx1ba

424 newself = tcls.super_transf('scale', self, xscale, yscale) 2 ctx1ba

425 dynkw['lscale'] = dynkw['lscale'] * xscale 2 ctx1ba

426 dynkw['rscale'] = dynkw['rscale'] * yscale 2 ctx1ba

427 return newself._clone(self.__class__, dynkw=dynkw) 2 ctx1ba