{"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":"### Idea","metadata":{}},{"cell_type":"markdown","source":"So the idea here is to train an autoencoder with only the Background signal and when presented with a Gravitation Wave, the autoencoder shouldn't be able to reconstruct it and hence should generate an unsually high reconstruction loss.","metadata":{}},{"cell_type":"code","source":"!pip install -qq gwpy ","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-10-06T10:18:58.627298Z","iopub.execute_input":"2021-10-06T10:18:58.628064Z","iopub.status.idle":"2021-10-06T10:19:18.443425Z","shell.execute_reply.started":"2021-10-06T10:18:58.627953Z","shell.execute_reply":"2021-10-06T10:19:18.442069Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport glob\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nsns.set()\n\nfrom sklearn.model_selection import train_test_split, StratifiedKFold\nfrom sklearn.metrics import roc_auc_score\n\nimport torch\nimport torch.nn.functional as F\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\n\nfrom pytorch_lightning import LightningModule, LightningDataModule, Trainer\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\n\nimport fastai\nfrom fastai import *\nfrom fastai.vision.all import *\n\nfrom gwpy.timeseries import TimeSeries\nfrom gwpy.plot import Plot\nimport numpy as np\nfrom scipy import signal\n\ndef seed_everything(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n    return seed\n    \n    \nSEED = 2704\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nseed_everything(SEED)","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:18.445306Z","iopub.execute_input":"2021-10-06T10:19:18.445643Z","iopub.status.idle":"2021-10-06T10:19:24.557424Z","shell.execute_reply.started":"2021-10-06T10:19:18.445606Z","shell.execute_reply":"2021-10-06T10:19:24.556025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To save time and resources, let's experiment with 1% of the data","metadata":{}},{"cell_type":"code","source":"%%time\n    \ndf = pd.read_csv(\"../input/g2net-gravitational-wave-detection/training_labels.csv\")\n# small_df = df\n\nsmall_df = df.sample(frac=0.01, random_state=SEED).reset_index(drop=True)\ntrain_df, test_df = train_test_split(small_df, test_size=0.2, random_state=SEED, shuffle=True, stratify=small_df['target'])\ntrain_df.shape, test_df.shape","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:24.560413Z","iopub.execute_input":"2021-10-06T10:19:24.560921Z","iopub.status.idle":"2021-10-06T10:19:25.112760Z","shell.execute_reply.started":"2021-10-06T10:19:24.560871Z","shell.execute_reply":"2021-10-06T10:19:25.111499Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To normalize the data, I used the statistics from: https://www.kaggle.com/mistag/mean-and-std-calculations-for-the-entire-dataset","metadata":{}},{"cell_type":"markdown","source":"### Data","metadata":{}},{"cell_type":"code","source":"import json\nwith open('../input/mean-and-std-calculations-for-the-entire-dataset/train_stats.json', 'r') as f:\n    train_stats = json.load(f)\n    \ntrain_mu, train_sigma = [], []\nfor item in train_stats['detector']:\n    train_mu += [item['mean']]\n    train_sigma += [item['std']]\n    \ntrain_mu, train_sigma = np.array(train_mu), np.array(train_sigma)\ntrain_mu, train_sigma","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:25.115075Z","iopub.execute_input":"2021-10-06T10:19:25.115550Z","iopub.status.idle":"2021-10-06T10:19:25.133213Z","shell.execute_reply.started":"2021-10-06T10:19:25.115504Z","shell.execute_reply":"2021-10-06T10:19:25.132246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def filters(array, sample_frequency=2048, lf=35, hf=350):\n    \"\"\" Apply preprocessing such as whitening and bandpass \"\"\"\n    strain = TimeSeries(array, sample_rate=int(sample_frequency))\n    # white_data = strain.whiten(fftlength=4, fduration=4)\n    # white_data = strain.whiten(window=(\"tukey\", 0.2))\n    white_data = strain\n    bp_data = white_data.bandpass(lf, hf)\n    return bp_data.value\n\nfrom scipy.ndimage import gaussian_filter1d\nfrom pathlib import Path\nINPUT_PATH = Path(\"../input/g2net-gravitational-wave-detection/\")\n\ndef load_wave(id_, mu, sigma, nc, folder='train'):\n    path = INPUT_PATH / folder / id_[0] / id_[1] / id_[2] / f\"{id_}.npy\"\n    waves = np.load(path).astype('float32').T\n    waves = ((waves - mu)/sigma).astype(np.float32)\n    # waves = gaussian_filter1d(waves, 0.5)\n    for idx in range(nc):\n        waves[:, idx] = filters(waves[:, idx])\n    return waves\n    \n\nclass Anomaly_Dataset(Dataset):\n    def __init__(self, df, nc=3):\n        self.df = df\n        self.nc = nc\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx, :]\n        waves = load_wave(row['id'], train_mu, train_sigma, self.nc)\n        return waves[:, :self.nc], row['target']\n\n    \n# Return only `wnd_size` sized random chunk background\nclass Background_Dataset(Dataset):\n    def __init__(self, df, wnd_size, nc=3):\n        self.df = df[df['target'] == 0].reset_index(drop=True)\n        self.nc = nc\n        self.wnd_size = wnd_size\n    def __len__(self):\n        return len(self.df)\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx, :]\n        waves = load_wave(row['id'], train_mu, train_sigma, self.nc)\n        \n        rnd_pos = np.random.randint(0, len(waves)-1-self.wnd_size)\n        waves = waves[rnd_pos:rnd_pos+self.wnd_size, :self.nc]\n        \n        return waves, waves","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:25.134686Z","iopub.execute_input":"2021-10-06T10:19:25.134993Z","iopub.status.idle":"2021-10-06T10:19:25.153346Z","shell.execute_reply.started":"2021-10-06T10:19:25.134963Z","shell.execute_reply":"2021-10-06T10:19:25.151717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Vis","metadata":{}},{"cell_type":"code","source":"nc = 3\nwnd_size = 128\nnum_samples = 4096\nn_wnd = num_samples // wnd_size\n# wnd_size = num_samples // n_wnd\n\nbds = Background_Dataset(train_df, wnd_size=wnd_size, nc=nc)\nrnd_idx = np.random.randint(len(bds))-1\nwaves = bds[rnd_idx][0]\nprint('rnd_idx:', rnd_idx, waves.shape)\nplt.figure(figsize=(20,4))\nfor _nc in range(nc):\n    plt.plot(waves[:, _nc], label=f'site{_nc}'); \nplt.legend()\n\ndl = DataLoader(bds, bs=4, shuffle=True)\nX, y = next(iter(dl))\nX.shape, y.shape","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:25.155341Z","iopub.execute_input":"2021-10-06T10:19:25.156146Z","iopub.status.idle":"2021-10-06T10:19:26.017793Z","shell.execute_reply.started":"2021-10-06T10:19:25.156089Z","shell.execute_reply":"2021-10-06T10:19:26.016292Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model","metadata":{}},{"cell_type":"markdown","source":"Yes, it's from the liverpool competition :p ","metadata":{}},{"cell_type":"code","source":"# from https://www.kaggle.com/hanjoonchoe/wavenet-lstm-pytorch-ignite-ver        \nclass Wave_Block(nn.Module):\n    \n    def __init__(self,in_channels,out_channels,dilation_rates):\n        super(Wave_Block,self).__init__()\n        self.num_rates = dilation_rates\n        self.convs = nn.ModuleList()\n        self.filter_convs = nn.ModuleList()\n        self.gate_convs = nn.ModuleList()\n        \n        self.convs.append(nn.Conv1d(in_channels,out_channels,kernel_size=1))\n        dilation_rates = [2**i for i in range(dilation_rates)]\n        for dilation_rate in dilation_rates:\n            self.filter_convs.append(nn.Conv1d(out_channels,out_channels,kernel_size=3,padding=dilation_rate,dilation=dilation_rate))\n            self.gate_convs.append(nn.Conv1d(out_channels,out_channels,kernel_size=3,padding=dilation_rate,dilation=dilation_rate))\n            self.convs.append(nn.Conv1d(out_channels,out_channels,kernel_size=1))\n            \n    def forward(self,x):\n        x = self.convs[0](x)\n        res = x\n        for i in range(self.num_rates):\n            x = torch.tanh(self.filter_convs[i](x)) * torch.sigmoid(self.gate_convs[i](x))\n            x = self.convs[i+1](x)\n            res = torch.add(res, x)\n        return res\n    \nclass Wavenet_denoiser(nn.Module):\n    def __init__(self, nc, hidden_dim=64, latent_dim=4):\n        super().__init__()\n        torch.cuda.empty_cache()\n        \n        self.conv1 = nn.Conv1d(nc, hidden_dim//2, kernel_size=3, padding=1)\n        self.encoder = nn.Sequential(\n            Wave_Block(hidden_dim//2, hidden_dim, 1),\n            nn.BatchNorm1d(hidden_dim),\n            nn.SiLU(),\n        )\n        self.rnn1 = nn.LSTM(input_size=hidden_dim, hidden_size=latent_dim, num_layers=1, batch_first=True, bidirectional=False)\n        self.rnn2 = nn.LSTM(input_size=latent_dim, hidden_size=hidden_dim, num_layers=1, batch_first=True, bidirectional=False)\n        \n        self.decoder = nn.Sequential(\n            Wave_Block(hidden_dim, hidden_dim//2, 1),\n            nn.BatchNorm1d(hidden_dim//2),\n            nn.SiLU(),\n        )\n        self.conv2 = nn.Conv1d(hidden_dim//2, nc, kernel_size=3, padding=1)\n            \n    def forward(self,x):\n        # ---- Encoder ----\n        x = x.permute(0, 2, 1)\n        x = self.conv1(x)\n        x = self.encoder(x)\n\n        # ---- Bottleneck ----\n        x = x.permute(0, 2, 1)\n        x, _ = self.rnn1(x)\n        \n        # ---- Decoder ----\n        x, _ = self.rnn2(x)\n        x = x.permute(0, 2, 1)\n        x = self.decoder(x)\n        \n        # ---- Output ----\n        x = self.conv2(x)\n        x = x.permute(0, 2, 1)        \n        return x","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:26.019394Z","iopub.execute_input":"2021-10-06T10:19:26.019714Z","iopub.status.idle":"2021-10-06T10:19:26.038334Z","shell.execute_reply.started":"2021-10-06T10:19:26.019684Z","shell.execute_reply":"2021-10-06T10:19:26.037610Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = Wavenet_denoiser(nc=nc).to(device)\nmodel","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:26.040304Z","iopub.execute_input":"2021-10-06T10:19:26.040809Z","iopub.status.idle":"2021-10-06T10:19:26.081368Z","shell.execute_reply.started":"2021-10-06T10:19:26.040776Z","shell.execute_reply":"2021-10-06T10:19:26.080274Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity check\nwith torch.no_grad():\n    print(X.shape, model(X.to(device)).shape)","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:26.083112Z","iopub.execute_input":"2021-10-06T10:19:26.083447Z","iopub.status.idle":"2021-10-06T10:19:26.242658Z","shell.execute_reply.started":"2021-10-06T10:19:26.083414Z","shell.execute_reply":"2021-10-06T10:19:26.241861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Fastai Learner","metadata":{}},{"cell_type":"code","source":"train_ds = Background_Dataset(train_df, wnd_size=wnd_size, nc=nc)\nval_ds = Background_Dataset(test_df, wnd_size=wnd_size, nc=nc)\nanomaly_ds = Anomaly_Dataset(test_df, nc=nc)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nbs = 64 if torch.cuda.is_available() else 4\n\ndef RMSELoss(yhat, y):\n    return torch.sqrt(torch.mean((yhat-y)**2))\n\ndls = DataLoaders.from_dsets(train_ds, val_ds, bs=bs, device=device)\nlearn = Learner(\n    dls, model, loss_func=RMSELoss, opt_func=Adam,\n    cbs=[fastai.callback.all.ShowGraphCallback(),fastai.callback.all.SaveModelCallback(fname='best'), fastai.callback.all.CSVLogger()]\n)","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:26.243908Z","iopub.execute_input":"2021-10-06T10:19:26.244400Z","iopub.status.idle":"2021-10-06T10:19:26.259121Z","shell.execute_reply.started":"2021-10-06T10:19:26.244357Z","shell.execute_reply":"2021-10-06T10:19:26.258114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learn.lr_find()","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:26.261070Z","iopub.execute_input":"2021-10-06T10:19:26.261610Z","iopub.status.idle":"2021-10-06T10:19:40.233154Z","shell.execute_reply.started":"2021-10-06T10:19:26.261564Z","shell.execute_reply":"2021-10-06T10:19:40.231662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learn.fit_one_cycle(n_epoch=200, lr_max=5e-3)\nlearn.save('dae')","metadata":{"execution":{"iopub.status.busy":"2021-10-06T10:19:40.235055Z","iopub.execute_input":"2021-10-06T10:19:40.235553Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import scipy\n\ndef moving_average(x, w=8):\n    return np.convolve(x, np.ones(w), 'valid') / w\n\n# https://www.kaggle.com/alexnitz/pycbc-making-images\ndef bandpass2rgb(data, f_range=(35,350), q_range=(16,32), q_max=10):\n    data = map(lambda x: TimeSeries(x, sample_rate=2048), data.T)\n    # Q-transform\n    data = map(lambda x: x.q_transform(qrange=q_range, frange=f_range, logf=True, whiten=False), data)\n    # Convert to RGB image\n    img = np.stack(list(data), axis = -1)\n    img = np.clip(img, 0, q_max)/q_max * 255\n    img = img.astype(np.uint8)\n    img = Image.fromarray(img).rotate(90, expand=1)\n    return img\n\ndef plot_denoised(data, folder='train', n_wnd=n_wnd, num_samples=num_samples, nc=nc, figsize=(20, 4)):\n    if isinstance(data, str):\n        waves = load_wave(data, train_mu, train_sigma, nc, folder=folder)\n    else:\n        waves = data\n    rgb = bandpass2rgb(waves)\n    \n    plt.figure(figsize=figsize)\n    plt.title('Q-Transformed')\n    plt.imshow(rgb)\n    \n    plt.figure(figsize=figsize)\n    plt.title('Original')\n    for _nc in range(nc):\n        plt.plot(waves[:, _nc], label=f'site{_nc}'); \n    plt.legend()\n    \n    wnd_size = num_samples // n_wnd\n    __waves = torch.from_numpy(waves).unsqueeze(dim=0).view(1*n_wnd, num_samples//n_wnd, nc)\n    with torch.no_grad():\n        print(__waves.shape)\n        pred = model(__waves.to(device))\n    raw_rmse = torch.sqrt((pred.cpu()-__waves)**2).reshape(num_samples, nc).mean(-1).numpy()\n    pred = pred.reshape(num_samples, nc).cpu().numpy()\n        \n    # \"\"\"\n    plt.figure(figsize=figsize)\n    plt.title(f'Reconstructed: RMSELoss {raw_rmse.mean():.05f}, maxRMSELoss {np.max(raw_rmse):.05f} @ wnd_idx:[{np.argmax(raw_rmse)+1}/{n_wnd}]')\n    for _nc in range(nc):\n        plt.plot(pred[:, _nc], label=f'site{_nc}');\n    # for i in range(n_wnd+1):\n    #     plt.axvline(x=i*wnd_size)\n    plt.legend()\n    # \"\"\"\n    \n    plt.figure(figsize=figsize)\n    plt.title('RMSELoss')\n    plt.plot(raw_rmse, label='raw')\n    \n    peaks, _ = scipy.signal.find_peaks(raw_rmse, width=2)\n    print('#peaks', len(peaks))\n    plt.plot(peaks, raw_rmse[peaks], \"x\")\n    plt.legend()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The model can also provide some localization for the Gravitation Waves.","metadata":{}},{"cell_type":"code","source":"# a very clean chirp:\nplot_denoised('0021f9dd71', folder='test', n_wnd=num_samples//1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## a not-so-clean chirp\nplot_denoised('000a5b6e5c', n_wnd=num_samples//4)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Some random samples","metadata":{}},{"cell_type":"code","source":"rnd_idx = np.random.randint(len(anomaly_ds))-1\nrnd_X, rnd_y = anomaly_ds[rnd_idx][0], anomaly_ds[rnd_idx][1]\nplot_denoised(rnd_X)\nprint('GW:', rnd_y==1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rnd_idx = np.random.randint(len(anomaly_ds))-1\nrnd_X, rnd_y = anomaly_ds[rnd_idx][0], anomaly_ds[rnd_idx][1]\nplot_denoised(rnd_X)\nprint('GW:', rnd_y==1)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}