From e49838587e96167ebeed40ea57fd09a91b334e8a Mon Sep 17 00:00:00 2001
From: Sergio Aguirre <djsergioaguirre@gmail.com>
Date: Sun, 1 Nov 2026 01:47:29 +0100
Subject: [PATCH] [FEATURE] add prony and invfreqz for IIR fits of measured
 responses

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 docs/pyttadoc/functions.rst |  3 ++-
 pytta/__init__.py           |  5 ++++-
 pytta/functions.py          | 40 +++++++++++++++++++++++++++++++++++++
 tests/test_prony.py         | 36 +++++++++++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_prony.py

diff --git a/docs/pyttadoc/functions.rst b/docs/pyttadoc/functions.rst
index add57c0..d6adee6 100644
--- a/docs/pyttadoc/functions.rst
+++ b/docs/pyttadoc/functions.rst
@@ -26,4 +26,5 @@ Functions
 .. autofunction:: pytta.plot_bars
 .. autofunction:: pytta.plot_spectrogram
 .. autofunction:: pytta.plot_waterfall
-
+.. autofunction:: pytta.prony
+.. autofunction:: pytta.invfreqz
diff --git a/pytta/__init__.py b/pytta/__init__.py
index d8740aa..a55c295 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
 
 # Must go on v0.1.0 (DEPRECATED)
 from .functions import fft_degree
@@ -146,6 +147,8 @@
            'resample',
            'corr_coef',
            'peak_time',
+           'prony',
+           'invfreqz',
            'save',
            'load',
            'weighting',
diff --git a/pytta/functions.py b/pytta/functions.py
index 9b39239..ae71a35 100644
--- a/pytta/functions.py
+++ b/pytta/functions.py
@@ -274,6 +274,46 @@ 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 fft_degree(*args,**kwargs):
     """
     DEPRECATED
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()