text_plain_output_1.png
from numba import njit, jit, cuda, guvectorize

from scipy.signal import butter, filtfilt, sosfiltfilt

from scipy.signal import cwt, ricker

from scipy.signal import spectrogram

import glob

import librosa

import matplotlib.pyplot as plt

import numpy as np # linear algebra

import pathlib

import torch

import torch

import glob
import pathlib
head = pathlib.Path('../input/g2net-gravitational-wave-detection')
train_files = sorted(glob.glob('../input/g2net-gravitational-wave-detection/train/*/*/*/*.npy'))

wave = np.load(train_files[0])

from numba import njit, jit, cuda, guvectorize
@njit(nogil=True)
def min_max_scaler(wave):
    for i in range(len(wave)):
        wave[i] = (wave[i] - min(wave[i])) / (max(wave[i]) - min(wave[i]))
        wave[i] = 2 * wave[i] - 1
    return wave

wave1 = min_max_scaler(wave)

from scipy.signal import butter, filtfilt, sosfiltfilt
T = 2
fs = 2048.0
cutoff = 2.5
nyq = 0.5 * fs
order = 3
n = int(T * fs)
normal_cutoff = cutoff / nyq

normal_cutoff = (21.83 / fs, 500 / fs)
def butter_bandpass_filter(data, normal_cutoff, fs, order=2):
    b, a = butter(order, normal_cutoff, btype='bandpass', analog=False)
    y = filtfilt(b, a, data)
    return y

def butter_bandpass_filter(data, low, high, fs, order):
    sos = butter(order, [low, high], btype='bandpass', output='sos', fs=fs)
    normalization = np.sqrt((high - low) / (fs / 2))
    return sosfiltfilt(sos, data) / normalization

def butter_lowpass_filter(data, normal_cutoff, fs, order):
    b, a = butter(order, normal_cutoff, btype='lowpass', analog=False)
    y = filtfilt(b, a, data)
    return y

data = torch.from_numpy(wave)
y = butter_bandpass_filter(data, 21.83, 500, fs, 4)

from scipy.signal import spectrogram
for i in range(len(wave)):
    f, t, Sxx = spectrogram(wave1[i], fs=10)
    plt.pcolormesh(t, f, Sxx, shading='gouraud')

def wrapper_plot(m):
    m()

stacked = []
for j in range(len(wave1)):
    melspec = librosa.feature.melspectrogram(wave1[j], sr=4096, n_mels=128, fmin=21.83, fmax=2048)
    melspec = librosa.power_to_db(melspec)
    melspec = melspec.transpose((1, 0))
    stacked.append(melspec)
image = np.vstack(stacked)

wave.shape

cwt(wave[0], signal.ricker, np.arange(1, 300))

class CWT(nn.Module):
    def __init__(self, widths, wavelet='ricker', channels=1, filter_len=2000):
        """PyTorch implementation of a continuous wavelet transform.
        Args:
            widths (iterable): The wavelet scales to use, e.g. np.arange(1, 33)
            wavelet (str, optional): Name of wavelet. Either "ricker" or "morlet".
            Defaults to "ricker".
            channels (int, optional): Number of audio channels in the input. Defaults to 3.
            filter_len (int, optional): Size of the wavelet filter bank. Set to
            the number of samples but can be smaller to save memory. Defaults to 2000.
        """
        super().__init__()
        self.widths = widths
        self.wavelet = getattr(self, wavelet)
        self.filter_len = filter_len
        self.channels = channels
        self.wavelet_bank = self._build_wavelet_bank()
    def ricker(self, points, a):
        A = 2 / (np.sqrt(3 * a) * np.pi ** 0.25)
        wsq = a ** 2
        vec = torch.arange(0, points) - (points - 1.0) / 2
        xsq = vec ** 2
        mod = 1 - xsq / wsq
        gauss = torch.exp(-xsq / (2 * wsq))
        total = A * mod * gauss
        return total
    def morlet(self, points, s):
        x = torch.arange(0, points) - (points - 1.0) / 2
        x = x / s
        wavelet = torch.exp(-x ** 2.0 / 2.0) * torch.cos(5.0 * x)
        output = np.sqrt(1 / s) * wavelet
        return output
    def cmorlet(self, points, s, wavelet_width=1, center_freq=1):
        x = torch.arange(0, points) - (points - 1.0) / 2
        x = x / s
        norm_constant = np.sqrt(np.pi * wavelet_width)
        exp_term = torch.exp(-x ** 2 / wavelet_width)
        kernel_base = exp_term / norm_constant
        kernel = kernel_base * torch.exp(1j * 2 * np.pi * center_freq * x)
        return kernel
    def _build_wavelet_bank(self):
        """This function builds a 2D wavelet filter using wavelets at different scales
        Returns:
            tensor: Tensor of shape (num_widths, 1, channels, filter_len)
        """
        wavelet_bank = [torch.conj(torch.flip(self.wavelet(self.filter_len, w), [-1])) for w in self.widths]
        wavelet_bank = torch.stack(wavelet_bank)
        wavelet_bank = wavelet_bank.view(wavelet_bank.shape[0], 1, 1, wavelet_bank.shape[1])
        wavelet_bank = torch.cat([wavelet_bank] * self.channels, 2)
        return wavelet_bank
    def forward(self, x):
        """Compute CWT arrays from a batch of multi-channel inputs
        Args:
            x (torch.tensor): Tensor of shape (batch_size, channels, time)
        Returns:
            torch.tensor: Tensor of shape (batch_size, channels, widths, time)
        """
        x = x.unsqueeze(1)
        if self.wavelet_bank.is_complex():
            wavelet_real = self.wavelet_bank.real.to(device=x.device, dtype=x.dtype)
            wavelet_imag = self.wavelet_bank.imag.to(device=x.device, dtype=x.dtype)
            output_real = nn.functional.conv2d(x, wavelet_real, padding='same')
            output_imag = nn.functional.conv2d(x, wavelet_imag, padding='same')
            output_real = torch.transpose(output_real, 1, 2)
            output_imag = torch.transpose(output_imag, 1, 2)
            return torch.complex(output_real, output_imag)
        else:
            self.wavelet_bank = self.wavelet_bank.to(device=x.device, dtype=x.dtype)
            output = nn.functional.conv2d(x, self.wavelet_bank, padding='same')
            return torch.transpose(output, 1, 2)

widths = np.arange(20, 120)
pycwt = CWT(widths, 'cmorlet', 3, 4096)

def apply_qtransform(waves, transform=CQT1992v2(sr=2048, fmin=21.83, fmax=1024, hop_length=64), cuda=False):
    waves = min_max_scaler(butter_bandpass_filter(waves, 21.83, 500, fs, 4))
    waves = torch.from_numpy(waves).float().view(1, 3, 4096)
    if cuda:
        waves = waves.cuda()
    image = pycwt(waves)
    return image
imgs = []
for i in range(10):
    img = apply_qtransform(wave, transform=CQT1992v2(sr=2048, fmin=21.83, n_bins=63, hop_length=64))
    imgs.append(img)

img[0, 0].shape