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

1# lsqfitgp/_fit.py 

2# 

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

21import warnings 6 ctx1feabcd

22import functools 6 ctx1feabcd

23import time 6 ctx1feabcd

24import textwrap 6 ctx1feabcd

25import datetime 6 ctx1feabcd

26 

27import gvar 6 ctx1feabcd

28import jax 6 ctx1feabcd

29from jax import numpy as jnp 6 ctx1feabcd

30import numpy 6 ctx1feabcd

31from scipy import optimize 6 ctx1feabcd

32from jax import tree_util 6 ctx1feabcd

33 

34from . import _GP 6 ctx1feabcd

35from . import _linalg 6 ctx1feabcd

36from . import _jaxext 6 ctx1feabcd

37from . import _gvarext 6 ctx1feabcd

38from . import _array 6 ctx1feabcd

39 

40# TODO the following token_ thing functionality may be provided by jax in the 

41# future, follow the developments 

42 

43@functools.singledispatch 6 ctx1feabcd

44def token_getter(x): 6 ctx1feabcd

45 return x 

46 

47@functools.singledispatch 6 ctx1feabcd

48def token_setter(x, token): 6 ctx1feabcd

49 return token 

50 

51@token_getter.register(jnp.ndarray) 6 ctx1feabcd

52@token_getter.register(numpy.ndarray) 6 ctx1feabcd

53def _(x): 6 ctx1feabcd

54 return x[x.ndim * (0,)] if x.size else x 6 ctx1feabcd

55 

56@token_setter.register(jnp.ndarray) 6 ctx1feabcd

57@token_setter.register(numpy.ndarray) 6 ctx1feabcd

58def _(x, token): 6 ctx1feabcd

59 x = jnp.asarray(x) 6 ctx1feabcd

60 return x.at[x.ndim * (0,)].set(token) if x.size else token 6 ctx1feabcd

61 

62def token_map_leaf(func, x): 6 ctx1feabcd

63 if isinstance(x, (jnp.ndarray, numpy.ndarray)): 63 ↛ 74line 63 didn't jump to line 74 because the condition on line 63 was always true6 ctx1feabcd

64 token = token_getter(x) 6 ctx1feabcd

65 @jax.custom_jvp 6 ctx1feabcd

66 def jaxfunc(token): 6 ctx1feabcd

67 return jax.pure_callback(func, token, token, vectorized=True) 6 ctx1feabcd

68 @jaxfunc.defjvp 6 ctx1feabcd

69 def _(p, t): 6 ctx1feabcd

70 return (jaxfunc(*p), *t) 6 ctx1feabcd

71 token = jaxfunc(token) 6 ctx1feabcd

72 return token_setter(x, token) 6 ctx1feabcd

73 else: 

74 token = token_getter(x) 

75 token = func(token) 

76 return token_setter(x, token) 

77 

78def token_map(func, x): 6 ctx1feabcd

79 return tree_util.tree_map(lambda x: token_map_leaf(func, x), x) 6 ctx1feabcd

80 

81class Logger: 6 ctx1feabcd

82 """ Class to manage a log. Can be used as superclass. Each line of the log 

83 has a verbosity level (an integer >= 0) and is printed only if this level is 

84 below a threshold. All lines are saved and the log can be retrieved. """ 

85 

86 def __init__(self, target_verbosity=0): 6 ctx1feabcd

87 """ set the threshold used to exclude log lines """ 

88 self._verbosity = target_verbosity 6 ctx1feabcd

89 self._loggedlines = [] 6 ctx1feabcd

90 

91 def _indent(self, text, level=0): 6 ctx1feabcd

92 """ indent a text by provided level or by global current level """ 

93 level = max(0, level + self.loglevel._level) 5 ctx1eabcd

94 prefix = 4 * level * ' ' 5 ctx1eabcd

95 return textwrap.indent(text, prefix) 5 ctx1eabcd

96 

97 def _select(self, verbosity, target_verbosity=None): 6 ctx1feabcd

