{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"},{"sourceId":345534,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":288699,"modelId":309457}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## UNet Model","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nimport numpy as np\n\ndef conv3x3(in_channels, out_channels, stride=1, padding=1, bias=True):    \n    return nn.Conv2d(\n        in_channels,\n        out_channels,\n        kernel_size=3,\n        stride=stride,\n        padding=padding,\n        bias=bias)\n\ndef conv1x1(in_channels, out_channels):\n    return nn.Conv2d(\n        in_channels,\n        out_channels,\n        kernel_size=1,\n        stride=1)\n\nclass CustomUNetRegression(nn.Module):\n    \"\"\"\n    Custom UNet architecture for transforming input of shape (batch, 5, 1000, 70)\n    to output of shape (batch, 1, 70, 70)\n    \"\"\"\n    def __init__(self, in_channels=5, out_channels=1):\n        super(CustomUNetRegression, self).__init__()\n        \n        # Initial feature extractor\n        self.conv_init = nn.Sequential(\n            conv3x3(in_channels, 64),\n            nn.ReLU(inplace=True),\n            conv3x3(64, 64),\n            nn.ReLU(inplace=True)\n        )\n        \n        # Downsampling path - progressively reduce spatial dimensions\n        # First downsampling block - reduce height by factor of ~2\n        self.down1 = nn.Sequential(\n            nn.MaxPool2d(kernel_size=(2, 1)),  # Pool only along height, preserving width\n            conv3x3(64, 128),\n            nn.ReLU(inplace=True),\n            conv3x3(128, 128),\n            nn.ReLU(inplace=True)\n        )\n        \n        # Second downsampling block - further reduce height\n        self.down2 = nn.Sequential(\n            nn.MaxPool2d(kernel_size=(2, 1)),  # Pool only along height\n            conv3x3(128, 256),\n            nn.ReLU(inplace=True),\n            conv3x3(256, 256),\n            nn.ReLU(inplace=True)\n        )\n        \n        # Third downsampling block - further reduce height\n        self.down3 = nn.Sequential(\n            nn.MaxPool2d(kernel_size=(2, 1)),  # Pool only along height\n            conv3x3(256, 512),\n            nn.ReLU(inplace=True),\n            conv3x3(512, 512),\n            nn.ReLU(inplace=True)\n        )\n        \n        # Final downsampling to reach target size (70,70)\n        # Need to reduce from ~125x70 to 70x70 \n        # We'll use a custom adaptive pooling layer\n        self.final_pool = nn.AdaptiveAvgPool2d((70, 70))\n        \n        # Final convolution to get to the desired output channels\n        self.final_conv = conv1x1(512, out_channels)\n        \n        # Initialize weights\n        self._initialize_weights()\n        \n    def _initialize_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                nn.init.xavier_normal_(m.weight)\n                if m.bias is not None:\n                    nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.constant_(m.weight, 1)\n                nn.init.constant_(m.bias, 0)\n    \n    def forward(self, x):\n        # Initial feature extraction\n        x = self.conv_init(x)  # -> (B, 64, 1000, 70)\n        \n        # Downsampling path\n        x = self.down1(x)      # -> (B, 128, 500, 70)\n        x = self.down2(x)      # -> (B, 256, 250, 70)\n        x = self.down3(x)      # -> (B, 512, 125, 70)\n        \n        # Adaptive pooling to get to target size\n        x = self.final_pool(x) # -> (B, 512, 70, 70)\n        \n        # Final convolution to get to the right number of channels\n        x = self.final_conv(x) # -> (B, 1, 70, 70)\n        \n        return x\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-19T07:29:39.430237Z","iopub.execute_input":"2025-04-19T07:29:39.430436Z","iopub.status.idle":"2025-04-19T07:29:43.840276Z","shell.execute_reply.started":"2025-04-19T07:29:39.430393Z","shell.execute_reply":"2025-04-19T07:29:43.839703Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Training","metadata":{}},{"cell_type":"markdown","source":"### Data Loading","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\n\n# # Find all model and data pairs\n# base_path = '/kaggle/input/waveform-inversion/train_samples'\n# dataset_types = ['CurveFault_A','CurveFault_B','CurveVel_A','CurveVel_B','FlatFault_A','FlatFault_B', 'FlatVel_A', 'FlatVel_B', 'Style_A', 'Style_B']\n\n# # Function to process a single dataset\n# def load_and_process_dataset(model_path, data_path):\n#     velocity = np.load(model_path)  # Shape: (500, 1, 70, 70)\n#     data = np.load(data_path)       # Shape: (500, 5, 1000, 70)\n    \n#     return velocity, data\n\n# # Lists to store all loaded data\n# all_velocity = []\n# all_seismic = []\n\n# # Load all available datasets\n# for dataset_type in dataset_types:\n#     dataset_path = os.path.join(base_path, dataset_type)\n    \n#     # Check if this folder structure exists\n#     if os.path.exists(os.path.join(dataset_path, 'model')) and os.path.exists(os.path.join(dataset_path, 'data')):\n#         model_files = [f for f in os.listdir(os.path.join(dataset_path, 'model')) if f.endswith('.npy')]\n#         data_files = [f for f in os.listdir(os.path.join(dataset_path, 'data')) if f.endswith('.npy')]\n        \n#         # Match model and data files by name\n#         for model_file in model_files:\n#             base_name = model_file.replace('model', 'data')\n#             data_file = base_name\n            \n#             if data_file in data_files:\n#                 model_path = os.path.join(dataset_path, 'model', model_file)\n#                 data_path = os.path.join(dataset_path, 'data', data_file)\n                \n#                 velocity, seismic = load_and_process_dataset(model_path, data_path)\n#                 all_velocity.append(velocity)\n#                 all_seismic.append(seismic)\n\n#     # elif \"Fault\" in dataset_path:\n#     #     for file in os.listdir(dataset_path):\n#     #         if \"seis\" in file[:4]:\n#     #             data_path = os.path.join(dataset_path,file)\n#     #             model_path = os.path.join(dataset_path, f\"vel{file.split('seis')[-1][:5]}.npy\")\n#     #             velocity, seismic = load_and_process_dataset(model_path, data_path)\n#     #             all_velocity.append(velocity)\n#     #             all_seismic.append(seismic)\n\n            \n\n\n# # Combine all loaded datasets\n# if all_velocity and all_seismic:\n#     velocity = np.concatenate(all_velocity, axis=0)\n#     seismic_data = np.concatenate(all_seismic, axis=0)\n    \n#     print('Combined velocity map size:', velocity.shape)\n#     print('Combined seismic data size:', seismic_data.shape)\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T07:29:43.841594Z","iopub.execute_input":"2025-04-19T07:29:43.841864Z","iopub.status.idle":"2025-04-19T07:29:43.846062Z","shell.execute_reply.started":"2025-04-19T07:29:43.841848Z","shell.execute_reply":"2025-04-19T07:29:43.845432Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Training ","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom torchvision import transforms\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.metrics import mean_squared_error, mean_absolute_error\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# # Custom dataset for 5-layer image input\n# class MultiLayerImageDataset(Dataset):\n#     def __init__(self, inputs, targets, transform=None):\n#         \"\"\"\n#         Args:\n#             inputs: List of input arrays with shape (N, 5, H, W) where 5 is the number of channels\n#             targets: List of target arrays with shape (N, 1, H, W)\n#             transform: Optional transform to be applied\n#         \"\"\"\n#         self.inputs = inputs\n#         self.targets = targets\n#         self.transform = transform\n\n#     def __len__(self):\n#         return len(self.inputs)\n\n#     def __getitem__(self, idx):\n#         input_img = self.inputs[idx]\n#         target_img = self.targets[idx]\n\n#         if self.transform:\n#             input_img = self.transform(input_img)\n#             target_img = self.transform(target_img)\n\n#         return input_img, target_img\n\n# # Convert numpy arrays to torch tensors\n# inputs = torch.FloatTensor(seismic_data)\n# targets = torch.FloatTensor(velocity)\n\n\n# # Split into train and validation sets (80% train, 20% validation)\n# num_samples = inputs.shape[0]\n# train_size = int(0.8 * num_samples)\n# val_size = num_samples - train_size\n\n# # Create random indices for splitting\n# indices = torch.randperm(num_samples)\n# train_indices = indices[:train_size]\n# val_indices = indices[train_size:]\n\n# train_inputs = inputs[train_indices]\n# train_targets = targets[train_indices]\n# val_inputs = inputs[val_indices]\n# val_targets = targets[val_indices]\n\n# print(f\"Training set: {train_inputs.shape}, {train_targets.shape}\")\n# print(f\"Validation set: {val_inputs.shape}, {val_targets.shape}\")\n\n# # Create datasets and dataloaders\n# batch_size = 16\n\n# train_dataset = MultiLayerImageDataset(train_inputs, train_targets)\n# val_dataset = MultiLayerImageDataset(val_inputs, val_targets)\n\n# train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n# val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n\n# # Initialize model (using your existing model)\n# model = CustomUNetRegression(in_channels=5, out_channels=1)\n# model.to(device)\n\n# # Define training function\n# def train_model(model, train_loader, val_loader, device, num_epochs=10, lr=0.001):\n#     criterion = torch.nn.MSELoss()\n#     optimizer = optim.Adam(model.parameters(), lr=lr)\n    \n#     train_losses = []\n#     val_losses = []\n    \n#     for epoch in range(num_epochs):\n#         # Training phase\n#         model.train()\n#         running_loss = 0.0\n        \n#         for inputs, targets in train_loader:\n#             inputs, targets = inputs.to(device), targets.to(device)\n            \n#             optimizer.zero_grad()\n#             outputs = model(inputs)\n#             loss = criterion(outputs, targets)\n#             loss.backward()\n#             optimizer.step()\n            \n#             running_loss += loss.item() * inputs.size(0)\n        \n#         epoch_train_loss = running_loss / len(train_loader.dataset)\n#         train_losses.append(epoch_train_loss)\n        \n#         # Validation phase\n#         model.eval()\n#         running_loss = 0.0\n        \n#         with torch.no_grad():\n#             for inputs, targets in val_loader:\n#                 inputs, targets = inputs.to(device), targets.to(device)\n#                 outputs = model(inputs)\n#                 loss = criterion(outputs, targets)\n#                 running_loss += loss.item() * inputs.size(0)\n        \n#         epoch_val_loss = running_loss / len(val_loader.dataset)\n#         val_losses.append(epoch_val_loss)\n        \n#         print(f'Epoch {epoch+1}/{num_epochs}, '\n#               f'Train Loss: {epoch_train_loss:.6f}, '\n#               f'Val Loss: {epoch_val_loss:.6f}')\n    \n#     return model, train_losses, val_losses\n\n# # Train model\n# optimizer = optim.Adam(model.parameters(), lr=0.001)\n# trained_model, train_losses, val_losses = train_model(\n#     model,\n#     train_loader,\n#     val_loader,\n#     device,\n#     num_epochs=10,\n#     lr=0.001\n# )\n\n# print(\"Training completed!\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T07:29:53.348490Z","iopub.execute_input":"2025-04-19T07:29:53.348762Z","iopub.status.idle":"2025-04-19T07:29:57.115217Z","shell.execute_reply.started":"2025-04-19T07:29:53.348743Z","shell.execute_reply":"2025-04-19T07:29:57.114596Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Training Results","metadata":{}},{"cell_type":"code","source":"# # Visualize training process\n# plt.figure(figsize=(10, 5))\n# plt.plot(train_losses, label='Training Loss')\n# plt.plot(val_losses, label='Validation Loss')\n# plt.xlabel('Epochs')\n# plt.ylabel('Loss')\n# plt.legend()\n# plt.title('Training and Validation Loss')\n# plt.savefig('loss_curves.png')\n# plt.show()\n\n# # Test with a sample image\n# with torch.no_grad():\n#     model.eval()\n#     sample_input = val_inputs[0:1].to(device)  # Get one sample and add batch dimension\n#     prediction = model(sample_input)\n    \n#     # Move tensors to CPU for visualization\n#     sample_input = sample_input.cpu()\n#     sample_target = val_targets[0:1].cpu()\n#     prediction = prediction.cpu()\n    \n#     # Visualize results\n#     plt.figure(figsize=(15, 10))\n    \n#     # Show some of the input channels\n#     for i in range(5):\n#         plt.subplot(2, 5, i+1)\n#         plt.title(f'Input (Channel {i})')\n#         plt.imshow(sample_input[0, i].numpy(), cmap='gray')\n    \n#     # Show target\n#     plt.subplot(2, 5, 6)\n#     plt.title('Ground Truth Velocity')\n#     plt.imshow(sample_target[0, 0].numpy(), cmap='viridis')\n    \n#     # Show prediction\n#     plt.subplot(2, 5, 7)\n#     plt.title('Predicted Velocity')\n#     plt.imshow(prediction[0, 0].numpy(), cmap='viridis')\n    \n#     # Show error map\n#     plt.subplot(2, 5, 8)\n#     plt.title('Error Map')\n#     error = np.abs(sample_target[0, 0].numpy() - prediction[0, 0].numpy())\n#     plt.imshow(error, cmap='hot')\n    \n#     plt.savefig('velocity_prediction_results.png')\n#     plt.show()\n    \n#     # Calculate and print metrics\n#     # Evaluate on the entire validation set\n#     all_predictions = []\n#     all_targets = []\n    \n#     for inputs, targets in val_loader:\n#         inputs = inputs.to(device)\n#         pred = model(inputs)\n#         pred = pred.cpu().numpy()\n#         targets = targets.numpy()\n        \n#         all_predictions.append(pred)\n#         all_targets.append(targets)\n    \n#     all_predictions = np.vstack(all_predictions)\n#     all_targets = np.vstack(all_targets)\n    \n#     mse = mean_squared_error(all_targets.flatten(), all_predictions.flatten())\n#     mae = mean_absolute_error(all_targets.flatten(), all_predictions.flatten())\n    \n#     print(f\"Validation MSE: {mse:.6f}\")\n#     print(f\"Validation MAE: {mae:.6f}\")\n    \n#     torch.save(model, 'velocity_model_full3.pth')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T07:15:25.922266Z","iopub.execute_input":"2025-04-19T07:15:25.922883Z","iopub.status.idle":"2025-04-19T07:15:50.301204Z","shell.execute_reply.started":"2025-04-19T07:15:25.922856Z","shell.execute_reply":"2025-04-19T07:15:50.300287Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prediction on test data","metadata":{}},{"cell_type":"code","source":"model = torch.load('/kaggle/input/updated-unet-model/pytorch/default/1/velocity_model_full3.pth')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T07:29:59.681030Z","iopub.execute_input":"2025-04-19T07:29:59.681732Z","iopub.status.idle":"2025-04-19T07:30:00.309880Z","shell.execute_reply.started":"2025-04-19T07:29:59.681706Z","shell.execute_reply":"2025-04-19T07:30:00.309300Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\nfrom pathlib import Path\nfrom tqdm import tqdm\nimport csv\n\ntest_files = list(Path('/kaggle/input/waveform-inversion/test').glob('*.npy'))\nlen(test_files)\n \nx_cols = [f'x_{i}' for i in range(1, 70, 2)]\nfieldnames = ['oid_ypos'] + x_cols\n\nclass 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        seismic_data = np.load(test_file)\n        return seismic_data, test_file.stem\n\nds = TestDataset(test_files)\ndl = DataLoader(ds, batch_size=8, num_workers=4, pin_memory=True)\n\n#Test\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 = torch.FloatTensor(inputs)\n        inputs = inputs.to(device)\n    \n        with torch.inference_mode():\n            # Model inference\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-19T07:30:02.738827Z","iopub.execute_input":"2025-04-19T07:30:02.739436Z","iopub.status.idle":"2025-04-19T07:30:52.805111Z","shell.execute_reply.started":"2025-04-19T07:30:02.739398Z","shell.execute_reply":"2025-04-19T07:30:52.804364Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}