Coverage for src/lsqfitgp/_fit.py: 84%

1# lsqfitgp/_fit.py 

2# 

3# Copyright (c) 2020, 2022, 2023, 2024, 2025, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

7# lsqfitgp is free software: you can redistribute it and/or modify 

8# it under the terms of the GNU General Public License as published by 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# lsqfitgp is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>. 

19 

20import datetime 3 ctx1cba

21import functools 3 ctx1cba

22import re 3 ctx1cba

23import textwrap 3 ctx1cba

24import time 3 ctx1cba

25import warnings 3 ctx1cba

26 

27import gvar 3 ctx1cba

28import jax 3 ctx1cba

29import jax.core 3 ctx1cba

30import numpy 3 ctx1cba

31from jax import jit as jax_jit # `jit` is a parameter of empbayes_fit 3 ctx1cba

32from jax import numpy as jnp 3 ctx1cba

33from jax import tree 3 ctx1cba

34from scipy import optimize 3 ctx1cba

35 

36from lsqfitgp import _GP, _array, _gvarext, _linalg 3 ctx1cba

37 

38 

39@functools.singledispatch 3 ctx1cba

40def token_getter(x): 3 ctx1cba

41 return x 

42 

43 

44@functools.singledispatch 3 ctx1cba

45def token_setter(x, token): # noqa: ARG001, the dispatch argument 3 ctx1cba

46 return token 

47 

48 

49@token_getter.register(jax.core.Tracer) 3 ctx1cba

50@token_getter.register(jnp.ndarray) 3 ctx1cba

51@token_getter.register(numpy.ndarray) 3 ctx1cba

52def _(x): 3 ctx1cba

53 return x[x.ndim * (0,)] if x.size else x 3 ctx1cba

54 

55 

56@token_setter.register(jax.core.Tracer) 3 ctx1cba

57@token_setter.register(jnp.ndarray) 3 ctx1cba

58@token_setter.register(numpy.ndarray) 3 ctx1cba

59def _(x, token): 3 ctx1cba

60 x = jnp.asarray(x) 3 ctx1cba

61 return x.at[x.ndim * (0,)].set(token) if x.size else token 3 ctx1cba

62 

63 

64def token_map_leaf(func, x): 3 ctx1cba

65 if isinstance(x, (jnp.ndarray, numpy.ndarray)): 65 ↛ 79line 65 didn't jump to line 79 because the condition on line 65 was always true3 ctx1cba

66 token = token_getter(x) 3 ctx1cba

67 

68 @jax.custom_jvp 3 ctx1cba

69 def jaxfunc(token): 3 ctx1cba

70 return jax.pure_callback(func, token, token, vmap_method='expand_dims') 3 ctx1cba

71 

72 @jaxfunc.defjvp 3 ctx1cba

73 def _(p, t): 3 ctx1cba

74 return (jaxfunc(*p), *t) 3 ctx1cba

75 

76 token = jaxfunc(token) 3 ctx1cba

77 return token_setter(x, token) 3 ctx1cba

78 else: 

79 token = token_getter(x) 

80 token = func(token) 

81 return token_setter(x, token) 

82 

83 

84def token_map(func, x): 3 ctx1cba

85 return tree.map(lambda x: token_map_leaf(func, x), x) 3 ctx1cba

86 

87 

88class Logger: 3 ctx1cba

89 """ 

90 Class to manage a log. 

91 

92 Can be used as superclass. Each line of the log has a verbosity level (an 

93 integer >= 0) and is printed only if this level is below a threshold. All 

94 lines are saved and the log can be retrieved. 

95 """ 

96 

97 def __init__(self, target_verbosity=0): 3 ctx1cba

98 """Set the threshold used to exclude log lines.""" 

99 self._verbosity = target_verbosity 3 ctx1cba

100 self._loggedlines = [] 3 ctx1cba

101 

102 def _indent(self, text, level=0): 3 ctx1cba

103 """Indent a text by provided level or by global current level.""" 

104 level = max(0, level + self.loglevel._level) 2 ctx1ba

105 prefix = 4 * level * ' ' 2 ctx1ba

106 return textwrap.indent(text, prefix) 2 ctx1ba

107 

108 def _select(self, verbosity, target_verbosity=None): 3 ctx1cba

