def test_stft_linearity(self):
     """Test the linearity property of the Fourier transform."""
     x = fmlin(128, 0.0, 0.2)[0]
     y = fmlin(128, 0.3, 0.5)[0]
     h = x + y
     tfr, _, _ = linear.ShortTimeFourierTransform(h).run()
     tfrx, _, _ = linear.ShortTimeFourierTransform(x).run()
     tfry, _, _ = linear.ShortTimeFourierTransform(y).run()
     np.testing.assert_allclose(tfr, tfrx + tfry)

 def test_ideal_tfr(self):
     """Test if the ideal TFR can be found using the instantaneous frequency
     laws."""
     _, iflaw1 = fmlin(128, 0.0, 0.2)
     _, iflaw2 = fmlin(128, 0.3, 0.5)
     iflaws = np.c_[iflaw1, iflaw2].T
     tfr, _, _ = pproc.ideal_tfr(iflaws)
     tfr[tfr == 1] = 255
     tfr = tfr.astype(np.uint8)
     hspace, angles, dists = hough_line(tfr)
     for x in hough_line_peaks(hspace, angles, dists):
         self.assertEqual(len(x), 2)

 def test_spectrogram_energy_conservation(self):
     """Test the energy conservation property of the spectrogram."""
     signal, _ = fmlin(128, 0.1, 0.4)
     window = kaiser(17, 3 * np.pi)
     tfr, ts, freqs = cohen.Spectrogram(signal, n_fbins=64, fwindow=window).run()
     e_sig = (np.abs(signal) ** 2).sum()
     self.assertAlmostEqual(tfr.sum().sum() / 64, e_sig)

 def test_stft_conjugation(self):
     x = fmlin(128, 0, 0.2)[0]
     h = np.conjugate(x)
     lhs, _, _ = linear.ShortTimeFourierTransform(h).run()
     rhs, _, _ = linear.ShortTimeFourierTransform(x[::-1]).run()
     rhs = np.conjugate(rhs)
     np.testing.assert_allclose(lhs, rhs)

 def test_spectrogram_linearity(self):
     signal, _ = fmlin(128, 0.1, 0.4)
     window = kaiser(17, 3 * np.pi)
     tfr1, _, _ = cohen.spectrogram(signal, n_fbins=64, window=window)
     tfr2, _, _ = cohen.spectrogram(signal * 2, n_fbins=64, window=window)
     x = np.sum(np.sum(tfr2))
     y = np.sum(np.sum(tfr1))
     self.assertEqual(x / y, 4)

 def test_stft_translation(self):
     """Test the time-shift property of the Fourier transform."""
     x = fmlin(128, 0.0, 0.2)[0]
     tfrx, _, freqs = linear.ShortTimeFourierTransform(x).run()
     x_shifted = np.roll(x, 64)
     tfrxs, _, _ = linear.ShortTimeFourierTransform(x_shifted).run()
     f_mul = np.exp(-1j * 2 * np.pi * 64 * freqs)
     np.testing.assert_allclose(tfrxs, f_mul * tfrx)

 def test_sigmerge(self):
     """Test merging of signals with a given SNR."""
     signal = fmlin(128)[0]
     noise = np.random.randn(128,)
     gamma = 0.1
     x = utils.sigmerge(signal, noise, gamma)
     h_est = np.linalg.norm(signal) / np.linalg.norm(noise) * 10 ** (-gamma / 20)
     x_hat = signal + h_est * noise
     np.testing.assert_allclose(x, x_hat)

 def test_sigmerge(self):
     """Test merging of signals with a given SNR."""
     signal = fmlin(128)[0]
     noise = np.random.randn(128,)
     noisy_signal = utils.sigmerge(signal, noise)
     gamma_estimate = np.sqrt(signal.var() / noise.var())
     np.testing.assert_allclose(noisy_signal,
                                signal + gamma_estimate * noise, rtol=1e-2,
                                atol=1e-2)

 def test_stft_modulation(self):
     """Test the modulation / frequency shifting property of STFT."""
     x = fmlin(128, 0.0, 0.2)[0]
     tfrx, _, _ = linear.ShortTimeFourierTransform(x).run()
     f_0 = 0.3
     h = np.exp(1j * 2 * np.pi * f_0 * np.arange(128)) * x
     tfrh, _, _ = linear.ShortTimeFourierTransform(h).run()
     tfrx_shifted = np.roll(tfrx, int(np.ceil(128 * 0.3)), axis=0)
     self.assert_tfr_equal(tfrx_shifted, tfrh, tol=0.95)

 def test_spectrogram_linearity(self):
     """Test the linearity property of the spectrogram."""
     signal, _ = fmlin(128, 0.1, 0.4)
     window = kaiser(17, 3 * np.pi)
     tfr1, _, _ = cohen.Spectrogram(signal, n_fbins=64,
                                    fwindow=window).run()
     tfr2, _, _ = cohen.Spectrogram(signal * 2, n_fbins=64,
                                    fwindow=window).run()
     x = np.sum(np.sum(tfr2))
     y = np.sum(np.sum(tfr1))
     self.assertEqual(x / y, 4)

 def test_spectrogram_time_invariance(self):
     """Test the time invariance property of the spectrogram."""
     signal, _ = fmlin(128, 0.1, 0.4)
     window = kaiser(17, 3 * np.pi)
     tfr, ts, freqs = cohen.Spectrogram(signal, n_fbins=64, fwindow=window).run()
     shift = 64
     timeshifted_signal = np.roll(signal, shift)
     timeshifted_tfr, _, _ = cohen.Spectrogram(timeshifted_signal, n_fbins=64,
                                         fwindow=window).run()
     rolled_tfr = np.roll(tfr, shift, axis=1)
     # the time invariance property holds mostly qualitatively. The shifted
     # TFR is not numerically indentical to the rolled TFR, having
     # differences at the edges; so clip with two TFRs where there are
     # discontinuities in the TFR.
     edge = 10
     xx = np.c_[timeshifted_tfr[:, edge:(shift - edge)],
                timeshifted_tfr[:, (shift + edge):-edge]]
     yy = np.c_[rolled_tfr[:, edge:(shift - edge)],
                rolled_tfr[:, (shift + edge):-edge]]
     np.testing.assert_allclose(xx, yy)

