From a934a5dda75308d4e6375e96591a67045becef39 Mon Sep 17 00:00:00 2001 From: Sergio Aguirre Date: Sat, 31 Oct 2026 21:53:58 +0100 Subject: [PATCH 1/2] [FEATURE] add group_delay, coherence, cepstrum and MLS group_delay(signal) returns the group delay of each channel in seconds, from the derivative of the spectrum, Re{DFT(n x[n]) / DFT(x[n])} / fs, which needs no phase unwrapping. coherence(x, y, winSize=1024) returns the magnitude squared coherence with scipy.signal.coherence (Welch, Hann, 50 % overlap). Channels are paired one to one, and a single channel is paired with all of the other. cepstrum(signal, kind='complex') returns the real or complex cepstrum. For the complex one the integer linear phase is removed first and returned, and icepstrum(cepstrumObj, delay) inverts it. The magnitude floor is the same as in minimum_phase. generate.mls(order, samplingRate) returns one period of a symmetrical MLS from scipy.signal.max_len_seq, and mls_deconvolution(response, order) recovers the impulse response by periodic cross-correlation normalized by L + 1. With more than one period the first is dropped and the others are averaged. Tests pin the group delay of a pure delay and of a first order allpass, the coherence of a noisy filtered signal against SNR / (1 + SNR), the cepstra of a minimum phase filter against the analytic series, the complex cepstrum round trip with its delay, the exact MLS autocorrelation and the deconvolution of a FIR with its known offset -sum(h) / (L + 1). Co-Authored-By: Claude Fable 5.1 --- docs/pyttadoc/functions.rst | 5 + pytta/__init__.py | 8 +- pytta/functions.py | 210 +++++++++++++++++++++++++++++++++++- pytta/generate.py | 33 ++++++ tests/test_dsp_utilities.py | 153 ++++++++++++++++++++++++++ 5 files changed, 407 insertions(+), 2 deletions(-) create mode 100644 tests/test_dsp_utilities.py diff --git a/docs/pyttadoc/functions.rst b/docs/pyttadoc/functions.rst index add57c0..3bbf3c6 100644 --- a/docs/pyttadoc/functions.rst +++ b/docs/pyttadoc/functions.rst @@ -16,6 +16,11 @@ Functions .. autofunction:: pytta.resample .. autofunction:: pytta.corr_coef .. autofunction:: pytta.peak_time +.. autofunction:: pytta.group_delay +.. autofunction:: pytta.coherence +.. autofunction:: pytta.cepstrum +.. autofunction:: pytta.icepstrum +.. autofunction:: pytta.mls_deconvolution .. autofunction:: pytta.save .. autofunction:: pytta.load .. autofunction:: pytta.weighting diff --git a/pytta/__init__.py b/pytta/__init__.py index d8740aa..efcee55 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,\ + group_delay, coherence, cepstrum, icepstrum, mls_deconvolution # Must go on v0.1.0 (DEPRECATED) from .functions import fft_degree @@ -146,6 +147,11 @@ 'resample', 'corr_coef', 'peak_time', + 'group_delay', + 'coherence', + 'cepstrum', + 'icepstrum', + 'mls_deconvolution', 'save', 'load', 'weighting', diff --git a/pytta/functions.py b/pytta/functions.py index 9b39239..55eb306 100644 --- a/pytta/functions.py +++ b/pytta/functions.py @@ -17,6 +17,11 @@ >>> pytta.corr_coef(signalObj1, signalObj2) >>> pytta.resample(signalObj, newSamplingRate) >>> pytta.peak_time(signalObj1, signalObj2, ..., signalObjN) + >>> pytta.group_delay(signalObj) + >>> pytta.coherence(signalObj1, signalObj2) + >>> pytta.cepstrum(signalObj) + >>> pytta.icepstrum(cepstrumObj, delay) + >>> pytta.mls_deconvolution(signalObj, order) >>> pytta.plot_time(signalObj1, signalObj2, ..., signalObjN) >>> pytta.plot_time_dB(signalObj1, signalObj2, ..., signalObjN) >>> pytta.plot_freq(signalObj1, signalObj2, ..., signalObjN) @@ -44,7 +49,7 @@ RecMeasure, PlayRecMeasure, FRFMeasure, \ Analysis, OctFilter from pytta.classes._base import ChannelsList, ChannelObj -from pytta.generate import measurement # TODO: Change to class instantiation. +from pytta.generate import measurement, mls # TODO: Change to class instantiation. from pytta import _h5utils as _h5 import copy as cp from warnings import warn @@ -274,6 +279,209 @@ def peak_time(signal): else: return peaks_time[0] + +def group_delay(signal): + """ + Group delay of each channel, -d(phase)/d(omega), in seconds. + + Computed with the derivative of the spectrum, which needs no phase + unwrapping (Oppenheim and Schafer, Discrete Time Signal Processing, + ch. 5 and 13): + + tau = Re{ DFT(n x[n]) / DFT(x[n]) } / fs + + The DFT is circular, so the signal should decay within its length. + Bins where the spectrum is zero give NaN. + + Parameters + ---------- + signal : SignalObj + + Returns + ------- + numpy.ndarray + Group delay with shape (numBins, numChannels), at signal.freqVector. + + >>> tau = pytta.group_delay(ir) + """ + if not isinstance(signal, SignalObj): + raise TypeError('signal must be a SignalObj.') + x = np.asarray(signal.timeSignal, dtype=float) + n = np.arange(signal.numSamples)[:, None] + X = np.fft.rfft(x, axis=0) + with np.errstate(divide='ignore', invalid='ignore'): + tau = np.real(np.fft.rfft(n * x, axis=0) / X) + tau[X == 0] = np.nan + return tau / signal.samplingRate + + +def coherence(x, y, winSize=1024): + """ + Magnitude squared coherence |Sxy|^2 / (Sxx Syy) between two SignalObj. + + Welch estimate from scipy.signal.coherence, Hann window of winSize + samples with 50 % overlap. Channels are paired one to one; a signal + with one channel is paired with every channel of the other. + + Parameters + ---------- + x, y : SignalObj + Same sampling rate and number of samples. + winSize : int, optional + Welch segment length in samples. The default is 1024. + + Returns + ------- + freq : numpy.ndarray + Frequencies in Hz, winSize // 2 + 1 bins. + gamma2 : numpy.ndarray + Coherence with shape (numBins, numChannels), between 0 and 1. + + >>> freq, gamma2 = pytta.coherence(excitation, response) + """ + if not isinstance(x, SignalObj) or not isinstance(y, SignalObj): + raise TypeError('x and y must be SignalObj.') + if x.samplingRate != y.samplingRate or x.numSamples != y.numSamples: + raise ValueError('x and y must have the same sampling rate and ' + 'number of samples.') + return ss.coherence(np.asarray(x.timeSignal, dtype=float), + np.asarray(y.timeSignal, dtype=float), + fs=x.samplingRate, window='hann', nperseg=winSize, + axis=0) + + +def cepstrum(signal, kind='complex'): + """ + Real or complex cepstrum of each channel (Oppenheim and Schafer, + Discrete Time Signal Processing, ch. 13). + + The real cepstrum is IDFT(log|X|). The complex cepstrum is + IDFT(log|X| + j arg X), with arg X unwrapped and its linear part, an + integer delay in samples, removed first and returned so icepstrum can + put it back. As in Oppenheim and Schafer, the sign of the gain is not + represented: the input should have a positive gain at 0 Hz. + + Both are computed on the signal length, so they are time aliased; + zero padding the signal before the call reduces it. Spectral zeros are + floored at -200 dB re the peak so the logarithm is finite, as in + minimum_phase. + + Parameters + ---------- + signal : SignalObj + kind : str, optional + 'complex' (default) or 'real'. + + Returns + ------- + SignalObj + The cepstrum, for kind='real'. + (SignalObj, numpy.ndarray) + The cepstrum and the delay removed from each channel, in samples, + for kind='complex'. + + >>> c = pytta.cepstrum(ir, 'real') + >>> chat, delay = pytta.cepstrum(ir) + >>> ir2 = pytta.icepstrum(chat, delay) + """ + if not isinstance(signal, SignalObj): + raise TypeError('signal must be a SignalObj.') + if kind not in ('complex', 'real'): + raise ValueError("kind must be 'complex' or 'real'.") + N = signal.numSamples + X = np.fft.rfft(np.asarray(signal.timeSignal, dtype=float), axis=0) + magnitude = np.abs(X) + magnitude = np.maximum(magnitude, 1e-10 * magnitude.max(axis=0)) + if kind == 'real': + return _cepstrum_signal(signal, np.fft.irfft(np.log(magnitude), n=N, + axis=0)) + w = 2 * np.pi * np.arange(X.shape[0])[:, None] / N + phase = np.unwrap(np.angle(X), axis=0) + delay = np.round(-phase[-1] / w[-1]).astype(int) + logX = np.log(magnitude) + 1j * (phase + w * delay) + return _cepstrum_signal(signal, np.fft.irfft(logX, n=N, axis=0)), delay + + +def icepstrum(cepstrumObj, delay=0): + """ + Inverse complex cepstrum: the signal whose cepstrum and delay, from + cepstrum(signal), are given. + + Parameters + ---------- + cepstrumObj : SignalObj + Complex cepstrum. + delay : int or array_like, optional + Delay in samples to put back, one per channel or one for all. + The default is 0. + + Returns + ------- + SignalObj + + >>> ir2 = pytta.icepstrum(*pytta.cepstrum(ir)) + """ + if not isinstance(cepstrumObj, SignalObj): + raise TypeError('cepstrumObj must be a SignalObj.') + N = cepstrumObj.numSamples + logX = np.fft.rfft(np.asarray(cepstrumObj.timeSignal, dtype=float), + axis=0) + w = 2 * np.pi * np.arange(logX.shape[0])[:, None] / N + X = np.exp(logX - 1j * w * np.asarray(delay)) + return _cepstrum_signal(cepstrumObj, np.fft.irfft(X, n=N, axis=0)) + + +def _cepstrum_signal(signal, data): + newSignal = SignalObj(data, 'time', samplingRate=signal.samplingRate, + signalType='energy') + newSignal.channels = cp.deepcopy(signal.channels) + return newSignal + + +def mls_deconvolution(response, order): + """ + Impulse response from the response to a periodic MLS, by periodic + cross-correlation with the sequence of pytta.generate.mls(order). + + The response must start with the excitation. With two or more periods, + the first one (the transient) is dropped and the others are averaged; + a single period is taken as already periodic. The correlation is + normalized by L + 1, L = 2^order - 1, which returns + h - sum(h) / (L + 1): the impulse response plus a constant of order + 1 / L (Rife and Vanderkooy 1989; Peltonen 2000, eq. 2.22). + + Parameters + ---------- + response : SignalObj + Recorded response, at least L samples per channel. + order : int + MLS order used for the excitation. + + Returns + ------- + SignalObj + Impulse response, L samples per channel. + + >>> ir = pytta.mls_deconvolution(recording, 16) + """ + if not isinstance(response, SignalObj): + raise TypeError('response must be a SignalObj.') + L = 2**order - 1 + periods = response.numSamples // L + if periods < 1: + raise ValueError('response is shorter than one MLS period.') + y = np.asarray(response.timeSignal[:periods * L], dtype=float) + y = y.reshape(periods, L, -1)[1 if periods > 1 else 0:].mean(axis=0) + z = mls(order).timeSignal[:, :1].astype(float) + ir = np.fft.irfft(np.fft.rfft(y, axis=0) + * np.conj(np.fft.rfft(z, axis=0)), n=L, axis=0) + newSignal = SignalObj(ir / (L + 1), 'time', + samplingRate=response.samplingRate, + signalType='energy') + newSignal.channels = cp.deepcopy(response.channels) + return newSignal + + def fft_degree(*args,**kwargs): """ DEPRECATED diff --git a/pytta/generate.py b/pytta/generate.py index 1979f54..5e762a8 100644 --- a/pytta/generate.py +++ b/pytta/generate.py @@ -19,6 +19,7 @@ >>> pytta.generate.sweep() >>> pytta.generate.random_noise() >>> pytta.generate.impulse() + >>> pytta.generate.mls() >>> pytta.generate.measurement() For further information see the specific function documentation @@ -504,6 +505,38 @@ def impulse(samplingRate=None, return impulseSignal +def mls(order=16, samplingRate=None): + """ + One period of a symmetrical maximum length sequence, L = 2^order - 1 + samples of +1 and -1. + + The binary sequence is scipy.signal.max_len_seq(order), mapped from + (0, 1) to (+1, -1). Its periodic autocorrelation is (L + 1) delta - 1 + (Rife and Vanderkooy 1989; Peltonen 2000, eq. 2.14). Play it + periodically and recover the impulse response with + pytta.mls_deconvolution(response, order). + + Parameters + ---------- + order : int, optional + Sequence order, 2 or more. The default is 16. + samplingRate : int, optional + The default is pytta.default.samplingRate. + + Returns + ------- + SignalObj + + >>> excitation = pytta.generate.mls(16) + """ + if samplingRate is None: + samplingRate = default.samplingRate + sequence = 1 - 2 * ss.max_len_seq(order)[0].astype(float) + return SignalObj(signalArray=sequence, domain='time', + samplingRate=samplingRate, + comment='MLS order {}'.format(order)) + + def measurement(kind='playrec', samplingRate=None, freqMin=None, diff --git a/tests/test_dsp_utilities.py b/tests/test_dsp_utilities.py new file mode 100644 index 0000000..27095f1 --- /dev/null +++ b/tests/test_dsp_utilities.py @@ -0,0 +1,153 @@ +""" +Tests for the DSP utilities group_delay, coherence, cepstrum, icepstrum, +generate.mls and mls_deconvolution. + +Each test pins one rule with an analytic value. SignalObj stores its samples +in float32, so the exact cases are checked to about 1e-6. +""" +import os +import unittest + +os.environ.setdefault('MPLBACKEND', 'Agg') + +import numpy as np +import scipy.signal as ss + +import pytta + +FS = 48000 +TOL32 = 1e-6 # rounding of the float32 SignalObj storage + + +class TestGroupDelay(unittest.TestCase): + + def test_pure_delay_and_first_order_allpass(self): + # A delay of k samples has group delay k / fs at every frequency. + # The allpass H(z) = (a + z^-1) / (1 + a z^-1) has group delay + # (1 - a^2) / (1 + 2 a cos(w) + a^2) samples. + N, k, a = 1024, 37, 0.8 + data = np.zeros((N, 2)) + data[k, 0] = 1 + impulse = np.zeros(N) + impulse[0] = 1 + data[:, 1] = ss.lfilter([a, 1], [1, a], impulse) + sig = pytta.SignalObj(data, 'time', FS) + tau = pytta.group_delay(sig) + w = 2 * np.pi * sig.freqVector / FS + expected = np.column_stack([np.full(w.size, k / FS), + (1 - a**2) / (1 + 2*a*np.cos(w) + a**2) / FS]) + err = np.max(np.abs(tau - expected), axis=0) * FS + print(f'\n max error: delay {err[0]:.1e}, allpass {err[1]:.1e} samples') + self.assertEqual(tau.shape, (N // 2 + 1, 2)) + self.assertLess(np.max(err), 1e-6) + + +class TestCoherence(unittest.TestCase): + + def test_noisy_output_and_identical_signals(self): + # y = h * x + n with white x (variance 1) and white n (variance s2): + # SNR(f) = |H(f)|^2 / s2 and gamma^2 = SNR / (1 + SNR). With 8192 + # Welch segments the bias is below 1e-3 and the random error has a + # standard deviation of about 0.004, so 0.02 is five of them. + # A signal against itself has gamma^2 = 1. + rng = np.random.default_rng(0) + N, s2, h = 2**19, 0.5, [1, 0.9] + x = rng.standard_normal(N) + y = ss.lfilter(h, 1, x) + np.sqrt(s2) * rng.standard_normal(N) + sigX = pytta.SignalObj(x, 'time', FS) + sigY = pytta.SignalObj(np.column_stack([y, x]), 'time', FS) + freq, gamma2 = pytta.coherence(sigX, sigY, winSize=128) + H2 = np.abs(ss.freqz(h, 1, worN=freq, fs=FS)[1])**2 + snr = H2 / s2 + err = np.max(np.abs(gamma2[:, 0] - snr / (1 + snr))) + errSelf = np.max(np.abs(gamma2[:, 1] - 1)) + print(f'\n max |error|: noisy {err:.1e}, identical {errSelf:.1e}') + self.assertLess(err, 0.02) + self.assertLess(errSelf, TOL32) + + +class TestCepstrum(unittest.TestCase): + + def test_known_minimum_phase_filter_matches_analytic_series(self): + # For x = (1 - a z^-1)(1 - b z^-1) with |a|, |b| < 1 the complex + # cepstrum is -(a^n + b^n) / n for n > 0 and zero for n <= 0 + # (Oppenheim and Schafer, DTSP 3rd ed., eq. 13.36). The real + # cepstrum is its even part. + N, a, b = 1024, 0.9, -0.6 + x = np.zeros(N) + x[:3] = np.poly([a, b]) + sig = pytta.SignalObj(x, 'time', FS) + n = np.arange(1, N // 2) + series = np.zeros(N) + series[n] = -(a**n + b**n) / n + real = pytta.cepstrum(sig, kind='real').timeSignal[:, 0] + cplx, delay = pytta.cepstrum(sig, kind='complex') + errReal = np.max(np.abs(real - (series + np.roll(series[::-1], 1)) / 2)) + errCplx = np.max(np.abs(cplx.timeSignal[:, 0] - series)) + print(f'\n max |error|: real {errReal:.1e}, complex {errCplx:.1e}') + self.assertLess(errReal, TOL32) + self.assertLess(errCplx, TOL32) + self.assertEqual(list(delay), [0]) + + def test_complex_cepstrum_round_trip_with_linear_phase(self): + # A mixed phase filter (one zero outside the unit circle) delayed by + # 7 samples: the linear phase removed is 7 samples of delay plus one + # per zero outside the circle, and icepstrum puts it back. + N = 1024 + x = np.zeros((N, 2)) + x[7:10, 0] = np.poly([0.8, -1.6]) + x[:, 1] = np.random.default_rng(1).standard_normal(N) \ + * np.exp(-np.arange(N) / 50) + x[:, 1] *= np.sign(x[:, 1].sum()) # positive gain at 0 Hz + sig = pytta.SignalObj(x, 'time', FS) + cep, delay = pytta.cepstrum(sig) + back = pytta.icepstrum(cep, delay) + err = np.max(np.abs(back.timeSignal - sig.timeSignal), axis=0) + print(f'\n delay {delay}, max |error| {err[0]:.1e}, {err[1]:.1e}') + self.assertEqual(delay[0], 8) + self.assertLess(np.max(err), TOL32) + + +class TestMLS(unittest.TestCase): + + def test_periodic_autocorrelation_is_exact(self): + # A symmetrical MLS of period L = 2^N - 1 has periodic + # autocorrelation (L + 1) delta - 1 (Peltonen 2000, eq. 2.14-2.15). + order = 10 + L = 2**order - 1 + s = pytta.generate.mls(order, samplingRate=FS) + z = s.timeSignal[:, 0].astype(int) + R = np.array([np.dot(z, np.roll(z, k)) for k in range(L)]) + expected = (L + 1) * (np.arange(L) == 0) - 1 + print(f'\n L = {s.numSamples}, max |R - expected| = ' + f'{np.max(np.abs(R - expected))}') + self.assertEqual(s.numSamples, L) + self.assertTrue(np.array_equal(R, expected)) + + def test_deconvolution_of_fir_excited_by_periodic_mls(self): + # Periodic cross-correlation normalized by L + 1 returns + # h - sum(h) / (L + 1) (Peltonen 2000, eq. 2.22): the FIR within + # 1 / L, with that exact constant offset. + order = 12 + L = 2**order - 1 + h = np.random.default_rng(2).standard_normal(40) \ + * np.exp(-np.arange(40) / 8) + z = pytta.generate.mls(order).timeSignal[:, 0] + y = ss.lfilter(h, 1, np.tile(z, 3)) + resp = pytta.SignalObj(np.column_stack([y, 2 * y]), 'time', FS) + ir = pytta.mls_deconvolution(resp, order) + hL = np.zeros(L) + hL[:h.size] = h + expected = hL - h.sum() / (L + 1) + errOffset = np.max(np.abs(ir.timeSignal[:, 0] - expected)) + errFir = np.max(np.abs(ir.timeSignal[:, 0] - hL)) + err2 = np.max(np.abs(ir.timeSignal[:, 1] - 2 * expected)) + print(f'\n max |ir - h| = {errFir:.1e} (1/L = {1/L:.1e}), ' + f'with offset {errOffset:.1e}, channel 2 {err2:.1e}') + self.assertEqual(ir.numSamples, L) + self.assertLess(errFir, 1 / L) + self.assertLess(max(errOffset, err2), TOL32) + + +if __name__ == '__main__': + unittest.main() From 010c0ec7305f7097264188fc8882b3ea3b4cdd66 Mon Sep 17 00:00:00 2001 From: Sergio Aguirre Date: Sun, 1 Nov 2026 01:39:30 +0100 Subject: [PATCH 2/2] [FEATURE] add plot_group_delay, plot_nyquist, plot_zplane and plot_polar Co-Authored-By: Claude Fable 5.1 --- docs/pyttadoc/functions.rst | 5 ++- pytta/__init__.py | 7 +++- pytta/functions.py | 81 +++++++++++++++++++++++++++++++++++++ tests/test_plot_more.py | 64 +++++++++++++++++++++++++++++ 4 files changed, 155 insertions(+), 2 deletions(-) create mode 100644 tests/test_plot_more.py diff --git a/docs/pyttadoc/functions.rst b/docs/pyttadoc/functions.rst index 3bbf3c6..429f026 100644 --- a/docs/pyttadoc/functions.rst +++ b/docs/pyttadoc/functions.rst @@ -31,4 +31,7 @@ Functions .. autofunction:: pytta.plot_bars .. autofunction:: pytta.plot_spectrogram .. autofunction:: pytta.plot_waterfall - +.. autofunction:: pytta.plot_group_delay +.. autofunction:: pytta.plot_nyquist +.. autofunction:: pytta.plot_zplane +.. autofunction:: pytta.plot_polar diff --git a/pytta/__init__.py b/pytta/__init__.py index efcee55..8964676 100644 --- a/pytta/__init__.py +++ b/pytta/__init__.py @@ -114,7 +114,8 @@ 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,\ - group_delay, coherence, cepstrum, icepstrum, mls_deconvolution + group_delay, coherence, cepstrum, icepstrum, mls_deconvolution,\ + plot_group_delay, plot_nyquist, plot_zplane, plot_polar # Must go on v0.1.0 (DEPRECATED) from .functions import fft_degree @@ -152,6 +153,10 @@ 'cepstrum', 'icepstrum', 'mls_deconvolution', + 'plot_group_delay', + 'plot_nyquist', + 'plot_zplane', + 'plot_polar', 'save', 'load', 'weighting', diff --git a/pytta/functions.py b/pytta/functions.py index 55eb306..de471f9 100644 --- a/pytta/functions.py +++ b/pytta/functions.py @@ -482,6 +482,87 @@ def mls_deconvolution(response, order): return newSignal +def plot_group_delay(*sigObjs, xLim=None, yLim=None, title=None): + """ + Group delay of every channel of the SignalObjs, in ms, on a logarithmic + frequency axis, as ita_plot_groupdelay. Returns the Figure. + + >>> pytta.plot_group_delay(ir1, ir2, xLim=[20, 20000]) + """ + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + for sig in sigObjs: + tau = group_delay(sig) + for c in range(tau.shape[1]): + ax.semilogx(sig.freqVector[1:], 1e3 * tau[1:, c]) + ax.set(xlabel='Frequency [Hz]', ylabel='Group delay [ms]', xlim=xLim, ylim=yLim, title=title) + ax.grid(True, which='both') + return fig + + +def plot_nyquist(*sigObjs, title=None): + """ + Imaginary against real part of the spectrum of every channel of the + SignalObjs, as ita_plot_nyquist. Returns the Figure. + """ + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + for sig in sigObjs: + X = np.asarray(sig.freqSignal) + for c in range(X.shape[1]): + ax.plot(X[:, c].real, X[:, c].imag) + ax.set(xlabel='Real part', ylabel='Imaginary part', title=title) + ax.set_aspect('equal', 'datalim') + ax.grid(True) + return fig + + +def plot_zplane(zeros, poles, title=None): + """ + Zeros (o) and poles (x) of a discrete-time system on the z plane, with the + unit circle, as ita_plot_zplane; scipy.signal.tf2zpk(b, a) gives them + from the coefficients. Returns the Figure. + + >>> pytta.plot_zplane(*scipy.signal.tf2zpk(b, a)[:2]) + """ + import matplotlib.pyplot as plt + zeros, poles = np.atleast_1d(zeros), np.atleast_1d(poles) + fig, ax = plt.subplots() + w = np.linspace(0, 2 * np.pi, 361) + ax.plot(np.cos(w), np.sin(w), 'k:', lw=0.8) + ax.plot(zeros.real, zeros.imag, 'o', mfc='none', label='zeros') + ax.plot(poles.real, poles.imag, 'x', label='poles') + ax.set(xlabel='Real part', ylabel='Imaginary part', title=title) + ax.set_aspect('equal', 'datalim') + ax.grid(True) + ax.legend() + return fig + + +def plot_polar(angles, levels, labels=None, rLim=None, title=None): + """ + Directivity in a polar plot, as ita_plot_polar: levels (dB) at the + angles (degrees, 0 to the front, counterclockwise), one column per curve, + closed at the first angle. Returns the Figure. + + >>> pytta.plot_polar(np.arange(0, 360, 15), L, labels=['1 kHz', '4 kHz']) + """ + import matplotlib.pyplot as plt + theta = np.deg2rad(np.r_[angles, angles[0]]) + levels = np.asarray(levels, dtype=float).reshape(len(angles), -1) + levels = np.vstack([levels, levels[:1]]) + fig, ax = plt.subplots(subplot_kw={'projection': 'polar'}) + ax.set_theta_zero_location('N') + for c in range(levels.shape[1]): + ax.plot(theta, levels[:, c], label=None if labels is None else labels[c]) + if rLim is not None: + ax.set_rlim(rLim) + if labels is not None: + ax.legend() + ax.set_title(title) + return fig + + def fft_degree(*args,**kwargs): """ DEPRECATED diff --git a/tests/test_plot_more.py b/tests/test_plot_more.py new file mode 100644 index 0000000..8722704 --- /dev/null +++ b/tests/test_plot_more.py @@ -0,0 +1,64 @@ +""" +Tests for plot_group_delay, plot_nyquist, plot_zplane and plot_polar of +pytta.functions, which follow ita_plot_groupdelay, ita_plot_nyquist, +ita_plot_zplane and ita_plot_polar of the ITA-Toolbox. Each test reads back +the plotted data and compares it with an analytic result. +""" +import os +import unittest + +import numpy as np +import scipy.signal as ss + +os.environ.setdefault('MPLBACKEND', 'Agg') + +import matplotlib.pyplot as plt +import pytta + +FS = 8000 + + +def delayed_impulse(d, n=1024): + label = 'x' + x = np.zeros(n) + x[d] = 1 + return pytta.SignalObj(x, 'time', FS) + + +class TestPlotMore(unittest.TestCase): + + def tearDown(self): + plt.close('all') + + def test_group_delay_of_a_pure_delay(self): + # A delay of 40 samples is 5 ms at every frequency. + y = pytta.plot_group_delay(delayed_impulse(40)).axes[0].lines[0].get_ydata() + print(f'\n group delay: {np.abs(y - 5).max():.1e} ms off 5 ms', end='') + self.assertLess(np.abs(y - 5).max(), 1e-9) + + def test_nyquist_of_a_pure_delay_is_a_circle(self): + sig = delayed_impulse(7) + line = pytta.plot_nyquist(sig).axes[0].lines[0] + # The spectrum of SignalObj scales the bins above DC alike, so they + # lie on one circle. + r = np.hypot(line.get_xdata(), line.get_ydata())[1:] + print(f'\n nyquist: radius spread {np.ptp(r) / r.mean():.1e}', end='') + self.assertLess(np.ptp(r) / r.mean(), 1e-6) + + def test_zplane_marks_the_poles_and_zeros(self): + z, p, _ = ss.tf2zpk(*ss.butter(4, 0.2)) + lines = pytta.plot_zplane(z, p).axes[0].lines + np.testing.assert_allclose(lines[1].get_xdata() + 1j * lines[1].get_ydata(), z) + np.testing.assert_allclose(lines[2].get_xdata() + 1j * lines[2].get_ydata(), p) + + def test_polar_curve_closes_at_the_front(self): + angles = np.arange(0, 360, 30) + levels = np.column_stack([np.cos(np.deg2rad(angles)), np.ones(12)]) + ax = pytta.plot_polar(angles, levels, labels=['a', 'b']).axes[0] + theta, r = ax.lines[0].get_xdata(), ax.lines[0].get_ydata() + self.assertEqual(ax.get_theta_offset(), np.pi / 2) + self.assertEqual((theta.size, r[-1]), (13, r[0])) + + +if __name__ == '__main__': + unittest.main()