forked from scipy/scipy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
_solve_toeplitz.pyx
128 lines (108 loc) · 3.79 KB
/
_solve_toeplitz.pyx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
# Author: Robert T. McGibbon, December 2014
#
# cython: boundscheck=False, wraparound=False, cdivision=True
from __future__ import absolute_import
from numpy import zeros, asarray, complex128, float64
from numpy.linalg import LinAlgError
from numpy cimport complex128_t, float64_t
cdef fused dz:
float64_t
complex128_t
def levinson(dz[::1] a, dz[::1] b):
"""Solve a linear Toeplitz system using Levinson recursion.
Parameters
----------
a : array, dtype=double or complex128, shape=(2n-1,)
The first column of the matrix in reverse order (without the diagonal)
followed by the first (see below)
b : array, dtype=double or complex128, shape=(n,)
The right hand side vector. Both a and b must have the same type
(double or complex128).
Notes
-----
For example, the 5x5 toeplitz matrix below should be represented as
the linear array ``a`` on the right ::
[ a0 a1 a2 a3 a4 ]
[ a-1 a0 a1 a2 a3 ]
[ a-2 a-1 a0 a1 a2 ] -> [a-4 a-3 a-2 a-1 a0 a1 a2 a3 a4]
[ a-3 a-2 a-1 a0 a1 ]
[ a-4 a-3 a-2 a-1 a0 ]
Returns
-------
x : arrray, shape=(n,)
The solution vector
reflection_coeff : array, shape=(n+1,)
Toeplitz reflection coefficients. When a is symmetric Toeplitz and
``b`` is ``a[n:]``, as in the solution of autoregressive systems,
then ``reflection_coeff`` also correspond to the partial
autocorrelation function.
"""
# Adapted from toeplitz.f90 by Alan Miller, accessed at
# http://jblevins.org/mirror/amiller/toeplitz.f90
# Released under a Public domain declaration.
if dz is float64_t:
dtype = float64
else:
dtype = complex128
cdef ssize_t n, m, j, nmj, k, m2
n = b.shape[0]
cdef dz x_num, g_num, h_num, x_den, g_den
cdef dz gj, gk, hj, hk, c1, c2
cdef dz[:] x = zeros(n, dtype=dtype) # result
cdef dz[:] g = zeros(n, dtype=dtype) # workspace
cdef dz[:] h = zeros(n, dtype=dtype) # workspace
cdef dz[:] reflection_coeff = zeros(n+1, dtype=dtype) # history
assert len(a) == (2*n) - 1
if a[n-1] == 0:
raise LinAlgError('Singular principal minor')
x[0] = b[0] / a[n-1]
reflection_coeff[0] = 1
reflection_coeff[1] = x[0]
if (n == 1):
return asarray(x), asarray(reflection_coeff)
g[0] = a[n-2] / a[n-1]
h[0] = a[n] / a[n-1]
for m in range(1, n):
# Compute numerator and denominator of x[m]
x_num = -b[m]
x_den = -a[n-1]
for j in range(m):
nmj = n + m - (j+1)
x_num = x_num + a[nmj] * x[j]
x_den = x_den + a[nmj] * g[m-j-1]
if x_den == 0:
raise LinAlgError('Singular principal minor')
x[m] = x_num / x_den
reflection_coeff[m+1] = x[m]
# Compute x
for j in range(m):
x[j] = x[j] - x[m] * g[m-j-1]
if m == n-1:
return asarray(x), asarray(reflection_coeff)
# Compute the numerator and denominator of g[m] and h[m]
g_num = -a[n-m-2]
h_num = -a[n+m]
g_den = -a[n-1]
for j in range(m):
g_num = g_num + a[n+j-m-1] * g[j]
h_num = h_num + a[n+m-j-1] * h[j]
g_den = g_den + a[n+j-m-1] * h[m-j-1]
if g_den == 0.0:
raise LinAlgError("Singular principal minor")
# Compute g and h
g[m] = g_num / g_den
h[m] = h_num / x_den
k = m - 1
m2 = (m + 1) >> 1
c1 = g[m]
c2 = h[m]
for j in range(m2):
gj = g[j]
gk = g[k]
hj = h[j]
hk = h[k]
g[j] = gj - (c1 * hk)
g[k] = gk - (c1 * hj)
h[j] = hj - (c2 * gk)
h[k] = hk - (c2 * gj)
k -= 1