109 if target_verbosity is None: 109 ↛ 111line 109 didn't jump to line 111 because the condition on line 109 was always true3 ctx1cba

110 target_verbosity = self._verbosity 3 ctx1cba

111 if isinstance(verbosity, int): 3 ctx1cba

112 return target_verbosity >= verbosity 3 ctx1cba

113 else: 

114 return target_verbosity in verbosity 3 ctx1cba

115 

116 def log(self, message, verbosity=1, *, level=0): 3 ctx1cba

117 """ 

118 Print and record a message. 

119 

120 Parameters 

121 ---------- 

122 message : str 

123 The message to print. A newline is added unconditionally. 

124 verbosity : int or set, default 1 

125 The verbosity level(s) at which the message is printed. If an 

126 integer, it's printed at all levels >= that integer. If a set, at 

127 the specified levels. 

128 level : int, default 0 

129 The indentation level of the message. 

130 """ 

131 if self._select(verbosity): 3 ctx1cba

132 print(self._indent(message, level)) # noqa: T201, the fit log 2 ctx1ba

133 self._loggedlines.append((message, verbosity, level + self.loglevel._level)) 3 ctx1cba

134 

135 def getlog(self, target_verbosity=None, *, base_level=0): 3 ctx1cba

136 """Return all logged line as a single string.""" 

137 return '\n'.join( 

138 self._indent(message, base_level + level) 

139 for message, verbosity, level in self._loggedlines 

140 if self._select(verbosity, target_verbosity) 

141 ) 

142 

143 class _LogLevel: 3 ctx1cba

144 """Shared context manager to indent messages.""" 

145 

146 _level = 0 3 ctx1cba

147 

148 @classmethod 3 ctx1cba

149 def __enter__(cls): 3 ctx1cba

150 cls._level += 1 3 ctx1cba

151 

152 @classmethod 3 ctx1cba

153 def __exit__(cls, *_): 3 ctx1cba

154 cls._level -= 1 3 ctx1cba

155 

156 loglevel = _LogLevel() 3 ctx1cba

157 

158 

159class empbayes_fit(Logger): 3 ctx1cba

160 SEPARATE_JAC = False 3 ctx1cba

161 

162 def __init__( 3 ctx1cba

163 self, 

164 hyperprior, 

165 gpfactory, 

166 data, 

167 *, 

168 raises=True, 

169 minkw={}, # noqa: B006, read only 

170 gpfactorykw={}, # noqa: B006, read only 

171 jit=True, 

172 method='gradient', 

173 initial='priormean', 

174 verbosity=0, 

175 covariance='auto', 

176 fix=None, 

177 mlkw={}, # noqa: B006, read only 

178 forward=False, 

179 additional_loss=None, 

180 ): 