98 if target_verbosity is None: 98 ↛ 100line 98 didn't jump to line 100 because the condition on line 98 was always true6 ctx1feabcd

99 target_verbosity = self._verbosity 6 ctx1feabcd

100 if isinstance(verbosity, int): 6 ctx1feabcd

101 return target_verbosity >= verbosity 6 ctx1feabcd

102 else: 

103 return target_verbosity in verbosity 6 ctx1feabcd

104 

105 def log(self, message, verbosity=1, *, level=0): 6 ctx1feabcd

106 """ 

107 Print and record a message. 

108 

109 Parameters 

110 ---------- 

111 message : str 

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

113 verbosity : int or set, default 1 

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

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

116 the specified levels. 

117 level : int, default 0 

118 The indentation level of the message. 

119 """ 

120 if self._select(verbosity): 6 ctx1feabcd

121 print(self._indent(message, level)) 5 ctx1eabcd

122 self._loggedlines.append((message, verbosity, level + self.loglevel._level)) 6 ctx1feabcd

123 

124 def getlog(self, target_verbosity=None, *, base_level=0): 6 ctx1feabcd

125 """ return all logged line as a single string """ 

126 return '\n'.join( 

127 self._indent(message, base_level + level) 

128 for message, verbosity, level in self._loggedlines 

129 if self._select(verbosity, target_verbosity) 

130 ) 

131 

132 class _LogLevel: 6 ctx1feabcd

133 """ shared context manager to indent messages """ 

134 

135 _level = 0 6 ctx1feabcd

136 

137 @classmethod 6 ctx1feabcd

138 def __enter__(cls): 6 ctx1feabcd

139 cls._level += 1 6 ctx1feabcd

140 

141 @classmethod 6 ctx1feabcd

142 def __exit__(cls, *_): 6 ctx1feabcd

143 cls._level -= 1 6 ctx1feabcd

144 

145 loglevel = _LogLevel() 6 ctx1feabcd

146 

147class empbayes_fit(Logger): 6 ctx1feabcd

148 

149 SEPARATE_JAC = False 6 ctx1feabcd

150 

151 def __init__( 6 ctx1feabcd

152 self, 

153 hyperprior, 

154 gpfactory, 

155 data, 

156 *, 

157 raises=True, 

158 minkw={}, 

159 gpfactorykw={}, 

160 jit=True, 

161 method='gradient', 

162 initial='priormean', 

163 verbosity=0, 

164 covariance='auto', 

165 fix=None, 

166 mlkw={}, 

167 forward=False, 

168 additional_loss=None, 

169 ): 

170 """ 

171  

172 Maximum a posteriori fit. 

173  

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

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

176 hyperparameters. 

177  

178 Parameters 

179 ---------- 

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

181 A collection of gvars representing the prior for the 

182 hyperparameters. 

183 gpfactory : callable 

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

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

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

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

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

189 data : dict, tuple or callable 

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

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

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

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

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

195 raises : bool, optional 

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

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

198 minkw : dict, optional 

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

200 values specified by `empbayes_fit`. 

201 gpfactorykw : dict, optional 

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

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

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

205 jit : bool 

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

207 method : str 

208 Minimization strategy. Options: 

209  

210 'nograd' 

211 Use a gradient-free method. 

212 'gradient' (default) 

213 Use a gradient-only method. 

214 'fisher' 

215 Use a Newton method with the Fisher information matrix plus 

216 the hyperprior precision matrix. 

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

218 Starting point for the minimization, matching the format of 

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

220  

221 'priormean' (default) 

222 Start from the hyperprior mean. 

223 'priorsample' 

224 Take a random sample from the hyperprior. 

225 verbosity : int 

226 An integer indicating how much information is printed on the 

227 terminal: 

228  

229 0 (default) 

230 No logging. 

231 1 

232 Minimal report. 

233 2 

234 Detailed report. 

235 3 

236 Log each iteration. 

237 4 

238 More detailed iteration log. 

239 5 

240 Print the current parameter values at each iteration. 

241 covariance : str 

242 Method to estimate the posterior covariance matrix of the 

243 hyperparameters: 

244  

245 'fisher' 

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

247 as precision matrix. 

248 'minhess' 

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

250 'none' 

251 Do not estimate the covariance matrix. 

252 'auto' (default) 

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

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

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

256 indicating which hyperparameters are kept fixed to their initial 

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

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

259 mlkw : dict 

260 Additional arguments passed to `GP.marginal_likelihood`. 

261 forward : bool, default False 

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

263 faster with a small number of parameters. 

264 additional_loss : callable, optional 

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

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

267 hyperparameters. 

268  

269 Attributes 

270 ---------- 

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

272 A collection of gvars representing the hyperparameters that 

273 maximize their posterior. These gvars do not track correlations 

274 with the hyperprior or the data. 

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

276 A copy of the hyperprior. 

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

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

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

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

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

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

283 Mean of ``p``. 

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

285 Covariance matrix of ``p``. 

286 minresult : scipy.optimize.OptimizeResult 

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

288 minargs : dict 

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

290 gpfactory : callable 

291 The ``gpfactory`` argument. 

292 gpfactorykw : dict 

293 The ``gpfactorykw`` argument. 

294 data : dict, tuple or callable 

295 The ``data`` argument. 

296 

297 Raises 

298 ------ 

299 RuntimeError 

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

301  

302 """ 

