Coverage for src/lsqfitgp/_kernels/_arma.py: 97%

1# lsqfitgp/_kernels/_arma.py 

2# 

3# Copyright (c) 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/>. 

19 

20import numpy 3 ctx1bca

21from jax import jit, lax 3 ctx1bca

22from jax import numpy as jnp 3 ctx1bca

23 

24from lsqfitgp import _jaxext 3 ctx1bca

25from lsqfitgp._Kernel import stationarykernel 3 ctx1bca

26from lsqfitgp._linalg import _toeplitz 3 ctx1bca

27 

28 

29# use positive delta because negative indices wrap around 

30@stationarykernel(derivable=False, maxdim=1, input='abs') 3 ctx1bca

31def MA(delta, w=None, norm=False): 3 ctx1bca

32 r""" 

33 Discrete moving average kernel. 

34 

35 .. math:: 

36 k(\Delta) = \sum_{k=|\Delta|}^{n-1} w_k w_{k-|\Delta|}, 

37 \quad \mathbf w = (w_0, \ldots, w_{n-1}). 

38 

39 The inputs must be integers. It is the autocovariance function of a moving 

40 average with weights :math:`\mathbf w` applied to white noise: 

41 

42 .. math:: 

43 k(i, j) &= \operatorname{Cov}[y_i, y_j], \\ 

44 y_i &= \sum_{k=0}^{n-1} w_k \epsilon_{i-k}, \\ 

45 \operatorname{Cov}[\epsilon_i,\epsilon_j] &= \delta_{ij}. 

46 

47 If ``norm=True``, the variance is normalized to 1, which amounts to 

48 normalizing :math:`\mathbf w` to unit length. 

49 """ 

50 w = jnp.asarray(w) 2 ctx1ba

51 assert w.ndim == 1 2 ctx1ba

52 if len(w): 2 ctx1ba

53 cov = jnp.convolve(w, w[::-1]) 2 ctx1ba

54 if norm: 2 ctx1ba

55 cov /= cov[len(w) - 1] 1 ctx1a

56 return cov.at[delta + len(w) - 1].get(mode='fill', fill_value=0) 2 ctx1ba

57 else: 

58 return jnp.zeros(delta.shape) 1 ctx1a

59 

60 

61@stationarykernel(derivable=False, maxdim=1, input='abs') 3 ctx1bca

62def _ARBase(delta, phi=None, gamma=None, maxlag=None, slnr=None, lnc=None, norm=False): 3 ctx1bca

