work in progress

Author: tomrunia

torch_wavelets/extractors.py

@@ -30,7 +30,7 @@
 
 class TemporalFilterBankBase(metaclass=ABCMeta):
 
-    def __init__(self, dt=1.0, dj=0.125, wavelet=Morlet(), unbias=False, signal_length=None):
+    def __init__(self, dt=1.0, dj=0.125, wavelet=Morlet(), unbias=False, signal_length=512):
 
         self._dt = dt
         self._dj = dj
@@ -47,16 +47,16 @@ def compute(self, signal):
         raise NotImplementedError
 
     def _init_filters(self):
-        filters = []
-        for i, scale in enumerate(self._scales):
+        filters = [None]*len(self.scales)
+        for scale_idx, scale in enumerate(self._scales):
             # number of points needed to capture wavelet
             M = 10 * scale / self.dt
             # times to use, centred at zero
             t = np.arange((-M + 1) / 2., (M + 1) / 2.) * dt
             if len(t) % 2 == 0: t = t[0:-1]  # requires odd filter size
             # sample wavelet and normalise
             norm = (self.dt / scale) ** .5
-            filters[i] = norm * self.wavelet(t, scale)
+            filters[scale_idx] = norm * self.wavelet(t, scale)
         return filters
 
     def compute_optimal_scales(self):
@@ -78,11 +78,11 @@ def func_to_solve(s):
             return self.fourier_period(s) - 2 * dt
         return scipy.optimize.fsolve(func_to_solve, 1)[0]
 
-    def power(self, signal):
+    def power(self, x):
         if self.unbias:
-            return (np.abs(self.compute(signal)).T ** 2 / self._scales).T
+            return (np.abs(self.compute(x)).T ** 2 / self.scales).T
         else:
-            return np.abs(self.compute(signal)) ** 2
+            return np.abs(self.compute(x)) ** 2
 
     @property
     def fourier_period(self):
@@ -136,13 +136,15 @@ def compute(self, x):
         num_examples = x.shape[0]
         output = np.zeros((num_examples, len(self.scales), x.shape[-1]), dtype=np.complex)
         for example_idx in range(num_examples):
-            output[example_idx] = self.compute_single(x[example_idx])
-        return np.squeeze(output, 0)
+            output[example_idx] = self._compute_single(x[example_idx])
+        if num_examples == 1:
+            output = output.squeeze(0)
+        return output
 
-    def compute_single(self, x):
+    def _compute_single(self, x):
         assert x.ndim == 1, 'input signal must have single dimension.'
         output = np.zeros((len(self.scales), len(x)), dtype=np.complex)
-        for scale_idx, filt in enumerate(self._filters)
+        for scale_idx, filt in enumerate(self._filters):
             output[scale_idx,:] = scipy.signal.fftconvolve(x, filt, mode='same')
         return output
 
@@ -199,13 +201,36 @@ def _get_padding(padding_type, kernel_size):
 
 if __name__ == "__main__":
 
-    dt = 1.0
-    dj = 0.125
-    wavelet = Morlet(w0=6)
+    import torch_wavelets.utils as utils
+    import matplotlib.pyplot as plt
+
+    fps = 20
+    dt  = 1.0/fps
+    dj  = 0.125
+    w0  = 6
     unbias = False
+    wavelet = Morlet()
+
+    t_min = 0
+    t_max = 10
+    t = np.linspace(t_min, t_max, (t_max-t_min)*fps)
+
+    batch_size = 12
+
+    # Generate a batch of sine waves with random frequency
+    random_frequencies = np.random.uniform(-0.5, 2.0, size=batch_size)
+    batch = np.asarray([np.sin(2*np.pi*f*t) for f in random_frequencies])
+
+    wa = TemporalFilterBankSciPy(dt, dj, wavelet, unbias)
+    power = wa.power(batch)
+
+    fig, ax = plt.subplots(3, 4, figsize=(16,8))
+    ax = ax.flatten()
+    for i in range(batch_size):
+        utils.plot_scalogram(power[i], wa.scales, t, ax=ax[i])
+        ax[i].axhline(1.0 / random_frequencies[i], lw=1, color='k')
+    plt.show()
 
-    cls_scipy = TemporalFilterBankSciPy(dt, dj, wavelet, unbias)
-    cls_torch = TemporalFilterBankTorch(dt, dj, wavelet, unbias)
 
 
 

torch_wavelets/utils.py

@@ -17,28 +17,51 @@
 from __future__ import print_function
 
 import numpy as np
-from torch_wavelets.wavelets import *
-
-
-
-def scale_distribution(wavelet, min_period, max_period, dj=0.125):
-    """
-    Given a minimum and maximum period compute a distribution of scales. The choice of
-    dj depends on the width in spectral space of the wavelet function.
-    For the Morlet, dj=0.5 is the largest that still adequately samples scale.
-    Smaller dj gives finer scale resolution.
-
-    :param wavelet: wavelet instance from (Morlet, Ricker, MexicanHat, Marr)
-    :param min_period: float, minimum period
-    :param max_period: float, maximum period
-    :param dj: float, scale sample density
-    :return: np.ndarray, containing a list of scales
-    """
-
-    assert isinstance(wavelet, (Morlet, Ricker, Marr, Mexican_hat))
-    scale_min = wavelet.scale_from_period(min_period)
-    scale_max = wavelet.scale_from_period(max_period)
-    num_scales = int(np.ceil((1.0/dj) * np.log2(scale_max/scale_min)))
-    j = np.arange(0, num_scales+1)
-    scales = scale_min*np.power(2, j*dj)
-    return scales
+import matplotlib.pyplot as plt
+
+# This makes the color map of same height as the image
+import matplotlib.ticker as ticker
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+
+def plot_scalogram(power, scales, t, normalize_columns=True, cmap=None, ax=None, scale_legend=True):
+
+    if not cmap:
+        cmap = plt.get_cmap("PuBu_r")
+
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    if normalize_columns:
+        power = power/np.max(power, axis=0)
+
+    T, S = np.meshgrid(t, scales)
+    cnt = ax.contourf(T, S, power, 100, cmap=cmap)
+
+    # Fix for saving as PDF (aliasing)
+    for c in cnt.collections:
+        c.set_edgecolor("face")
+
+    ax.set_yscale('log')
+    ax.set_ylabel("Scale (Log Scale)")
+    ax.set_xlabel("Time (s)")
+    ax.set_title("Wavelet Power Spectrum")
+
+    if scale_legend:
+
+        def format_axes_label(x, pos):
+            return "{:.2f}".format(x)
+
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(
+            cnt, cax=cax, ticks=[np.min(power), 0, np.max(power)],
+            format=ticker.FuncFormatter(format_axes_label))
+
+    return ax
+
+
+
+
+
+