303 

304 Logger.__init__(self, verbosity) 6 ctx1feabcd

305 del verbosity 6 ctx1feabcd

306 self.log('**** call lsqfitgp.empbayes_fit ****') 6 ctx1feabcd

307 

308 assert callable(gpfactory) 6 ctx1feabcd

309 

310 # analyze the hyperprior 

311 hpinitial, hpunflat = self._parse_hyperprior(hyperprior, initial, fix) 6 ctx1feabcd

312 del hyperprior, initial, fix 6 ctx1feabcd

313 

314 # analyze data 

315 data, cachedargs = self._parse_data(data) 6 ctx1feabcd

316 

317 # define functions 

318 timer, functions = self._prepare_functions( 6 ctx1feabcd

319 gpfactory=gpfactory, gpfactorykw=gpfactorykw, data=data, 

320 cachedargs=cachedargs, hpunflat=hpunflat, mlkw=mlkw, jit=jit, 

321 forward=forward, additional_loss=additional_loss, 

322 ) 

323 del gpfactory, gpfactorykw, data, cachedargs, mlkw, forward, additional_loss 6 ctx1feabcd

324 

325 # prepare minimizer arguments 

326 minargs = self._prepare_minargs(method, functions, hpinitial) 6 ctx1feabcd

327 

328 # set up callback to time and log iterations 

329 callback = self._Callback(self, functions, timer, hpunflat) 6 ctx1feabcd

330 minargs.update(callback=callback) 6 ctx1feabcd

331 

332 # check invalid argument before running minimizer 

333 if not covariance in ('auto', 'fisher', 'minhess', 'none'): 333 ↛ 334line 333 didn't jump to line 334 because the condition on line 333 was never true6 ctx1feabcd

334 raise KeyError(covariance) 

335 

336 # add user arguments and minimize 

337 minargs.update(minkw) 6 ctx1feabcd

338 self.log(f'minimizer method {minargs["method"]!r}', 2) 6 ctx1feabcd

339 total = time.perf_counter() 6 ctx1feabcd

340 result = optimize.minimize(**minargs) 6 ctx1feabcd

341 

342 # check the minimization was successful 

343 self._check_success(result, raises) 6 ctx1feabcd

344 

345 # compute posterior covariance of the hyperparameters 

346 cov = self._posterior_covariance(method, covariance, result, functions['fisher']) 6 ctx1feabcd

347 

348 # log total timings and function calls 

349 total = time.perf_counter() - total 6 ctx1feabcd

350 self._log_totals(total, timer, callback, jit, functions) 6 ctx1feabcd

351 

352 # join posterior mean and covariance matrix 

353 uresult = gvar.gvar(result.x, cov) 6 ctx1feabcd

354 