181 """ 

182 Maximum a posteriori fit. 

183 

184 Maximizes the marginal likelihood of the data with a Gaussian process 

185 model that depends on hyperparameters, multiplied by a prior on the 

186 hyperparameters. 

187 

188 Parameters 

189 ---------- 

190 hyperprior : scalar, array or dictionary of scalars/arrays 

191 A collection of gvars representing the prior for the 

192 hyperparameters. 

193 gpfactory : callable 

194 A function with signature gpfactory(hyperparams) -> GP object. The 

195 argument ``hyperparams`` has the same structure of the 

196 `empbayes_fit` argument ``hyperprior``. gpfactory must be 

197 JAX-friendly, i.e., use `jax.numpy` and `jax.scipy` instead of plain 

198 `numpy`/`scipy` and avoid assignments to arrays. 

199 data : dict, tuple or callable 

200 Dictionary of data that is passed to `GP.marginal_likelihood` on 

201 the GP object returned by ``gpfactory``. If a tuple, it contains the 

202 first two arguments to `GP.marginal_likelihood`. If a callable, it 

203 is called with the same arguments of ``gpfactory`` and must return 

204 the argument(s) for `GP.marginal_likelihood`. 

205 raises : bool, optional 

206 If True (default), raise an error when the minimization fails. 

207 Otherwise, use the last point of the minimization as result. 

208 minkw : dict, optional 

209 Keyword arguments passed to `scipy.optimize.minimize`, overwrites 

210 values specified by `empbayes_fit`. 

211 gpfactorykw : dict, optional 

212 Keyword arguments passed to ``gpfactory``, and also to ``data`` if 

213 it is a callable. If ``jit``, ``gpfactorykw`` crosses a `jax.jit` 

214 boundary, so it must contain objects understandable by `jax`. 

215 jit : bool 

216 If True (default), use `jax.jit` to compile the minimization target. 

217 method : str 

218 Minimization strategy. Options: 

219 

220 'nograd' 

221 Use a gradient-free method. 

222 'gradient' (default) 

223 Use a gradient-only method. 

224 'fisher' 

225 Use a Newton method with the Fisher information matrix plus 

226 the hyperprior precision matrix. 

227 initial : str, scalar, array, dictionary of scalars/arrays 

228 Starting point for the minimization, matching the format of 

229 ``hyperprior``, or one of the following options: 

230 

231 'priormean' (default) 

232 Start from the hyperprior mean. 

233 'priorsample' 

234 Take a random sample from the hyperprior. 

235 verbosity : int 

236 An integer indicating how much information is printed on the 

237 terminal: 

238 

239 0 (default) 

240 No logging. 

241 1 

242 Minimal report. 

243 2 

244 Detailed report. 

245 3 

246 Log each iteration. 

247 4 

248 More detailed iteration log. 

249 5 

250 Print the current parameter values at each iteration. 

251 covariance : str 

252 Method to estimate the posterior covariance matrix of the 

253 hyperparameters: 

254 

255 'fisher' 

256 Use the Fisher information in the MAP, plus the prior precision, 

257 as precision matrix. 

258 'minhess' 

259 Use the hessian estimate of the minimizer as precision matrix. 

260 'none' 

261 Do not estimate the covariance matrix. 

262 'auto' (default) 

263 ``'minhess'`` if applicable, ``'none'`` otherwise. 

264 fix : scalar, array or dictionary of scalars/arrays 

265 A set of booleans, with the same format as ``hyperprior``, 

266 indicating which hyperparameters are kept fixed to their initial 

267 value. Scalars and arrays are broadcasted to the shape of 

268 ``hyperprior``. If a dictionary, missing keys are treated as False. 

269 mlkw : dict 

270 Additional arguments passed to `GP.marginal_likelihood`. 

271 forward : bool, default False 

272 Use forward instead of backward derivatives. Typically, forward is 

273 faster with a small number of parameters. 

274 additional_loss : callable, optional 

275 A function with signature ``additional_loss(hyperparams) -> float`` 

276 which is added to the minus log marginal posterior of the 

277 hyperparameters. 

278 

279 Attributes 

280 ---------- 

281 p : scalar, array or dictionary of scalars/arrays 

282 A collection of gvars representing the hyperparameters that 

283 maximize their posterior. These gvars do not track correlations 

284 with the hyperprior or the data. 

285 prior : scalar, array or dictionary of scalars/arrays 

286 A copy of the hyperprior. 

287 initial : scalar, array or dictionary of scalars/arrays 

288 Starting point of the minimization, with the same format as ``p``. 

289 fix : scalar, array or dictionary of scalars/arrays 

290 A set of booleans, with the same format as ``p``, indicating which 

291 parameters were kept fixed to the values in ``initial``. 

292 pmean : scalar, array or dictionary of scalars/arrays 

293 Mean of ``p``. 

294 pcov : scalar, array or dictionary of scalars/arrays 

295 Covariance matrix of ``p``. 

296 minresult : scipy.optimize.OptimizeResult 

297 The result object returned by `scipy.optimize.minimize`. 

298 minargs : dict 

299 The arguments passed to `scipy.optimize.minimize`. 

300 gpfactory : callable 

301 The ``gpfactory`` argument. 

302 gpfactorykw : dict 

303 The ``gpfactorykw`` argument. 

304 data : dict, tuple or callable 

305 The ``data`` argument. 

306 

307 Raises 

308 ------ 

309 RuntimeError 

310 The minimization failed and ``raises`` is True. 

311 """ 

312 Logger.__init__(self, verbosity) 3 ctx1cba

313 del verbosity 3 ctx1cba

314 self.log('**** call lsqfitgp.empbayes_fit ****') 3 ctx1cba

315 