63 r""" 

64 Discrete autoregressive kernel. 

65 

66 You have to specify one and only one of the sets of parameters 

67 ``phi+maxlag``, ``gamma+maxlag``, ``slnr+lnc``. 

68 

69 Parameters 

70 ---------- 

71 phi : (p,) real 

72 The autoregressive coefficients at lag 1...p. 

73 gamma : (p + 1,) real 

74 The autocovariance at lag 0...p. 

75 maxlag : int 

76 The maximum lag that the kernel will be evaluated on. If the actual 

77 inputs produce higher lags, the missing values are filled with ``nan``. 

78 slnr : (nr,) real 

79 The real roots of the characteristic polynomial, expressed in the 

80 following way: ``sign(slnr)`` is the sign of the root, and 

81 ``abs(slnr)`` is the natural logarithm of the absolute value. 

82 lnc : (nc,) complex 

83 The natural logarithm of the complex roots of the characteristic 

84 polynomial (:math:`\log z = \log|z| + i\arg z`), where each root 

85 also stands for its paired conjugate. 

86 

87 In `slnr` and `lnc`, the multiplicity of a root is expressed by 

88 repeating the root in the array (not necessarily next to each other). 

89 Only exact repetition counts; very close yet distinct roots are treated 

90 as separate and lead to numerical instability, in particular complex 

91 roots very close to the real line. An exactly real complex root behaves 

92 like a pair of identical real roots. Two complex roots also count as 

93 equal if conjugate, and the argument is standardized to :math:`[0, 

94 2\pi)`. 

95 norm : bool, default False 

96 If True, normalize the autocovariance to be 1 at lag 0. If False, 

97 normalize such that the variance of the generating noise is 1, or use 

98 the user-provided normalization if `gamma` is specified. 

99 

100 Notes 

101 ----- 

102 This is the covariance function of a stationary autoregressive process, 

103 which is defined recursively as 

104 

105 .. math:: 

106 y_i = \sum_{k=1}^p \phi_k y_{i-k} + \epsilon_i, 

107 

108 where :math:`\epsilon_i` is white noise, i.e., 

109 :math:`\operatorname{Cov}[\epsilon_i, \epsilon_j] = \delta_{ij}`. The 

110 length :math:`p` of the vector of coefficients :math:`\boldsymbol\phi` 

111 is the "order" of the process. 

112 

113 The covariance function can be expressed in two ways. First as the same 

114 recursion defining the process: 

115 

116 .. math:: 

117 \gamma_m = \sum_{k=1}^p \phi_k \gamma_{m-k} + \delta_{m0}, 

118 

119 where :math:`\gamma_m \equiv \operatorname{Cov}[y_i, y_{i+m}]`. This is 

120 called "Yule-Walker equation." Second, as a linear combination of mixed 

121 power-exponentials: 

122 

123 .. math:: 

124 \gamma_m = \sum_{j=1}^n 

125 \sum_{l=1}^{\mu_j} 

126 a_{jl} |m|^{l-1} x_j^{-|m|}, 

127 

128 where :math:`x_j` and :math:`\mu_j` are the (complex) roots and 

129 corresponding multiplicities of the "characteristic polynomial" 

130 

131 .. math:: 

132 P(x) = 1 - \sum_{k=1}^p \phi_k x^k, 

133 

134 and the :math:`a_{jl}` are uniquely determined complex coefficients. The 

135 :math:`\boldsymbol\phi` vector is valid iff :math:`|x_j|>1, \forall j`. 

136 

137 There are three alternative parametrization for this kernel. 

138 

139 If you specify `phi`, the first terms of the covariance are computed 

140 solving the Yule-Walker equation, and then evolved up to `maxlag`. It 

141 is necessary to specify `maxlag` instead of letting the code figure it out 

142 from the actual inputs for technical reasons. 

143 

144 Likewise, if you specify `gamma`, the coefficients are obtained with 

145 Yule-Walker and then used to evolve the covariance. The only difference is 

146 that the normalization can be different: starting from `phi`, the variance 

147 of the generating noise :math:`\epsilon` is fixed to 1, while giving 

148 `gamma` directly implies an arbitrary value. 

149 

150 Instead, if you specify the roots with `slnr` and `lnc`, the coefficients 

151 are obtained from the polynomial defined in terms of the roots, and then 

152 the amplitudes :math:`a_{jl}` are computed by solving a linear system with 

153 the covariance (from YW) as RHS. Finally, the full covariance function is 

154 evaluated with the analytical expression. 

155 

156 The reasons for using the logarithm are that 1) in practice the roots are 

157 tipically close to 1, so the logarithm is numerically more accurate, and 2) 

158 the logarithm is readily interpretable as the inverse of the correlation 

159 length. 

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

161 cond = ( 2 ctx1ba

162 ( 

163 phi is not None 

164 and maxlag is not None 

165 and gamma is None 

166 and slnr is None 

167 and lnc is None 

168 ) 

169 or ( 

170 phi is None 

171 and maxlag is not None 

172 and gamma is not None 

173 and slnr is None 

174 and lnc is None 

175 ) 

176 or ( 

177 phi is None 

178 and maxlag is None 

179 and gamma is None 

180 and slnr is not None 

181 and lnc is not None 

182 ) 

183 ) 

184 if not cond: 2 ctx1ba

185 msg = 'invalid set of specified parameters' 1 ctx1a

186 raise ValueError(msg) 1 ctx1a

187 

188 if phi is None and gamma is None: 2 ctx1ba

189 return _ar_with_roots(delta, slnr, lnc, norm) 2 ctx1ba

190 else: 

191 return _ar_with_phigamma(delta, phi, gamma, maxlag, norm) 2 ctx1ba

192 

193 

194def _ar_with_phigamma(delta, phi, gamma, maxlag, norm): 3 ctx1bca

195 if phi is None: 2 ctx1ba

196 phi = AR.phi_from_gamma(gamma) 2 ctx1ba

197 if gamma is None: 2 ctx1ba

198 gamma = AR.gamma_from_phi(phi) 2 ctx1ba

199 if norm: 2 ctx1ba

200 gamma = gamma / gamma[0] 2 ctx1ba

201 acf = AR.extend_gamma(gamma, phi, maxlag + 1 - len(gamma)) 2 ctx1ba

202 return acf.at[delta].get(mode='fill', fill_value=jnp.nan) 2 ctx1ba

203 

204 

205def _yule_walker(gamma): 3 ctx1bca

206 """ 

