diff --git a/docs/pyttadoc/functions.rst b/docs/pyttadoc/functions.rst index add57c0..45e1190 100644 --- a/docs/pyttadoc/functions.rst +++ b/docs/pyttadoc/functions.rst @@ -26,4 +26,6 @@ Functions .. autofunction:: pytta.plot_bars .. autofunction:: pytta.plot_spectrogram .. autofunction:: pytta.plot_waterfall - +.. autofunction:: pytta.prony +.. autofunction:: pytta.invfreqz +.. autofunction:: pytta.extrapolate diff --git a/pytta/__init__.py b/pytta/__init__.py index d8740aa..116b08e 100644 --- a/pytta/__init__.py +++ b/pytta/__init__.py @@ -113,7 +113,8 @@ list_devices, print_devices, get_device_from_user,\ fft_convolve, find_delay, corr_coef, resample,\ peak_time, save, load, plot_time, plot_time_dB, plot_freq,\ - plot_bars, plot_spectrogram, plot_waterfall, SPL + plot_bars, plot_spectrogram, plot_waterfall, SPL,\ + prony, invfreqz, extrapolate # Must go on v0.1.0 (DEPRECATED) from .functions import fft_degree @@ -146,6 +147,9 @@ 'resample', 'corr_coef', 'peak_time', + 'prony', + 'invfreqz', + 'extrapolate', 'save', 'load', 'weighting', diff --git a/pytta/functions.py b/pytta/functions.py index 9b39239..ca8fe74 100644 --- a/pytta/functions.py +++ b/pytta/functions.py @@ -274,6 +274,72 @@ def peak_time(signal): else: return peaks_time[0] +def prony(h, nb, na): + """ + Coefficients b (nb + 1) and a (na + 1, a[0] = 1) of the IIR filter whose + impulse response best matches h, by Prony's method (Parks and Burrus, + Digital Filter Design, 1987), as the time domain of ita_prony_analysis: + the denominator fits the samples after nb by least squares, and the + numerator then matches the first nb + 1 samples exactly. + scipy.signal.tf2zpk(b, a) gives the zeros and poles. + + >>> b, a = pytta.prony(ir.timeSignal[:, 0], 8, 8) + """ + from scipy.linalg import toeplitz + h = np.asarray(h) + if h.size <= max(nb, na): + h = np.pad(h, (0, max(nb, na) + 1 - h.size)) + H = toeplitz(h, np.r_[h[0], np.zeros(na)]) + a = np.r_[1, -np.linalg.lstsq(H[nb + 1:, 1:], H[nb + 1:, 0], rcond=None)[0]] + return H[:nb + 1] @ a, a + + +def invfreqz(H, w, nb, na, weights=None): + """ + Real coefficients b (nb + 1) and a (na + 1, a[0] = 1) of the IIR filter + whose frequency response fits H at the normalized frequencies w (rad per + sample), by the linear equation error method of Levi (Complex curve + fitting, IRE Trans. Automatic Control 4(1), 1959), minimizing + sum weights |B(w) - H(w) A(w)|^2, as MATLAB's invfreqz without + iterations and the frequency domain of ita_prony_analysis. + + >>> b, a = pytta.invfreqz(H, 2 * np.pi * f / fs, 8, 8) + """ + H, w = np.asarray(H), np.asarray(w) + E = np.exp(-1j * np.outer(w, np.arange(max(nb, na) + 1))) + M = np.hstack([E[:, :nb + 1], -H[:, None] * E[:, 1:na + 1]]) + sw = np.sqrt(np.ones(w.size) if weights is None else np.asarray(weights))[:, None] + x = np.linalg.lstsq(np.vstack([(sw * M).real, (sw * M).imag]), + np.r_[(sw[:, 0] * H).real, (sw[:, 0] * H).imag], rcond=None)[0] + return x[:nb + 1], np.r_[1, x[nb + 1:]] + + +def extrapolate(signal, freqRange, order=4): + """ + Spectrum outside freqRange = [f1, f2] (Hz) replaced by IIR models of + order order fitted with invfreqz near each end, as ita_extrapolate: + the low model on f1 + (0.01 to 0.2)(f2 - f1), used below f1, and the + high model on f1 + (0.9 to 0.99)(f2 - f1), used above f2, the bands + of the ITA-Toolbox. A pure delay should be removed first (time_shift), + since it is not a low-order model. + """ + fs = signal.samplingRate + x = np.asarray(signal.timeSignal, dtype=float) + n = x.shape[0] + H = np.fft.rfft(x, axis=0) + f = np.fft.rfftfreq(n, 1 / fs) + w = 2 * np.pi * f / fs + f1, f2 = freqRange + out = H.copy() + for edge, (lo, hi) in (('low', (0.01, 0.2)), ('high', (0.9, 0.99))): + band = (f >= f1 + lo * (f2 - f1)) & (f <= f1 + hi * (f2 - f1)) + outside = f < f1 if edge == 'low' else f > f2 + for c in range(H.shape[1]): + b, a = invfreqz(H[band, c], w[band], order, order) + out[outside, c] = ss.freqz(b, a, w[outside])[1] + return SignalObj(np.fft.irfft(out, n, axis=0), 'time', fs) + + def fft_degree(*args,**kwargs): """ DEPRECATED diff --git a/tests/test_extrapolate.py b/tests/test_extrapolate.py new file mode 100644 index 0000000..2f76fb9 --- /dev/null +++ b/tests/test_extrapolate.py @@ -0,0 +1,40 @@ +""" +Tests for extrapolate of pytta.functions, which follows ita_extrapolate of +the ITA-Toolbox with IIR fits by invfreqz in place of MATLAB's rationalfit. +""" +import unittest + +import numpy as np +import scipy.signal as ss + +import pytta + +FS = 48000 + + +class TestExtrapolate(unittest.TestCase): + + def test_rational_response_is_extrapolated_exactly(self): + # A band-pass of order 4 known only between 300 Hz and 3 kHz (the + # rest of the spectrum replaced by noise) comes back outside that + # range from models of the same order, within what the float32 + # samples of SignalObj allow. + label = 'x' + n = 2 ** 14 + b, a = ss.butter(2, [100, 5000], 'bandpass', fs=FS) + h = ss.lfilter(b, a, np.eye(1, n)[0]) + H = np.fft.rfft(h) + f = np.fft.rfftfreq(n, 1 / FS) + outside = (f < 300) | (f > 3000) + rng = np.random.default_rng(7) + corrupted = H.copy() + corrupted[outside] = rng.standard_normal(outside.sum()) + 1j * rng.standard_normal(outside.sum()) + y = pytta.extrapolate(pytta.SignalObj(np.fft.irfft(corrupted, n), 'time', FS), [300, 3000], 4) + Y = np.fft.rfft(np.asarray(y.timeSignal[:, 0], dtype=float)) + e = np.abs(Y[outside] - H[outside]).max() / np.abs(H).max() + print(f'\n extrapolation outside 300 Hz to 3 kHz: {e:.1e}', end='') + self.assertLess(e, 1e-3) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_prony.py b/tests/test_prony.py new file mode 100644 index 0000000..21de996 --- /dev/null +++ b/tests/test_prony.py @@ -0,0 +1,36 @@ +""" +Tests for prony and invfreqz of pytta.functions, the time and frequency +domains of ita_prony_analysis of the ITA-Toolbox. Each test recovers the +coefficients of a known filter. +""" +import unittest + +import numpy as np +import scipy.signal as ss + +import pytta + + +class TestProny(unittest.TestCase): + + def test_prony_recovers_an_iir_filter(self): + # Impulse response of a filter with 3 zeros and 5 poles. + a = ss.cheby1(5, 1, 0.3)[1] + b = np.poly([0.9, -0.5, 0.2]) + h = ss.lfilter(b, a, np.r_[1, np.zeros(199)]) + bp, ap = pytta.prony(h, 3, 5) + e = max(np.abs(bp - b).max() / np.abs(b).max(), np.abs(ap - a).max() / np.abs(a).max()) + print(f'\n prony: {e:.1e}', end='') + self.assertLess(e, 1e-9) + + def test_invfreqz_recovers_an_iir_filter(self): + b, a = ss.butter(4, 0.3) + w, H = ss.freqz(b, a, 512) + bi, ai = pytta.invfreqz(H, w, 4, 4) + e = max(np.abs(bi - b).max() / np.abs(b).max(), np.abs(ai - a).max()) + print(f'\n invfreqz: {e:.1e}', end='') + self.assertLess(e, 1e-9) + + +if __name__ == '__main__': + unittest.main()