316 assert callable(gpfactory) 3 ctx1cba

317 

318 # analyze the hyperprior 

319 hpinitial, hpunflat = self._parse_hyperprior(hyperprior, initial, fix) 3 ctx1cba

320 del hyperprior, initial, fix 3 ctx1cba

321 

322 # analyze data 

323 data, cachedargs = self._parse_data(data) 3 ctx1cba

324 

325 # define functions 

326 timer, functions = self._prepare_functions( 3 ctx1cba

327 gpfactory=gpfactory, 

328 gpfactorykw=gpfactorykw, 

329 data=data, 

330 cachedargs=cachedargs, 

331 hpunflat=hpunflat, 

332 mlkw=mlkw, 

333 jit=jit, 

334 forward=forward, 

335 additional_loss=additional_loss, 

336 ) 

337 del gpfactory, gpfactorykw, data, cachedargs, mlkw, forward, additional_loss 3 ctx1cba

338 

339 # prepare minimizer arguments 

340 minargs = self._prepare_minargs(method, functions, hpinitial) 3 ctx1cba

341 

342 # set up callback to time and log iterations 

343 callback = self._Callback(self, functions, timer, hpunflat) 3 ctx1cba

344 minargs.update(callback=callback) 3 ctx1cba

345 

346 # check invalid argument before running minimizer 

347 if not covariance in ('auto', 'fisher', 'minhess', 'none'): 347 ↛ 348line 347 didn't jump to line 348 because the condition on line 347 was never true3 ctx1cba

348 raise KeyError(covariance) 

349 

350 # add user arguments and minimize 

351 minargs.update(minkw) 3 ctx1cba

352 self.log(f'minimizer method {minargs["method"]!r}', 2) 3 ctx1cba

353 total = time.perf_counter() 3 ctx1cba

354 result = optimize.minimize(**minargs) 3 ctx1cba

355 

356 # check the minimization was successful 

357 self._check_success(result, raises) 3 ctx1cba

358 

359 # compute posterior covariance of the hyperparameters 

360 cov = self._posterior_covariance( 3 ctx1cba

361 method, covariance, result, functions['fisher'] 

362 ) 

363 

364 # log total timings and function calls 

365 total = time.perf_counter() - total 3 ctx1cba

366 self._log_totals(total, timer, callback, jit, functions) 3 ctx1cba

367 

368 ##### temporary fix for gplepage/gvar#50 ##### 

369 cov = numpy.array(cov, order='C') 3 ctx1cba

370 ############################################## 

371 

372 # join posterior mean and covariance matrix 

373 uresult = gvar.gvar(result.x, cov) 3 ctx1cba

374 

375 # set attributes 

376 self.p = gvar.gvar(hpunflat(uresult)) 3 ctx1cba

377 self.pmean = gvar.mean(self.p) 3 ctx1cba

378 self.pcov = gvar.evalcov(self.p) 3 ctx1cba

379 self.minresult = result 3 ctx1cba

380 self.minargs = minargs 3 ctx1cba

381 

382 # tabulate hyperparameter prior and posterior 

383 if self._verbosity >= 2: 3 ctx1cba

384 self.log( 2 ctx1ba

385 _gvarext.tabulate_together( 

386 self.prior, self.p, headers=['param', 'prior', 'posterior'] 

387 ) 

388 ) 

389 

390 self.log('**** exit lsqfitgp.empbayes_fit ****') 3 ctx1cba

391 

392 class _CountCalls: 3 ctx1cba

393 """Wrap a callable to count calls.""" 

394 

395 def __init__(self, func): 3 ctx1cba

396 self._func = func 3 ctx1cba

397 self._total = 0 3 ctx1cba

398 self._partial = 0 3 ctx1cba

399 functools.update_wrapper(self, func) 3 ctx1cba

400 

401 def __call__(self, *args, **kw): 3 ctx1cba

402 self._total += 1 3 ctx1cba

403 self._partial += 1 3 ctx1cba

404 return self._func(*args, **kw) 3 ctx1cba

405 

406 def partial(self): 3 ctx1cba

407 """Return the partial counter and reset it.""" 

408 result = self._partial 3 ctx1cba

409 self._partial = 0 3 ctx1cba

410 return result 3 ctx1cba

411 

412 def total(self): 3 ctx1cba

