Coverage for src/lsqfitgp/_linalg/_toeplitz.py: 96%
136 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-04 11:43 +0000
1# lsqfitgp/_linalg/_toeplitz.py
2#
3# Copyright (c) 2020, 2022, 2023, 2026, Giacomo Petrillo
4#
5# This file is part of lsqfitgp.
6#
7# lsqfitgp is free software: you can redistribute it and/or modify
8# it under the terms of the GNU General Public License as published by
9# the Free Software Foundation, either version 3 of the License, or
10# (at your option) any later version.
11#
12# lsqfitgp is distributed in the hope that it will be useful,
13# but WITHOUT ANY WARRANTY; without even the implied warranty of
14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
15# GNU General Public License for more details.
16#
17# You should have received a copy of the GNU General Public License
18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>.
20import numpy 1bca
21from jax import jit 1bca
22from jax import numpy as jnp 1bca
24from lsqfitgp._linalg import _seqalg 1bca
27class SymSchur(_seqalg.Producer): 1bca
28 """
29 Cholesky decomposition of a symmetric Toeplitz matrix.
31 Adapted from TOEPLITZ_CHOLESKY by John Burkardt (LGPL license)
32 http://people.sc.fsu.edu/~jburkardt/py_src/toeplitz_cholesky/toeplitz_cholesky.html
33 """
35 def __init__(self, t): 1bca
36 """`t` = first row of the matrix."""
37 t = jnp.asarray(t) 1a
38 assert t.ndim == 1 1a
39 # assert t[0] > 0, '1-th leading minor is not positive definite'
40 self.t = t 1a
42 inputs = () 1bca
44 def init(self, n): 1bca
45 t = self.t 1a
46 del self.t 1a
47 assert len(t) == n 1a
48 norm = t[0] 1a
49 t = t / norm 1a
50 self.g = jnp.stack([t, t]) 1a
51 self.snorm = jnp.sqrt(norm) 1a
53 def iter_out(self, i): # noqa: ARG002, interface method 1bca
54 """`i`-th column of Cholesky factor L."""
55 return self.g[0, :] * self.snorm 1a
57 def iter(self, i): 1bca
58 g = self.g 1a
59 g = g.at[0, :].set(jnp.roll(g[0, :], 1)) 1a
60 g = g.at[:, 0].set(0).at[:, i - 1].set(0) 1a
61 # assert g[0, i] > 0, 'what??'
62 rho = -g[1, i] / g[0, i] 1a
63 # assert abs(rho) < 1, f'{i+1}-th leading minor is not positive definite'
64 gamma = jnp.sqrt((1 - rho) * (1 + rho)) 1a
65 self.g = (g + g[::-1] * rho) / gamma 1a
68class SymLevinson(_seqalg.Producer): 1bca
69 """
70 Cholesky decomposition of the *inverse* of a symmetric Toeplitz matrix.
72 Adapted form SuperGauss/DurbinLevinson.h (GPL license)
73 https://cran.r-project.org/package=SuperGauss
75 Note: Schur should be more accurate than Levinson
76 """
78 def __init__(self, t): 1bca
79 """`t` = first row of the matrix."""
80 t = jnp.asarray(t, float) 1ba
81 assert t.ndim == 1 1ba
82 # assert t[0] > 0, '1-th leading minor is not positive definite'
83 self.t = t 1ba
85 inputs = () 1bca
87 def init(self, n): 1bca
88 self.phi1 = jnp.zeros(n) 1ba
89 self.phi2 = jnp.zeros(n) 1ba
90 self.nu = self.t[0] 1ba
91 self.tlag = jnp.roll(self.t, -1) 1ba
92 del self.t 1ba
94 def iter_out(self, i): 1bca
95 """`i`-th row of L^-1."""
96 return -self.phi2.at[i].set(-1) / jnp.sqrt(self.nu) 1ba
98 def iter(self, i): 1bca
99 phi1 = self.phi1 1ba
100 phi2 = self.phi2 1ba
101 nu = self.nu 1ba
102 tlag = self.tlag 1ba
104 pi = i - 1 1ba
105 rp = phi2 @ tlag 1ba
106 phi1 = phi1.at[pi].set((tlag[pi] - rp) / nu) 1ba
107 phi1 = phi1 - phi1[pi] * phi2 1ba
108 # assert abs(phi1[pi]) < 1, f'{i+1}-th leading minor is not positive definite'
109 nu = nu * (1 - phi1[pi]) * (1 + phi1[pi]) 1ba
110 phi2 = jnp.roll(phi1[::-1], i) 1ba
112 self.phi1 = phi1 1ba
113 self.phi2 = phi2 1ba
114 self.nu = nu 1ba
117@jit 1bca
118def chol(t): 1bca
119 _, out = _seqalg.sequential_algorithm(len(t), [SymSchur(t), _seqalg.Stack(0)]) 1a
120 return out.T 1a
123@jit 1bca
124def chol_solve(t, *bs): 1bca
125 ops = [SymSchur(t)] + [_seqalg.SolveTriLowerColByFull(0, b) for b in bs] 1a
126 out = _seqalg.sequential_algorithm(len(t), ops) 1a
127 return out[1] if len(bs) == 1 else out[1:] 1a
130@jit 1bca
131def chol_matmul(t, b): 1bca
132 ops = [SymSchur(t), _seqalg.Rows(b), _seqalg.MatMulColByRow(0, 1)] 1a
133 _, _, out = _seqalg.sequential_algorithm(len(t), ops) 1a
134 return out 1a
137@jit 1bca
138def chol_transp_matmul(t, b): 1bca
139 ops = [SymSchur(t), _seqalg.MatMulRowByFull(0, b), _seqalg.Stack(1)]
140 _, _, out = _seqalg.sequential_algorithm(len(t), ops)
141 return out
144@jit 1bca
145def logdet(t): 1bca
146 _, out = _seqalg.sequential_algorithm(len(t), [SymSchur(t), _seqalg.SumLogDiag(0)]) 1a
147 return 2 * out 1a
150@jit 1bca
151def solve(t, b): 1bca
152 ops = [SymLevinson(t), _seqalg.MatMulRowByFull(0, b), _seqalg.MatMulColByRow(0, 1)] 1ba
153 _, _, out = _seqalg.sequential_algorithm(len(t), ops) 1ba
154 return out 1ba
157@jit 1bca
158def chol_transp_solve(t, b): 1bca
159 ops = [SymLevinson(t), _seqalg.Rows(b), _seqalg.MatMulColByRow(0, 1)]
160 _, _, out = _seqalg.sequential_algorithm(len(t), ops)
161 return out
164def chol_solve_numpy(t, b, diageps=None): 1bca
165 """
166 Solve a linear system for the cholesky factor of a symmetric Toeplitz matrix.
168 The algorithm is:
170 t[0] += diageps
171 m = toeplitz(t)
172 l = chol(m)
173 return solve(l, b)
175 Numpy object arrays are supported. Broadcasts like matmul.
177 Parameters
178 ----------
179 t : (..., n) array
180 The first row or column of the matrix.
181 b : (..., n, m) or (n,) array
182 The right hand side of the linear system.
183 diageps : scalar, optional
184 Term added to the diagonal elements of the matrix for regularization.
186 Returns
187 -------
188 x : (..., n, m) or (..., n) array
189 The solution ``solve(l, b)``.
190 """
191 t = numpy.array(t, subok=True) 1a
192 n = t.shape[-1] 1a
194 b = numpy.asanyarray(b) 1a
195 vec = b.ndim < 2 1a
196 if vec: 1a
197 b = b[:, None] 1a
198 assert b.shape[-2] == n 1a
200 t = t.astype(numpy.result_type(t, 0.1), copy=False) 1a
201 b = b.astype(numpy.result_type(b, 0.1), copy=False) 1a
203 if n == 0: 1a
204 shape = numpy.broadcast_shapes(t.shape[:-1], b.shape[:-2]) 1a
205 shape += (n,) if vec else b.shape[-2:] 1a
206 dtype = numpy.result_type(t.dtype, b.dtype) 1a
207 return numpy.empty(shape, dtype) 1a
209 if diageps is not None: 1a
210 t[..., 0] += diageps 1a
212 if numpy.any(t[..., 0] <= 0): 1a
213 msg = '1-th leading minor is not positive definite' 1a
214 raise numpy.linalg.LinAlgError(msg) 1a
216 norm = numpy.copy(t[..., 0, None], subok=True) 1a
217 t /= norm 1a
218 invLb = numpy.copy(numpy.broadcast_arrays(b, t[..., None])[0], subok=True) 1a
219 prevLi = t 1a
220 g = numpy.stack([numpy.roll(t, 1, -1), t], -2) 1a
222 for i in range(1, n): 1a
223 assert numpy.all(g[..., 0, i] > 0) 1a
224 rho = -g[..., 1, i, None, None] / g[..., 0, i, None, None] 1a
226 if numpy.any(numpy.abs(rho) >= 1): 1a
227 msg = f'{i + 1}-th leading minor is not positive definite' 1a
228 raise numpy.linalg.LinAlgError(msg) 1a
230 gamma = numpy.sqrt((1 - rho) * (1 + rho)) 1a
231 g[..., :, i:] += g[..., ::-1, i:] * rho 1a
232 g[..., :, i:] /= gamma 1a
233 Li = g[..., 0, i:] # i-th column of L from row i 1a
234 invLb[..., i:, :] -= invLb[..., i - 1, None, :] * prevLi[..., i:, None] 1a
235 invLb[..., i, :] /= Li[..., 0, None] 1a
236 prevLi[..., i:] = Li 1a
237 g[..., 0, i:] = numpy.roll(g[..., 0, i:], 1, -1) 1a
239 invLb /= numpy.sqrt(norm[..., None]) 1a
240 if vec: 1a
241 invLb = numpy.squeeze(invLb, -1) 1a
242 return invLb 1a
245def eigv_bound(t): 1bca
246 """
247 Bound the eigenvalues of a symmetric Toeplitz matrix.
249 Parameters
250 ----------
251 t : array
252 The first row of the matrix.
254 Returns
255 -------
256 m : scalar
257 Any eigenvalue `v` of the matrix satisfies `|v| <= m`.
258 """
259 s = jnp.abs(t) 1a
260 c = jnp.cumsum(s) 1a
261 d = c + c[::-1] - s[0] 1a
262 return jnp.max(d) 1a