207 Compute the autoregressive coefficients from the autocovariance. 

208 

209 `gamma` = autocovariance at lag 0...p 

210 output: autoregressive coefficients at lag 1...p. 

211 """ 

212 gamma = jnp.asarray(gamma) 2 ctx1ba

213 assert gamma.ndim == 1 2 ctx1ba

214 t = gamma[:-1] 2 ctx1ba

215 b = gamma[1:] 2 ctx1ba

216 if t.size: 2 ctx1ba

217 return _toeplitz.solve(t, b) 2 ctx1ba

218 else: 

219 return jnp.empty(0) 1 ctx1a

220 

221 

222def _yule_walker_inv_mat(phi): 3 ctx1bca

223 phi = jnp.asarray(phi) 2 ctx1ba

224 assert phi.ndim == 1 2 ctx1ba

225 p = len(phi) 2 ctx1ba

226 m = jnp.arange(p + 1)[:, None] # rows 2 ctx1ba

227 n = m.T # columns 2 ctx1ba

228 phi = jnp.pad(phi, (1, 1)) 2 ctx1ba

229 kp = jnp.clip(m + n, 0, p + 1) 2 ctx1ba

230 km = jnp.clip(m - n, 0, p + 1) 2 ctx1ba

231 return jnp.eye(p + 1) - (phi[kp] + phi[km]) / jnp.where(n, 1, 2) 2 ctx1ba

232 

233 

234def _yule_walker_inv(phi): 3 ctx1bca

235 """ 

236 Compute the autocovariance from the autoregressive coefficients. 

237 

238 `phi` = autoregressive coefficients at lag 1...p 

239 output: autocovariance at lag 0...p, assuming driving noise has sdev 1. 

240 """ 

241 a = _yule_walker_inv_mat(phi) 2 ctx1ba

242 b = jnp.zeros(len(a)).at[0].set(1) 2 ctx1ba

243 # gamma = _pseudo_solve(a, b) 

244 return jnp.linalg.solve(a, b) 2 ctx1ba

245 

246 

247def _ar_evolve(phi, start, noise): 3 ctx1bca

248 """ 

249 Evolve an autoregressive process. 

250 

251 `phi` = autoregressive coefficients at lag 1...p 

252 start = first p values of the process (increasing time) 

253 noise = n noise values added at each step 

254 output: n new process values. 

