{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Initial Data Loading and Visualization\n\nI begin by importing essential libraries: numpy for numerical operations, matplotlib.pyplot for visualization, and os for file handling. I define paths to the training data and velocity models from the datasets, listing all .npy files in both directories. To explore the data, I load one example of seismic data and its corresponding velocity map, printing their shapes and value ranges to verify their structure—seismic data should have dimensions like (batch_size, num_sources, time_steps, num_receivers), and the velocity map should be (batch_size, 1, height, width). I visualize the seismic waveform for one source and receiver over time to understand its temporal behavior, and plot the velocity map as a 2D heatmap after removing the batch dimension, ensuring it’s a (70, 70) grid, to inspect the spatial distribution of velocities.","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nimport os\n\n# List of all datasets to analyze\ndatasets = [\n    'FlatVel_A',\n    'FlatVel_B',\n    'CurveFault_A',\n    'CurveFault_B',\n    'CurveVel_A',\n    'FlatFault_A',\n    'FlatFault_B',\n    'Style_A'\n]\n\n# Function to load one example from a dataset\ndef load_example(dataset_name):\n    base_path = f'/kaggle/input/waveform-inversion/train_samples/{dataset_name}'\n    \n    # Check if the dataset has 'data' and 'model' subdirectories (Vel and Style families)\n    if os.path.exists(os.path.join(base_path, 'data')) and os.path.exists(os.path.join(base_path, 'model')):\n        # Vel/Style family: FlatVel_A, FlatVel_B, CurveVel_A, Style_A\n        data_dir = os.path.join(base_path, 'data')\n        model_dir = os.path.join(base_path, 'model')\n        \n        # Get list of data and model files\n        data_files = sorted([f for f in os.listdir(data_dir) if f.endswith('.npy')])\n        model_files = sorted([f for f in os.listdir(model_dir) if f.endswith('.npy')])\n        \n        if not data_files or not model_files:\n            print(f\"No .npy files found in {dataset_name} (data: {len(data_files)}, model: {len(model_files)}).\")\n            return None, None\n        \n        # Load the first file\n        example_data_path = os.path.join(data_dir, data_files[0])\n        example_model_path = os.path.join(model_dir, model_files[0])\n    else:\n        # Fault family: CurveFault_A, CurveFault_B, FlatFault_A, FlatFault_B\n        all_files = sorted([f for f in os.listdir(base_path) if f.endswith('.npy')])\n        \n        # Separate files into seismic data (seis*.npy) and velocity maps (vel*.npy)\n        data_files = [f for f in all_files if f.startswith('seis')]  # Changed from 'seis_' to 'seis'\n        model_files = [f for f in all_files if f.startswith('vel')]  # Changed from 'vel_' to 'vel'\n        \n        if not data_files or not model_files:\n            print(f\"Could not separate files into seismic data and velocity maps in {dataset_name}.\")\n            print(f\"Seismic files: {data_files}\")\n            print(f\"Velocity files: {model_files}\")\n            return None, None\n        \n        # Load the first pair\n        example_data_path = os.path.join(base_path, data_files[0])\n        example_model_path = os.path.join(base_path, model_files[0])\n    \n    # Load the data\n    seismic_data = np.load(example_data_path)\n    velocity_map = np.load(example_model_path)\n    \n    return seismic_data, velocity_map\n\n# Analyze each dataset\nfor dataset_name in datasets:\n    print(f\"\\n=== Analyzing dataset: {dataset_name} ===\")\n    \n    # Load example data\n    seismic_data, velocity_map = load_example(dataset_name)\n    \n    if seismic_data is None or velocity_map is None:\n        print(f\"Skipping {dataset_name} due to loading error.\")\n        continue\n    \n    # Verify shapes to ensure correct separation\n    expected_seismic_shape = (500, 5, 1000, 70)\n    expected_velocity_shape = (500, 1, 70, 70)\n    if seismic_data.shape != expected_seismic_shape or velocity_map.shape != expected_velocity_shape:\n        print(f\"Unexpected shapes in {dataset_name}:\")\n        print(\"Seismic data shape:\", seismic_data.shape)\n        print(\"Velocity map shape:\", velocity_map.shape)\n        print(\"Skipping due to shape mismatch.\")\n        continue\n    \n    # Print shapes and min/max values\n    print(\"Seismic data shape:\", seismic_data.shape)\n    print(\"Velocity map shape:\", velocity_map.shape)\n    print(\"Seismic data min/max:\", seismic_data.min(), seismic_data.max())\n    print(\"Velocity map min/max:\", velocity_map.min(), velocity_map.max())\n    \n    # Visualize seismic data (one source, one receiver over time)\n    plt.figure(figsize=(12, 4))\n    plt.plot(seismic_data[0, 0, :, 0], label=\"Source 0, Receiver 0\")\n    plt.title(f\"Example Seismic Waveform ({dataset_name})\")\n    plt.xlabel(\"Time Steps\")\n    plt.ylabel(\"Amplitude\")\n    plt.legend()\n    plt.show()\n    \n    # Prepare velocity map for visualization by removing batch dimension\n    velocity_map_2d = np.squeeze(velocity_map[0])  # Shape should be (70, 70)\n    print(\"Velocity map 2D shape after squeeze:\", velocity_map_2d.shape)\n    \n    # Visualize velocity map\n    plt.figure(figsize=(8, 6))\n    plt.imshow(velocity_map_2d, cmap='viridis', aspect='equal')\n    plt.colorbar(label=\"Velocity (m/s)\")\n    plt.title(f\"Example Velocity Map ({dataset_name})\")\n    plt.xlabel(\"Width (x)\")\n    plt.ylabel(\"Height (y)\")\n    plt.show()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-13T04:23:06.452179Z","iopub.execute_input":"2025-05-13T04:23:06.452426Z","iopub.status.idle":"2025-05-13T04:23:43.175888Z","shell.execute_reply.started":"2025-05-13T04:23:06.452402Z","shell.execute_reply":"2025-05-13T04:23:43.174918Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Data Normalization\n\nAfter loading, I note the shapes and value ranges: seismic data is (500, 5, 1000, 70) and velocity maps are (500, 1, 70, 70), with seismic values ranging from -25.7 to 49.97 and velocity values from 1501 to 4500 m/s. I normalize the seismic data to the range [-1, 1] using min-max scaling per sample to prepare it for the neural network, and scale the velocity maps to [0, 1] for consistency, printing the new min/max values to confirm the normalization. Since the seismic data has 1000 time steps (matching the updated time_steps_to_keep=1000), no cropping is needed, and the shape remains (500, 5, 1000, 70). I visualize the normalized seismic waveform for one source and receiver to ensure the signal’s behavior is preserved.","metadata":{}},{"cell_type":"code","source":"# Define paths to training data (using FlatVel_A as an example)\ndata_dir = '/kaggle/input/waveform-inversion/train_samples/FlatVel_A/data'\nmodel_dir = '/kaggle/input/waveform-inversion/train_samples/FlatVel_A/model'\n\n# Get list of files\ndata_files = sorted([f for f in os.listdir(data_dir) if f.endswith('.npy')])\nmodel_files = sorted([f for f in os.listdir(model_dir) if f.endswith('.npy')])\n\n# Load one example seismic data and velocity map\nexample_data_path = os.path.join(data_dir, data_files[0])  # First seismic data file\nexample_model_path = os.path.join(model_dir, model_files[0])  # Corresponding velocity map\nseismic_data = np.load(example_data_path)\nvelocity_map = np.load(example_model_path)\n\n# Print shapes and basic info\nprint(\"Seismic data shape:\", seismic_data.shape)  # Expected: (500, 5, 1000, 70)\nprint(\"Velocity map shape:\", velocity_map.shape)  # Expected: (500, 1, 70, 70)\nprint(\"Seismic data min/max:\", seismic_data.min(), seismic_data.max())\nprint(\"Velocity map min/max:\", velocity_map.min(), velocity_map.max())\n\n# Define global min/max for normalization (based on analysis of all datasets)\nglobal_seismic_min = -27.16  # Minimum across all datasets\nglobal_seismic_max = 55.31   # Maximum across all datasets\nglobal_velocity_min = 1500.0  # Minimum across all datasets\nglobal_velocity_max = 4500.0  # Maximum across all datasets\n\n# Normalize seismic data to [-1, 1] using global min/max\nseismic_data_normalized = 2 * (seismic_data - global_seismic_min) / (global_seismic_max - global_seismic_min) - 1\n\n# Normalize velocity map to [0, 1] using global min/max\nvelocity_map_normalized = (velocity_map - global_velocity_min) / (global_velocity_max - global_velocity_min)\n\n# Print new min/max to verify normalization\nprint(\"Normalized seismic data min/max:\", seismic_data_normalized.min(), seismic_data_normalized.max())\nprint(\"Normalized velocity map min/max:\", velocity_map_normalized.min(), velocity_map_normalized.max())\n\n# Visualize seismic data (one source, one receiver over time) - using all 1000 steps\nplt.figure(figsize=(12, 4))\nplt.plot(seismic_data_normalized[0, 0, :, 0], label=\"Source 0, Receiver 0\")\nplt.title(\"Seismic Waveform (All 1000 Steps)\")\nplt.xlabel(\"Time Steps\")\nplt.ylabel(\"Normalized Amplitude\")\nplt.legend()\nplt.show()\n\n# Prepare velocity map for visualization by removing batch dimension\nvelocity_map_2d = np.squeeze(velocity_map_normalized[0])  # Remove singleton dimensions\nprint(\"Velocity map 2D shape after squeeze:\", velocity_map_2d.shape)  # Should be (70, 70)\n\n# Visualize velocity map\nplt.figure(figsize=(8, 6))\nplt.imshow(velocity_map_2d, cmap='viridis', aspect='equal')\nplt.colorbar(label=\"Normalized Velocity\")\nplt.title(\"Example Velocity Map\")\nplt.xlabel(\"Width (x)\")\nplt.ylabel(\"Height (y)\")\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T04:24:14.457696Z","iopub.execute_input":"2025-05-13T04:24:14.458088Z","iopub.status.idle":"2025-05-13T04:24:16.330258Z","shell.execute_reply.started":"2025-05-13T04:24:14.458056Z","shell.execute_reply":"2025-05-13T04:24:16.329168Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Creating a Custom Dataset and DataLoader\n\nWith the normalized seismic data ranging from -1 to 1 and velocity maps from 0 to 1, and the seismic data shape confirmed as (500, 5, 1000, 70), I set up the data pipeline for training. I import PyTorch libraries and define a SeismicDataset class to handle loading and preprocessing of seismic data and velocity maps from multiple directories (e.g., FlatVel_A and FlatVel_B). In the dataset, I collect all .npy files, normalize the data (seismic to [-1, 1] and velocity to [0, 1]), and convert it to PyTorch tensors, ensuring the velocity map is squeezed to (height, width). I then create a DataLoader with a batch size of 64, enabling shuffling and multi-threading (with num_workers=2 for efficiency). To verify, I load one batch and print the shapes, expecting (batch_size, 5, 1000, 70) for the seismic batch and (batch_size, 70, 70) for the velocity batch.","metadata":{}},{"cell_type":"code","source":"import torch\nimport os\nfrom torch.utils.data import Dataset, DataLoader, Subset\n\n# Define global min/max values for normalization (same as in analysis step)\nglobal_seismic_min = -27.16  # Minimum across all datasets\nglobal_seismic_max = 55.31   # Maximum across all datasets\nglobal_velocity_min = 1500.0  # Minimum across all datasets\nglobal_velocity_max = 4500.0  # Maximum across all datasets\n\n# Custom Dataset with augmentation\nclass SeismicDataset(Dataset):\n    def __init__(self, data_dirs, model_dirs, time_steps_to_keep=1000, augment=True):\n        self.data_files = []\n        self.model_files = []\n        \n        # Process each pair of directories\n        for data_dir, model_dir in zip(data_dirs, model_dirs):\n            # Check if the dataset has 'data' and 'model' subdirectories (Vel/Style family)\n            if os.path.exists(os.path.join(data_dir, 'data')) and os.path.exists(os.path.join(model_dir, 'model')):\n                # Vel/Style family: FlatVel_A, FlatVel_B, CurveVel_A, Style_A\n                data_files = sorted([f for f in os.listdir(os.path.join(data_dir, 'data')) if f.endswith('.npy')])\n                model_files = sorted([f for f in os.listdir(os.path.join(model_dir, 'model')) if f.endswith('.npy')])\n                self.data_files.extend([os.path.join(data_dir, 'data', f) for f in data_files])\n                self.model_files.extend([os.path.join(model_dir, 'model', f) for f in model_files])\n            else:\n                # Fault family: CurveFault_A, CurveFault_B, FlatFault_A, FlatFault_B\n                all_files = sorted([f for f in os.listdir(data_dir) if f.endswith('.npy')])\n                # Separate files into seismic data (seis*.npy) and velocity maps (vel*.npy)\n                data_files = [f for f in all_files if f.startswith('seis')]\n                model_files = [f for f in all_files if f.startswith('vel')]\n                \n                if len(data_files) != len(model_files):\n                    print(f\"Mismatch in number of data and model files in {data_dir}: {len(data_files)} data, {len(model_files)} model\")\n                    continue\n                \n                self.data_files.extend([os.path.join(data_dir, f) for f in data_files])\n                self.model_files.extend([os.path.join(data_dir, f) for f in model_files])\n        \n        # Verify that the number of data and model files matches\n        assert len(self.data_files) == len(self.model_files), f\"Mismatch between data and model files: {len(self.data_files)} data files, {len(self.model_files)} model files\"\n        self.time_steps_to_keep = time_steps_to_keep\n        self.augment = augment\n    \n    def __len__(self):\n        return len(self.data_files) * 500\n    \n    def __getitem__(self, idx):\n        file_idx = idx // 500\n        sample_idx = idx % 500\n        \n        seismic_data = np.load(self.data_files[file_idx])[sample_idx]\n        velocity_map = np.load(self.model_files[file_idx])[sample_idx]\n        \n        # Normalize seismic data to [-1, 1] using global min/max\n        seismic_data = 2 * (seismic_data - global_seismic_min) / (global_seismic_max - global_seismic_min) - 1\n        \n        # Normalize velocity map to [0, 1] using global min/max\n        velocity_map = (velocity_map - global_velocity_min) / (global_velocity_max - global_velocity_min)\n        \n        # Crop time steps for seismic data (using all 1000 steps)\n        seismic_data = seismic_data[:, :self.time_steps_to_keep, :]\n        \n        # Apply augmentation if enabled\n        if self.augment and np.random.rand() > 0.5:\n            # Add random noise\n            noise = np.random.normal(0, 0.01, seismic_data.shape)\n            seismic_data = seismic_data + noise\n            # Random amplitude scaling\n            scale = np.random.uniform(0.9, 1.1)\n            seismic_data = seismic_data * scale\n            # Clip to ensure values stay in [-1, 1]\n            seismic_data = np.clip(seismic_data, -1, 1)\n        \n        velocity_map = np.squeeze(velocity_map)\n        \n        seismic_data = torch.FloatTensor(seismic_data)\n        velocity_map = torch.FloatTensor(velocity_map)\n        \n        return seismic_data, velocity_map\n\n# Define directories for training data (include all sets)\ntrain_data_dirs = [\n    '/kaggle/input/waveform-inversion/train_samples/FlatVel_A',\n    '/kaggle/input/waveform-inversion/train_samples/FlatVel_B',\n    '/kaggle/input/waveform-inversion/train_samples/CurveFault_A',\n    '/kaggle/input/waveform-inversion/train_samples/CurveFault_B',\n    '/kaggle/input/waveform-inversion/train_samples/CurveVel_A',\n    '/kaggle/input/waveform-inversion/train_samples/FlatFault_A',\n    '/kaggle/input/waveform-inversion/train_samples/FlatFault_B',\n    '/kaggle/input/waveform-inversion/train_samples/Style_A'\n]\ntrain_model_dirs = [\n    '/kaggle/input/waveform-inversion/train_samples/FlatVel_A',\n    '/kaggle/input/waveform-inversion/train_samples/FlatVel_B',\n    '/kaggle/input/waveform-inversion/train_samples/CurveFault_A',\n    '/kaggle/input/waveform-inversion/train_samples/CurveFault_B',\n    '/kaggle/input/waveform-inversion/train_samples/CurveVel_A',\n    '/kaggle/input/waveform-inversion/train_samples/FlatFault_A',\n    '/kaggle/input/waveform-inversion/train_samples/FlatFault_B',\n    '/kaggle/input/waveform-inversion/train_samples/Style_A'\n]\n\n# Create dataset with updated time_steps_to_keep\ntime_steps_to_keep = 1000  # Use all 1000 time steps\ntrain_dataset = SeismicDataset(train_data_dirs, train_model_dirs, time_steps_to_keep=time_steps_to_keep, augment=True)\n\n# Split dataset into train and validation\ndataset_size = len(train_dataset)\nindices = list(range(dataset_size))\nnp.random.shuffle(indices)\ntrain_split = int(0.8 * dataset_size)\ntrain_indices = indices[:train_split]\nval_indices = indices[train_split:]\n\ntrain_subset = Subset(train_dataset, train_indices)\nval_subset = Subset(train_dataset, val_indices)\n\ntrain_dataloader = DataLoader(train_subset, batch_size=32, shuffle=True, num_workers=0)\nval_dataloader = DataLoader(val_subset, batch_size=32, shuffle=False, num_workers=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T04:24:23.845123Z","iopub.execute_input":"2025-05-13T04:24:23.845472Z","iopub.status.idle":"2025-05-13T04:24:27.706529Z","shell.execute_reply.started":"2025-05-13T04:24:23.845449Z","shell.execute_reply":"2025-05-13T04:24:27.705478Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Defining the InversionNet Model\n\nHaving confirmed the batch shapes—seismic data at (64, 5, 1000, 70) and velocity maps at (64, 70, 70)—I define the InversionNet model for training. I create building blocks: ConvBlock for convolutional layers with batch normalization, LeakyReLU, and optional dropout; DeconvBlock for upsampling with target sizes; and ConvBlock_Tanh for the final layer with a Tanh activation. The InversionNet follows a U-Net-like structure, with an encoder downsizing the input from (5, 1000, 70) to (512, 1, 1) through convolutional blocks, and a decoder upsampling to (1, 70, 70) using deconvolutional blocks with dynamic target_size. I incorporate skip connections at multiple levels, adjusting their sizes with convolutional and upsampling layers to match the decoder, ensuring feature preservation. The final output is a velocity map prediction, matching the target shape (70, 70).","metadata":{}},{"cell_type":"code","source":"import torch.nn as nn\nimport torch.nn.functional as F\nfrom math import ceil\n\nclass ConvBlock(nn.Module):\n    def __init__(self, in_fea, out_fea, kernel_size=3, stride=1, padding=1, norm='bn', relu_slop=0.2, dropout=None):\n        super(ConvBlock, self).__init__()\n        layers = [\n            nn.Conv2d(in_channels=in_fea, out_channels=out_fea, kernel_size=kernel_size, stride=stride, padding=padding)\n        ]\n        if norm == 'bn':\n            layers.append(nn.BatchNorm2d(out_fea))\n        layers.append(nn.LeakyReLU(relu_slop, inplace=True))\n        if dropout:\n            layers.append(nn.Dropout2d(dropout))\n        self.layers = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.layers(x)\n\nclass ConvBlock_Tanh(nn.Module):\n    def __init__(self, in_fea, out_fea, kernel_size=3, stride=1, padding=1, norm='bn'):\n        super(ConvBlock_Tanh, self).__init__()\n        layers = [\n            nn.Conv2d(in_channels=in_fea, out_channels=out_fea, kernel_size=kernel_size, stride=stride, padding=padding)\n        ]\n        if norm == 'bn':\n            layers.append(nn.BatchNorm2d(out_fea))\n        layers.append(nn.Tanh())\n        self.layers = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.layers(x)\n\nclass DeconvBlock(nn.Module):\n    def __init__(self, in_fea, out_fea, target_size=None, norm='bn'):\n        super(DeconvBlock, self).__init__()\n        self.target_size = target_size  # (height, width) to upsample to\n        self.conv = nn.Conv2d(in_channels=in_fea, out_channels=out_fea, kernel_size=3, stride=1, padding=1)\n        self.norm = nn.BatchNorm2d(out_fea) if norm == 'bn' else None\n        self.relu = nn.LeakyReLU(0.2, inplace=True)\n\n    def forward(self, x):\n        if self.target_size is not None:\n            x = nn.functional.interpolate(x, size=self.target_size, mode='bilinear', align_corners=True)\n        x = self.conv(x)\n        if self.norm is not None:\n            x = self.norm(x)\n        x = self.relu(x)\n        return x\n\nclass InversionNet(nn.Module):\n    def __init__(self, dim1=32, dim2=64, dim3=128, dim4=256, dim5=512, sample_spatial=1.0):\n        super(InversionNet, self).__init__()\n        # Энкодер\n        self.convblock1 = ConvBlock(5, dim1, kernel_size=(7, 1), stride=(2, 1), padding=(3, 0))  # (5, 1000, 70) -> (32, 500, 70)\n        self.convblock2_1 = ConvBlock(dim1, dim2, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (32, 500, 70) -> (64, 250, 70)\n        self.convblock2_2 = ConvBlock(dim2, dim2, kernel_size=(3, 1), padding=(1, 0))  # (64, 250, 70) -> (64, 250, 70)\n        self.convblock3_1 = ConvBlock(dim2, dim2, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (64, 250, 70) -> (64, 125, 70)\n        self.convblock3_2 = ConvBlock(dim2, dim2, kernel_size=(3, 1), padding=(1, 0))  # (64, 125, 70) -> (64, 125, 70)\n        self.convblock4_1 = ConvBlock(dim2, dim3, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (64, 125, 70) -> (128, 62, 70)\n        self.convblock4_2 = ConvBlock(dim3, dim3, kernel_size=(3, 1), padding=(1, 0))  # (128, 62, 70) -> (128, 62, 70)\n        self.convblock5_1 = ConvBlock(dim3, dim3, stride=2)  # (128, 62, 70) -> (128, 31, 35)\n        self.convblock5_2 = ConvBlock(dim3, dim3, dropout=0.3)  # (128, 31, 35) -> (128, 31, 35)\n        self.convblock6_1 = ConvBlock(dim3, dim4, stride=2)  # (128, 31, 35) -> (256, 15, 17)\n        self.convblock6_2 = ConvBlock(dim4, dim4, dropout=0.3)  # (256, 15, 17) -> (256, 15, 17)\n        self.convblock7_1 = ConvBlock(dim4, dim4, stride=2)  # (256, 15, 17) -> (256, 7, 8)\n        self.convblock7_2 = ConvBlock(dim4, dim4, dropout=0.3)  # (256, 7, 8) -> (256, 7, 8)\n        self.convblock8 = ConvBlock(dim4, dim5, kernel_size=(3, ceil(70 * sample_spatial / 8)), padding=0)  # (256, 7, 8) -> (512, 1, 1)\n\n        # Дополнительные слои для подгонки размеров skip-соединений\n        self.skip1_adjust = ConvBlock(dim1, dim1, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (32, 500, 70) -> (32, 250, 70)\n        self.skip1_adjust2 = ConvBlock(dim1, dim1, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (32, 250, 70) -> (32, 125, 70)\n        self.skip1_adjust3 = ConvBlock(dim1, dim1, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (32, 125, 70) -> (32, 62, 70)\n        self.skip1_adjust4 = nn.Conv2d(dim1, dim1, kernel_size=(3, 1), stride=(2, 1), padding=(1, 0))  # (32, 62, 70) -> (32, 31, 70)\n        self.skip1_adjust5 = DeconvBlock(dim1, dim1, target_size=(40, 40))  # (32, 31, 70) -> (32, 40, 40)\n\n        self.skip4_adjust = ConvBlock(dim3, dim2, kernel_size=(3, 1), stride=(1, 2), padding=(1, 0))  # (128, 62, 70) -> (64, 62, 35)\n        self.skip4_adjust2 = DeconvBlock(dim2, dim2, target_size=(20, 20))  # (64, 62, 35) -> (64, 20, 20)\n\n        self.skip5_adjust = ConvBlock(dim3, dim3, kernel_size=(3, 1), stride=(1, 2), padding=(1, 0))  # (128, 31, 35) -> (128, 31, 17)\n        self.skip5_adjust2 = DeconvBlock(dim3, dim3, target_size=(10, 10))  # (128, 31, 17) -> (128, 10, 10)\n\n        self.skip6_adjust = ConvBlock(dim4, dim4, kernel_size=(3, 1), stride=(1, 2), padding=(1, 0))  # (256, 15, 17) -> (256, 15, 8)\n        self.skip6_adjust2 = DeconvBlock(dim4, dim4, target_size=(5, 5))  # (256, 15, 8) -> (256, 5, 5)\n\n        # Декодер\n        self.deconv1_1 = DeconvBlock(dim5, dim5, target_size=(5, 5))  # (512, 1, 1) -> (512, 5, 5)\n        self.deconv1_2 = ConvBlock(dim5, dim5)  # (512, 5, 5) -> (512, 5, 5)\n        self.deconv2_1 = DeconvBlock(dim5 + dim4, dim4, target_size=(10, 10))  # (512+256, 5, 5) -> (256, 10, 10)\n        self.deconv2_2 = ConvBlock(dim4, dim4)  # (256, 10, 10) -> (256, 10, 10)\n        self.deconv3_1 = DeconvBlock(dim4 + dim3, dim3, target_size=(20, 20))  # (256+128, 10, 10) -> (128, 20, 20)\n        self.deconv3_2 = ConvBlock(dim3, dim3)  # (128, 20, 20) -> (128, 20, 20)\n        self.deconv4_1 = DeconvBlock(dim3 + dim2, dim2, target_size=(40, 40))  # (128+64, 20, 20) -> (64, 40, 40)\n        self.deconv4_2 = ConvBlock(dim2, dim2)  # (64, 40, 40) -> (64, 40, 40)\n        self.deconv5_1 = DeconvBlock(dim2 + dim1, dim1, target_size=(70, 70))  # (64+32, 40, 40) -> (32, 70, 70)\n        self.deconv5_2 = ConvBlock(dim1, dim1)  # (32, 70, 70) -> (32, 70, 70)\n        self.deconv6 = ConvBlock_Tanh(dim1, 1)  # (32, 70, 70) -> (1, 70, 70)\n\n    def forward(self, x):\n        # Энкодер\n        x1 = self.convblock1(x)  # (32, 500, 70)\n        x2 = self.convblock2_1(x1)  # (64, 250, 70)\n        x2 = self.convblock2_2(x2)  # (64, 250, 70)\n        x3 = self.convblock3_1(x2)  # (64, 125, 70)\n        x3 = self.convblock3_2(x3)  # (64, 125, 70)\n        x4 = self.convblock4_1(x3)  # (128, 62, 70)\n        x4 = self.convblock4_2(x4)  # (128, 62, 70)\n        x5 = self.convblock5_1(x4)  # (128, 31, 35)\n        x5 = self.convblock5_2(x5)  # (128, 31, 35)\n        x6 = self.convblock6_1(x5)  # (256, 15, 17)\n        x6 = self.convblock6_2(x6)  # (256, 15, 17)\n        x7 = self.convblock7_1(x6)  # (256, 7, 8)\n        x7 = self.convblock7_2(x7)  # (256, 7, 8)\n        x = self.convblock8(x7)  # (512, 1, 1)\n\n        # Подгонка размеров для skip-соединений\n        x1_skip = self.skip1_adjust(x1)  # (32, 500, 70) -> (32, 250, 70)\n        x1_skip = self.skip1_adjust2(x1_skip)  # (32, 250, 70) -> (32, 125, 70)\n        x1_skip = self.skip1_adjust3(x1_skip)  # (32, 125, 70) -> (32, 62, 70)\n        x1_skip = self.skip1_adjust4(x1_skip)  # (32, 62, 70) -> (32, 31, 70)\n        x1_skip = self.skip1_adjust5(x1_skip)  # (32, 31, 70) -> (32, 40, 40)\n\n        x4_skip = self.skip4_adjust(x4)  # (128, 62, 70) -> (64, 62, 35)\n        x4_skip = self.skip4_adjust2(x4_skip)  # (64, 62, 35) -> (64, 20, 20)\n\n        x5_skip = self.skip5_adjust(x5)  # (128, 31, 35) -> (128, 31, 17)\n        x5_skip = self.skip5_adjust2(x5_skip)  # (128, 31, 17) -> (128, 10, 10)\n\n        x6_skip = self.skip6_adjust(x6)  # (256, 15, 17) -> (256, 15, 8)\n        x6_skip = self.skip6_adjust2(x6_skip)  # (256, 15, 8) -> (256, 5, 5)\n\n        # Декодер с skip-соединениями\n        x = self.deconv1_1(x)  # (512, 1, 1) -> (512, 5, 5)\n        x = self.deconv1_2(x)  # (512, 5, 5)\n        x = torch.cat([x, x6_skip], dim=1)  # (512+256, 5, 5)\n        x = self.deconv2_1(x)  # (256, 10, 10)\n        x = self.deconv2_2(x)  # (256, 10, 10)\n        x = torch.cat([x, x5_skip], dim=1)  # (256+128, 10, 10)\n        x = self.deconv3_1(x)  # (128, 20, 20)\n        x = self.deconv3_2(x)  # (128, 20, 20)\n        x = torch.cat([x, x4_skip], dim=1)  # (128+64, 20, 20)\n        x = self.deconv4_1(x)  # (64, 40, 40)\n        x = self.deconv4_2(x)  # (64, 40, 40)\n        x = torch.cat([x, x1_skip], dim=1)  # (64+32, 40, 40)\n        x = self.deconv5_1(x)  # (32, 70, 70)\n        x = self.deconv5_2(x)  # (32, 70, 70)\n        x = self.deconv6(x)  # (1, 70, 70)\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T04:35:13.985040Z","iopub.execute_input":"2025-05-13T04:35:13.985383Z","iopub.status.idle":"2025-05-13T04:35:14.011175Z","shell.execute_reply.started":"2025-05-13T04:35:13.985359Z","shell.execute_reply":"2025-05-13T04:35:14.009960Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Training the InversionNet Model\n\nWith the dataset prepared, I split it into training and validation sets (80% train, 20% validation) by shuffling indices and creating subsets, then set up DataLoaders for both. I initialize the InversionNet model, move it to the appropriate device (GPU if available), and define the training setup: I use L1Loss as the criterion, Adam optimizer with a learning rate of 0.0005 and weight decay of 5e-4 for regularization, and a CosineAnnealingLR scheduler with T_max=25. I train the model for 25 epochs, computing the training loss by denormalizing model outputs from [-1, 1] to [0, 1] to match the target velocity maps, and validate after each epoch, tracking both L1 and gradient losses. I plot the training and validation losses to monitor convergence, achieving a best Val L1 Loss of 0.0824 on the 22nd epoch, and visualize a sample prediction against its ground truth from the validation set to qualitatively assess performance.","metadata":{}},{"cell_type":"code","source":"import torch.optim as optim\nfrom torch.utils.data import Subset\nimport matplotlib.pyplot as plt\n\n# Split train_dataset into train and validation sets (80% train, 20% validation)\ndataset_size = len(train_dataset)\nindices = list(range(dataset_size))\nnp.random.shuffle(indices)\ntrain_split = int(0.8 * dataset_size)\ntrain_indices = indices[:train_split]\nval_indices = indices[train_split:]\n\ntrain_subset = Subset(train_dataset, train_indices)\nval_subset = Subset(train_dataset, val_indices)\n\ntrain_dataloader = DataLoader(train_subset, batch_size=64, shuffle=True, num_workers=2)\nval_dataloader = DataLoader(val_subset, batch_size=64, shuffle=False, num_workers=2)\n\n# Gradient loss function\ndef gradient_loss(pred, target):\n    pred_dx = pred[:, :, :, 1:] - pred[:, :, :, :-1]\n    pred_dy = pred[:, :, 1:, :] - pred[:, :, :-1, :]\n    target_dx = target[:, :, :, 1:] - target[:, :, :, :-1]\n    target_dy = target[:, :, 1:, :] - target[:, :, :-1, :]\n    loss_dx = torch.mean(torch.abs(pred_dx - target_dx))\n    loss_dy = torch.mean(torch.abs(pred_dy - target_dy))\n    return (loss_dx + loss_dy) / 2\n\nmodel = InversionNet()\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\nmodel = model.to(device)\n\ncriterion = nn.L1Loss()\noptimizer = optim.Adam(model.parameters(), lr=0.0005, weight_decay=5e-4)  # Уменьшенный lr\nnum_epochs = 25\nscheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)  # T_max = num_epochs\n\ntrain_losses = []\nval_losses = []\ntrain_grad_losses = []\nval_grad_losses = []\nbest_val_loss = float('inf')\n\nfor epoch in range(num_epochs):\n    model.train()\n    running_loss = 0.0\n    running_grad_loss = 0.0\n    for seismic_batch, velocity_batch in train_dataloader:\n        seismic_batch = seismic_batch.to(device)\n        velocity_batch = velocity_batch.to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(seismic_batch)\n        outputs = (outputs + 1) / 2\n        target = velocity_batch.unsqueeze(1)\n        \n        l1_loss = criterion(outputs, target)\n        grad_loss = gradient_loss(outputs, target)\n        loss = l1_loss + 0.1 * grad_loss\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Градиентный клиппинг\n        optimizer.step()\n        \n        running_loss += l1_loss.item() * seismic_batch.size(0)\n        running_grad_loss += grad_loss.item() * seismic_batch.size(0)\n    \n    epoch_train_loss = running_loss / len(train_dataloader.dataset)\n    epoch_grad_loss = running_grad_loss / len(train_dataloader.dataset)\n    train_losses.append(epoch_train_loss)\n    train_grad_losses.append(epoch_grad_loss)\n    print(f\"Epoch {epoch+1}/{num_epochs}, Train L1 Loss: {epoch_train_loss:.4f}, Train Grad Loss: {epoch_grad_loss:.4f}\")\n    \n    model.eval()\n    val_loss = 0.0\n    val_grad_loss = 0.0\n    with torch.no_grad():\n        for seismic_batch, velocity_batch in val_dataloader:\n            seismic_batch = seismic_batch.to(device)\n            velocity_batch = velocity_batch.to(device)\n            outputs = model(seismic_batch)\n            outputs = (outputs + 1) / 2\n            target = velocity_batch.unsqueeze(1)\n            \n            l1_loss = criterion(outputs, target)\n            grad_loss = gradient_loss(outputs, target)\n            loss = l1_loss + 0.1 * grad_loss\n            \n            val_loss += l1_loss.item() * seismic_batch.size(0)\n            val_grad_loss += grad_loss.item() * seismic_batch.size(0)\n    \n    epoch_val_loss = val_loss / len(val_dataloader.dataset)\n    epoch_val_grad_loss = val_grad_loss / len(val_dataloader.dataset)\n    val_losses.append(epoch_val_loss)\n    val_grad_losses.append(epoch_val_grad_loss)\n    print(f\"Epoch {epoch+1}/{num_epochs}, Val L1 Loss: {epoch_val_loss:.4f}, Val Grad Loss: {epoch_val_grad_loss:.4f}\")\n    \n    scheduler.step()\n    \n    # Save the best model\n    if epoch_val_loss < best_val_loss:\n        best_val_loss = epoch_val_loss\n        torch.save({\n            'epoch': epoch + 1,\n            'model_state_dict': model.state_dict(),\n            'val_loss': epoch_val_loss,\n        }, \"/kaggle/working/model_best.pth\")\n        print(f\"Best model saved at epoch {epoch+1} with Val Loss: {epoch_val_loss:.4f}\")\n\n# Plot losses\nplt.figure(figsize=(10, 5))\nplt.plot(train_losses, label=\"Train L1 Loss\")\nplt.plot(val_losses, label=\"Val L1 Loss\")\nplt.plot(train_grad_losses, label=\"Train Grad Loss\")\nplt.plot(val_grad_losses, label=\"Val Grad Loss\")\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Loss\")\nplt.title(\"Training and Validation Losses\")\nplt.legend()\nplt.show()\n\n# Visualize a sample prediction\nmodel.eval()\nwith torch.no_grad():\n    sample_seismic, sample_velocity = next(iter(val_dataloader))\n    sample_seismic = sample_seismic.to(device)\n    sample_velocity = sample_velocity.to(device)\n    output = model(sample_seismic)\n    output = (output + 1) / 2\n    \n    plt.figure(figsize=(12, 5))\n    plt.subplot(1, 2, 1)\n    plt.imshow(output[0, 0].cpu().numpy(), cmap='viridis', aspect='equal')\n    plt.colorbar(label=\"Predicted Velocity (normalized)\")\n    plt.title(\"Predicted Velocity Map\")\n    plt.xlabel(\"Width (x)\")\n    plt.ylabel(\"Height (y)\")\n    \n    plt.subplot(1, 2, 2)\n    plt.imshow(sample_velocity[0].cpu().numpy(), cmap='viridis', aspect='equal')\n    plt.colorbar(label=\"Ground Truth Velocity (normalized)\")\n    plt.title(\"Ground Truth Velocity Map\")\n    plt.xlabel(\"Width (x)\")\n    plt.ylabel(\"Height (y)\")\n    \n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T04:54:06.317896Z","iopub.execute_input":"2025-05-13T04:54:06.318273Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Generating the Submission File\n\nAfter training, with a best validation loss of 0.0824 on the 22nd epoch, I prepare the submission. I define a TestDataset to load and preprocess the test seismic data, normalizing it to [-1, 1] as done during training, with a shape of (5, 1000, 70). I set the InversionNet model to evaluation mode, move it to the appropriate device, and load the best weights from /kaggle/working/model_best.pth. Using a DataLoader, I process the test data in batches of 64, make predictions, denormalize the outputs from [-1, 1] to [0, 1], and extract odd-numbered columns. I denormalize the predictions to the velocity range (1501 to 4500 m/s) using the scale factor 2999.0 and offset 1501.0, then format the results into a submission DataFrame with oid_ypos and x_{col} columns, saving it as /kaggle/working/submission.csv for evaluation on the platform.","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nimport numpy as np\nimport pandas as pd\nimport os\nfrom glob import glob\n\n# Custom Dataset for test data\n\nclass TestDataset(Dataset):\n    def __init__(self, test_dir):\n        self.test_files = sorted(glob(os.path.join(test_dir, \"*.npy\")))\n    \n    def __len__(self):\n        return len(self.test_files)\n    \n    def __getitem__(self, idx):\n        file_path = self.test_files[idx]\n        seismic_data = np.load(file_path)  # Shape: (5, 1000, 70)\n        # Crop to (5, 1000, 70) by taking the first 1000 time steps\n        seismic_data = seismic_data[:, :1000, :]\n        seismic_data = torch.tensor(seismic_data, dtype=torch.float32)\n        # Normalize the seismic data (same as training)\n        seismic_data = 2 * (seismic_data - seismic_data.min()) / (seismic_data.max() - seismic_data.min()) - 1\n        filename = os.path.basename(file_path).replace(\".npy\", \"\")\n        return seismic_data, filename\n\n# Set device\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = InversionNet()\nmodel = model.to(device)\n\n# Load the best weights\ncheckpoint_path = \"/kaggle/working/model_best.pth\"\ncheckpoint = torch.load(checkpoint_path, map_location=device)\nmodel.load_state_dict(checkpoint['model_state_dict'])\nmodel.eval()  # Switch to evaluation mode\nprint(f\"Model loaded from best checkpoint with Val Loss = {checkpoint['val_loss']:.4f}\")\n\n# Create test dataset and DataLoader\ntest_dir = \"/kaggle/input/waveform-inversion/test\"\ntest_dataset = TestDataset(test_dir)\ntest_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)\nprint(f\"Number of test samples: {len(test_dataset)}\")\n\n# Make predictions\nall_predictions = []\nall_filenames = []\n\nwith torch.no_grad():\n    for seismic_batch, filenames in test_dataloader:\n        seismic_batch = seismic_batch.to(device)\n        # Forward pass\n        outputs = model(seismic_batch)  # Shape: (batch_size, 1, 70, 70)\n        outputs = (outputs + 1) / 2  # Denormalize from [-1, 1] to [0, 1]\n        all_predictions.append(outputs.cpu().numpy())\n        all_filenames.extend(filenames)\n\n# Concatenate all predictions\nall_predictions = np.concatenate(all_predictions, axis=0)  # Shape: (num_samples, 1, 70, 70)\nall_predictions = all_predictions.squeeze(1)  # Shape: (num_samples, 70, 70)\nprint(f\"Predictions shape: {all_predictions.shape}\")\n\n# Extract odd-numbered columns\nodd_columns = np.arange(1, 70, 2)  # [1, 3, 5, ..., 69]\npredictions_odd = all_predictions[:, :, odd_columns]  # Shape: (num_samples, 70, 35)\nprint(f\"Predictions (odd columns) shape: {predictions_odd.shape}\")\n\n# Denormalize predictions to m/s\npredictions_denorm = predictions_odd * 2999.0 + 1501.0  # Shape: (num_samples, 70, 35)\nprint(f\"Denormalized predictions shape: {predictions_denorm.shape}\")\n\n# Prepare submission file\nsubmission_rows = []\nnum_samples = len(test_dataset)\n\nfor sample_idx in range(num_samples):\n    filename = all_filenames[sample_idx]\n    for row in range(70):\n        # Form the id in the format filename_y_row\n        sample_id = f\"{filename}_y_{row}\"\n        # Get values for odd-numbered columns\n        values = predictions_denorm[sample_idx, row, :]  # Shape: (35,)\n        row_data = [sample_id] + values.tolist()\n        submission_rows.append(row_data)\n\n# Create DataFrame\ncolumns = ['oid_ypos'] + [f\"x_{col}\" for col in odd_columns]\nsubmission_df = pd.DataFrame(submission_rows, columns=columns)\n\n# Verify submission format\nprint(\"Submission DataFrame shape:\", submission_df.shape)\nprint(\"Submission DataFrame columns:\", submission_df.columns.tolist())\nprint(submission_df.head())\n\n# Compare with sample submission\nsample_submission = pd.read_csv(\"/kaggle/input/waveform-inversion/sample_submission.csv\")\nprint(\"\\nSample submission shape:\", sample_submission.shape)\nprint(\"Sample submission columns:\", sample_submission.columns.tolist())\nprint(sample_submission.head())\n\n# Save to CSV with the correct path\nsubmission_df.to_csv(\"/kaggle/working/submission.csv\", index=False)\nprint(\"Submission file saved as /kaggle/working/submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-13T11:49:58.592814Z","iopub.execute_input":"2025-05-13T11:49:58.593215Z","iopub.status.idle":"2025-05-13T12:22:58.994661Z","shell.execute_reply.started":"2025-05-13T11:49:58.593182Z","shell.execute_reply":"2025-05-13T12:22:58.993702Z"}},"outputs":[],"execution_count":null}]}