413 """Return the total number of calls.""" 

414 return self._total 3 ctx1cba

415 

416 @staticmethod 3 ctx1cba

417 def fmtcalls(method, functions): 3 ctx1cba

418 """ 

419 Format summary of number of calls. 

420 

421 method : str 

422 functions: dict[str, _CountCalls]. 

423 """ 

424 

425 def counts(): 3 ctx1cba

426 for name, func in functions.items(): 3 ctx1cba

427 if count := getattr(func, method)(): 3 ctx1cba

428 yield f'{name} {count}' 3 ctx1cba

429 

430 return ', '.join(counts()) 3 ctx1cba

431 

432 class _Timer: 3 ctx1cba

433 """Object to time likelihood computations.""" 

434 

435 def __init__(self): 3 ctx1cba

436 self.totals = {} 3 ctx1cba

437 self.partials = {} 3 ctx1cba

438 self._last_start = False 3 ctx1cba

439 

440 def start(self, token): 3 ctx1cba

441 return token_map(self._start, token) 3 ctx1cba

442 

443 def _start(self, token): 3 ctx1cba

444 self.stamp = time.perf_counter() 3 ctx1cba

445 self.counter = 0 3 ctx1cba

446 assert not self._last_start # forbid consecutive start() calls 3 ctx1cba

447 self._last_start = True 3 ctx1cba

448 return token 3 ctx1cba

449 

450 def reset(self): 3 ctx1cba

451 self.partials = {} 3 ctx1cba

452 

453 def partial(self, token): 3 ctx1cba

454 return token_map(self._partial, token) 3 ctx1cba

455 

456 def _partial(self, token): 3 ctx1cba

457 now = time.perf_counter() 3 ctx1cba

458 delta = now - self.stamp 3 ctx1cba

459 self.partials[self.counter] = self.partials.get(self.counter, 0) + delta 3 ctx1cba

460 self.totals[self.counter] = self.totals.get(self.counter, 0) + delta 3 ctx1cba

461 self.stamp = now 3 ctx1cba

462 self.counter += 1 3 ctx1cba

463 self._last_start = False 3 ctx1cba

464 return token 3 ctx1cba

465 

466 def _parse_hyperprior(self, hyperprior, initial, fix): 3 ctx1cba

467 

468 # check fix against hyperprior and fill missing values 

469 hyperprior = self._copyasarrayorbufferdict(hyperprior) 3 ctx1cba

470 self._check_no_redundant_keys(hyperprior) 3 ctx1cba

471 fix = self._parse_fix(hyperprior, fix) 3 ctx1cba

472 flatfix = self._flatview(fix) 3 ctx1cba

473 

474 # extract distribution of free hyperparameters 

475 flathp = self._flatview(hyperprior) 3 ctx1cba

476 freehp = flathp[~flatfix] 3 ctx1cba

477 mean = gvar.mean(freehp) 3 ctx1cba

478 cov = gvar.evalcov(freehp) 3 ctx1cba

479 dec = _linalg.Chol(cov) 3 ctx1cba

480 assert dec.n == freehp.size 3 ctx1cba

481 self.log(f'{freehp.size}/{flathp.size} free hyperparameters', 2) 3 ctx1cba

482 

483 # determine starting point for minimization 

484 initial = self._parse_initial(hyperprior, initial, dec) 3 ctx1cba

485 flatinitial = self._flatview(initial) 3 ctx1cba

486 x0 = dec.pinv_correlate(flatinitial[~flatfix] - mean) 3 ctx1cba

487 

488 # make function to correlate, add fixed values, and reshape to original 

489 # format 

490 (fixed_indices,) = jnp.nonzero(flatfix) 3 ctx1cba

491 (unfixed_indices,) = jnp.nonzero(~flatfix) 3 ctx1cba

492 fixed_values = jnp.asarray(flatinitial[flatfix]) 3 ctx1cba

493 

494 def unflat(x): 3 ctx1cba

495 assert x.ndim == 1 3 ctx1cba

496 if x.dtype == object: 3 ctx1cba

497 jac, indices = _gvarext.jacobian(x) 3 ctx1cba

498 xmean = mean + dec.correlate(gvar.mean(x)) 3 ctx1cba