Instantaneous frequency and group delay are very closely related. The former is
the frequency of a signal at a given instant, and the latter is the time delay
of frequency components. As this example shows, they coincide with each other
for a given signal when the time bandwidth product of the signal is
sufficiently high.

"""

import numpy as np
import matplotlib.pyplot as plt
from tftb.generators import amgauss, fmlin
from tftb.processing import loctime, locfreq, inst_freq, group_delay

time_instants = np.arange(2, 256)
sig1 = amgauss(256, 128, 90) * fmlin(256)[0]
tm, T1 = loctime(sig1)
fm, B1 = locfreq(sig1)
ifr1 = inst_freq(sig1, time_instants)[0]
f1 = np.linspace(0, 0.5 - 1.0 / 256, 256)
gd1 = group_delay(sig1, f1)

plt.subplot(211)
plt.plot(time_instants, ifr1, '*', label='inst_freq')
plt.plot(gd1, f1, '-', label='group delay')
plt.xlim(0, 256)
plt.grid(True)
plt.legend()
plt.title("Time-Bandwidth product: {0}".format(T1 * B1))
plt.xlabel('Time')
plt.ylabel('Normalized Frequency')

            points = np.arange(-min([lg, signal.shape[0] - ti - tau]),
                               min([lg, ti - 1 - tau]) + 1)
            g2 = twindow[lg + points]
            g2 = g2 / np.sum(g2)
            R = np.sum(g2 * signal[ti + tau - points - 1] * np.conj(signal[ti - tau - points - 1]))
            tfr[1 + tau, icol] = fwindow[lh + tau + 1] * R
            R = np.sum(g2 * signal[ti - tau - points - 1] * np.conj(signal[ti + tau - points - 1]))
            tfr[freq_bins - tau - 1, icol] = fwindow[lh - tau + 1] * R
        tau = np.round(freq_bins / 2.0)
        if (ti <= signal.shape[0] - tau) and (ti >= tau + 1) and (tau <= lh):
            points = np.arange(-min([lg, signal.shape[0] - ti - tau]),
                               min([lg, ti - 1 - tau]) + 1)
            g2 = twindow[lg + 1 + points]
            g2 = g2 / np.sum(g2)
            _x = np.sum(g2 * signal[ti + tau - points] * np.conj(signal[ti - tau - points]))
            _x *= fwindow[lh + tau + 1]
            _y = np.sum(g2 * signal[ti - tau - points] * np.conj(signal[ti + tau - points]))
            _y *= fwindow[lh - tau + 1]
            tfr[tau, icol] = (_x + _y) * 0.5
    tfr = np.fft.fft(tfr, axis=0)
    return np.real(tfr)


if __name__ == '__main__':
    from tftb.generators import fmlin
    sig = fmlin(128, 0.1, 0.4)[0]
    spec = PseudoWignerVilleDistribution(sig)
    tfr, _, _ = spec.run()
    from scipy.io import savemat
    savemat("/tmp/foo.mat", dict(tfr2=tfr))

# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2015 jaidev <[email protected]>
#
# Distributed under terms of the MIT license.

"""
Example showing use of Hough transform on a Wigner-Ville distribution.
"""

import numpy as np
from tftb.generators import noisecg, sigmerge, fmlin
from tftb.processing.cohen import WignerVilleDistribution
from tftb.processing.postprocessing import hough_transform
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

N = 64
sig = sigmerge(fmlin(N, 0, 0.3)[0], noisecg(N), 1)
tfr, _, _ = WignerVilleDistribution(sig).run()

ht, rho, theta = hough_transform(tfr, N, N)
theta, rho = np.meshgrid(theta, rho)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot_wireframe(theta, rho, ht)
ax.set_xlabel('Theta')
ax.set_ylabel('Rho')
plt.show()

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2015 jaidev <[email protected]>
#
# Distributed under terms of the MIT license.

"""