255 """ 

256 phi = jnp.asarray(phi) 2 ctx1ba

257 start = jnp.asarray(start) 2 ctx1ba

258 noise = jnp.asarray(noise) 2 ctx1ba

259 assert phi.ndim == 1 2 ctx1ba

260 assert phi.shape == start.shape 2 ctx1ba

261 assert noise.ndim == 1 2 ctx1ba

262 return _ar_evolve_jit(phi, start, noise) 2 ctx1ba

263 

264 

265@jit 3 ctx1bca

266def _ar_evolve_jit(phi, start, noise): 3 ctx1bca

267 

268 def f(carry, eps): 2 ctx1ba

269 vals, cc, roll = carry 2 ctx1ba

270 phi = lax.dynamic_slice(cc, [vals.size - roll], [vals.size]) 2 ctx1ba

271 nextval = phi @ vals + eps 2 ctx1ba

272 if vals.size: 2 ctx1ba

273 vals = vals.at[roll].set(nextval) 2 ctx1ba

274 # maybe for some weird reason like alignment, actual rolling would 

275 # be faster. whatever 

276 roll = (roll + 1) % vals.size 2 ctx1ba

277 return (vals, cc, roll), nextval 2 ctx1ba

278 

279 cc = jnp.concatenate([phi, phi])[::-1] 2 ctx1ba

280 _, ev = lax.scan(f, (start, cc, 0), noise, unroll=16) 2 ctx1ba

281 return ev 2 ctx1ba

282 

283 

284def _ar_with_roots(delta, slnr, lnc, norm): 3 ctx1bca

285 phi = AR.phi_from_roots(slnr, lnc) # <---- weak 2 ctx1ba

286 gamma = AR.gamma_from_phi(phi) # <---- point 2 ctx1ba

287 if norm: 2 ctx1ba

288 gamma /= gamma[0] 2 ctx1ba

289 ampl = AR.ampl_from_roots(slnr, lnc, gamma) 2 ctx1ba

290 return AR.cov_from_ampl(slnr, lnc, ampl, delta) 2 ctx1ba

291 

292 # Is numerical integration of the spectrum a feasible way to get the 

293 # covariance? The roots correspond to peaks, and they get very high as the 

294 # roots get close to 1. But I know where the peaks are in advance => nope 

295 # because the e^iwx oscillates arbitrarily fast. However maybe I can compute 

296 # the first p terms, which solves my current problem with gamma. I guess I 

297 # just have to use a multiple of p of quadrature points. The spectrum 

298 # oscillates too but only up to mode p. The total calculation cost is then 

299 # O(p^2), better than current O(p^3). See Hamilton (1994, p. 155). 

300 

301 # Other solution (Hamilton p. 319): the covariance should be equal to the 

302 # impulse response, so I can get gamma from phi by an evolution starting 

303 # from zeros. => Nope, it's equal only for AR(1). 

304 

305 # condition for phi: in the region phi >= 0, it must be sum(phi) <= 1 

306 # (Hamilton p. 659). 

307 

308 # p = phi.size 

309 # yw = _yule_walker_inv_mat(phi) 

310 # b = jnp.zeros(p + 1).at[0].set(1) 

311 # ampl = jnp.linalg.solve(yw @ mat, b) 

312 # lag = delta if delta.ndim else delta[None] 

313 # acf = _gamma_from_ampl_matmul(slnr, lnc, lag, ampl) 

314 # if norm: 

315 # acf0 = _gamma_from_ampl_matmul(slnr, lnc, jnp.array([0]), ampl) 

316 # acf /= acf0 

317 # return acf if delta.ndim else acf.squeeze(0) 

318 

319 

320def _pseudo_solve(a, b): 3 ctx1bca

321 # this is less accurate than jnp.linalg.solve 

322 u, s, vh = jnp.linalg.svd(a) 2 ctx1ba

323 eps = jnp.finfo(a.dtype).eps 2 ctx1ba

324 s0 = s[0] if s.size else 0 2 ctx1ba

325 invs = jnp.where(s < s0 * eps * len(a), 0, 1 / s) 2 ctx1ba

326 return jnp.einsum('ij,j,jk,k', vh.conj().T, invs, u.conj().T, b) 2 ctx1ba

327 

328 

329@jit 3 ctx1bca

330def _gamma_from_ampl_matmul(slnr, lnc, lag, ampl, lagnorm=None): 3 ctx1bca

331 

332 vec = ampl.ndim == 1 2 ctx1ba

333 if vec: 2 ctx1ba

334 ampl = ampl[:, None] 2 ctx1ba

335 p = slnr.size + 2 * lnc.size 2 ctx1ba

336 assert ampl.shape[-2] == p + 1 2 ctx1ba

337 if lagnorm is None: 337 ↛ 340line 337 didn't jump to line 340 because the condition on line 337 was always true2 ctx1ba

338 lagnorm = p 2 ctx1ba

339 

340 def logcol(root, lag, llag, repeat): 2 ctx1ba

341 return -root * lag + jnp.where(repeat, repeat * llag, 0) 2 ctx1ba

342 

343 def lognorm(root, repeat, lagnorm): 2 ctx1ba

344 maxnorm = jnp.where(repeat, repeat * (-1 + jnp.log(repeat / root)), 0) 2 ctx1ba

345 defnorm = logcol(root, lagnorm, jnp.log(lagnorm), repeat) 2 ctx1ba

346 maxloc = repeat / root 2 ctx1ba

347 return jnp.where(maxloc <= lagnorm, maxnorm, defnorm) 2 ctx1ba

348 

349 # roots at infinity 

350 col = jnp.where(lag, 0, 1) 2 ctx1ba

351 out = col[..., :, None] * ampl[..., 0, :] 2 ctx1ba

352 

353 # real roots 

354 llag = jnp.log(lag) 2 ctx1ba

355 val = (jnp.nan, 0, out, slnr, lag, llag, lagnorm) 2 ctx1ba

356 

357 def loop(i, val): 2 ctx1ba

358 prevroot, repeat, out, slnr, lag, llag, lagnorm = val 2 ctx1ba

359 root = slnr[i] 2 ctx1ba

360 repeat = jnp.where(root == prevroot, repeat + 1, 0) 2 ctx1ba

361 prevroot = root 2 ctx1ba

362 sign = jnp.sign(root) ** lag 2 ctx1ba

363 aroot = jnp.abs(root) 2 ctx1ba

364 lcol = logcol(aroot, lag, llag, repeat) 2 ctx1ba

365 norm = lognorm(aroot, repeat, lagnorm) 2 ctx1ba

366 col = sign * jnp.exp(lcol - norm) 2 ctx1ba

367 out += col[..., :, None] * ampl[..., 1 + i, :] 2 ctx1ba

368 return prevroot, repeat, out, slnr, lag, llag, lagnorm 2 ctx1ba

369 

370 if slnr.size: 2 ctx1ba

371 _, _, out, _, _, _, _ = lax.fori_loop(0, slnr.size, loop, val) 2 ctx1ba

372 

373 # complex roots 

374 val = (jnp.nan, 0, out, lnc, lag, llag, lagnorm) 2 ctx1ba

375 

376 def loop(i, val): 2 ctx1ba

377 prevroot, repeat, out, lnc, lag, llag, lagnorm = val 2 ctx1ba

378 root = lnc[i] 2 ctx1ba

379 repeat = jnp.where(root == prevroot, repeat + 1, 0) 2 ctx1ba

380 prevroot = root 2 ctx1ba

381 lcol = logcol(root, lag, llag, repeat) 2 ctx1ba

382 norm = lognorm(root.real, repeat, lagnorm) 2 ctx1ba

383 col = jnp.exp(lcol - norm) 2 ctx1ba

384 idx = 1 + slnr.size + 2 * i 2 ctx1ba

385 out += col.real[..., :, None] * ampl[..., idx, :] 2 ctx1ba

386 

387 # real complex root = a pair of identical real roots 

388 repeat = jnp.where(root.imag, repeat, repeat + 1) 2 ctx1ba

389 col1 = jnp.where(root.imag, -col.imag, col.real * lag) 2 ctx1ba

390 out += col1[..., :, None] * ampl[..., idx + 1, :] 2 ctx1ba

391 

392 return prevroot, repeat, out, lnc, lag, llag, lagnorm 2 ctx1ba

393 

394 if lnc.size: 2 ctx1ba

395 _, _, out, _, _, _, _ = lax.fori_loop(0, lnc.size, loop, val) 2 ctx1ba

396 

397 if vec: 2 ctx1ba

398 out = out.squeeze(-1) 2 ctx1ba

399 

400 return out 2 ctx1ba

401 

402 

403class AR(_ARBase): 3 ctx1bca

404 __doc__ = _ARBase.__doc__ 3 ctx1bca

405 

406 @classmethod 3 ctx1bca

407 def phi_from_gamma(cls, gamma): 3 ctx1bca

408 """ 