499 xjac = dec.correlate(jac) 3 ctx1cba

500 x = _gvarext.from_jacobian(xmean, xjac, indices) 3 ctx1cba

501 y = numpy.empty(flatfix.size, x.dtype) 3 ctx1cba

502 numpy.put(y, unfixed_indices, x) 3 ctx1cba

503 numpy.put(y, fixed_indices, fixed_values) 3 ctx1cba

504 else: 

505 x = mean + dec.correlate(x) 3 ctx1cba

506 y = jnp.empty(flatfix.size, x.dtype) 3 ctx1cba

507 y = y.at[unfixed_indices].set(x) 3 ctx1cba

508 y = y.at[fixed_indices].set(fixed_values) 3 ctx1cba

509 return self._unflatview(y, hyperprior) 3 ctx1cba

510 

511 self.prior = hyperprior 3 ctx1cba

512 return x0, unflat 3 ctx1cba

513 

514 @staticmethod 3 ctx1cba

515 def _check_no_redundant_keys(hyperprior): 3 ctx1cba

516 if not hasattr(hyperprior, 'keys'): 3 ctx1cba

517 return 1 ctx1a

518 for k in hyperprior: 3 ctx1cba

519 m = hyperprior.extension_pattern.match(k) 3 ctx1cba

520 if m and m.group(1) in hyperprior.invfcn: 3 ctx1cba

521 altk = m.group(2) 3 ctx1cba

522 if altk in hyperprior: 522 ↛ 523line 522 didn't jump to line 523 because the condition on line 522 was never true3 ctx1cba

523 msg = f'duplicate keys {altk!r} and {k!r} in hyperprior' 

524 raise ValueError(msg) 

525 

526 def _parse_fix(self, hyperprior, fix): 3 ctx1cba

527 

528 if fix is None: 528 ↛ 536line 528 didn't jump to line 536 because the condition on line 528 was always true3 ctx1cba

529 if hasattr(hyperprior, 'keys'): 3 ctx1cba

530 fix = gvar.BufferDict( 3 ctx1cba

531 hyperprior, buf=numpy.zeros(hyperprior.size, bool) 

532 ) 

533 else: 

534 fix = numpy.zeros(hyperprior.shape, bool) 1 ctx1a

535 else: 

536 fix = self._copyasarrayorbufferdict(fix) 

537 if hasattr(fix, 'keys'): 

538 assert hasattr(hyperprior, 'keys'), ( 

539 'fix is dictionary but hyperprior is array' 

540 ) 

541 assert all(hyperprior.has_dictkey(k) for k in fix), ( 

542 'some keys in fix are missing in hyperprior' 

543 ) 

544 newfix = {} 

545 for k, v in hyperprior.items(): 

546 key = None 

547 m = hyperprior.extension_pattern.match(k) 

548 if m and m.group(1) in hyperprior.invfcn: 

549 altk = m.group(2) 

550 if altk in fix: 

551 assert k not in fix, ( 

552 f'duplicate keys {k!r} and {altk!r} in fix' 

553 ) 

554 key = altk 

555 if key is None and k in fix: 

556 key = k 

557 if key is None: 

558 elem = numpy.zeros(v.shape, bool) 

559 else: 

560 elem = numpy.broadcast_to(fix[key], v.shape) 

561 newfix[k] = elem 

562 fix = gvar.BufferDict(newfix, dtype=bool) 

563 else: 

564 assert not hasattr(hyperprior, 'keys'), ( 

565 'fix is array but hyperprior is dictionary' 

566 ) 

567 fix = numpy.broadcast_to(fix, hyperprior.shape).astype(bool) 

568 

569 self.fix = fix 3 ctx1cba

570 return fix 3 ctx1cba

571 

572 def _parse_initial(self, hyperprior, initial, dec): 3 ctx1cba

573 

574 if not isinstance(initial, str): 574 ↛ 575line 574 didn't jump to line 575 because the condition on line 574 was never true3 ctx1cba

575 self.log('start from provided point', 2) 

576 initial = self._copyasarrayorbufferdict(initial) 

577 if hasattr(hyperprior, 'keys'): 

578 assert hasattr(initial, 'keys'), ( 

579 'hyperprior is dictionary but initial is array' 

580 ) 

