301 lines
10 KiB
Python
301 lines
10 KiB
Python
import numpy as np
|
|
from numpy.testing import (assert_allclose, assert_equal,
|
|
assert_almost_equal, assert_array_equal,
|
|
assert_array_almost_equal)
|
|
|
|
from scipy.ndimage import convolve1d
|
|
|
|
from scipy.signal import savgol_coeffs, savgol_filter
|
|
from scipy.signal._savitzky_golay import _polyder
|
|
|
|
|
|
def check_polyder(p, m, expected):
|
|
dp = _polyder(p, m)
|
|
assert_array_equal(dp, expected)
|
|
|
|
|
|
def test_polyder():
|
|
cases = [
|
|
([5], 0, [5]),
|
|
([5], 1, [0]),
|
|
([3, 2, 1], 0, [3, 2, 1]),
|
|
([3, 2, 1], 1, [6, 2]),
|
|
([3, 2, 1], 2, [6]),
|
|
([3, 2, 1], 3, [0]),
|
|
([[3, 2, 1], [5, 6, 7]], 0, [[3, 2, 1], [5, 6, 7]]),
|
|
([[3, 2, 1], [5, 6, 7]], 1, [[6, 2], [10, 6]]),
|
|
([[3, 2, 1], [5, 6, 7]], 2, [[6], [10]]),
|
|
([[3, 2, 1], [5, 6, 7]], 3, [[0], [0]]),
|
|
]
|
|
for p, m, expected in cases:
|
|
check_polyder(np.array(p).T, m, np.array(expected).T)
|
|
|
|
|
|
#--------------------------------------------------------------------
|
|
# savgol_coeffs tests
|
|
#--------------------------------------------------------------------
|
|
|
|
def alt_sg_coeffs(window_length, polyorder, pos):
|
|
"""This is an alternative implementation of the SG coefficients.
|
|
|
|
It uses numpy.polyfit and numpy.polyval. The results should be
|
|
equivalent to those of savgol_coeffs(), but this implementation
|
|
is slower.
|
|
|
|
window_length should be odd.
|
|
|
|
"""
|
|
if pos is None:
|
|
pos = window_length // 2
|
|
t = np.arange(window_length)
|
|
unit = (t == pos).astype(int)
|
|
h = np.polyval(np.polyfit(t, unit, polyorder), t)
|
|
return h
|
|
|
|
|
|
def test_sg_coeffs_trivial():
|
|
# Test a trivial case of savgol_coeffs: polyorder = window_length - 1
|
|
h = savgol_coeffs(1, 0)
|
|
assert_allclose(h, [1])
|
|
|
|
h = savgol_coeffs(3, 2)
|
|
assert_allclose(h, [0, 1, 0], atol=1e-10)
|
|
|
|
h = savgol_coeffs(5, 4)
|
|
assert_allclose(h, [0, 0, 1, 0, 0], atol=1e-10)
|
|
|
|
h = savgol_coeffs(5, 4, pos=1)
|
|
assert_allclose(h, [0, 0, 0, 1, 0], atol=1e-10)
|
|
|
|
h = savgol_coeffs(5, 4, pos=1, use='dot')
|
|
assert_allclose(h, [0, 1, 0, 0, 0], atol=1e-10)
|
|
|
|
|
|
def compare_coeffs_to_alt(window_length, order):
|
|
# For the given window_length and order, compare the results
|
|
# of savgol_coeffs and alt_sg_coeffs for pos from 0 to window_length - 1.
|
|
# Also include pos=None.
|
|
for pos in [None] + list(range(window_length)):
|
|
h1 = savgol_coeffs(window_length, order, pos=pos, use='dot')
|
|
h2 = alt_sg_coeffs(window_length, order, pos=pos)
|
|
assert_allclose(h1, h2, atol=1e-10,
|
|
err_msg=("window_length = %d, order = %d, pos = %s" %
|
|
(window_length, order, pos)))
|
|
|
|
|
|
def test_sg_coeffs_compare():
|
|
# Compare savgol_coeffs() to alt_sg_coeffs().
|
|
for window_length in range(1, 8, 2):
|
|
for order in range(window_length):
|
|
compare_coeffs_to_alt(window_length, order)
|
|
|
|
|
|
def test_sg_coeffs_exact():
|
|
polyorder = 4
|
|
window_length = 9
|
|
halflen = window_length // 2
|
|
|
|
x = np.linspace(0, 21, 43)
|
|
delta = x[1] - x[0]
|
|
|
|
# The data is a cubic polynomial. We'll use an order 4
|
|
# SG filter, so the filtered values should equal the input data
|
|
# (except within half window_length of the edges).
|
|
y = 0.5 * x ** 3 - x
|
|
h = savgol_coeffs(window_length, polyorder)
|
|
y0 = convolve1d(y, h)
|
|
assert_allclose(y0[halflen:-halflen], y[halflen:-halflen])
|
|
|
|
# Check the same input, but use deriv=1. dy is the exact result.
|
|
dy = 1.5 * x ** 2 - 1
|
|
h = savgol_coeffs(window_length, polyorder, deriv=1, delta=delta)
|
|
y1 = convolve1d(y, h)
|
|
assert_allclose(y1[halflen:-halflen], dy[halflen:-halflen])
|
|
|
|
# Check the same input, but use deriv=2. d2y is the exact result.
|
|
d2y = 3.0 * x
|
|
h = savgol_coeffs(window_length, polyorder, deriv=2, delta=delta)
|
|
y2 = convolve1d(y, h)
|
|
assert_allclose(y2[halflen:-halflen], d2y[halflen:-halflen])
|
|
|
|
|
|
def test_sg_coeffs_deriv():
|
|
# The data in `x` is a sampled parabola, so using savgol_coeffs with an
|
|
# order 2 or higher polynomial should give exact results.
|
|
i = np.array([-2.0, 0.0, 2.0, 4.0, 6.0])
|
|
x = i ** 2 / 4
|
|
dx = i / 2
|
|
d2x = np.full_like(i, 0.5)
|
|
for pos in range(x.size):
|
|
coeffs0 = savgol_coeffs(5, 3, pos=pos, delta=2.0, use='dot')
|
|
assert_allclose(coeffs0.dot(x), x[pos], atol=1e-10)
|
|
coeffs1 = savgol_coeffs(5, 3, pos=pos, delta=2.0, use='dot', deriv=1)
|
|
assert_allclose(coeffs1.dot(x), dx[pos], atol=1e-10)
|
|
coeffs2 = savgol_coeffs(5, 3, pos=pos, delta=2.0, use='dot', deriv=2)
|
|
assert_allclose(coeffs2.dot(x), d2x[pos], atol=1e-10)
|
|
|
|
|
|
def test_sg_coeffs_deriv_gt_polyorder():
|
|
"""
|
|
If deriv > polyorder, the coefficients should be all 0.
|
|
This is a regression test for a bug where, e.g.,
|
|
savgol_coeffs(5, polyorder=1, deriv=2)
|
|
raised an error.
|
|
"""
|
|
coeffs = savgol_coeffs(5, polyorder=1, deriv=2)
|
|
assert_array_equal(coeffs, np.zeros(5))
|
|
coeffs = savgol_coeffs(7, polyorder=4, deriv=6)
|
|
assert_array_equal(coeffs, np.zeros(7))
|
|
|
|
|
|
def test_sg_coeffs_large():
|
|
# Test that for large values of window_length and polyorder the array of
|
|
# coefficients returned is symmetric. The aim is to ensure that
|
|
# no potential numeric overflow occurs.
|
|
coeffs0 = savgol_coeffs(31, 9)
|
|
assert_array_almost_equal(coeffs0, coeffs0[::-1])
|
|
coeffs1 = savgol_coeffs(31, 9, deriv=1)
|
|
assert_array_almost_equal(coeffs1, -coeffs1[::-1])
|
|
|
|
|
|
#--------------------------------------------------------------------
|
|
# savgol_filter tests
|
|
#--------------------------------------------------------------------
|
|
|
|
|
|
def test_sg_filter_trivial():
|
|
""" Test some trivial edge cases for savgol_filter()."""
|
|
x = np.array([1.0])
|
|
y = savgol_filter(x, 1, 0)
|
|
assert_equal(y, [1.0])
|
|
|
|
# Input is a single value. With a window length of 3 and polyorder 1,
|
|
# the value in y is from the straight-line fit of (-1,0), (0,3) and
|
|
# (1, 0) at 0. This is just the average of the three values, hence 1.0.
|
|
x = np.array([3.0])
|
|
y = savgol_filter(x, 3, 1, mode='constant')
|
|
assert_almost_equal(y, [1.0], decimal=15)
|
|
|
|
x = np.array([3.0])
|
|
y = savgol_filter(x, 3, 1, mode='nearest')
|
|
assert_almost_equal(y, [3.0], decimal=15)
|
|
|
|
x = np.array([1.0] * 3)
|
|
y = savgol_filter(x, 3, 1, mode='wrap')
|
|
assert_almost_equal(y, [1.0, 1.0, 1.0], decimal=15)
|
|
|
|
|
|
def test_sg_filter_basic():
|
|
# Some basic test cases for savgol_filter().
|
|
x = np.array([1.0, 2.0, 1.0])
|
|
y = savgol_filter(x, 3, 1, mode='constant')
|
|
assert_allclose(y, [1.0, 4.0 / 3, 1.0])
|
|
|
|
y = savgol_filter(x, 3, 1, mode='mirror')
|
|
assert_allclose(y, [5.0 / 3, 4.0 / 3, 5.0 / 3])
|
|
|
|
y = savgol_filter(x, 3, 1, mode='wrap')
|
|
assert_allclose(y, [4.0 / 3, 4.0 / 3, 4.0 / 3])
|
|
|
|
|
|
def test_sg_filter_2d():
|
|
x = np.array([[1.0, 2.0, 1.0],
|
|
[2.0, 4.0, 2.0]])
|
|
expected = np.array([[1.0, 4.0 / 3, 1.0],
|
|
[2.0, 8.0 / 3, 2.0]])
|
|
y = savgol_filter(x, 3, 1, mode='constant')
|
|
assert_allclose(y, expected)
|
|
|
|
y = savgol_filter(x.T, 3, 1, mode='constant', axis=0)
|
|
assert_allclose(y, expected.T)
|
|
|
|
|
|
def test_sg_filter_interp_edges():
|
|
# Another test with low degree polynomial data, for which we can easily
|
|
# give the exact results. In this test, we use mode='interp', so
|
|
# savgol_filter should match the exact solution for the entire data set,
|
|
# including the edges.
|
|
t = np.linspace(-5, 5, 21)
|
|
delta = t[1] - t[0]
|
|
# Polynomial test data.
|
|
x = np.array([t,
|
|
3 * t ** 2,
|
|
t ** 3 - t])
|
|
dx = np.array([np.ones_like(t),
|
|
6 * t,
|
|
3 * t ** 2 - 1.0])
|
|
d2x = np.array([np.zeros_like(t),
|
|
np.full_like(t, 6),
|
|
6 * t])
|
|
|
|
window_length = 7
|
|
|
|
y = savgol_filter(x, window_length, 3, axis=-1, mode='interp')
|
|
assert_allclose(y, x, atol=1e-12)
|
|
|
|
y1 = savgol_filter(x, window_length, 3, axis=-1, mode='interp',
|
|
deriv=1, delta=delta)
|
|
assert_allclose(y1, dx, atol=1e-12)
|
|
|
|
y2 = savgol_filter(x, window_length, 3, axis=-1, mode='interp',
|
|
deriv=2, delta=delta)
|
|
assert_allclose(y2, d2x, atol=1e-12)
|
|
|
|
# Transpose everything, and test again with axis=0.
|
|
|
|
x = x.T
|
|
dx = dx.T
|
|
d2x = d2x.T
|
|
|
|
y = savgol_filter(x, window_length, 3, axis=0, mode='interp')
|
|
assert_allclose(y, x, atol=1e-12)
|
|
|
|
y1 = savgol_filter(x, window_length, 3, axis=0, mode='interp',
|
|
deriv=1, delta=delta)
|
|
assert_allclose(y1, dx, atol=1e-12)
|
|
|
|
y2 = savgol_filter(x, window_length, 3, axis=0, mode='interp',
|
|
deriv=2, delta=delta)
|
|
assert_allclose(y2, d2x, atol=1e-12)
|
|
|
|
|
|
def test_sg_filter_interp_edges_3d():
|
|
# Test mode='interp' with a 3-D array.
|
|
t = np.linspace(-5, 5, 21)
|
|
delta = t[1] - t[0]
|
|
x1 = np.array([t, -t])
|
|
x2 = np.array([t ** 2, 3 * t ** 2 + 5])
|
|
x3 = np.array([t ** 3, 2 * t ** 3 + t ** 2 - 0.5 * t])
|
|
dx1 = np.array([np.ones_like(t), -np.ones_like(t)])
|
|
dx2 = np.array([2 * t, 6 * t])
|
|
dx3 = np.array([3 * t ** 2, 6 * t ** 2 + 2 * t - 0.5])
|
|
|
|
# z has shape (3, 2, 21)
|
|
z = np.array([x1, x2, x3])
|
|
dz = np.array([dx1, dx2, dx3])
|
|
|
|
y = savgol_filter(z, 7, 3, axis=-1, mode='interp', delta=delta)
|
|
assert_allclose(y, z, atol=1e-10)
|
|
|
|
dy = savgol_filter(z, 7, 3, axis=-1, mode='interp', deriv=1, delta=delta)
|
|
assert_allclose(dy, dz, atol=1e-10)
|
|
|
|
# z has shape (3, 21, 2)
|
|
z = np.array([x1.T, x2.T, x3.T])
|
|
dz = np.array([dx1.T, dx2.T, dx3.T])
|
|
|
|
y = savgol_filter(z, 7, 3, axis=1, mode='interp', delta=delta)
|
|
assert_allclose(y, z, atol=1e-10)
|
|
|
|
dy = savgol_filter(z, 7, 3, axis=1, mode='interp', deriv=1, delta=delta)
|
|
assert_allclose(dy, dz, atol=1e-10)
|
|
|
|
# z has shape (21, 3, 2)
|
|
z = z.swapaxes(0, 1).copy()
|
|
dz = dz.swapaxes(0, 1).copy()
|
|
|
|
y = savgol_filter(z, 7, 3, axis=0, mode='interp', delta=delta)
|
|
assert_allclose(y, z, atol=1e-10)
|
|
|
|
dy = savgol_filter(z, 7, 3, axis=0, mode='interp', deriv=1, delta=delta)
|
|
assert_allclose(dy, dz, atol=1e-10)
|