[WIP] Adding NUFFT and alternative mother wavelet support

jazzhands/base_wavelet.py

@@ -0,0 +1,146 @@
+import numpy as np
+
+class BaseWavelet:
+    """
+    Base class used to provide the functionality for the mother wavelets used
+    in the wavelet transform
+
+    Parameters
+    ----------
+    c : float
+        Decay rate of Gaussian envelope of wavelet. If not given defaults
+        to 0.0125.
+
+    w0 : float
+        Non-dimensional frequency, as given in eqn. 1 of Torrence and Compo 
+        (1998). It is related to the parameter c.
+
+    """
+    def __init__(self, c=None, w0=None):
+
+        if (c is None) and (w0 is None):
+            print("No resolution factor given. Setting c=0.0125")
+            self.c = 0.0125
+            # Omega0 as defined in Torrence and Compo
+            self.w0 = 1/np.sqrt(2*c)
+        elif (c is not None) and (w0 is None):
+            self.c = c
+            self.w0 = 1/np.sqrt(2*c)
+        elif (c is None) and (w0 is not None):
+            self.w0 = w0
+            self.c = 1/(2*w0**2)
+        else:
+            raise ValueError("Either c or w0 needs to be defined, "
+                             "but not both!")
+
+    def weights(self, time, omega, tau, c):
+        """
+        Weighting function for each point at a given omega and tau; (5-3) in Foster (1996).
+
+        Parameters
+        ----------
+        time : array-like
+            times of observations
+
+        omega : float
+            angular frequency in radians per unit time.
+
+        tau : float
+            time shift in same units as t
+
+        c : float
+            Decay constant of the Gaussian envelope for the wavelet
+
+        Returns
+        -------
+        array-like
+            Statistical weights of data points
+
+        """
+        raise NotImplementedError("Weights function has not been implemented"
+                                  " for chosen mother wavelet.")
+
+    def construct_wavelet(self, time, omega, tau, scale, domain='time'):
+        """
+        Construction function for a given mother wavelet. This can be done
+        either in the time domain, or the frequency domain.
+
+        Parameters
+        ----------
+        time : array-like
+            times of observations
+
+        omega : float
+            angular frequency in radians per unit time.
+
+        tau : float
+            time shift in same units as t
+
+        scale : float
+            scale ranges of the wavelet
+
+        Returns
+        -------
+        array-like
+            Mother wavelet evaluated either at given time stamps or frequencies
+
+        """
+        if domain == 'time':
+            return self._construct_wavelet_time(time, tau, scale)
+        elif domain == 'frequency':
+            return self._construct_wavelet_frequency(omega, tau, scale)
+        else:
+            # TODO this is just a placeholder - need to raise correct exception!
+            raise ValueError("Incorrect domain keyword set. Must be either time"
+                             " or frequency.")
+
+    def _construct_wavelet_time(self, time, tau, scale):
+        """
+        Construction function for a given mother wavelet in the time domain.
+
+        Parameters
+        ----------
+        time : array-like
+            times of observations
+
+        tau : float
+            time shift in same units as t
+
+        scale : float
+            scale ranges of the wavelet
+
+        Returns
+        -------
+        array-like
+            Mother wavelet evaluated either at given time stamps
+
+        """
+        raise NotImplementedError("Current mother wavelet has no function "
+                                  "implementing its construction in the time "
+                                  "domain.")
+
+    def _construct_wavelet_frequency(self, omega, tau, scale):
+        """
+        Construction function for a given mother wavelet in the frequency 
+        domain.
+
+        Parameters
+        ----------
+        omega : float
+            angular frequency in radians per unit time.
+
+        tau : float
+            time shift in same units as t
+
+        scale : float
+            scale ranges of the wavelet
+
+        Returns
+        -------
+        array-like
+            Mother wavelet evaluated either at given time stamps
+
+        """
+        raise NotImplementedError("Current mother wavelet has no function "
+                                  "implementing its construction in the "
+                                  "frequency domain.")

jazzhands/mother_wavelets.py

@@ -0,0 +1,105 @@
+import numpy as np
+import scipy
+import scipy.signal
+from scipy.special import factorial
+
+from .base_wavelet import BaseWavelet
+
+class Morlet(BaseWavelet):
+
+    def __init__(self, *args, **kwargs):
+        # Initialise parent:
+        super().__init__(*args, **kwargs)
+
+    def weights(self, time, omega, tau, c):
+        """
+        Weighting function for each point at a given omega and tau; (5-3) in Foster (1996).
+
+        Parameters
+        ----------
+        time : array-like
+            times of observations
+
+        omega : float
+            angular frequency in radians per unit time.
+
+        tau : float
+            time shift in same units as t
+
+        c : float
+            Decay constant of the Gaussian envelope for the wavelet
+
+        Returns
+        -------
+        array-like
+            Statistical weights of data points
+
+        """
+        weights = np.exp(-c * np.power(omega * (time - tau), 2.0))
+        return weights / np.sum(weights)
+
+    def _construct_wavelet_time(self, time, tau, scale):
+        """
+        Construction function for a given mother wavelet in the time domain.
+
+        Parameters
+        ----------
+        time : array-like
+            times of observations
+
+        tau : float
+            time shift in same units as t
+
+        scale : float
+            scale ranges of the wavelet
+
+        Returns
+        -------
+        array-like
+            Mother wavelet evaluated either at given time stamps
+
+        """
+        # TODO: Depends on whether scale is really used, but ignore for the
+        # moment
+        # Compute eta following Torrence and Compo (1998)
+        eta = (time - tau)/scale
+        output = np.exp(1j*eta) * np.exp(-self.c * eta**2)
+        return output
+
+    def _construct_wavelet_frequency(self, omega, tau, scale):
+        """
+        Construction function for a given mother wavelet in the frequency 
+        domain.
+
+        Parameters
+        ----------
+        omega : float
+            angular frequency in radians per unit time.
+
+        tau : float
+            time shift in same units as t
+
+        scale : float
+            scale ranges of the wavelet
+
+        Returns
+        -------
+        array-like
+            Mother wavelet evaluated either at given time stamps
+
+        """
+        # TODO: Again this is dependent on whether the actual scale is used or
+        # not. Will come back to this later.
+        x = omega * scale # Assuming scale is given
+        # Mock heaviside function
+        Heaviside = np.zeros_like(omega)
+        Heaviside[omega > 0] = 1
+        # Need to check that this is consistent with the formulation from 
+        # Foster (1996)
+        return Hw * np.exp((-x - self.w0) **2 / 2) 
+
+
+    def __call__(self, *args, **kwargs):
+        # Call method constructs the wavelet with given arguments and keyword
+        # arguments
+        return self.construct_wavelet(**args, **kwargs)