Coverage for src/lsqfitgp/_linalg/_decomp.py: 85%
249 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-04 11:43 +0000
1# lsqfitgp/_linalg/_decomp.py
2#
3# Copyright (c) 2023, 2024, 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/>.
20"""
21Matrix decompositions.
23Copy-pasted from the notes:
252023-02-14
26==========
28My current decomposition system is a mess. I can't take reverse gradients.
29I can't straightforwardly implement optimized algorithms that compute together
30likelihood, gradient, and fisher. Jax patterns break down unpredictably. I
31have to redesign it from scratch.
33Guidelines and requirements:
35 - Sufficient modularity to implement composite decompositions (Woodbury,
36 Block)
38 - Does not mess up jax in any way
40 - Caches decompositions
42 - Favors optimizing together the likelihood and its derivatives
44Operations (in the following I indicate with lowercase inputs which are
45typically vectors or tall matrices, and uppercase inputs which are typically
46large matrices, since optimization requires taking it into account):
48 pinv_bilinear(A, r) => A'K⁺r (for the posterior mean)
49 pinv_bilinear_robj(A, r) same but r can be gvars
50 ginv_quad(A) => A'K⁻A (for the posterior covariance)
51 I want the pseudoinverse for the mean because the data may not be
52 in the span and I want to project it orthogonally, while for the
53 covariance I expect A and K to come from a pos def matrix so they are
54 coherent
55 ginv_diagquad(A) => diag(A'K⁻A) (for the posterior variance)
56 minus_log_normal_density(
57 r: 1d array, # the residuals (data - prior mean)
58 dr_vjp: callable, # x -> x_i ∂r_i/∂p_j, gradrev and fishvec
59 dK_vjp: callable, # x -> x_ij ∂K_ij/∂p_k, gradrev and fishvec
60 dr_jvp: callable, # x -> ∂r_i/∂r_j x_j, fishvec
61 dK_jvp: callable, # x -> ∂K_ij/∂p_k x_k, fishvec
62 dr: 2d array, # ∂r_i/∂p_j, gradfwd and fisher
63 dK: 3d array, # ∂K_ij/∂p_k, gradfwd and fisher
64 vec: 1d array, # input vector of fishvec, same size as params
65 value: bool,
66 gradrev: bool,
67 gradfwd: bool,
68 fisher: bool,
69 fishvec: bool,
70 )
71 This computes on request
72 value: 1/2 tr(KK⁺) log 2π
73 + 1/2 tr(I-KK⁺) log 2π
74 + 1/2 log pdet K
75 + 1/2 tr(I-KK⁺) log ε
76 + 1/2 r'(K⁺+(I-KK⁺)/ε)r
77 gradrev,
78 gradfwd: 1/2 tr(K⁺dK)
79 + r'(K⁺+(I-KK⁺)/ε) dr
80 - 1/2 r'(K⁺+2(I-KK⁺)/ε)dKK⁺r
81 fisher: 1/2 tr(K⁺dK(K⁺+2(I-KK⁺)/ε)d'K)
82 - 2 tr(K⁺dK(I-KK⁺)d'KK⁺)
83 + dr'(K⁺+(I-KK⁺)/ε)d'r
84 fishvec: fisher matrix times vec
85 There should be options for omitting the pieces with ε. I also need a
86 way to make densities with different values of ε comparable with each
87 other (may not be possible, if it is, it probably requires a history of
88 ranks and ε). gradfwd/rev form K⁺ explicitly to compute tr(K⁺dK) for
89 efficiency.
90 correlate(x)
91 Zx where K = ZZ'.
92 back_correlate(X):
93 Z'X, this is used by Sandwich and Woodbury.
95Since I also want to compute the Student density, I could split
96minus_log_normal_density's return value into logdet and quad. The gradient
97splits nicely between the two terms, but I have to redo the calculation of the
98Fisher matrix for the Student distribution. Alternatively, I could use the the
99Normal Fisher. => See Lange et al. (1989, app. B). => I think I can split the
100gradient and Fisher matrix too.
1022023-03-07
103==========
105To compute a Fisher-vector product when there are many parameters, do
107 tr(K+ dK K+ dK) v =
108 = K_vjp(K+ K_jvp(v) K+)
109"""
111import abc 1cba
112import functools 1cba
114import jax 1cba
115import numpy 1cba
116from jax import lax 1cba
117from jax import numpy as jnp 1cba
118from jax.scipy import linalg as jlinalg 1cba
120from lsqfitgp import _jaxext 1cba
121from lsqfitgp._linalg import _pytree 1cba
124class Decomposition(_pytree.AutoPyTree, abc.ABC): 1cba
125 """Abstract base class for decompositions of positive semidefinite matrices."""
127 @abc.abstractmethod 1cba
128 def __init__(self, *args, **kw): 1cba
129 """Decompose the input matrix."""
130 pass
132 @abc.abstractmethod 1cba
133 def matrix(self): 1cba
134 """Return the input matrix."""
135 pass
137 @abc.abstractmethod 1cba
138 def ginv_linear(self, X): 1cba
139 """Compute K⁻X."""
140 pass
142 @abc.abstractmethod 1cba
143 def pinv_bilinear(self, A, r): 1cba
144 """Compute A'K⁺r."""
145 pass
147 @abc.abstractmethod 1cba
148 def pinv_bilinear_robj(self, A, r): 1cba
149 """Compute A'K⁺r, where r can be an array of objects."""
150 pass
152 @abc.abstractmethod 1cba
153 def ginv_quad(self, A): 1cba
154 """Compute A'K⁻A."""
155 pass
157 @abc.abstractmethod 1cba
158 def ginv_diagquad(self, A): 1cba
159 """Compute diag(A'K⁻A)."""
160 pass
162 @abc.abstractmethod 1cba
163 def correlate(self, x): 1cba
164 """Compute Zx where K = ZZ'."""
165 pass
167 @abc.abstractmethod 1cba
168 def back_correlate(self, X): 1cba
169 """Compute Z'X."""
170 pass
172 @abc.abstractmethod 1cba
173 def pinv_correlate(self, x): 1cba
174 """Compute Z⁺x."""
175 pass
177 @abc.abstractmethod 1cba
178 def minus_log_normal_density( 1cba
179 self,
180 r,
181 *,
182 dr_vjp=None,
183 dK_vjp=None,
184 dr_jvp_vec=None,
185 dK_jvp_vec=None,
186 dr=None,
187 dK=None,
188 value=False,
189 gradrev=False,
190 gradfwd=False,
191 fisher=False,
192 fishvec=False,
193 ):
194 """
195 Compute minus log a Normal density and its derivatives.
197 The covariance matrix of the Normal is K.
199 If an input derivative is not specified, it is assumed to be zero.
201 Parameters
202 ----------
203 r : 1d array
204 The residuals (value - mean)
205 dr_vjp : callable
206 x -> x_i ∂r_i/∂p_j, for gradrev and fishvec
207 dK_vjp : callable
208 x -> x_ij ∂K_ij/∂p_k, for gradrev and fishvec
209 dr_jvp_vec : 1d array
210 ∂r_i/∂r_j vec_j, for fishvec
211 dK_jvp_vec : 2d array
212 ∂K_ij/∂p_k vec_k, for fishvec
213 dr : 2d array
214 ∂r_i/∂p_j for gradfwd and fisher
215 dK : 3d array
216 ∂K_ij/∂p_k, for gradfwd and fisher
217 value : bool
218 gradrev : bool
219 gradfwd : bool
220 fisher : bool
221 fishvec : bool
222 These parameters indicate which of the return values to compute.
223 Default all False.
225 Returns
226 -------
227 value : 1/2 tr(KK⁺) log 2π
228 + 1/2 tr(I-KK⁺) log 2π
229 + 1/2 log pdet K
230 + 1/2 tr(I-KK⁺) log ε
231 + 1/2 r'(K⁺+(I-KK⁺)/ε)r
232 gradrev,
233 gradfwd : 1/2 tr(K⁺dK)
234 + r'(K⁺+(I-KK⁺)/ε) dr
235 - 1/2 r'(K⁺+2(I-KK⁺)/ε)dKK⁺r
236 fisher : 1/2 tr(K⁺dK(K⁺+2(I-KK⁺)/ε)d'K)
237 - 2 tr(K⁺dK(I-KK⁺)d'KK⁺)
238 + dr'(K⁺+(I-KK⁺)/ε)d'r
239 fishvec : fisher matrix @ vec
240 """ # noqa: DOC202
241 pass
243 def _parseeps(self, K, epsrel, epsabs, maxeigv=None): 1cba
244 """Determine eps from input arguments."""
245 machine_eps = jnp.finfo(_jaxext.float_type(K)).eps 1cba
246 if epsrel == 'auto': 1cba
247 epsrel = len(K) * machine_eps 1cba
248 if epsabs == 'auto': 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true1cba
249 epsabs = machine_eps
250 if maxeigv is None: 250 ↛ 252line 250 didn't jump to line 252 because the condition on line 250 was always true1cba
251 maxeigv = eigval_bound(K) 1cba
252 self._eps = epsrel * maxeigv + epsabs 1cba
253 return self._eps 1cba
255 @property 1cba
256 def eps(self): 1cba
257 """The threshold below which eigenvalues are too small to be determined."""
258 return self._eps 1a
260 @property 1cba
261 @abc.abstractmethod 1cba
262 def n(self): 1cba
263 """Number of rows/columns of the matrix."""
264 pass
266 @property 1cba
267 @abc.abstractmethod 1cba
268 def m(self): 1cba
269 """Number of columns of Z."""
270 pass
272 def ginv(self): 1cba
273 """Compute K⁻."""
274 return self.ginv_quad(jnp.eye(self.n)) 1a
277def solve_triangular_python(a, b, *, lower=False): 1cba
278 """
279 Pure python implementation of scipy.linalg.solve_triangular.
281 For when a or b are object arrays.
282 """
283 a = numpy.asarray(a) 1cba
284 x = numpy.copy(b) 1cba
286 vec = x.ndim < 2 1cba
287 if vec: 287 ↛ 290line 287 didn't jump to line 290 because the condition on line 287 was always true1cba
288 x = x[:, None] 1cba
290 n = a.shape[-1] 1cba
291 assert x.shape[-2] == n 1cba
293 if not lower: 293 ↛ 294line 293 didn't jump to line 294 because the condition on line 293 was never true1cba
294 a = a[..., ::-1, ::-1]
295 x = x[..., ::-1, :]
297 x[..., 0, :] /= a[..., 0, 0, None] 1cba
298 for i in range(1, n): 1cba
299 x[..., i:, :] -= x[..., None, i - 1, :] * a[..., i:, i - 1, None] 1cba
300 x[..., i, :] /= a[..., i, i, None] 1cba
302 if not lower: 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true1cba
303 x = x[..., ::-1, :]
305 if vec: 305 ↛ 307line 305 didn't jump to line 307 because the condition on line 305 was always true1cba
306 x = numpy.squeeze(x, -1) 1cba
307 return x 1cba
310def solve_triangular_batched(a, b, *, lower=False): 1cba
311 """Version of jax.scipy.linalg.solve_triangular that batches matmul-like."""
312 a = jnp.asarray(a) 1a
313 b = jnp.asarray(b) 1a
314 vec = b.ndim < 2 1a
315 if vec: 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true1a
316 b = b[:, None]
318 batch_shape = jnp.broadcast_shapes(a.shape[:-2], b.shape[:-2]) 1a
319 a_shape = batch_shape + a.shape[-2:] 1a
320 b_shape = batch_shape + b.shape[-2:] 1a
321 result = lax.linalg.triangular_solve( 1a
322 jnp.broadcast_to(a, a_shape),
323 jnp.broadcast_to(b, b_shape),
324 left_side=True,
325 lower=lower,
326 )
327 assert result.shape == b_shape 1a
329 if vec: 329 ↛ 330line 329 didn't jump to line 330 because the condition on line 329 was never true1a
330 result = result.squeeze(-1)
331 return result 1a
334def solve_batched(a, b, **kw): 1cba
335 """Version of jax.scipy.linalg.solve that batches matmul-like."""
336 a = jnp.asarray(a) 1a
337 b = jnp.asarray(b) 1a
338 vec = b.ndim < 2 1a
339 if vec: 339 ↛ 340line 339 didn't jump to line 340 because the condition on line 339 was never true1a
340 b = b[:, None]
342 @functools.partial(jnp.vectorize, signature='(i,j),(j,k)->(i,k)') 1a
343 def solve_batched(a, b): 1a
344 return jlinalg.solve(a, b, **kw) 1a
346 result = solve_batched(a, b) 1a
348 if vec: 348 ↛ 349line 348 didn't jump to line 349 because the condition on line 348 was never true1a
349 result = result.squeeze(-1)
350 return result 1a
353def eigval_bound(K): 1cba
354 """
355 Upper bound on the largest magnitude eigenvalue of the matrix.
357 The bound is from Gershgorin's theorem.
358 """
359 return jnp.max(jnp.sum(jnp.abs(K), axis=1)) 1cba
362def diag_scale_pow2(K): 1cba
363 """Compute a vector s of powers of 2 such that diag(K / outer(s, s)) ~ 1."""
364 d = jnp.diag(K) 1cba
365 return jnp.where(d, 2 ** jnp.rint(0.5 * jnp.log2(d)), 1) 1cba
367 # Golub and Van Loan (2013) say this is not a totally general heuristic
370def transpose(x): 1cba
371 """
372 Swap the last two axes of array x.
374 Corresponds to matrix tranposition with the broadcasting convention of
375 matmul.
376 """
377 if x.ndim < 2:
378 return x
379 elif isinstance(x, jnp.ndarray):
380 return jnp.swapaxes(x, -2, -1)
381 else:
382 # need to support numpy because this function is used with gvars
383 return numpy.swapaxes(x, -2, -1)
386class Chol(Decomposition): 1cba
387 """
388 Cholesky decomposition.
390 The matrix is regularized adding a small multiple of the identity.
391 """
393 def __init__(self, K, *, epsrel='auto', epsabs=0): 1cba
394 # K <- K + Iε
395 # K = LL'
396 self._K = K 1cba
397 s = diag_scale_pow2(K) 1cba
398 K = K / s / s[:, None] 1cba
399 eps = self._parseeps(K, epsrel, epsabs) 1cba
400 K = K.at[jnp.diag_indices_from(K)].add(eps) 1cba
401 L = jlinalg.cholesky(K, lower=True) 1cba
402 with _jaxext.skipifabstract(): 1cba
403 if not jnp.all(jnp.isfinite(L)): 1cba
404 msg = 'cholesky decomposition not finite, probably matrix not pos def numerically' 1a
405 raise numpy.linalg.LinAlgError(msg) 1a
406 self._L = L * s[:, None] 1cba
407 self._eps = eps * jnp.min(s * s) 1cba
409 def matrix(self): 1cba
410 return self._K 1cba
412 def ginv_linear(self, X): 1cba
413 # = K⁻¹X
414 # K⁻¹ = L'⁻¹L⁻¹
415 # K⁻¹X = L'⁻¹(L⁻¹X)
416 invLX = jlinalg.solve_triangular(self._L, X, lower=True) 1a
417 return jlinalg.solve_triangular(self._L.T, invLX, lower=False) 1a
419 def pinv_bilinear(self, A, r): 1cba
420 # = A'K⁻¹r = A'L'⁻¹L⁻¹r = (L⁻¹A)'(L⁻¹r)
421 invLr = jlinalg.solve_triangular(self._L, r, lower=True) 1cba
422 invLA = jlinalg.solve_triangular(self._L, A, lower=True) 1cba
423 return invLA.T @ invLr 1cba
425 def pinv_bilinear_robj(self, A, r): 1cba
426 # = A'K⁻¹r
427 invLr = solve_triangular_python(self._L, r, lower=True) 1cba
428 invLA = jlinalg.solve_triangular(self._L, A, lower=True) 1cba
429 return numpy.asarray(invLA).T @ invLr 1cba
431 def ginv_quad(self, A): 1cba
432 # = A'K⁻¹A = A'K⁻¹A = A'L'⁻¹L⁻¹A = (L⁻¹A)'(L⁻¹A)
433 invLA = jlinalg.solve_triangular(self._L, A, lower=True) 1cba
434 return invLA.T @ invLA 1cba
436 def ginv_diagquad(self, A): 1cba
437 # = diag(A'K⁻¹A)
438 # X = L⁻¹A
439 # diag(A'K⁻¹A)_i = diag(X'X)_i = ∑_j X'_ij X_ji = ∑_j X_ji X_ji
440 invLA = jlinalg.solve_triangular(self._L, A, lower=True) 1a
441 return jnp.einsum('ji,ji->i', invLA, invLA) 1a
443 def correlate(self, x): 1cba
444 # = Lx
445 return self._L @ x 1cba
447 def back_correlate(self, X): 1cba
448 # = L'X
449 return self._L.T @ X 1a
451 def pinv_correlate(self, x): 1cba
452 # = L⁻¹x
453 return jlinalg.solve_triangular(self._L, x, lower=True) 1cba
455 def minus_log_normal_density( # noqa: C901, PLR0915 1cba
456 self,
457 r, # 1d array, the residuals (data - prior mean)
458 *,
459 dr_vjp=None, # callable, x -> x_i ∂r_i/∂p_j, gradrev and fishvec
460 dK_vjp=None, # callable, x -> x_ij ∂K_ij/∂p_k, gradrev and fishvec
461 dr_jvp_vec=None, # 1d array, ∂r_i/∂r_j v_j, fishvec
462 dK_jvp_vec=None, # 2d array, ∂K_ij/∂p_k v_k, fishvec
463 dr=None, # 2d array, ∂r_i/∂p_j, gradfwd and fisher
464 dK=None, # 3d array, ∂K_ij/∂p_k, gradfwd and fisher
465 value=False,
466 gradrev=False,
467 gradfwd=False,
468 fisher=False,
469 fishvec=False,
470 ):
472 L = self._L 1cba
474 out = {} 1cba
476 # compute shared factors
477 grad = (gradrev and (dK_vjp is not None or dr_vjp is not None)) or ( 1cba
478 gradfwd and (dK is not None or dr is not None)
479 )
480 if value or grad: 1cba
481 invLr = jlinalg.solve_triangular(L, r, lower=True) 1cba
482 if grad: 1cba
483 invKr = jlinalg.solve_triangular(L.T, invLr, lower=False) 1cba
484 if (gradrev and dK_vjp is not None) or (gradfwd and dK is not None): 1cba
485 invL = jlinalg.solve_triangular(L, jnp.eye(len(L)), lower=True) 1cba
486 invK = invL.T @ invL 1cba
488 if value: 1cba
489 # = 1/2 n log 2π
490 # + 1/2 log det K
491 # + 1/2 r'K⁻¹r
492 # K = LL'
493 # K⁻¹ = L'⁻¹L⁻¹
494 # det K = (det L)² =
495 # = (∏_i L_ii)²
496 # r'K⁻¹r = r'L'⁻¹L⁻¹r =
497 # = (L⁻¹r)'(L⁻¹r)
498 out['value'] = ( 1cba
499 1
500 / 2
501 * (
502 len(L) * jnp.log(2 * jnp.pi)
503 + 2 * jnp.sum(jnp.log(jnp.diag(L)))
504 + invLr @ invLr
505 )
506 )
507 else:
508 out['value'] = None 1a
510 if gradrev: 1cba
511 # = 1/2 tr(K⁻¹dK)
512 # + r'K⁻¹dr
513 # - 1/2 r'K⁻¹dKK⁻¹r
514 # tr(K⁻¹dK) = K⁻¹_ij dK_ji =
515 # = K⁻¹_ij dK_ij =
516 # = dK_vjp(K⁻¹)
517 # r'K⁻¹dr = r_i K⁻¹_ij dr_j =
518 # = (K⁻¹r)_j dr_j =
519 # = dr_vjp(K⁻¹r)
520 # r'K⁻¹dKK⁻¹r = r_i K⁻¹_ij dK_jl K⁻¹_lm r_m =
521 # = (K⁻¹r)_j dK_jl (K⁻¹r)_l =
522 # = dK_vjp((K⁻¹r) ⊗ (K⁻¹r))
523 out['gradrev'] = 0 1cba
524 if dK_vjp is not None: 524 ↛ 528line 524 didn't jump to line 528 because the condition on line 524 was always true1cba
525 tr_invK_dK = dK_vjp(invK) 1cba
526 r_invK_dK_invK_r = dK_vjp(jnp.outer(invKr, invKr)) 1cba
527 out['gradrev'] += 1 / 2 * (tr_invK_dK - r_invK_dK_invK_r) 1cba
528 if dr_vjp is not None: 528 ↛ 534line 528 didn't jump to line 534 because the condition on line 528 was always true1cba
529 r_invK_dr = dr_vjp(invKr) 1cba
530 out['gradrev'] += r_invK_dr 1cba
531 else:
532 out['gradrev'] = None 1ba
534 if gradfwd: 1cba
535 # = 1/2 tr(K⁻¹dK)
536 # + r'K⁻¹dr
537 # - 1/2 r'K⁻¹dKK⁻¹r
538 # tr(K⁻¹dK)_k = K⁻¹_ij dK_ijk
539 # r'K⁻¹dr = (K⁻¹r)'dr
540 # (r'K⁻¹dKK⁻¹r)_k = r_i K⁻¹_ij dK_jlk K⁻¹_lm r_m =
541 # = (K⁻¹r)_j dK_jlk (K⁻¹r)_l
542 out['gradfwd'] = 0 1ba
543 if dK is not None: 543 ↛ 547line 543 didn't jump to line 547 because the condition on line 543 was always true1ba
544 tr_invK_dK = jnp.einsum('ij,ijk->k', invK, dK) 1ba
545 r_invK_dK_invK_r = jnp.einsum('i,ijk,j->k', invKr, dK, invKr) 1ba
546 out['gradfwd'] += 1 / 2 * (tr_invK_dK - r_invK_dK_invK_r) 1ba
547 if dr is not None: 547 ↛ 553line 547 didn't jump to line 553 because the condition on line 547 was always true1ba
548 r_invK_dr = invKr @ dr 1ba
549 out['gradfwd'] += r_invK_dr 1ba
550 else:
551 out['gradfwd'] = None 1cba
553 if fisher: 1cba
554 # = 1/2 tr(K⁻¹dKK⁻¹d'K)
555 # + dr'K⁻¹d'r
556 # tr(K⁻¹dKK⁻¹d'K)_ij = tr(L'⁻¹L⁻¹dKL'⁻¹L⁻¹d'K)_ij =
557 # = tr(L⁻¹dKL'⁻¹L⁻¹d'KL'⁻¹)_ij =
558 # = (L⁻¹dKL'⁻¹)_kli (L⁻¹dKL'⁻¹)_klj
559 # (L⁻¹dKL'⁻¹)_ijk = L⁻¹_il dK_lmk L'⁻¹_mj =
560 # = L⁻¹_il L⁻¹_jm dK_lmk
561 # (dr'K⁻¹d'r)_kq = dr'_k L'⁻¹L⁻¹dr_q =
562 # = (L⁻¹dr_k)_i (L⁻¹dr_q)_i
563 out['fisher'] = 0 1a
564 if dK is not None: 564 ↛ 575line 564 didn't jump to line 575 because the condition on line 564 was always true1a
565 invL_dK = solve_triangular_batched( 1a
566 L, jnp.moveaxis(dK, 2, 0), lower=True
567 ) # kim: L⁻¹_il dK_lmk
568 invL_dK_invL = solve_triangular_batched( 1a
569 L, jnp.swapaxes(invL_dK, 1, 2), lower=True
570 ) # kji: L⁻¹_jm (L⁻¹_il dK_lmk)
571 tr_invK_dK_invK_dK = jnp.einsum( 1a
572 'kij,qij->kq', invL_dK_invL, invL_dK_invL
573 )
574 out['fisher'] += 1 / 2 * tr_invK_dK_invK_dK 1a
575 if dr is not None: 575 ↛ 582line 575 didn't jump to line 582 because the condition on line 575 was always true1a
576 invLdr = jlinalg.solve_triangular(L, dr, lower=True) 1a
577 dr_invK_dr = invLdr.T @ invLdr 1a
578 out['fisher'] += dr_invK_dr 1a
579 else:
580 out['fisher'] = None 1cba
582 if fishvec: 1cba
583 # = 1/2 tr(K⁻¹dKK⁻¹d'K) v
584 # + dr'K⁻¹d'r v
585 # tr(K⁻¹dKK⁻¹d'K) v = K_vjp(K⁻¹K_jvp(v)K⁻¹) =
586 # = K_vjp(L'⁻¹L⁻¹ K_jvp(v) L'⁻¹L⁻¹)
587 # dr'K⁻¹d'r v = dr'K⁻¹dr_jvp(v) =
588 # = dr_vjp(K⁻¹dr_jvp(v)) =
589 # = dr_vjp(L'⁻¹L⁻¹ dr_jvp(v))
590 out['fishvec'] = 0 1a
591 if not (dK_jvp_vec is None and dK_vjp is None): 591 ↛ 600line 591 didn't jump to line 600 because the condition on line 591 was always true1a
592 invL_dKv = jlinalg.solve_triangular(L, dK_jvp_vec, lower=True) 1a
593 invK_dKv = jlinalg.solve_triangular(L.T, invL_dKv, lower=False) 1a
594 invL_dKv_invK = jlinalg.solve_triangular(L, invK_dKv.T, lower=True) 1a
595 invK_dKv_invK = jlinalg.solve_triangular( 1a
596 L.T, invL_dKv_invK, lower=False
597 )
598 tr_invK_dK_invK_dK_v = dK_vjp(invK_dKv_invK) # ty: ignore[call-non-callable] 1a
599 out['fishvec'] += 1 / 2 * tr_invK_dK_invK_dK_v 1a
600 if not (dr_jvp_vec is None and dr_vjp is None): 600 ↛ 608line 600 didn't jump to line 608 because the condition on line 600 was always true1a
601 invL_drv = jlinalg.solve_triangular(L, dr_jvp_vec, lower=True) 1a
602 invK_drv = jlinalg.solve_triangular(L.T, invL_drv, lower=False) 1a
603 dr_invK_drv_v = dr_vjp(invK_drv) # ty: ignore[call-non-callable] 1a
604 out['fishvec'] += dr_invK_drv_v 1a
605 else:
606 out['fishvec'] = None 1cba
608 return tuple(out.values()) 1cba
610 @classmethod 1cba
611 def make_derivs( 1cba
612 cls,
613 K_fun,
614 r_fun,
615 primal,
616 *,
617 args=(),
618 kw={}, # noqa: B006, read only
619 vec=None,
620 value=False, # noqa: ARG003, no derivatives needed for the value
621 gradrev=False,
622 gradfwd=False,
623 fisher=False,
624 fishvec=False,
625 ):
626 """
627 Prepare arguments for `minus_log_normal_density`.
629 Parameters
630 ----------
631 K_fun : callable
632 Function with signature ``f(primal, *args, **kw)`` that produces the
633 `K` init argument.
634 r_fun : callable
635 Function with signature ``f(primal, *args, **kw)`` that produces the
636 `r` `minus_log_normal_density` argument.
637 primal : 1d array
638 The first argument to `K_fun` and `r_fun`.
639 args : tuple
640 Additional positional arguments to `K_fun` and `r_fun`.
641 kw : dict
642 Keyword arguments to `K_fun` and `r_fun`.
643 vec : 1d array
644 A tangent vector to compute the jacobian-vector products.
645 value : bool
646 Argument to `minus_log_normal_density`, used to determine which
647 derivatives are needed.
648 gradrev : bool
649 Like `value`.
650 gradfwd : bool
651 Like `value`.
652 fisher : bool
653 Like `value`.
654 fishvec : bool
655 Like `value`.
657 Returns
658 -------
659 K : 2d array
660 Output of `K_fun`.
661 r : 1d array
662 Output of `r_fun`.
663 out : dict
664 Dictionary with derivative arguments to `minus_log_normal_density`.
665 """
666 partial = lambda f: lambda x: f(x, *args, **kw) 1a
667 K_fun = partial(K_fun) 1a
668 r_fun = partial(r_fun) 1a
670 out = {} 1a
672 if gradrev or fishvec: 1a
673 K, dK_vjp = jax.vjp(K_fun, primal) 1a
674 r, dr_vjp = jax.vjp(r_fun, primal) 1a
675 out['dK_vjp'] = lambda x: dK_vjp(x)[0] 1a
676 out['dr_vjp'] = lambda x: dr_vjp(x)[0] 1a
677 else:
678 K = K_fun(primal) 1a
679 r = r_fun(primal) 1a
680 if fishvec: 1a
681 _, out['dK_jvp_vec'] = jax.jvp(K_fun, (primal,), (vec,)) 1a
682 _, out['dr_jvp_vec'] = jax.jvp(r_fun, (primal,), (vec,)) 1a
683 if gradfwd or fisher: 1a
684 out['dK'] = jax.jacfwd(K_fun)(primal) 1a
685 out['dr'] = jax.jacfwd(r_fun)(primal) 1a
687 return K, r, out 1a
689 @property 1cba
690 def n(self): 1cba
691 return len(self._L) 1cba
693 m = n 1cba