diff --git a/docs/pyttadoc/functions.rst b/docs/pyttadoc/functions.rst index add57c0..429f026 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 @@ -26,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 d8740aa..8964676 100644 --- a/pytta/__init__.py +++ b/pytta/__init__.py @@ -113,7 +113,9 @@ 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,\ + plot_group_delay, plot_nyquist, plot_zplane, plot_polar # Must go on v0.1.0 (DEPRECATED) from .functions import fft_degree @@ -146,6 +148,15 @@ 'resample', 'corr_coef', 'peak_time', + 'group_delay', + 'coherence', + '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 9b39239..de471f9 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,290 @@ 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 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/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() 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()