"""

from tftb.generators import amgauss, fmlin
import numpy as np
import matplotlib.pyplot as plt

z = amgauss(128, 50.0, 30.0) * fmlin(128, 0.05, 0.3, 50)[0]
plt.plot(np.real(z))
plt.xlim(0, 128)
plt.grid()
plt.title('Linear Frequency Modulation')
plt.show()

# Distributed under terms of the MIT license.

"""
Example from section 2.7 of the tutorial.
Short time Fourier transform of a multi-component nonstationary signal.
"""

from tftb.generators import fmlin
from tftb.processing.linear import ShortTimeFourierTransform
import matplotlib.pyplot as plt
from scipy.signal import hamming
import numpy as np
from mpl_toolkits.axes_grid1 import make_axes_locatable

N = 128
x1, _ = fmlin(N, 0, 0.2)
x2, _ = fmlin(N, 0.3, 0.5)
x = x1 + x2

n_fbins = 128
window = hamming(33)
tfr, _, _ = ShortTimeFourierTransform(x, timestamps=None, n_fbins=n_fbins,
                                      fwindow=window).run()
tfr = tfr[:64, :]
threshold = np.amax(np.abs(tfr)) * 0.05
tfr[np.abs(tfr) <= threshold] = 0.0 + 1j * 0.0
tfr = np.abs(tfr) ** 2
t = np.arange(tfr.shape[1])
f = np.linspace(0, 0.5, tfr.shape[0])

T, F = np.meshgrid(t, f)

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2015 jaidev <[email protected]>
#
# Distributed under terms of the MIT license.

"""
Examples from section 3.4.1 of the tutorial.
"""

from tftb.generators import fmlin
from tftb.processing.cohen import Spectrogram
import numpy as np
import matplotlib.pyplot as plt

sig = fmlin(128, 0, 0.3)[0] + fmlin(128, 0.2, 0.5)[0]
window = np.exp(np.log(0.005) * np.linspace(-1, 1, 63) ** 2)
spec = Spectrogram(sig, fwindow=window, n_fbins=128)
spec.run()
spec.plot(show_tf=True, cmap=plt.cm.gray)

 def test_spectrogram_reality(self):
     """Test the reality property of the spectrogram."""
     signal, _ = fmlin(128, 0.1, 0.4)
     window = kaiser(17, 3 * np.pi)
     tfr, _, _ = cohen.Spectrogram(signal, n_fbins=64, fwindow=window).run()
     self.assertTrue(np.all(np.isreal(tfr)))

#
# Copyright © 2015 jaidev <[email protected]>
#
# Distributed under terms of the MIT license.

"""
==================================================
Linear Frequency and Gaussian Amplitude Modulation
==================================================

Generate a mono-component nonstationary signal with linear frequency
modulation and Gaussian amplitude modulation.

"""

from tftb.generators import fmlin, amgauss
from numpy import real
import matplotlib.pyplot as plt


fm, _ = fmlin(256)
am = amgauss(256)
signal = fm * am
plt.plot(real(signal))
plt.xlabel("Time")
plt.ylabel("Real part")
plt.title("Linear Frequency, Gaussian Amplitude")
plt.xlim(0, 256)
plt.grid()
plt.show()

 def test_spectrogram_non_negativity(self):
     """Test that the spectrogram is non negative."""
     signal, _ = fmlin(128, 0.1, 0.4)
     window = kaiser(17, 3 * np.pi)
     tfr, _, _ = cohen.Spectrogram(signal, n_fbins=64, fwindow=window).run()
     self.assertTrue(np.all(tfr >= 0))