[FEATURE] add extrapolate, IIR extrapolation of a spectrum beyond its range

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

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',

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

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()

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()