355 # set attributes 

356 self.p = gvar.gvar(hpunflat(uresult)) 6 ctx1feabcd

357 self.pmean = gvar.mean(self.p) 6 ctx1feabcd

358 self.pcov = gvar.evalcov(self.p) 6 ctx1feabcd

359 self.minresult = result 6 ctx1feabcd

360 self.minargs = minargs 6 ctx1feabcd

361 

362 # tabulate hyperparameter prior and posterior 

363 if self._verbosity >= 2: 6 ctx1feabcd

364 self.log(_gvarext.tabulate_together( 5 ctx1eabcd

365 self.prior, self.p, 

366 headers=['param', 'prior', 'posterior'], 

367 )) # TODO replace tabulate_toegether with something more flexible I 

368 # can use for the callback as well. Maybe import TextMatrix from 

369 # miscpy. 

370 # TODO print the transformed parameters 

371 

372 self.log('**** exit lsqfitgp.empbayes_fit ****') 6 ctx1feabcd

373 

374 class _CountCalls: 6 ctx1feabcd

375 """ wrap a callable to count calls """ 

376 

377 def __init__(self, func): 6 ctx1feabcd

378 self._func = func 6 ctx1feabcd

379 self._total = 0 6 ctx1feabcd

380 self._partial = 0 6 ctx1feabcd

381 functools.update_wrapper(self, func) 6 ctx1feabcd

382 

383 def __call__(self, *args, **kw): 6 ctx1feabcd

384 self._total += 1 6 ctx1feabcd

385 self._partial += 1 6 ctx1feabcd

386 return self._func(*args, **kw) 6 ctx1feabcd

387 

388 def partial(self): 6 ctx1feabcd

389 """ return the partial counter and reset it """ 

390 result = self._partial 6 ctx1feabcd

391 self._partial = 0 6 ctx1feabcd

392 return result 6 ctx1feabcd

393 

394 def total(self): 6 ctx1feabcd

395 """ return the total number of calls """ 

396 return self._total 6 ctx1feabcd

397 

398 @staticmethod 6 ctx1feabcd

399 def fmtcalls(method, functions): 6 ctx1feabcd

400 """ 

401 format summary of number of calls 

402 method : str 

403 functions: dict[str, _CountCalls] 

404 """ 

405 def counts(): 6 ctx1feabcd

406 for name, func in functions.items(): 6 ctx1feabcd

407 if count := getattr(func, method)(): 6 ctx1feabcd

408 yield f'{name} {count}' 6 ctx1feabcd

409 return ', '.join(counts()) 6 ctx1feabcd

410 

411 class _Timer: 6 ctx1feabcd

412 """ object to time likelihood computations """ 

413 

414 def __init__(self): 6 ctx1feabcd

415 self.totals = {} 6 ctx1feabcd

416 self.partials = {} 6 ctx1feabcd

417 self._last_start = False 6 ctx1feabcd

418 

419 def start(self, token): 6 ctx1feabcd

420 return token_map(self._start, token) 6 ctx1feabcd

421 

422 def _start(self, token): 6 ctx1feabcd

423 self.stamp = time.perf_counter() 4 ctx1abcd

424 self.counter = 0 4 ctx1abcd

425 assert not self._last_start # forbid consecutive start() calls 4 ctx1abcd

426 self._last_start = True 4 ctx1abcd

427 return token 4 ctx1abcd

428 

429 def reset(self): 6 ctx1feabcd

430 self.partials = {} 6 ctx1feabcd

431 

432 def partial(self, token): 6 ctx1feabcd

433 return token_map(self._partial, token) 6 ctx1feabcd

434 

435 def _partial(self, token): 6 ctx1feabcd

436 now = time.perf_counter() 4 ctx1abcd

437 delta = now - self.stamp 4 ctx1abcd

438 self.partials[self.counter] = self.partials.get(self.counter, 0) + delta 4 ctx1abcd

439 self.totals[self.counter] = self.totals.get(self.counter, 0) + delta 4 ctx1abcd

440 self.stamp = now 4 ctx1abcd