409 Determine the autoregressive coefficients from the covariance. 

410 

411 Parameters 

412 ---------- 

413 gamma : (p + 1,) array 

414 The autocovariance at lag 0...p. 

415 

416 Returns 

417 ------- 

418 phi : (p,) array 

419 The autoregressive coefficients at lag 1...p. 

420 """ 

421 gamma = cls._process_gamma(gamma) 2 ctx1ba

422 return _yule_walker(gamma) 2 ctx1ba

423 

424 @classmethod 3 ctx1bca

425 def gamma_from_phi(cls, phi): 3 ctx1bca

426 """ 

427 Determine the covariance from the autoregressive coefficients. 

428 

429 Parameters 

430 ---------- 

431 phi : (p,) array 

432 The autoregressive coefficients at lag 1...p. 

433 

434 Returns 

435 ------- 

436 gamma : (p + 1,) array 

437 The autocovariance at lag 0...p. The normalization is 

438 with noise variance 1. 

439 

440 Notes 

441 ----- 

442 The result is wildly inaccurate for roots with high multiplicity and/or 

443 close to 1. 

444 """ 

445 phi = cls._process_phi(phi) 2 ctx1ba

446 return _yule_walker_inv(phi) 2 ctx1ba

447 

448 @classmethod 3 ctx1bca

449 def extend_gamma(cls, gamma, phi, n): 3 ctx1bca

450 """ 