581 assert set(hyperprior.keys()) == set(initial.keys()) 

582 assert all(hyperprior[k].shape == initial[k].shape for k in hyperprior) 

583 else: 

584 assert not hasattr(initial, 'keys'), ( 

585 'hyperprior is array but initial is dictionary' 

586 ) 

587 assert hyperprior.shape == initial.shape 

588 

589 elif initial == 'priormean': 589 ↛ 593line 589 didn't jump to line 593 because the condition on line 589 was always true3 ctx1cba

590 self.log('start from prior mean', 2) 3 ctx1cba

591 initial = gvar.mean(hyperprior) 3 ctx1cba

592 

593 elif initial == 'priorsample': 

594 self.log('start from a random sample from the prior', 2) 

595 if dec.n < hyperprior.size: 

596 flathp = self._flatview(hyperprior) 

597 cov = gvar.evalcov(flathp) 

598 fulldec = _linalg.Chol(cov) 

599 else: 

600 fulldec = dec 

601 iid = numpy.random.randn(fulldec.m) # noqa: NPY002, uses the global numpy state on purpose 

602 flatinitial = numpy.asarray(fulldec.correlate(iid)) 

603 initial = self._unflatview(flatinitial, hyperprior) 

604 

605 else: 

606 raise KeyError(initial) 

607 

608 self.initial = initial 3 ctx1cba

609 return initial 3 ctx1cba

610 

611 def _parse_data(self, data): 3 ctx1cba

612 

613 self.data = data 3 ctx1cba

614 if isinstance(data, tuple) and len(data) == 1: 3 ctx1cba

615 (data,) = data 1 ctx1a

616 

617 if callable(data): 3 ctx1cba

618 self.log('data is callable', 2) 2 ctx1ba

619 cachedargs = None 2 ctx1ba

620 elif isinstance(data, tuple): 3 ctx1cba

621 self.log('data errors provided separately', 2) 1 ctx1a

622 assert len(data) == 2 1 ctx1a

623 cachedargs = data 1 ctx1a

624 elif (gdata := self._copyasarrayorbufferdict(data)).dtype == object: 3 ctx1cba

625 self.log('data has errors as gvars', 2) 2 ctx1ba

626 data = gvar.gvar(gdata) 2 ctx1ba

627 # convert to gvar because non-gvars in the array would upset 

628 # gvar.mean and gvar.evalcov 

629 cachedargs = (gvar.mean(data), gvar.evalcov(data)) 2 ctx1ba

630 else: 

631 self.log('data has no errors', 2) 3 ctx1cba

632 cachedargs = (data,) 3 ctx1cba

633 

634 return data, cachedargs 3 ctx1cba

635 

636 def _prepare_functions( # noqa: C901, PLR0915 3 ctx1cba

637 self, 

638 *, 

639 gpfactory, 

640 gpfactorykw, 

641 data, 

642 cachedargs, 

643 hpunflat, 

644 mlkw, 

645 jit, 

646 forward, 

647 additional_loss, 

648 ): 

649 

650 timer = self._Timer() 3 ctx1cba

651 firstcall = [None] 3 ctx1cba

652 

653 def make_decomp(p, **kw): 3 ctx1cba

654 """Decomposition of the prior covariance and data.""" 

655 # start timer and convert hypers to user format 

656 p = timer.start(p) 3 ctx1cba

657 hp = hpunflat(p) 3 ctx1cba

658 

659 # create GP object 

660 gp = gpfactory(hp, **kw) 3 ctx1cba

661 assert isinstance(gp, _GP.GP) 3 ctx1cba

662 

663 # extract data 

664 if cachedargs: 3 ctx1cba

665 args = cachedargs 3 ctx1cba

666 else: 

667 args = data(hp, **kw) 2 ctx1ba

668 if not isinstance(args, tuple): 2 ctx1ba

669 args = (args,) 2 ctx1ba

670 

671 # decompose covariance matrix and flatten data 

672 decomp, r = gp._prior_decomp(*args, covtransf=timer.partial, **mlkw) 3 ctx1cba

673 r = r.astype(float) # int data upsets jax 3 ctx1cba

674 

675 # log number of datapoints 

676 if firstcall: 3 ctx1cba

677 # it is convenient to do here because the data is flattened. 

678 # works under jit since the first call is tracing 