441 self.counter += 1 4 ctx1abcd

442 self._last_start = False 4 ctx1abcd

443 return token 4 ctx1abcd

444 

445 def _parse_hyperprior(self, hyperprior, initial, fix): 6 ctx1feabcd

446 

447 # check fix against hyperprior and fill missing values 

448 hyperprior = self._copyasarrayorbufferdict(hyperprior) 6 ctx1feabcd

449 self._check_no_redundant_keys(hyperprior) 6 ctx1feabcd

450 fix = self._parse_fix(hyperprior, fix) 6 ctx1feabcd

451 flatfix = self._flatview(fix) 6 ctx1feabcd

452 

453 # extract distribution of free hyperparameters 

454 flathp = self._flatview(hyperprior) 6 ctx1feabcd

455 freehp = flathp[~flatfix] 6 ctx1feabcd

456 mean = gvar.mean(freehp) 6 ctx1feabcd

457 cov = gvar.evalcov(freehp) # TODO use evalcov_blocks 6 ctx1feabcd

458 dec = _linalg.Chol(cov) 6 ctx1feabcd

459 assert dec.n == freehp.size 6 ctx1feabcd

460 self.log(f'{freehp.size}/{flathp.size} free hyperparameters', 2) 6 ctx1feabcd

461 

462 # determine starting point for minimization 

463 initial = self._parse_initial(hyperprior, initial, dec) 6 ctx1feabcd

464 flatinitial = self._flatview(initial) 6 ctx1feabcd

465 x0 = dec.pinv_correlate(flatinitial[~flatfix] - mean) 6 ctx1feabcd

466 # TODO for initial = 'priormean', x0 is zero, skip decorrelate 

467 # for initial = 'priorsample', x0 is iid normal, but I have to sync 

468 # it with the user-exposed unflattened initial in _parse_initial 

469 

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

471 # format 

472 fixed_indices, = jnp.nonzero(flatfix) 6 ctx1feabcd

473 unfixed_indices, = jnp.nonzero(~flatfix) 6 ctx1feabcd

474 fixed_values = jnp.asarray(flatinitial[flatfix]) 6 ctx1feabcd

475 def unflat(x): 6 ctx1feabcd

476 assert x.ndim == 1 6 ctx1feabcd

477 if x.dtype == object: 6 ctx1feabcd

478 jac, indices = _gvarext.jacobian(x) 6 ctx1feabcd

479 xmean = mean + dec.correlate(gvar.mean(x)) 6 ctx1feabcd

480 xjac = dec.correlate(jac) 6 ctx1feabcd

481 x = _gvarext.from_jacobian(xmean, xjac, indices) 6 ctx1feabcd

482 y = numpy.empty(flatfix.size, x.dtype) 6 ctx1feabcd

483 numpy.put(y, unfixed_indices, x) 6 ctx1feabcd

484 numpy.put(y, fixed_indices, fixed_values) 6 ctx1feabcd

485 else: 

486 x = mean + dec.correlate(x) 6 ctx1feabcd

487 y = jnp.empty(flatfix.size, x.dtype) 6 ctx1feabcd

488 y = y.at[unfixed_indices].set(x) 6 ctx1feabcd

489 y = y.at[fixed_indices].set(fixed_values) 6 ctx1feabcd

490 return self._unflatview(y, hyperprior) 6 ctx1feabcd

491 

492 self.prior = hyperprior 6 ctx1feabcd

493 return x0, unflat 6 ctx1feabcd

494 

495 @staticmethod 6 ctx1feabcd

496 def _check_no_redundant_keys(hyperprior): 6 ctx1feabcd

497 if not hasattr(hyperprior, 'keys'): 6 ctx1feabcd

498 return 4 ctx1abcd

499 for k in hyperprior: 6 ctx1feabcd

500 m = hyperprior.extension_pattern.match(k) 6 ctx1feabcd

501 if m and m.group(1) in hyperprior.invfcn: 6 ctx1feabcd

502 altk = m.group(2) 6 ctx1feabcd