451 Extend values of the covariance function to higher lags. 

452 

453 Parameters 

454 ---------- 

455 gamma : (m,) array 

456 The autocovariance at lag q-m+1...q, with q >= 0 and m >= p + 1. 

457 phi : (p,) array 

458 The autoregressive coefficients at lag 1...p. 

459 n : int 

460 The number of new values to generate. 

461 

462 Returns 

463 ------- 

464 ext : (m + n,) array 

465 The autocovariance at lag q-m+1...q+n. 

466 """ 

467 gamma = cls._process_gamma(gamma) 2 ctx1ba

468 phi = cls._process_phi(phi) 2 ctx1ba

469 assert gamma.size > phi.size 2 ctx1ba

470 ext = _ar_evolve( 2 ctx1ba

471 phi, gamma[len(gamma) - len(phi) :], jnp.broadcast_to(0.0, (n,)) 

472 ) 

473 return jnp.concatenate([gamma, ext]) 2 ctx1ba

474 

475 @classmethod 3 ctx1bca

476 def phi_from_roots(cls, slnr, lnc): 3 ctx1bca

477 r""" 

478 Determine the AR coefficients from the roots of the characteristic polynomial. 

479 

480 Parameters 

481 ---------- 

482 slnr : (nr,) real 

483 The real roots of the characteristic polynomial, expressed in the 

484 following way: ``sign(slnr)`` is the sign of the root, and 

485 ``abs(slnr)`` is the natural logarithm of the absolute value. 

486 lnc : (nc,) complex 

487 The natural logarithm of the complex roots of the characteristic 

488 polynomial (:math:`\log z = \log|z| + i\arg z`), where each root 

489 also stands for its paired conjugate. 

490 

491 Returns 

492 ------- 

493 phi : (p,) real 

494 The autoregressive coefficients at lag 1...p, with p = nr + 2 nc. 

