labels for the frequency axis.

Author: v0lta

examples/wavelet_packet_chirp_analysis/chirp.png

@@ -11,24 +11,36 @@
 t = np.linspace(0, 2, int(2//(1/fs)))
 w = np.sin(256*np.pi*t**2)
 
-wavelet = pywt.Wavelet("sym6")
+wavelet = pywt.Wavelet("sym8")
 wp = WaveletPacket(
     data=torch.from_numpy(w.astype(np.float32)), wavelet=wavelet, mode="boundary"
 )
-nodes = wp.get_level(5)
+level = 5
+nodes = wp.get_level(level)
 np_lst = []
 for node in nodes:
     np_lst.append(wp[node])
 viz = np.stack(np_lst).squeeze()
 
+n = list(range(int(np.power(2, level))))
+freqs = (fs/2)*(n/(np.power(2, level)))
+
+xticks = list(range(viz.shape[-1]))[::6]
+xlabels = np.round(np.linspace(min(t), max(t), viz.shape[-1]), 2)[::6]
+
 fig, axs = plt.subplots(2)
 axs[0].plot(t, w)
 axs[0].set_title("Analyzed signal")
-axs[0].set_xlabel("t [s]")
+axs[0].set_xlabel("time [s]")
+axs[0].set_ylabel("magnitude")
 
-axs[1].set_title("Wavelet analysis")
+axs[1].set_title("Wavelet packet analysis")
 axs[1].imshow(np.abs(viz))
-axs[1].set_xlabel("time")
-axs[1].set_ylabel("frequency")
+axs[1].set_xlabel("time [s]")
+axs[1].set_xticks(xticks)
+axs[1].set_xticklabels(xlabels)
+axs[1].set_ylabel("frequency [Hz]")
+axs[1].set_yticks(n[::4])
+axs[1].set_yticklabels(freqs[::4])
 axs[1].invert_yaxis()
 plt.show()

examples/wavelet_packet_chirp_analysis/chirp_analysis.py

@@ -8,6 +8,7 @@
 
 import pywt
 import torch
+import numpy as np
 
 from ._util import Wavelet, _as_wavelet
 from .conv_transform import wavedec, waverec
@@ -22,6 +23,23 @@
     BaseDict = collections.UserDict
 
 
+def _wpfreq(fs: float, level: int) -> np.ndarray:
+    """Compute the frequencies for a fully decomposed single dimensional
+       packet tree. The packet transform linearly subdivides all frequencies
+       from zero up to the Nyquist frequency.
+
+    Args:
+        fs (float): The sampling frequency
+        level (int): The decomposition level
+
+    Returns:
+        np.ndarray: The frequency bins of the packets in frequency order.
+    """
+    n = list(range(int(np.power(2., level))))
+    freqs = (fs/2.)*(n/(np.power(2., level)))
+    return freqs
+
+
 class WaveletPacket(BaseDict):
     """One dimensional wavelet packets."""