503 if altk in hyperprior: 503 ↛ 504line 503 didn't jump to line 504 because the condition on line 503 was never true6 ctx1feabcd

504 raise ValueError(f'duplicate keys {altk!r} and {k!r} in hyperprior') 

505 

506 def _parse_fix(self, hyperprior, fix): 6 ctx1feabcd

507 

508 if fix is None: 508 ↛ 514line 508 didn't jump to line 514 because the condition on line 508 was always true6 ctx1feabcd

509 if hasattr(hyperprior, 'keys'): 6 ctx1feabcd

510 fix = gvar.BufferDict(hyperprior, buf=numpy.zeros(hyperprior.size, bool)) 6 ctx1feabcd

511 else: 

512 fix = numpy.zeros(hyperprior.shape, bool) 4 ctx1abcd

513 else: 

514 fix = self._copyasarrayorbufferdict(fix) 

515 if hasattr(fix, 'keys'): 

516 assert hasattr(hyperprior, 'keys'), 'fix is dictionary but hyperprior is array' 

517 assert all(hyperprior.has_dictkey(k) for k in fix), 'some keys in fix are missing in hyperprior' 

518 newfix = {} 

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

520 key = None 

521 m = hyperprior.extension_pattern.match(k) 

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

523 altk = m.group(2) 

524 if altk in fix: 

525 assert k not in fix, f'duplicate keys {k!r} and {altk!r} in fix' 

526 key = altk 

527 if key is None and k in fix: 

528 key = k 

529 if key is None: 

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

531 else: 

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

533 newfix[k] = elem 

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

535 else: 

536 assert not hasattr(hyperprior, 'keys'), 'fix is array but hyperprior is dictionary' 

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

538 

539 self.fix = fix 6 ctx1feabcd

540 return fix 6 ctx1feabcd

541 

542 def _parse_initial(self, hyperprior, initial, dec): 6 ctx1feabcd

543 

544 if not isinstance(initial, str): 544 ↛ 545line 544 didn't jump to line 545 because the condition on line 544 was never true6 ctx1feabcd

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

546 initial = self._copyasarrayorbufferdict(initial) 

547 if hasattr(hyperprior, 'keys'): 

548 assert hasattr(initial, 'keys'), 'hyperprior is dictionary but initial is array' 

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

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

551 else: 

552 assert not hasattr(initial, 'keys'), 'hyperprior is array but initial is dictionary' 

553 assert hyperprior.shape == initial.shape 

554 

555 elif initial == 'priormean': 555 ↛ 559line 555 didn't jump to line 559 because the condition on line 555 was always true6 ctx1feabcd

556 self.log('start from prior mean', 2) 6 ctx1feabcd

557 initial = gvar.mean(hyperprior) 6 ctx1feabcd

558 

559 elif initial == 'priorsample': 

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

561 if dec.n < hyperprior.size: 

562 flathp = self._flatview(hyperprior) 

563 cov = gvar.evalcov(flathp) # TODO use evalcov_blocks 

564 fulldec = _linalg.Chol(cov) 

565 else: 

566 fulldec = dec 

567 iid = numpy.random.randn(fulldec.m) 

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

569 initial = self._unflatview(flatinitial, hyperprior) 

570 

571 else: 

572 raise KeyError(initial) 

573 

574 self.initial = initial 6 ctx1feabcd

575 return initial 6 ctx1feabcd

576 

577 def _parse_data(self, data): 6 ctx1feabcd

578 

579 self.data = data 6 ctx1feabcd

580 if isinstance(data, tuple) and len(data) == 1: 6 ctx1feabcd

581 data, = data 4 ctx1abcd

582 

583 if callable(data): 6 ctx1feabcd

584 self.log('data is callable', 2) 5 ctx1eabcd

585 cachedargs = None 5 ctx1eabcd

586 elif isinstance(data, tuple): 6 ctx1feabcd

587 self.log('data errors provided separately', 2) 4 ctx1abcd

588 assert len(data) == 2 4 ctx1abcd

589 cachedargs = data 4 ctx1abcd