495 """ 

496 slnr, lnc = cls._process_roots(slnr, lnc) 2 ctx1ba

497 r = jnp.copysign(jnp.exp(-jnp.abs(slnr)), slnr) # works with +/-0 2 ctx1ba

498 c = jnp.exp(-lnc) 2 ctx1ba

499 

500 # minus sign in the exponentials to do 1/z, the poly output is already 

501 # reversed 

502 

503 roots = jnp.concatenate([r, c, c.conj()]).sort() # <-- polyroots sorts 2 ctx1ba

504 coef = jnp.atleast_1d(jnp.poly(roots)) 2 ctx1ba

505 

506 if coef.size: 506 ↛ 510line 506 didn't jump to line 510 because the condition on line 506 was always true2 ctx1ba

507 with _jaxext.skipifabstract(): 2 ctx1ba

508 numpy.testing.assert_equal(coef[0].item(), 1) 2 ctx1ba

509 numpy.testing.assert_allclose(jnp.imag(coef), 0, rtol=0, atol=1e-4) 2 ctx1ba

510 return -coef.real[1:] 2 ctx1ba

511 

512 @classmethod 3 ctx1bca

513 def ampl_from_roots(cls, slnr, lnc, gamma): 3 ctx1bca

514 slnr, lnc = cls._process_roots(slnr, lnc) 2 ctx1ba

515 gamma = cls._process_gamma(gamma) 2 ctx1ba

516 assert gamma.size == 1 + slnr.size + 2 * lnc.size 2 ctx1ba

517 lag = jnp.arange(gamma.size) 2 ctx1ba

518 mat = _gamma_from_ampl_matmul(slnr, lnc, lag, jnp.eye(gamma.size)) 2 ctx1ba

519 # return jnp.linalg.solve(mat, gamma) 

520 return _pseudo_solve(mat, gamma) 2 ctx1ba

521 

522 @classmethod 3 ctx1bca

523 def cov_from_ampl(cls, slnr, lnc, ampl, lag): 3 ctx1bca

524 slnr, lnc = cls._process_roots(slnr, lnc) 2 ctx1ba

525 ampl = cls._process_ampl(ampl) 2 ctx1ba

526 assert ampl.size == 1 + slnr.size + 2 * lnc.size 2 ctx1ba

527 lag = cls._process_lag(lag) 2 ctx1ba

528 scalar = lag.ndim == 0 2 ctx1ba

529 if scalar: 529 ↛ 530line 529 didn't jump to line 530 because the condition on line 529 was never true2 ctx1ba

530 lag = lag[None] 

531 acf = _gamma_from_ampl_matmul(slnr, lnc, lag, ampl) 2 ctx1ba

532 return acf.squeeze(0) if scalar else acf 2 ctx1ba

533 

534 @classmethod 3 ctx1bca

535 def inverse_roots_from_phi(cls, phi): 3 ctx1bca

536 phi = cls._process_phi(phi) 

537 poly = jnp.concatenate([jnp.ones(1), -phi]) 

538 return jnp.roots(poly, strip_zeros=False) 

539 

540 @staticmethod 3 ctx1bca

541 def _process_roots(slnr, lnc): 3 ctx1bca

542 slnr = jnp.asarray(slnr, float).sort() 2 ctx1ba

543 lnc = jnp.asarray(lnc, complex) 2 ctx1ba

544 assert slnr.ndim == lnc.ndim == 1 2 ctx1ba

545 imag = jnp.abs(lnc.imag) % (2 * jnp.pi) 2 ctx1ba

546 imag = jnp.where(imag > jnp.pi, 2 * jnp.pi - imag, imag) 2 ctx1ba

547 lnc = lnc.real + 1j * imag 2 ctx1ba

548 lnc = lnc.sort() 2 ctx1ba

549 return slnr, lnc 2 ctx1ba

550 

551 @staticmethod 3 ctx1bca

552 def _process_gamma(gamma): 3 ctx1bca

553 gamma = jnp.asarray(gamma, float) 2 ctx1ba

554 assert gamma.ndim == 1 2 ctx1ba

555 assert gamma.size >= 1 2 ctx1ba

556 return gamma 2 ctx1ba

557 

558 @staticmethod 3 ctx1bca

559 def _process_phi(phi): 3 ctx1bca

560 phi = jnp.asarray(phi, float) 2 ctx1ba

561 assert phi.ndim == 1 2 ctx1ba

562 return phi 2 ctx1ba

563 

564 @staticmethod 3 ctx1bca

565 def _process_ampl(ampl): 3 ctx1bca

566 ampl = jnp.asarray(ampl, float) 2 ctx1ba

567 assert ampl.ndim == 1 2 ctx1ba

568 assert ampl.size >= 1 2 ctx1ba

569 return ampl 2 ctx1ba

570 

571 @staticmethod 3 ctx1bca

572 def _process_lag(lag): 3 ctx1bca

573 lag = jnp.asarray(lag) 2 ctx1ba

574 assert jnp.issubdtype(lag.dtype, jnp.integer) 2 ctx1ba

575 return lag.astype(int) 2 ctx1ba