{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:31:58.853202Z","iopub.execute_input":"2025-04-10T00:31:58.853528Z","iopub.status.idle":"2025-04-10T00:32:02.944387Z","shell.execute_reply.started":"2025-04-10T00:31:58.853501Z","shell.execute_reply":"2025-04-10T00:32:02.943661Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Find files to load and create Dataset","metadata":{}},{"cell_type":"code","source":"all_inputs = [\n    f\n    for f in\n    Path('/kaggle/input/waveform-inversion/train_samples').rglob('*.npy')\n    if ('seis' in f.stem) or ('data' in f.stem)\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:02.945471Z","iopub.execute_input":"2025-04-10T00:32:02.945917Z","iopub.status.idle":"2025-04-10T00:32:03.047684Z","shell.execute_reply.started":"2025-04-10T00:32:02.945892Z","shell.execute_reply":"2025-04-10T00:32:03.046935Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def inputs_files_to_output_files(input_files):\n    return [\n        Path(str(f).replace('seis', 'vel').replace('data', 'model'))\n        for f in input_files\n    ]\n\nall_outputs = inputs_files_to_output_files(all_inputs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:03.049138Z","iopub.execute_input":"2025-04-10T00:32:03.049440Z","iopub.status.idle":"2025-04-10T00:32:03.053899Z","shell.execute_reply.started":"2025-04-10T00:32:03.049413Z","shell.execute_reply":"2025-04-10T00:32:03.052970Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"assert all(f.exists() for f in all_outputs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:03.055164Z","iopub.execute_input":"2025-04-10T00:32:03.055381Z","iopub.status.idle":"2025-04-10T00:32:03.068342Z","shell.execute_reply.started":"2025-04-10T00:32:03.055362Z","shell.execute_reply":"2025-04-10T00:32:03.067475Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_inputs = [all_inputs[i] for i in range(0, len(all_inputs), 2)] # Sample every two\nvalid_inputs = [f for f in all_inputs if not f in train_inputs]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:03.069365Z","iopub.execute_input":"2025-04-10T00:32:03.069643Z","iopub.status.idle":"2025-04-10T00:32:03.082200Z","shell.execute_reply.started":"2025-04-10T00:32:03.069620Z","shell.execute_reply":"2025-04-10T00:32:03.081301Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_outputs = inputs_files_to_output_files(train_inputs)\nvalid_outputs = inputs_files_to_output_files(valid_inputs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:03.082944Z","iopub.execute_input":"2025-04-10T00:32:03.083200Z","iopub.status.idle":"2025-04-10T00:32:03.096540Z","shell.execute_reply.started":"2025-04-10T00:32:03.083179Z","shell.execute_reply":"2025-04-10T00:32:03.095611Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SeismicDataset(Dataset):\n    def __init__(self, inputs_files, output_files, n_examples_per_file=500):\n        assert len(inputs_files) == len(output_files)\n        self.inputs_files = inputs_files\n        self.output_files = output_files\n        self.n_examples_per_file = n_examples_per_file\n\n    def __len__(self):\n        return len(self.inputs_files) * self.n_examples_per_file\n\n    def __getitem__(self, idx):\n        # Calculate file offset and sample offset within file\n        file_idx = idx // self.n_examples_per_file\n        sample_idx = idx % self.n_examples_per_file\n\n        X = np.load(self.inputs_files[file_idx], mmap_mode='r')\n        y = np.load(self.output_files[file_idx], mmap_mode='r')\n\n        try:\n            return X[sample_idx].copy(), y[sample_idx].copy()\n        finally:\n            del X, y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:03.097443Z","iopub.execute_input":"2025-04-10T00:32:03.097751Z","iopub.status.idle":"2025-04-10T00:32:03.106448Z","shell.execute_reply.started":"2025-04-10T00:32:03.097722Z","shell.execute_reply":"2025-04-10T00:32:03.105611Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dstrain = SeismicDataset(train_inputs, train_outputs)\ndltrain = DataLoader(dstrain, batch_size=64, shuffle=True, pin_memory=True, drop_last=True, num_workers=4, persistent_workers=True)\n\ndsvalid = SeismicDataset(valid_inputs, valid_outputs)\ndlvalid = DataLoader(dsvalid, batch_size=64, shuffle=False, pin_memory=True, drop_last=False, num_workers=4, persistent_workers=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:03.107265Z","iopub.execute_input":"2025-04-10T00:32:03.107561Z","iopub.status.idle":"2025-04-10T00:32:03.121061Z","shell.execute_reply.started":"2025-04-10T00:32:03.107533Z","shell.execute_reply":"2025-04-10T00:32:03.120180Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## DumbNet","metadata":{}},{"cell_type":"code","source":"class DumbNet(nn.Module):\n    '''DumbNet is just a MLP Model, with a avg-pool first to reduze input size'''\n    def __init__(self, pool_size=(8, 2), input_size=5 * 1000 * 70, hidden_size=70 * 70, output_size=70 * 70):\n        super().__init__()\n\n        self.pool = nn.AvgPool2d(kernel_size=pool_size)\n\n        self.model = nn.Sequential(\n            nn.Linear(input_size // (pool_size[0] * pool_size[1]), hidden_size),\n            \n            nn.ReLU(),\n            nn.Dropout(0.25),\n            nn.Linear(hidden_size, hidden_size),\n\n            nn.ReLU(),\n            nn.Dropout(0.25),\n            nn.Linear(hidden_size, output_size),\n        )\n\n    def forward(self, x):\n        bs = x.shape[0]\n\n        # We apply a pool to reduze input size\n        x_pool = self.pool(x)\n\n        #Model is just a\n        out = self.model(x_pool.view(bs, -1))\n\n        return out.view(bs, 1, 70, 70) * 1000 + 1500","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:42.022780Z","iopub.execute_input":"2025-04-10T00:32:42.023188Z","iopub.status.idle":"2025-04-10T00:32:42.030033Z","shell.execute_reply.started":"2025-04-10T00:32:42.023148Z","shell.execute_reply":"2025-04-10T00:32:42.029170Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\ndevice","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:42.521142Z","iopub.execute_input":"2025-04-10T00:32:42.521452Z","iopub.status.idle":"2025-04-10T00:32:42.527322Z","shell.execute_reply.started":"2025-04-10T00:32:42.521428Z","shell.execute_reply":"2025-04-10T00:32:42.526540Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = DumbNet().to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:42.876952Z","iopub.execute_input":"2025-04-10T00:32:42.877286Z","iopub.status.idle":"2025-04-10T00:32:44.330929Z","shell.execute_reply.started":"2025-04-10T00:32:42.877261Z","shell.execute_reply":"2025-04-10T00:32:44.329987Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train loop","metadata":{}},{"cell_type":"code","source":"criterion = nn.L1Loss()\noptim = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-3)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:44.332063Z","iopub.execute_input":"2025-04-10T00:32:44.332382Z","iopub.status.idle":"2025-04-10T00:32:44.336875Z","shell.execute_reply.started":"2025-04-10T00:32:44.332352Z","shell.execute_reply":"2025-04-10T00:32:44.336020Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"n_epochs = 50\n\nhistory = []\n\nfor epoch in range(1, n_epochs+1):\n    print(f'[{epoch:02d}] Begin train')\n\n    # Train\n    model.train()\n    train_losses = []\n    for inputs, targets in tqdm(dltrain, desc='train', leave=False):\n        inputs = inputs.to(device)\n        targets = targets.to(device)\n        \n        optim.zero_grad()\n        \n        outputs = model(inputs)\n        \n        loss = criterion(outputs, targets)\n        \n        loss.backward()\n        \n        optim.step()\n\n        train_losses.append(loss.item())\n\n    print('Train loss: {:.5f}'.format( np.mean(train_losses) ))\n\n    # Valid\n    model.eval()\n    valid_losses = []\n    for inputs, targets in tqdm(dlvalid, desc='valid', leave=False):\n        inputs = inputs.to(device)\n        targets = targets.to(device)\n\n        with torch.inference_mode():\n            outputs = model(inputs)\n        \n        loss = criterion(outputs, targets)\n\n        valid_losses.append(loss.item())\n    \n    print('Valid loss: {:.5f}'.format( np.mean(valid_losses)) )\n    history.append({\n        'train': np.mean(train_losses),\n        'valid': np.mean(valid_losses)\n    })\n\n    # Plot last result\n    if epoch % 4 == 0:\n        y = targets[0, 0].detach().cpu()\n        y_pred = outputs[0, 0].detach().cpu()\n        \n        fig, ax = plt.subplots(1, 2, figsize=(5, 2.5))\n        fig.suptitle(f'Epoch {epoch} | Valid: {np.mean(valid_losses):.5f}')\n        ax[0].imshow(y)\n        ax[1].imshow(y_pred)\n        plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:45.445342Z","iopub.execute_input":"2025-04-10T00:32:45.445652Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.DataFrame(history).plot();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.188061Z","iopub.status.idle":"2025-04-10T00:32:13.188413Z","shell.execute_reply":"2025-04-10T00:32:13.188294Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Predict test","metadata":{}},{"cell_type":"code","source":"import csv  # Use \"low-level\" CSV to save memory on predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.189716Z","iopub.status.idle":"2025-04-10T00:32:13.190227Z","shell.execute_reply":"2025-04-10T00:32:13.190002Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\ntest_files = list(Path('/kaggle/input/waveform-inversion/test').glob('*.npy'))\nlen(test_files)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.191022Z","iopub.status.idle":"2025-04-10T00:32:13.191396Z","shell.execute_reply":"2025-04-10T00:32:13.191238Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"x_cols = [f'x_{i}' for i in range(1, 70, 2)]\nfieldnames = ['oid_ypos'] + x_cols","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.192877Z","iopub.status.idle":"2025-04-10T00:32:13.193415Z","shell.execute_reply":"2025-04-10T00:32:13.193153Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, test_files):\n        self.test_files = test_files\n\n\n    def __len__(self):\n        return len(self.test_files)\n\n\n    def __getitem__(self, i):\n        test_file = self.test_files[i]\n\n        return np.load(test_file), test_file.stem","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.195123Z","iopub.status.idle":"2025-04-10T00:32:13.195612Z","shell.execute_reply":"2025-04-10T00:32:13.195393Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ds = TestDataset(test_files)\ndl = DataLoader(ds, batch_size=8, num_workers=4, pin_memory=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.196329Z","iopub.status.idle":"2025-04-10T00:32:13.196707Z","shell.execute_reply":"2025-04-10T00:32:13.196525Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train\nmodel.eval()\nwith open('submission.csv', 'wt', newline='') as csvfile:\n    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n    writer.writeheader()\n    \n    for inputs, oids_test in tqdm(dl, desc='test'):\n        inputs = inputs.to(device)\n        with torch.inference_mode():\n            outputs = model(inputs)\n\n        y_preds = outputs[:, 0].cpu().numpy()\n        \n        for y_pred, oid_test in zip(y_preds, oids_test):\n            for y_pos in range(70):\n                row = dict(\n                    zip(\n                        x_cols,\n                        [y_pred[y_pos, x_pos] for x_pos in range(1, 70, 2)]\n                    )\n                )\n                row['oid_ypos'] = f\"{oid_test}_y_{y_pos}\"\n            \n                writer.writerow(row)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-10T00:32:13.197484Z","iopub.status.idle":"2025-04-10T00:32:13.197929Z","shell.execute_reply":"2025-04-10T00:32:13.197746Z"}},"outputs":[],"execution_count":null}]}