590 elif (gdata := self._copyasarrayorbufferdict(data)).dtype == object: 6 ctx1feabcd

591 self.log('data has errors as gvars', 2) 5 ctx1eabcd

592 data = gvar.gvar(gdata) 5 ctx1eabcd

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

594 # gvar.mean and gvar.evalcov 

595 cachedargs = (gvar.mean(data), gvar.evalcov(data)) 5 ctx1eabcd

596 else: 

597 self.log('data has no errors', 2) 6 ctx1feabcd

598 cachedargs = (data,) 6 ctx1feabcd

599 

600 return data, cachedargs 6 ctx1feabcd

601 

602 def _prepare_functions(self, *, gpfactory, gpfactorykw, data, cachedargs, 6 ctx1feabcd

603 hpunflat, mlkw, jit, forward, additional_loss): 

604 

605 timer = self._Timer() 6 ctx1feabcd

606 firstcall = [None] 6 ctx1feabcd

607 

608 def make_decomp(p, **kw): 6 ctx1feabcd

609 """ decomposition of the prior covariance and data """ 

610 

611 # start timer and convert hypers to user format 

612 p = timer.start(p) 6 ctx1feabcd

613 hp = hpunflat(p) 6 ctx1feabcd

614 

615 # create GP object 

616 gp = gpfactory(hp, **kw) 6 ctx1feabcd

617 assert isinstance(gp, _GP.GP) 6 ctx1feabcd

618 

619 # extract data 

620 if cachedargs: 6 ctx1feabcd

621 args = cachedargs 6 ctx1feabcd

622 else: 

623 args = data(hp, **kw) 5 ctx1eabcd

624 if not isinstance(args, tuple): 5 ctx1eabcd

625 args = (args,) 5 ctx1eabcd

626 

627 # decompose covariance matrix and flatten data 

628 decomp, r = gp._prior_decomp(*args, covtransf=timer.partial, **mlkw) 6 ctx1feabcd

629 r = r.astype(float) # int data upsets jax 6 ctx1feabcd

630 

631 # log number of datapoints 

632 if firstcall: 6 ctx1feabcd

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

634 # works under jit since the first call is tracing 

635 firstcall.pop() 6 ctx1feabcd

636 xdtype = gp._get_x_dtype() 6 ctx1feabcd

637 nd = '?' if xdtype is None else _array._nd(xdtype) 6 ctx1feabcd

638 self.log(f'{r.size} datapoints, {nd} covariates') 6 ctx1feabcd

639 

640 # compute user loss 

641 if additional_loss is None: 6 ctx1feabcd

642 loss = 0. 6 ctx1feabcd

643 else: 

644 loss = additional_loss(hp) 5 ctx1eabcd

645 

646 # split timer and return decomposition 

647 return timer.partial(decomp), r, loss 6 ctx1feabcd

648 # TODO what's the correct way of checkpointing r? 

649 

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

651 def wrap(func): 6 ctx1feabcd

652 if jit: 652 ↛ 654line 652 didn't jump to line 654 because the condition on line 652 was always true6 ctx1feabcd

653 func = jax.jit(func) 6 ctx1feabcd

654 func = functools.partial(func, **gpfactorykw) 6 ctx1feabcd

655 return self._CountCalls(func) 6 ctx1feabcd

656 if jit: 656 ↛ 660line 656 didn't jump to line 660 because the condition on line 656 was always true6 ctx1feabcd

657 self.log('compile functions with jax jit', 2) 6 ctx1feabcd

658 

659 # log derivation method 

660 modename = 'forward' if forward else 'reverse' 6 ctx1feabcd

661 self.log(f'{modename}-mode autodiff (if used)', 2) 6 ctx1feabcd

662 

663 # TODO time the derivatives separately => maybe I need a custom 

664 # derivative rule for timer token acknoledgement? 

665 

666 def prior(p): 6 ctx1feabcd

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

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

669 return 1/2 * (len(p) * jnp.log(2 * jnp.pi) + p @ p) 6 ctx1feabcd

