{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CWT in Pytorch that actually works in GPUs with no overhead\n\nsay thanks to [datasaurus](https://www.kaggle.com/anjum48) and like the original [notebook](https://www.kaggle.com/anjum48/continuous-wavelet-transform-cwt-in-pytorch).","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport scipy.optimize as optimize\nimport torch\nimport torch.nn as nn","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Morlet(object):\n    def __init__(self, w0=6):\n        \"\"\"w0 is the nondimensional frequency constant. If this is\n        set too low then the wavelet does not sample very well: a\n        value over 5 should be ok; Terrence and Compo set it to 6.\n        \"\"\"\n        self.w0 = w0\n        if w0 == 6:\n            # value of C_d from TC98\n            self.C_d = 0.776\n\n    def __call__(self, *args, **kwargs):\n        return self.time(*args, **kwargs)\n\n    def time(self, t, s=1.0, complete=True):\n        \"\"\"\n        Complex Morlet wavelet, centred at zero.\n        Parameters\n        ----------\n        t : float\n            Time. If s is not specified, this can be used as the\n            non-dimensional time t/s.\n        s : float\n            Scaling factor. Default is 1.\n        complete : bool\n            Whether to use the complete or the standard version.\n        Returns\n        -------\n        out : complex\n            Value of the Morlet wavelet at the given time\n        See Also\n        --------\n        scipy.signal.gausspulse\n        Notes\n        -----\n        The standard version::\n            pi**-0.25 * exp(1j*w*x) * exp(-0.5*(x**2))\n        This commonly used wavelet is often referred to simply as the\n        Morlet wavelet.  Note that this simplified version can cause\n        admissibility problems at low values of `w`.\n        The complete version::\n            pi**-0.25 * (exp(1j*w*x) - exp(-0.5*(w**2))) * exp(-0.5*(x**2))\n        The complete version of the Morlet wavelet, with a correction\n        term to improve admissibility. For `w` greater than 5, the\n        correction term is negligible.\n        Note that the energy of the return wavelet is not normalised\n        according to `s`.\n        The fundamental frequency of this wavelet in Hz is given\n        by ``f = 2*s*w*r / M`` where r is the sampling rate.\n        \"\"\"\n        w = self.w0\n\n        x = t / s\n\n        output = np.exp(1j * w * x)\n\n        if complete:\n            output -= np.exp(-0.5 * (w ** 2))\n\n        output *= np.exp(-0.5 * (x ** 2)) * np.pi ** (-0.25)\n\n        return output\n\n    # Fourier wavelengths\n    def fourier_period(self, s):\n        \"\"\"Equivalent Fourier period of Morlet\"\"\"\n        return 4 * np.pi * s / (self.w0 + (2 + self.w0 ** 2) ** 0.5)\n\n    def scale_from_period(self, period):\n        \"\"\"\n        Compute the scale from the fourier period.\n        Returns the scale\n        \"\"\"\n        # Solve 4 * np.pi * scale / (w0 + (2 + w0 ** 2) ** .5)\n        #  for s to obtain this formula\n        coeff = np.sqrt(self.w0 * self.w0 + 2)\n        return (period * (coeff + self.w0)) / (4.0 * np.pi)\n\n    # Frequency representation\n    def frequency(self, w, s=1.0):\n        \"\"\"Frequency representation of Morlet.\n        Parameters\n        ----------\n        w : float\n            Angular frequency. If `s` is not specified, i.e. set to 1,\n            this can be used as the non-dimensional angular\n            frequency w * s.\n        s : float\n            Scaling factor. Default is 1.\n        Returns\n        -------\n        out : complex\n            Value of the Morlet wavelet at the given frequency\n        \"\"\"\n        x = w * s\n        # Heaviside mock\n        Hw = np.array(w)\n        Hw[w <= 0] = 0\n        Hw[w > 0] = 1\n        return np.pi ** -0.25 * Hw * np.exp((-((x - self.w0) ** 2)) / 2)\n\n    def coi(self, s):\n        \"\"\"The e folding time for the autocorrelation of wavelet\n        power at each scale, i.e. the timescale over which an edge\n        effect decays by a factor of 1/e^2.\n        This can be worked out analytically by solving\n            |Y_0(T)|^2 / |Y_0(0)|^2 = 1 / e^2\n        \"\"\"\n        return 2 ** 0.5 * s\n\n\nclass CWT(nn.Module):\n    def __init__(\n        self,\n        dj=0.0625,\n        dt=1 / 2048,\n        fmin: int = 20,\n        fmax: int = 500,\n        output_format=\"Magnitude\",\n        trainable=False,\n        padding=\"same\",\n        signal_length=4096,\n        channels=3\n    ):\n        super().__init__()\n        self.wavelet = Morlet()\n\n        self.dt = dt\n        self.dj = dj\n        self.fmin = fmin\n        self.fmax = fmax\n        self.output_format = output_format\n        self.trainable = trainable  # TODO make kernel a trainable parameter\n        self.stride = (1, 1)  # Strides > 1 not yet supported\n        self.padding = padding\n\n        self.signal_length = signal_length\n        self.channels = channels\n        \n        scale_minimum = self.compute_minimum_scale()\n        scales = self.compute_optimal_scales(scale_minimum)\n        kernel = self.build_wavelet_bank(scales)\n    \n        if kernel.is_complex():\n            self.register_buffer(\"kernel_real\", kernel.real, persistent=False)\n            self.register_buffer(\"kernel_imag\", kernel.imag, persistent=False)\n        else:\n            self.register_buffer(\"kernel\", kernel, persistent=False)\n\n    def compute_optimal_scales(self, scale_minimum):\n        \"\"\"\n        Determines the optimal scale distribution (see. Torrence & Combo, Eq. 9-10).\n        :return: np.ndarray, collection of scales\n        \"\"\"\n        if self.signal_length is None:\n            raise ValueError(\"Please specify signal_length before computing optimal scales.\")\n        J = int((1 / self.dj) * np.log2(self.signal_length * self.dt / scale_minimum))\n        scales = scale_minimum * 2 ** (self.dj * np.arange(0, J + 1))\n\n        # Remove high and low frequencies\n        frequencies = np.array([1 / self.wavelet.fourier_period(s) for s in scales])\n        if self.fmin:\n            frequencies = frequencies[frequencies >= self.fmin]\n            scales = scales[slice(0, len(frequencies))]\n        if self.fmax:\n            frequencies = frequencies[frequencies <= self.fmax]\n            scales = scales[slice(len(scales) - len(frequencies), len(scales))]\n        return scales\n\n    def compute_minimum_scale(self):\n        \"\"\"\n        Choose s0 so that the equivalent Fourier period is 2 * dt.\n        See Torrence & Combo Sections 3f and 3h.\n        :return: float, minimum scale level\n        \"\"\"\n\n        def func_to_solve(s):\n            return self.wavelet.fourier_period(s) - 2 * self.dt\n\n        return optimize.fsolve(func_to_solve, 1)[0]\n\n    def build_filters(self, scales):\n        filters = []\n        for scale_idx, scale in enumerate(scales):\n            # Number of points needed to capture wavelet\n            M = 10 * scale / self.dt\n            # Times to use, centred at zero\n            t = torch.arange((-M + 1) / 2.0, (M + 1) / 2.0) * self.dt\n            if len(t) % 2 == 0:\n                t = t[0:-1]  # requires odd filter size\n            # Sample wavelet and normalise\n            norm = (self.dt / scale) ** 0.5\n            filter_ = norm * self.wavelet(t, scale)\n            filters.append(torch.conj(torch.flip(filter_, [-1])))\n\n        filters = self.pad_filters(filters)\n        return filters\n\n    def pad_filters(self, filters):\n        filter_len = filters[-1].shape[0]\n        padded_filters = []\n\n        for f in filters:\n            pad = (filter_len - f.shape[0]) // 2\n            padded_filters.append(nn.functional.pad(f, (pad, pad)))\n\n        return padded_filters\n\n    def build_wavelet_bank(self, scales):\n        \"\"\"This function builds a 2D wavelet filter using wavelets at different scales\n\n        Returns:\n            tensor: Tensor of shape (num_widths, 1, channels, filter_len)\n        \"\"\"\n\n        filters = self.build_filters(scales)\n        wavelet_bank = torch.stack(filters)\n        wavelet_bank = wavelet_bank.view(wavelet_bank.shape[0], 1, 1, wavelet_bank.shape[1])\n\n        return wavelet_bank\n\n    def forward(self, x):\n        \"\"\"Compute CWT arrays from a batch of multi-channel inputs\n\n        Args:\n            x (torch.tensor): Tensor of shape (batch_size, channels, time)\n\n        Returns:\n            torch.tensor: Tensor of shape (batch_size, channels, widths, time)\n        \"\"\"\n\n        x = x.unsqueeze(1)\n\n        if not hasattr(self, 'kernel'):\n            output_real = nn.functional.conv2d(\n                x, self.kernel_real, padding=self.padding, stride=self.stride\n            )\n            output_imag = nn.functional.conv2d(\n                x, self.kernel_imag, padding=self.padding, stride=self.stride\n            )\n            output_real = torch.transpose(output_real, 1, 2)\n            output_imag = torch.transpose(output_imag, 1, 2)\n\n            if self.output_format == \"Magnitude\":\n                return torch.sqrt(output_real ** 2 + output_imag ** 2).contiguous()\n            else:\n                return torch.stack([output_real, output_imag], -1).contiguous()\n\n        else:\n\n            output = nn.functional.conv2d(x, self.kernel, padding=self.padding, stride=self.stride)\n            return torch.transpose(output, 1, 2).contiguous()\n","metadata":{},"execution_count":null,"outputs":[]}]}