{"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":"code","source":"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install torch --no-deps --upgrade --quiet\n!pip install timm --no-deps --quiet","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:37.748119Z","iopub.execute_input":"2022-06-19T03:29:37.749359Z","iopub.status.idle":"2022-06-19T03:29:41.987346Z","shell.execute_reply.started":"2022-06-19T03:29:37.749317Z","shell.execute_reply":"2022-06-19T03:29:41.986124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport pywt\nimport torch\nimport torch.nn as nn\nfrom scipy import signal\nfrom scipy import optimize\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom IPython.display import Image\nfrom timm.models.layers.conv2d_same import conv2d_same\n\nINPUT_PATH = Path(\"../input/g2net-gravitational-wave-detection/\")","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:41.98972Z","iopub.execute_input":"2022-06-19T03:29:41.990072Z","iopub.status.idle":"2022-06-19T03:29:41.99671Z","shell.execute_reply.started":"2022-06-19T03:29:41.990037Z","shell.execute_reply":"2022-06-19T03:29:41.995798Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"t = np.linspace(-1, 1, 200, endpoint=False)\nsig = np.cos(2 * np.pi * 7 * t) + signal.gausspulse(t - 0.4, fc=2)\n\nplt.plot(t, sig);","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:41.99817Z","iopub.execute_input":"2022-06-19T03:29:41.998531Z","iopub.status.idle":"2022-06-19T03:29:42.213678Z","shell.execute_reply.started":"2022-06-19T03:29:41.998502Z","shell.execute_reply":"2022-06-19T03:29:42.212745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"widths = np.arange(1, 31)\n","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:42.214716Z","iopub.execute_input":"2022-06-19T03:29:42.214993Z","iopub.status.idle":"2022-06-19T03:29:42.218931Z","shell.execute_reply.started":"2022-06-19T03:29:42.214968Z","shell.execute_reply":"2022-06-19T03:29:42.21822Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"coefs, freqs = pywt.cwt(sig, widths, \"mexh\")\n\nplt.imshow(\n    coefs,\n    extent=[-1, 1, 1, 31],\n    cmap=\"PRGn\",\n    aspect=\"auto\",\n    vmax=abs(coefs).max(),\n    vmin=-abs(coefs).max(),\n);","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:42.221622Z","iopub.execute_input":"2022-06-19T03:29:42.221937Z","iopub.status.idle":"2022-06-19T03:29:42.447985Z","shell.execute_reply.started":"2022-06-19T03:29:42.221908Z","shell.execute_reply":"2022-06-19T03:29:42.447358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#SciPy\ncwtmatr = signal.cwt(sig, signal.ricker, widths)\n\nplt.imshow(\n    cwtmatr,\n    extent=[-1, 1, 1, 31],\n    cmap=\"PRGn\",\n    aspect=\"auto\",\n    vmax=abs(cwtmatr).max(),\n    vmin=-abs(cwtmatr).max(),\n);","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:42.449004Z","iopub.execute_input":"2022-06-19T03:29:42.449772Z","iopub.status.idle":"2022-06-19T03:29:42.674266Z","shell.execute_reply.started":"2022-06-19T03:29:42.44974Z","shell.execute_reply":"2022-06-19T03:29:42.673455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Transformaciones complejas\nLo ideal es crear una imagen que relacione frecuencia y tiempo. Una forma de hacerlo es usar una transformación compleja y tomar la magnitud de las partes reales e imaginarias. A continuación usaremos el complejo wavelet de Morlet en SciPy que se ve así:","metadata":{}},{"cell_type":"code","source":"Image(url='https://upload.wikimedia.org/wikipedia/commons/thumb/4/4f/Wavelet_Cmor.svg/1920px-Wavelet_Cmor.svg.png', width=500)","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:42.675789Z","iopub.execute_input":"2022-06-19T03:29:42.67647Z","iopub.status.idle":"2022-06-19T03:29:42.683364Z","shell.execute_reply.started":"2022-06-19T03:29:42.676427Z","shell.execute_reply":"2022-06-19T03:29:42.682403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cwtmatr = signal.cwt(sig, signal.morlet2, widths)\n\nreal, imag = np.real(cwtmatr), np.imag(cwtmatr)\n\n# Plot the real & imaginary components after CWT using a complex Morlet wavelet\nfig, axes = plt.subplots(ncols=2, figsize=(20, 5))\naxes[0].imshow(\n    real,\n    extent=[-1, 1, 1, 31],\n    cmap=\"PRGn\",\n    aspect=\"auto\",\n    vmax=abs(real).max(),\n    vmin=-abs(real).max(),\n)\naxes[1].imshow(\n    imag,\n    extent=[-1, 1, 1, 31],\n    cmap=\"PRGn\",\n    aspect=\"auto\",\n    vmax=abs(imag).max(),\n    vmin=-abs(imag).max(),\n);\n","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:42.684457Z","iopub.execute_input":"2022-06-19T03:29:42.685124Z","iopub.status.idle":"2022-06-19T03:29:43.121689Z","shell.execute_reply.started":"2022-06-19T03:29:42.685079Z","shell.execute_reply":"2022-06-19T03:29:43.120615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot the magnitude with the raw signal\nmagn = np.absolute(cwtmatr)\nprint(magn.shape)\nfig, axes = plt.subplots(nrows=2, figsize=(10, 10))\naxes[0].imshow(\n    magn,\n    extent=[-1, 1, 1, 31],\n    aspect=\"auto\",\n)\naxes[1].plot(t, sig);\n(30, 200)","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:43.123101Z","iopub.execute_input":"2022-06-19T03:29:43.123689Z","iopub.status.idle":"2022-06-19T03:29:43.503501Z","shell.execute_reply.started":"2022-06-19T03:29:43.123653Z","shell.execute_reply":"2022-06-19T03:29:43.50236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"varía la matriz de magnitud cuando cambia la frecuencia de la señal bruta.\n\nImplementación de PyTorch\nLa clase siguiente implementa la wavelet de Morlet","metadata":{}},{"cell_type":"code","source":"# From https://github.com/tomrunia/PyTorchWavelets/blob/master/wavelets_pytorch/wavelets.py\nclass 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        wavelet=Morlet(),\n        fmin: int = 20,\n        fmax: int = 500,\n        output_format=\"Magnitude\",\n        trainable=False,\n        hop_length: int = 1,\n    ):\n        super().__init__()\n        self.wavelet = wavelet\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, hop_length)\n        # self.padding = 0  # \"same\"\n\n        self._scale_minimum = self.compute_minimum_scale()\n\n        self.signal_length = None\n        self._channels = None\n\n        self._scales = None\n        self._kernel = None\n        self._kernel_real = None\n        self._kernel_imag = None\n\n    def compute_optimal_scales(self):\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(\n                \"Please specify signal_length before computing optimal scales.\"\n            )\n        J = int(\n            (1 / self.dj) * np.log2(self.signal_length * self.dt / self._scale_minimum)\n        )\n        scales = self._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[0 : len(frequencies)]\n        if self.fmax:\n            frequencies = frequencies[frequencies <= self.fmax]\n            scales = scales[len(scales) - len(frequencies) : len(scales)]\n\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        dt = self.dt\n\n        def func_to_solve(s):\n            return self.wavelet.fourier_period(s) - 2 * dt\n\n        return optimize.fsolve(func_to_solve, 1)[0]\n\n    def _build_filters(self):\n        self._filters = []\n        for scale_idx, scale in enumerate(self._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            self._filters.append(torch.conj(torch.flip(filter_, [-1])))\n\n        self._pad_filters()\n\n    def _pad_filters(self):\n        filter_len = self._filters[-1].shape[0]\n        padded_filters = []\n\n        for f in self._filters:\n            pad = (filter_len - f.shape[0]) // 2\n            padded_filters.append(nn.functional.pad(f, (pad, pad)))\n\n        self._filters = padded_filters\n\n    def _build_wavelet_bank(self):\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        self._build_filters()\n        wavelet_bank = torch.stack(self._filters)\n        wavelet_bank = wavelet_bank.view(\n            wavelet_bank.shape[0], 1, 1, wavelet_bank.shape[1]\n        )\n        # See comment by tez6c32\n        # https://www.kaggle.com/anjum48/continuous-wavelet-transform-cwt-in-pytorch/comments#1499878\n        # wavelet_bank = torch.cat([wavelet_bank] * self.channels, 2)\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        if self.signal_length is None:\n            self.signal_length = x.shape[-1]\n            self.channels = x.shape[-2]\n            self._scales = self.compute_optimal_scales()\n            self._kernel = self._build_wavelet_bank()\n\n            if self._kernel.is_complex():\n                self._kernel_real = self._kernel.real\n                self._kernel_imag = self._kernel.imag\n\n        x = x.unsqueeze(1)\n\n        if self._kernel.is_complex():\n            if (\n                x.dtype != self._kernel_real.dtype\n                or x.device != self._kernel_real.device\n            ):\n                self._kernel_real = self._kernel_real.to(device=x.device, dtype=x.dtype)\n                self._kernel_imag = self._kernel_imag.to(device=x.device, dtype=x.dtype)\n\n            # Strides > 1 not yet supported for \"same\" padding\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 = conv2d_same(x, self._kernel_real, stride=self.stride)\n            output_imag = conv2d_same(x, self._kernel_imag, stride=self.stride)\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)\n            else:\n                return torch.stack([output_real, output_imag], -1)\n\n        else:\n            if x.device != self._kernel.device:\n                self._kernel = self._kernel.to(device=x.device, dtype=x.dtype)\n\n            # output = nn.functional.conv2d(\n            #     x, self._kernel, padding=self.padding, stride=self.stride\n            # )\n            output = conv2d_same(x, self._kernel, stride=self.stride)\n            return torch.transpose(output, 1, 2)","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:43.505311Z","iopub.execute_input":"2022-06-19T03:29:43.505713Z","iopub.status.idle":"2022-06-19T03:29:43.552375Z","shell.execute_reply.started":"2022-06-19T03:29:43.50568Z","shell.execute_reply":"2022-06-19T03:29:43.551148Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test with the complex Morlet wavelet and plot the magnitude\nsig_pt = torch.tensor(sig, dtype=torch.float32)\nsig_pt = torch.stack([sig_pt] * 3)  # 3 channels\nsig_pt = torch.stack([sig_pt] * 32)  # Batch size of 32\nprint(sig_pt.shape)\n\n\npycwt = CWT(dt=1/400)\nout = pycwt(sig_pt)\nprint(out.shape)\n\nplt.imshow(out[0, 0].numpy(), aspect=\"auto\");\n\n","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:43.553898Z","iopub.execute_input":"2022-06-19T03:29:43.554302Z","iopub.status.idle":"2022-06-19T03:29:43.810884Z","shell.execute_reply.started":"2022-06-19T03:29:43.554269Z","shell.execute_reply":"2022-06-19T03:29:43.809979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Prueba en datos de GW","metadata":{}},{"cell_type":"code","source":"def load_file(id_, folder=\"train\"):\n    path = INPUT_PATH / folder / id_[0] / id_[1] / id_[2] / f\"{id_}.npy\"\n    waves = np.load(path)\n    # return waves / np.max(waves, axis=1).reshape(3, 1)\n    return waves / np.max(waves)\n\n# Cell 33 of https://www.gw-openscience.org/LVT151012data/LOSC_Event_tutorial_LVT151012.html\n# https://scipy-cookbook.readthedocs.io/items/ButterworthBandpass.html\ndef apply_bandpass(x, lf=25, hf=500, order=4, sr=2048):\n    sos = signal.butter(order, [lf, hf], btype=\"bandpass\", output=\"sos\", fs=sr)\n    normalization = np.sqrt((hf - lf) / (sr / 2))\n    return signal.sosfiltfilt(sos, x) / normalization","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:43.813039Z","iopub.execute_input":"2022-06-19T03:29:43.81352Z","iopub.status.idle":"2022-06-19T03:29:43.821919Z","shell.execute_reply.started":"2022-06-19T03:29:43.813476Z","shell.execute_reply":"2022-06-19T03:29:43.820605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x = load_file(\"000a5b6e5c\")\nx *= signal.tukey(4096, 0.2)\nx = apply_bandpass(x, 35, 500)\nx_ten = torch.tensor(x, dtype=torch.float32).view(1, 3, 4096)\n\npycwt = CWT(fmin=20, fmax=500, hop_length=8, dj=0.125/8)\nout = pycwt(x_ten)\nprint(out.shape)\n\nplt.imshow(out[0, 2].numpy(), aspect=\"auto\");","metadata":{"execution":{"iopub.status.busy":"2022-06-19T03:29:43.823774Z","iopub.execute_input":"2022-06-19T03:29:43.824421Z","iopub.status.idle":"2022-06-19T03:29:43.898415Z","shell.execute_reply.started":"2022-06-19T03:29:43.824374Z","shell.execute_reply":"2022-06-19T03:29:43.897406Z"},"trusted":true},"execution_count":null,"outputs":[]}]}