[FEATURE] add psd, interpolate_spectrum and extrapolate_low

docs/pyttadoc/functions.rst

@@ -26,4 +26,6 @@ Functions
 .. autofunction:: pytta.plot_bars
 .. autofunction:: pytta.plot_spectrogram
 .. autofunction:: pytta.plot_waterfall
-
+.. autofunction:: pytta.psd
+.. autofunction:: pytta.interpolate_spectrum
+.. autofunction:: pytta.extrapolate_low

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,\
+    psd, interpolate_spectrum, extrapolate_low
 
 # Must go on v0.1.0 (DEPRECATED)
 from .functions import fft_degree
@@ -146,6 +147,9 @@
            'resample',
            'corr_coef',
            'peak_time',
+           'psd',
+           'interpolate_spectrum',
+           'extrapolate_low',
            'save',
            'load',
            'weighting',

pytta/functions.py

@@ -23,6 +23,9 @@
     >>> pytta.plot_bars(signalObj1, signalObj2, ..., signalObjN)
     >>> pytta.save(fileName, obj1, ..., objN)
     >>> pytta.load(fileName)
+    >>> pytta.psd(signalObj)
+    >>> pytta.interpolate_spectrum(signalObj, nSamples)
+    >>> pytta.extrapolate_low(signalObj, freqRange)
 
 For further information, check the function specific documentation.
 
@@ -274,6 +277,110 @@ def peak_time(signal):
     else:
         return peaks_time[0]
 
