Skip to content

Commit

Permalink
Pythranized version of _solve_toeplitz.pyx
Browse files Browse the repository at this point in the history
  • Loading branch information
serge-sans-paille committed Jun 15, 2023
1 parent 4d74135 commit 1f44e65
Show file tree
Hide file tree
Showing 3 changed files with 142 additions and 10 deletions.
113 changes: 113 additions & 0 deletions scipy/linalg/_solve_toeplitz.py
@@ -0,0 +1,113 @@
# Author: Robert T. McGibbon, December 2014

from numpy import zeros, sum

# FIXME: from numpy.linalg import LinAlgError when it's supported by Pythran
LinAlgError = RuntimeError

# FIXME: use '@' once pythran supports '@' without linking to blas.
def _dot(x, y):
return sum(x * y)


# pythran export levinson(float64 [], float64[])
# pythran export levinson(complex128 [], complex128[])
def levinson(a, 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.

dtype = a.dtype

(n,) = b.shape
x = zeros(n, dtype=dtype) # result
g = zeros(n, dtype=dtype) # workspace
h = zeros(n, dtype=dtype) # workspace
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 x, 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 = _dot(a[n + m - 1 : n - 1 : -1], x[:m]) - b[m]
x_den = _dot(a[n + m - 1 : n - 1 : -1], g[m - 1 :: -1]) - a[n - 1]
if x_den == 0:
raise LinAlgError("Singular principal minor")
x[m] = x_num / x_den
reflection_coeff[m + 1] = x[m]

# Compute x
x[:m] -= x[m] * g[m - 1 :: -1]
if m == n - 1:
return x, reflection_coeff

# Compute the numerator and denominator of g[m] and h[m]
g_num = _dot(a[n - m - 1 : n - 1], g[:m]) - a[n - m - 2]
h_num = _dot(a[n + m - 1 : n - 1 : -1], h[:m]) - a[n + m]
g_den = _dot(a[n - m - 1 : n - 1], h[m - 1 :: -1]) - a[n - 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) // 2
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
33 changes: 25 additions & 8 deletions scipy/linalg/meson.build
Expand Up @@ -176,14 +176,31 @@ py3.extension_module('_interpolative',
)

# _solve_toeplitz
py3.extension_module('_solve_toeplitz',
linalg_init_cython_gen.process('_solve_toeplitz.pyx'),
c_args: cython_c_args,
dependencies: np_dep,
link_args: version_link_args,
install: true,
subdir: 'scipy/linalg'
)
if use_pythran
_solve_toeplitz_cpp = custom_target('_solve_toeplitz',
output: ['_solve_toeplitz.cpp'],
input: '_solve_toeplitz.py',
command: [pythran, '-E', '@INPUT@', '-o', '@OUTDIR@/_solve_toeplitz.cpp']
)

py3.extension_module('_solve_toeplitz',
[_solve_toeplitz_cpp],
cpp_args: cpp_args_pythran,
dependencies: [pythran_dep, np_dep],
link_args: version_link_args,
install: true,
subdir: 'scipy/linalg'
)
else
py3.extension_module('_solve_toeplitz',
linalg_init_cython_gen.process('_solve_toeplitz.pyx'),
c_args: cython_c_args,
dependencies: np_dep,
link_args: version_link_args,
install: true,
subdir: 'scipy/linalg'
)
endif

# _matfuncs_sqrtm_triu:
py3.extension_module('_matfuncs_sqrtm_triu',
Expand Down
6 changes: 4 additions & 2 deletions scipy/linalg/tests/test_solve_toeplitz.py
Expand Up @@ -56,6 +56,8 @@ def test_native_list_arguments():
desired = solve(toeplitz(c, r=r), y)
assert_allclose(actual, desired)

# FIXME: use np.linalg.LinAlgError instead once pythran supports it
LinAlgErrors = (np.linalg.LinAlgError, RuntimeError)

def test_zero_diag_error():
# The Levinson-Durbin implementation fails when the diagonal is zero.
Expand All @@ -65,7 +67,7 @@ def test_zero_diag_error():
r = random.randn(n)
y = random.randn(n)
c[0] = 0
assert_raises(np.linalg.LinAlgError,
assert_raises(LinAlgErrors,
solve_toeplitz, (c, r), b=y)


Expand All @@ -75,7 +77,7 @@ def test_wikipedia_counterexample():
random = np.random.RandomState(1234)
c = [2, 2, 1]
y = random.randn(3)
assert_raises(np.linalg.LinAlgError, solve_toeplitz, c, b=y)
assert_raises(LinAlgErrors, solve_toeplitz, c, b=y)


def test_reflection_coeffs():
Expand Down

0 comments on commit 1f44e65

Please sign in to comment.