670 

671 def grad_prior(p): 6 ctx1feabcd

672 return p 6 ctx1feabcd

673 

674 def fisher_prior(p): 6 ctx1feabcd

675 return jnp.eye(len(p)) 4 ctx1abcd

676 

677 @wrap 6 ctx1feabcd

678 def fun(p, **kw): 6 ctx1feabcd

679 """ minus log marginal posterior of the hyperparameters (not 

680 normalized) """ 

681 decomp, r, loss = make_decomp(p, **kw) 4 ctx1abcd

682 cond, _, _, _, _ = decomp.minus_log_normal_density(r, value=True) 4 ctx1abcd

683 post = cond + prior(p) + loss 4 ctx1abcd

684 # TODO what's the correct way of checkpointing prior and loss? 

685 return timer.partial(post) 4 ctx1abcd

686 

687 def make_gradfwd_fisher_args(p, **kw): 6 ctx1feabcd

688 def make_decomp_tee(p): 5 ctx1eabcd

689 decomp, r, loss = make_decomp(p, **kw) 5 ctx1eabcd

690 return (decomp.matrix(), r, loss), (decomp, r, loss) 5 ctx1eabcd

691 (dK, dr, grad_loss), (decomp, r, loss) = jax.jacfwd(make_decomp_tee, has_aux=True)(p) 5 ctx1eabcd

692 lkw = dict(dK=dK, dr=dr) 5 ctx1eabcd

693 return decomp, r, lkw, loss, grad_loss 5 ctx1eabcd

694 

695 def make_gradrev_args(p, **kw): 6 ctx1feabcd

696 def make_decomp_loss(p): 6 ctx1feabcd

697 def make_decomp_r(p): 6 ctx1feabcd

698 def make_decomp_K(p): 6 ctx1feabcd

699 decomp, r, loss = make_decomp(p, **kw) 6 ctx1feabcd

700 return decomp.matrix(), (decomp, r, loss) 6 ctx1feabcd

701 _, dK_vjp, (decomp, r, loss) = jax.vjp(make_decomp_K, p, has_aux=True) 6 ctx1feabcd

702 return r, (decomp, r, dK_vjp, loss) 6 ctx1feabcd

703 _, dr_vjp, (decomp, r, dK_vjp, loss) = jax.vjp(make_decomp_r, p, has_aux=True) 6 ctx1feabcd

704 return loss, (decomp, r, dK_vjp, dr_vjp, loss) 6 ctx1feabcd

705 grad_loss, (decomp, r, dK_vjp, dr_vjp, loss) = jax.grad(make_decomp_loss, has_aux=True)(p) 6 ctx1feabcd

706 unpack = lambda f: lambda x: f(x)[0] 6 ctx1feabcd

707 dK_vjp = unpack(dK_vjp) 6 ctx1feabcd

708 dr_vjp = unpack(dr_vjp) 6 ctx1feabcd

709 lkw = dict(dK_vjp=dK_vjp, dr_vjp=dr_vjp) 6 ctx1feabcd

710 return decomp, r, lkw, loss, grad_loss 6 ctx1feabcd

711 

712 def make_jac_args(p, **kw): 6 ctx1feabcd

713 if forward: 6 ctx1feabcd

714 out = make_gradfwd_fisher_args(p, **kw) 5 ctx1eabcd

715 out[2].update(gradfwd=True) # out[2] is lkw 5 ctx1eabcd

716 else: 

717 out = make_gradrev_args(p, **kw) 6 ctx1feabcd

718 out[2].update(gradrev=True) 6 ctx1feabcd

719 return out 6 ctx1feabcd

720 

721 @wrap 6 ctx1feabcd

722 def fun_and_jac(p, **kw): 6 ctx1feabcd

723 """ fun and its gradient """ 

724 decomp, r, lkw, loss, grad_loss = make_jac_args(p, **kw) 6 ctx1feabcd

725 cond, gradrev, gradfwd, _, _ = decomp.minus_log_normal_density(r, value=True, **lkw) 6 ctx1feabcd