679 firstcall.pop() 3 ctx1cba

680 xdtype = gp._get_x_dtype() 3 ctx1cba

681 nd = '?' if xdtype is None else _array._nd(xdtype) 3 ctx1cba

682 self.log(f'{r.size} datapoints, {nd} covariates') 3 ctx1cba

683 

684 # compute user loss 

685 if additional_loss is None: 3 ctx1cba

686 loss = 0.0 3 ctx1cba

687 else: 

688 loss = additional_loss(hp) 2 ctx1ba

689 

690 # split timer and return decomposition 

691 return timer.partial(decomp), r, loss 3 ctx1cba

692 

693 # define wrapper to collect call stats, pass user args, compile 

694 def wrap(func): 3 ctx1cba

695 if jit: 695 ↛ 697line 695 didn't jump to line 697 because the condition on line 695 was always true3 ctx1cba

696 func = jax_jit(func) 3 ctx1cba

697 func = functools.partial(func, **gpfactorykw) 3 ctx1cba

698 return self._CountCalls(func) 3 ctx1cba

699 

700 if jit: 700 ↛ 704line 700 didn't jump to line 704 because the condition on line 700 was always true3 ctx1cba

701 self.log('compile functions with jax jit', 2) 3 ctx1cba

702 

703 # log derivation method 

704 modename = 'forward' if forward else 'reverse' 3 ctx1cba

705 self.log(f'{modename}-mode autodiff (if used)', 2) 3 ctx1cba

706 

707 def prior(p): 3 ctx1cba

708 # the marginal prior of the hyperparameters is a Normal with 

709 # identity covariance matrix because p is transformed to make it so 

710 return 1 / 2 * (len(p) * jnp.log(2 * jnp.pi) + p @ p) 3 ctx1cba

711 

712 def grad_prior(p): 3 ctx1cba

713 return p 3 ctx1cba

714 

715 def fisher_prior(p): 3 ctx1cba

716 return jnp.eye(len(p)) 1 ctx1a

717 

718 @wrap 3 ctx1cba

719 def fun(p, **kw): 3 ctx1cba

720 """Minus log marginal posterior of the hyperparameters (not normalized).""" 

721 decomp, r, loss = make_decomp(p, **kw) 1 ctx1a

722 cond, _, _, _, _ = decomp.minus_log_normal_density(r, value=True) 1 ctx1a

723 post = cond + prior(p) + loss 1 ctx1a

724 return timer.partial(post) 1 ctx1a

725 

726 def make_gradfwd_fisher_args(p, **kw): 3 ctx1cba

727 def make_decomp_tee(p): 2 ctx1ba

728 decomp, r, loss = make_decomp(p, **kw) 2 ctx1ba

729 return (decomp.matrix(), r, loss), (decomp, r, loss) 2 ctx1ba

730 

731 (dK, dr, grad_loss), (decomp, r, loss) = jax.jacfwd( 2 ctx1ba

732 make_decomp_tee, has_aux=True 

733 )(p) 

734 lkw = dict(dK=dK, dr=dr) 2 ctx1ba

735 return decomp, r, lkw, loss, grad_loss 2 ctx1ba

736 

737 def make_gradrev_args(p, **kw): 3 ctx1cba

738 def make_decomp_loss(p): 3 ctx1cba

739 def make_decomp_r(p): 3 ctx1cba

740 def make_decomp_K(p): 3 ctx1cba

741 decomp, r, loss = make_decomp(p, **kw) 3 ctx1cba

742 return decomp.matrix(), (decomp, r, loss) 3 ctx1cba

743 

744 _, dK_vjp, (decomp, r, loss) = jax.vjp( 3 ctx1cba

745 make_decomp_K, p, has_aux=True 

746 ) 

747 return r, (decomp, r, dK_vjp, loss) 3 ctx1cba

748 

749 _, dr_vjp, (decomp, r, dK_vjp, loss) = jax.vjp( 3 ctx1cba

750 make_decomp_r, p, has_aux=True 

751 ) 

752 return loss, (decomp, r, dK_vjp, dr_vjp, loss) 3 ctx1cba

753 

754 grad_loss, (decomp, r, dK_vjp, dr_vjp, loss) = jax.grad( 3 ctx1cba