+def psd(signal, signal2=None, block_size=None, fft_size=None):
+    """
+    Power spectral density of each channel, or the cross power spectral
+    density with the same channel of signal2, as ita_psd: Welch averaging
+    with a symmetric Hann window of block_size samples (the whole signal by
+    default) and half-block overlap, one-sided, in units^2/Hz, with the
+    convention of MATLAB cpsd, X conj(Y) (the conjugate of
+    scipy.signal.csd).
+
+    Returns
+    -------
+    (numpy.ndarray, numpy.ndarray)
+        Frequencies (Hz) and the density, (bins, channels).
+
+    >>> freq, P = pytta.psd(noise, block_size=1024)
+    """
+    x = np.asarray(signal.timeSignal, dtype=float)
+    y = x if signal2 is None else np.asarray(signal2.timeSignal, dtype=float)
+    if y.shape != x.shape:
+        raise ValueError('both signals need the same length and channels.')
+    block = x.shape[0] if block_size is None else int(round(block_size))
+    block += block % 2
+    if block > x.shape[0]:
+        x = np.vstack([x, np.zeros((block - x.shape[0], x.shape[1]))])
+        y = np.vstack([y, np.zeros((block - y.shape[0], y.shape[1]))])
+    nfft = block if fft_size is None else int(fft_size)
+    window = np.hanning(block + 2)[1:-1]          # MATLAB hanning(block)
+    f, P = ss.csd(x, y, fs=signal.samplingRate, window=window, noverlap=block // 2,
+                  nfft=nfft, axis=0, detrend=False, scaling='density')
+    return f, np.conj(P) if signal2 is not None else P.real    # MATLAB cpsd: X conj(Y)
+
+
+def interpolate_spectrum(signal, n_samples, method='spline', abs_phase=False):
+    """
+    signal with its spectrum interpolated onto the bins of n_samples, as
+    ita_interpolate_spk: the DFT is interpolated between DC and the highest
+    frequency (spline, as MATLAB interp1, or linear), on its real and
+    imaginary parts or (abs_phase=True) on magnitude and unwrapped phase.
+    Meant for energy signals such as impulse responses.
+
+    >>> longer = pytta.interpolate_spectrum(ir, 2 ** 16)
+    """
+    from scipy.interpolate import CubicSpline
+    x = np.asarray(signal.timeSignal, dtype=float)
+    X = np.fft.rfft(x, axis=0)
+    old = np.linspace(0, 1, X.shape[0])
+    new = np.linspace(0, 1, n_samples // 2 + 1)
+    if method == 'spline':
+        interp = lambda v: CubicSpline(old, v, axis=0)(new)
+    elif method == 'linear':
+        interp = lambda v: np.column_stack([np.interp(new, old, v[:, c]) for c in range(v.shape[1])])
+    else:
+        raise ValueError("method must be 'spline' or 'linear'.")
+    if abs_phase:
+        Y = interp(np.abs(X)) * np.exp(1j * interp(np.unwrap(np.angle(X), axis=0)))
+    else:
+        Y = interp(X.real) + 1j * interp(X.imag)
+    newSignal = SignalObj(np.fft.irfft(Y, n_samples, axis=0), 'time',
+                          samplingRate=signal.samplingRate, signalType=signal.signalType)
+    newSignal.channels = cp.deepcopy(signal.channels)
+    return newSignal
+
+
+def extrapolate_low(signal, freq_range, slope=0.):
+    """
+    Spectrum below freq_range replaced by an extrapolation, as the lower
+    part of ita_spk_extrapolate: the mean magnitude of freq_range, changed
+    by slope dB per decade towards low frequencies, and a linear phase from
+    the mean group delay of freq_range, joined to the phase at its lower
+    end.
+
+    Parameters
+    ----------
+    freq_range : (float, float)
+        Frequencies (Hz) whose magnitude and group delay are averaged.
+    slope : float
+        dB per decade below freq_range. Default 0.
+
+    >>> fixed = pytta.extrapolate_low(ir, (80, 160))
+    """
+    x = np.asarray(signal.timeSignal, dtype=float)
+    n = x.shape[0]
+    X = np.fft.rfft(x, axis=0)
+    f = np.fft.rfftfreq(n, 1 / signal.samplingRate)
+    lo = int(np.flatnonzero(f > min(freq_range))[0])
+    hi = int(np.flatnonzero(f > max(freq_range))[0])
+    df = f[1]
+    band = X[lo:hi + 1]
+    gdelay = -np.diff(np.unwrap(np.angle(X), axis=0), axis=0) / (2 * np.pi * df)
+    tau = np.mean(gdelay[lo:hi + 1], axis=0)
+    fl = f[:lo + 1, None]
+    phase = -2 * np.pi * tau * fl
+    phase = phase - (phase[-1] - np.angle(band[0]))
+    with np.errstate(divide='ignore', invalid='ignore'):
+        shape = 10 ** (slope * (np.log10(fl) - np.log10(fl[-1])) / 20)
+    shape[0] = 0 if slope > 0 else shape[min(1, lo)]   # DC: limit, or held for slope <= 0
+    X[:lo + 1] = np.mean(np.abs(band), axis=0) * shape * np.exp(1j * phase)
+    X[0] = X[0].real
+    newSignal = SignalObj(np.fft.irfft(X, n, axis=0), 'time',
+                          samplingRate=signal.samplingRate, signalType=signal.signalType)
+    newSignal.channels = cp.deepcopy(signal.channels)
+    return newSignal
+
+
 def fft_degree(*args,**kwargs):
     """
     DEPRECATED

tests/golden/make_psd_matlab.m

@@ -0,0 +1,12 @@
+% Golden data for tests/test_spectrum.py: MATLAB cpsd as ita_psd of the
+% ITA-Toolbox calls it (symmetric Hann window of the block size, half-block
+% overlap, one-sided density). The noise is rand with the Mersenne twister,
+% which numpy.random.RandomState(1).random_sample rebuilds.
+fs = 8000; n = 4096; block = 512;
+rng(1, 'twister');
+x = rand(n, 1) - 0.5;
+y = filter([1 0.5 0.25], 1, x);
+w = hanning(block);
+Pxx = cpsd(x, x, w, block/2, block, fs);
+Pxy = cpsd(x, y, w, block/2, block, fs);
+save('psd_matlab.mat', 'fs', 'block', 'Pxx', 'Pxy');

tests/test_spectrum.py

@@ -0,0 +1,80 @@
+"""
+Tests for the spectrum utilities of pytta.functions that follow the
+ITA-Toolbox (ita_psd, ita_interpolate_spk, ita_spk_extrapolate). Each test
+pins one rule.
+
+The golden data come from golden/make_psd_matlab.m: MATLAB cpsd called as
+ita_psd calls it; the noise is rebuilt with numpy.random.RandomState.
+"""
+import os
+import unittest
+
+import numpy as np
+import scipy.io as sio
+import scipy.signal as ss
+
+os.environ.setdefault('MPLBACKEND', 'Agg')
+
+import pytta
+
+GOLDEN = os.path.join(os.path.dirname(__file__), 'golden')
+
+
+class TestSpectrum(unittest.TestCase):
+
+    def test_psd_matches_matlab_cpsd_and_parseval(self):
+        label = 'psd'
+        d = sio.loadmat(os.path.join(GOLDEN, 'psd_matlab.mat'), squeeze_me=True)
+        fs, block = int(d['fs']), int(d['block'])
+        x = np.random.RandomState(1).random_sample(4096) - 0.5
+        y = ss.lfilter([1, 0.5, 0.25], 1, x)
+        sx, sy = (pytta.SignalObj(v, 'time', fs) for v in (x, y))
+        # SignalObj stores float32: compare with the float32 samples
+        x32, y32 = sx.timeSignal[:, 0].astype(float), sy.timeSignal[:, 0].astype(float)
+        f, Pxx = pytta.psd(sx, block_size=block)
+        _, Pxy = pytta.psd(sx, sy, block_size=block)
+        e_xx = np.abs(Pxx[:, 0] / d['Pxx'] - 1).max()
+        e_xy = np.abs(Pxy[:, 0] / d['Pxy'] - 1).max()
+        power = np.sum(Pxx[:, 0]) * (f[1] - f[0])
+        print(f'\n  PSD {e_xx:.1e}, CPSD {e_xy:.1e} relative to MATLAB; integral {power:.5f} '
+              f'against the variance {np.var(x32):.5f}')
+        self.assertLess(e_xx, 1e-5)
+        self.assertLess(e_xy, 1e-5)
+        self.assertAlmostEqual(power / np.var(x32), 1, delta=0.02)
+
+    def test_interpolated_spectrum_follows_the_true_response(self):
+        # h[n] = a^n has H = 1 / (1 - a exp(-j w)); its 256-point DFT,
+        # interpolated onto the bins of 1024 points, follows H there.
+        label = 'interp'
+        fs, a = 48000, 0.5
+        h = a ** np.arange(256)
+        sig = pytta.SignalObj(h, 'time', fs)
+        fine = pytta.interpolate_spectrum(sig, 1024)
+        H = np.fft.rfft(fine.timeSignal[:, 0].astype(float))
+        w = 2 * np.pi * np.arange(H.size) / 1024
+        err = np.abs(H - 1 / (1 - a * np.exp(-1j * w))).max()
+        same = pytta.interpolate_spectrum(sig, 256).timeSignal[:, 0]
+        print(f'\n  max |H - H_true| on the finer grid {err:.1e}; same length {np.abs(same - h).max():.1e}')
+        self.assertLess(err, 1e-3)
+        self.assertLess(np.abs(same - h).max(), 1e-6)
+
+    def test_low_extrapolation(self):
+        # A pure delay has a flat magnitude and the group delay of the
+        # delay, so its extrapolation reproduces it; a slope of 6 dB per
+        # decade leaves the magnitude 6 dB lower a decade below the range.
+        label = 'extrap'
+        fs, n, k = 1000, 1000, 25
+        x = np.zeros(n)
+        x[k] = 1
+        sig = pytta.SignalObj(x, 'time', fs)
+        same = pytta.extrapolate_low(sig, (100, 200)).timeSignal[:, 0]
+        sloped = np.abs(np.fft.rfft(pytta.extrapolate_low(sig, (100, 200), slope=6).timeSignal[:, 0]))
+        ratio = 20 * np.log10(sloped[10] / sloped[100])
+        print(f'\n  pure delay kept within {np.abs(same - x).max():.1e}; '
+              f'{ratio:.4f} dB from 100 Hz to 10 Hz')
+        self.assertLess(np.abs(same - x).max(), 1e-6)
+        self.assertAlmostEqual(ratio, -6, delta=1e-3)
+
+
+if __name__ == '__main__':
+    unittest.main()