hht_example.py

import numpy as np
import pylab as plt
from scipy.signal import hilbert

from PyEMD import EMD


def instant_phase(imfs):
    """Extract analytical signal through Hilbert Transform."""
    analytic_signal = hilbert(imfs)  # Apply Hilbert transform to each row
    phase = np.unwrap(np.angle(analytic_signal))  # Compute angle between img and real
    return phase


# Define signal
t = np.linspace(0, 1, 200)
dt = t[1] - t[0]

sin = lambda x, p: np.sin(2 * np.pi * x * t + p)
S = 3 * sin(18, 0.2) * (t - 0.2) ** 2
S += 5 * sin(11, 2.7)
S += 3 * sin(14, 1.6)
S += 1 * np.sin(4 * 2 * np.pi * (t - 0.8) ** 2)
S += t**2.1 - t

# Compute IMFs with EMD
emd = EMD()
imfs = emd(S, t)

# Extract instantaneous phases and frequencies using Hilbert transform
instant_phases = instant_phase(imfs)
instant_freqs = np.diff(instant_phases) / (2 * np.pi * dt)

# Create a figure consisting of 3 panels which from the top are the input signal, IMFs and instantaneous frequencies
fig, axes = plt.subplots(3, figsize=(12, 12))

# The top panel shows the input signal
ax = axes[0]
ax.plot(t, S)
ax.set_ylabel("Amplitude [arb. u.]")
ax.set_title("Input signal")

# The middle panel shows all IMFs
ax = axes[1]
for num, imf in enumerate(imfs):
    ax.plot(t, imf, label="IMF %s" % (num + 1))

# Label the figure
ax.legend()
ax.set_ylabel("Amplitude [arb. u.]")
ax.set_title("IMFs")

# The bottom panel shows all instantaneous frequencies
ax = axes[2]
for num, instant_freq in enumerate(instant_freqs):
    ax.plot(t[:-1], instant_freq, label="IMF %s" % (num + 1))

# Label the figure
ax.legend()
ax.set_xlabel("Time [s]")
ax.set_ylabel("Inst. Freq. [Hz]")
ax.set_title("Huang-Hilbert Transform")

plt.tight_layout()
plt.savefig("hht_example", dpi=120)
plt.show()