{"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"}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Yale/UNC-CH - Geophysical Waveform Inversion Starter Notebook\n\n## Competition Goal:\nPredict subsurface velocity maps from seismic waveform data using machine learning, potentially guided by physics.\n\n## Approach:\nThis notebook implements a baseline approach using a U-Net architecture, a common choice for image-to-image translation tasks.\n\n## Key Challenges & Considerations:\n*   **Physics-Guided ML:** This notebook uses a data-driven U-Net. Integrating physics (e.g., wave equation loss) is encouraged by the competition but not implemented here.\n*   **Data Size:** Uses the provided Kaggle dataset samples. For better results, using the full OpenFWI dataset is recommended.\n*   **Input Shape:** Handling the 4D input seismic data `(batch, sources, time, receivers)` to predict a 2D map `(height, width)` requires careful preprocessing or model design. This notebook uses a simplified approach.\n*   **Evaluation:** Mean Absolute Error (MAE).\n*   **Submission:** Specific format requiring only odd-indexed columns (`x_1, x_3, ...`) for each predicted row (`y_pos`).","metadata":{"_uuid":"71c05cff-fb89-49fc-bed4-a3560330d483","_cell_guid":"ac9a4033-7a83-4014-9f48-f7a7f2cf0d62","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport glob\nimport gc\nfrom tqdm.notebook import tqdm\n\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import train_test_split","metadata":{"_uuid":"1102f962-e4a2-4fcb-9861-93ec6efec611","_cell_guid":"ef2c761a-171a-470e-ae5c-8bb1ebd6eeff","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:23:38.497545Z","iopub.execute_input":"2025-04-10T05:23:38.497812Z","iopub.status.idle":"2025-04-10T05:23:44.547671Z","shell.execute_reply.started":"2025-04-10T05:23:38.497792Z","shell.execute_reply":"2025-04-10T05:23:44.547093Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Configuration","metadata":{"_uuid":"1560f512-6dfb-40ce-b500-635bc6720af8","_cell_guid":"ff4ea1a0-4be2-4365-8194-a68d59d37c1c","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# --- Configuration ---\nCOMPETITION_NAME = \"waveform-inversion\"\nBASE_PATH = f\"/kaggle/input/{COMPETITION_NAME}\"\nTRAIN_SAMPLES_PATH = os.path.join(BASE_PATH, \"train_samples\")\nTEST_PATH = os.path.join(BASE_PATH, \"test\")\nSAMPLE_SUB_PATH = os.path.join(BASE_PATH, \"sample_submission.csv\")\n\n# Model & Training Params (EXAMPLES - need tuning!)\nSEED = 42\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nBATCH_SIZE = 16 # Adjust based on GPU memory\nIMG_HEIGHT = 70 # Assuming velocity map height from typical data/desc\nIMG_WIDTH = 70  # Assuming velocity map width (consistent with submission x_69)\nN_INPUT_CHANNELS = 1 # Placeholder - how to represent 4D seismic data as input channels?\nN_OUTPUT_CHANNELS = 1 # Predicting a single velocity map\nEPOCHS = 20 # Start small, increase for real training\nLEARNING_RATE = 1e-4\nVALIDATION_SPLIT = 0.1 # Use 10% of training data for validation\n\n# Submission Params\nSUBMISSION_ODD_COLS_ONLY = True\nN_SUBMISSION_COLS = 70 # Number of columns in the full velocity map (0 to 69)\n\nprint(f\"Using device: {DEVICE}\")\nprint(f\"Base path: {BASE_PATH}\")\nprint(f\"Image dimensions (H, W): ({IMG_HEIGHT}, {IMG_WIDTH})\") # Verify these dimensions!\n\n# Set seed for reproducibility\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(SEED)\n    torch.cuda.manual_seed_all(SEED)\ntorch.backends.cudnn.deterministic = True\ntorch.backends.cudnn.benchmark = False","metadata":{"_uuid":"36b1d9f9-3671-4a07-9ef4-816fc018ba2f","_cell_guid":"0085fa2c-6b57-4b25-8d43-6bccb21525f5","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:23:44.548352Z","iopub.execute_input":"2025-04-10T05:23:44.548655Z","iopub.status.idle":"2025-04-10T05:23:44.638973Z","shell.execute_reply.started":"2025-04-10T05:23:44.548638Z","shell.execute_reply":"2025-04-10T05:23:44.638256Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Load Data & Metadata\n\nWe need to find pairs of seismic data files and velocity map files. The naming convention differs between dataset families (Vel/Style vs. Fault). We'll create a DataFrame to manage file paths.\n\n**Note:** This only scans the provided `train_samples`. For a real submission, adapt this to use the full OpenFWI dataset if downloaded separately.","metadata":{"_uuid":"14e01667-6a23-4b4f-a824-675cec61cf9d","_cell_guid":"6a6119ad-7a67-4fec-83db-499a6a5deaed","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"def get_train_files_df(train_path):\n    all_files = []\n\n    # --- Vel and Style families ---\n    # data/*.npy paired with model/*.npy\n    data_files_vel_style = sorted(glob.glob(os.path.join(train_path, \"data\", \"*.npy\")))\n    model_files_vel_style = sorted(glob.glob(os.path.join(train_path, \"model\", \"*.npy\")))\n\n    if len(data_files_vel_style) == len(model_files_vel_style) and len(data_files_vel_style) > 0:\n        print(f\"Found {len(data_files_vel_style)} Vel/Style data/model pairs.\")\n        # Basic check: Assume files correspond based on sorting by name parts\n        # Example: data1.npy -> model1.npy\n        file_map = {}\n        for f in data_files_vel_style:\n            basename = os.path.basename(f)\n            key = basename.replace('data','').replace('.npy','')\n            file_map[key] = {'data': f}\n        for f in model_files_vel_style:\n            basename = os.path.basename(f)\n            key = basename.replace('model','').replace('.npy','')\n            if key in file_map:\n                 file_map[key]['model'] = f\n\n        for key, paths in file_map.items():\n             if 'data' in paths and 'model' in paths:\n                 all_files.append({'data_path': paths['data'], 'model_path': paths['model'], 'family': 'Vel/Style'})\n             else:\n                 print(f\"Warning: Missing pair for key {key} in Vel/Style\")\n\n    else:\n        print(\"Mismatch or no Vel/Style files found.\")\n        if len(data_files_vel_style) != len(model_files_vel_style):\n             print(f\"Warning: Mismatch in Vel/Style file counts: {len(data_files_vel_style)} data vs {len(model_files_vel_style)} model files.\")\n\n\n    # --- Fault family ---\n    # seis_{n}_1_{i}.npy paired with vel_{n}_1_{i}.npy\n    data_files_fault = sorted(glob.glob(os.path.join(train_path, \"seis_*.npy\")))\n    model_files_fault = sorted(glob.glob(os.path.join(train_path, \"vel_*.npy\")))\n\n    if len(data_files_fault) > 0:\n        print(f\"Found {len(data_files_fault)} Fault seismic files and {len(model_files_fault)} Fault velocity files.\")\n        # Create a map based on the common part of the filename\n        file_map_fault = {}\n        for f in data_files_fault:\n            basename = os.path.basename(f)\n            key = basename.replace('seis_','').replace('.npy','') # e.g., \"10_1_0\"\n            file_map_fault[key] = {'data': f}\n        for f in model_files_fault:\n             basename = os.path.basename(f)\n             key = basename.replace('vel_','').replace('.npy','') # e.g., \"10_1_0\"\n             if key in file_map_fault:\n                 file_map_fault[key]['model'] = f\n             else:\n                  # This case might happen if vel files exist without corresponding seis files\n                  pass # Or add logging\n\n        for key, paths in file_map_fault.items():\n             if 'data' in paths and 'model' in paths:\n                 all_files.append({'data_path': paths['data'], 'model_path': paths['model'], 'family': 'Fault'})\n             else:\n                 print(f\"Warning: Missing pair for key {key} in Fault family\")\n\n    else:\n         print(\"No Fault files found.\")\n\n\n    if not all_files:\n        print(\"ERROR: No training file pairs were found. Check paths and file structures.\")\n        return pd.DataFrame()\n\n    df = pd.DataFrame(all_files)\n    print(f\"Total training file pairs found: {len(df)}\")\n    return df\n\ntrain_df = get_train_files_df(TRAIN_SAMPLES_PATH)\ndisplay(train_df.head())\n\n# --- Get Test Files ---\ntest_files = sorted(glob.glob(os.path.join(TEST_PATH, \"*.npy\")))\ntest_oids = [os.path.basename(f).replace('.npy', '') for f in test_files]\nprint(f\"\\nFound {len(test_files)} test files.\")\n# print(test_oids[:5]) # Example test oids","metadata":{"_uuid":"182f48cf-40c4-413d-a1d6-233ac36b8558","_cell_guid":"a8ed5ba9-ff8c-4511-99c9-2887f2aef88a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:23:44.639830Z","iopub.execute_input":"2025-04-10T05:23:44.640034Z","iopub.status.idle":"2025-04-10T05:23:45.202474Z","shell.execute_reply.started":"2025-04-10T05:23:44.640018Z","shell.execute_reply":"2025-04-10T05:23:45.201620Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Dataset and DataLoader\n\nCreate a custom PyTorch Dataset to load seismic data and velocity maps. This involves loading `.npy` files, which contain batches of samples. We'll iterate through the samples within each file.\n\n**Preprocessing:**\n*   **Input (Seismic):** Normalization/Standardization is crucial. The 4D shape `(batch, sources, time, receivers)` needs processing to fit a typical 2D CNN input `(batch, channels, height, width)`. This is a critical step needing careful consideration. **Here, we implement a placeholder:** *select the first source and treat time steps as input channels*. This is likely suboptimal and needs refinement based on domain knowledge or experimentation (e.g., averaging sources, using FFT, using 3D CNNs initially, etc.).\n*   **Output (Velocity Map):** Normalization is also important. Min-Max scaling to [0, 1] or [-1, 1] is common. We need to store the original range for inverse transformation during prediction.","metadata":{"_uuid":"ba0febd1-e017-416a-a295-51666f0e3602","_cell_guid":"b28b0435-1b70-42cf-84cd-4eab56b0d81b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# Placeholder: Store min/max velocity values for potential normalization\n# These might be estimated from the training data or known physical bounds\nVELOCITY_MIN = 1500 # Example value (m/s) - **MUST BE ADJUSTED BASED ON DATA**\nVELOCITY_MAX = 4500 # Example value (m/s) - **MUST BE ADJUSTED BASED ON DATA**\n\ndef preprocess_input(seismic_data_4d):\n    \"\"\"\n    Placeholder function to process 4D seismic data into a format suitable for the U-Net.\n    Input: numpy array (batch_size, num_sources, time_steps, num_receivers)\n    Output: torch tensor (batch_size, N_INPUT_CHANNELS, height, width) - requires careful design!\n\n    Current simple strategy:\n    1. Select the first source. -> (batch, time, receivers)\n    2. Treat time_steps as channels? Or maybe use Conv1D first?\n    3. Normalize.\n    4. Reshape/Interpolate if needed to match expected H, W for U-Net?\n\n    THIS IS A MAJOR SIMPLIFICATION AND LIKELY NEEDS SIGNIFICANT IMPROVEMENT.\n    \"\"\"\n    # Example: Select first source, keep time and receivers\n    # Shape: (batch_size, time_steps, num_receivers)\n    processed_data = seismic_data_4d[:, 0, :, :]\n\n    # Normalize (example: standardize per sample)\n    mean = np.mean(processed_data, axis=(1, 2), keepdims=True)\n    std = np.std(processed_data, axis=(1, 2), keepdims=True)\n    processed_data = (processed_data - mean) / (std + 1e-6) # Add epsilon for stability\n\n    # Reshape/Adapt to U-Net input (e.g., treat time as channels or use specific layers)\n    # Assuming IMG_HEIGHT=time_steps, IMG_WIDTH=num_receivers FOR THIS PLACEHOLDER\n    # This assumption is likely INCORRECT and depends heavily on the actual data dimensions\n    # and the chosen model architecture.\n    if processed_data.shape[1] != IMG_HEIGHT or processed_data.shape[2] != IMG_WIDTH:\n         # This part needs real implementation - maybe interpolation, padding, or different architecture\n         # For now, we'll just add a channel dimension assuming shapes match (WHICH THEY PROBABLY DON'T)\n         print(f\"Warning: Input shape mismatch {processed_data.shape} vs expected ({IMG_HEIGHT}, {IMG_WIDTH}). Requires proper handling.\")\n         # As a fallback placeholder, let's just take a slice or resize crudely if possible.\n         # This is highly problematic and just for code structure demonstration.\n         # Let's assume we can reshape/select to get (batch, N_INPUT_CHANNELS, H, W)\n         # Example: adding a channel dim - this assumes T=H, R=W\n         if N_INPUT_CHANNELS == 1:\n             processed_data = processed_data[:, np.newaxis, :, :] # Add channel dim\n             # If shapes *still* don't match H, W, resizing/cropping/padding needed here\n             # This step is complex and data-dependent.\n         else:\n             # Handle multiple input channels (e.g. if time is treated as channels)\n             # processed_data = processed_data.transpose(0, 1, 2) # Example, needs correct logic\n             raise NotImplementedError(\"Input channel handling > 1 needs specific implementation\")\n\n\n    # Ensure final shape matches (batch, N_INPUT_CHANNELS, IMG_HEIGHT, IMG_WIDTH)\n    # Add crude resizing/padding if shapes don't match (VERY basic placeholder)\n    current_h, current_w = processed_data.shape[2], processed_data.shape[3]\n    if current_h != IMG_HEIGHT or current_w != IMG_WIDTH:\n        # Use torch functional interpolate (requires tensor)\n        temp_tensor = torch.tensor(processed_data, dtype=torch.float32)\n        temp_tensor = nn.functional.interpolate(temp_tensor, size=(IMG_HEIGHT, IMG_WIDTH), mode='bilinear', align_corners=False)\n        processed_data = temp_tensor.numpy()\n        print(f\"Resized input from {(current_h, current_w)} to {(IMG_HEIGHT, IMG_WIDTH)}\")\n\n\n    return torch.tensor(processed_data, dtype=torch.float32)\n\n\ndef preprocess_output(velocity_map_3d):\n    \"\"\"\n    Process 3D velocity map (batch, H, W) for model target.\n    Normalize and add channel dimension.\n    \"\"\"\n    # Add channel dim: (batch, H, W) -> (batch, 1, H, W)\n    velocity_map_4d = velocity_map_3d[:, np.newaxis, :, :]\n\n    # Normalize (Example: Min-Max scaling to [0, 1])\n    normalized_map = (velocity_map_4d - VELOCITY_MIN) / (VELOCITY_MAX - VELOCITY_MIN)\n    normalized_map = np.clip(normalized_map, 0, 1) # Ensure values are within [0, 1]\n\n    return torch.tensor(normalized_map, dtype=torch.float32)\n\ndef postprocess_output(prediction_tensor):\n    \"\"\"\n    Inverse transform the model's normalized prediction back to original velocity scale.\n    Input: torch tensor (batch, 1, H, W)\n    Output: numpy array (batch, H, W)\n    \"\"\"\n    prediction = prediction_tensor.detach().cpu().numpy()\n    # Denormalize (from [0, 1] back to original scale)\n    velocity_map = prediction * (VELOCITY_MAX - VELOCITY_MIN) + VELOCITY_MIN\n    # Remove channel dim: (batch, 1, H, W) -> (batch, H, W)\n    velocity_map = velocity_map.squeeze(1)\n    return velocity_map\n\n\nclass WaveformDataset(Dataset):\n    def __init__(self, df, file_indices, transform_input=None, transform_output=None):\n        self.df = df\n        self.file_indices = file_indices # Indices of files in df to use for this dataset split\n        self.transform_input = transform_input\n        self.transform_output = transform_output\n\n        self.samples = []\n        print(f\"Loading data pointers for {len(self.file_indices)} files...\")\n        for idx in tqdm(self.file_indices):\n            row = self.df.iloc[idx]\n            try:\n                # Load the entire file content once\n                seismic_data_batch = np.load(row['data_path'])\n                velocity_map_batch = np.load(row['model_path'])\n\n                # Verify shapes (example: check first sample)\n                # print(f\"Loaded seismic shape: {seismic_data_batch.shape}, velocity shape: {velocity_map_batch.shape}\")\n                # Expected: (batch_size, num_sources, time_steps, num_receivers) and (batch_size, height, width)\n\n                num_samples_in_file = seismic_data_batch.shape[0]\n                if num_samples_in_file != velocity_map_batch.shape[0]:\n                     print(f\"Warning: Mismatch in batch size for file {idx}. Skipping.\")\n                     continue\n\n                for i in range(num_samples_in_file):\n                    # Store pointers or indices to individual samples within the loaded batches\n                    # Option 1: Store indices (memory efficient if files loaded lazily)\n                    self.samples.append({'file_idx': idx, 'sample_idx': i})\n                    # Option 2: Store actual data (simpler but uses more RAM if all loaded upfront)\n                    # self.samples.append({'seismic': seismic_data_batch[i], 'velocity': velocity_map_batch[i]})\n\n            except Exception as e:\n                print(f\"Error loading file index {idx} ({row['data_path']}/{row['model_path']}): {e}\")\n\n        print(f\"Total individual samples found: {len(self.samples)}\")\n        # If using Option 2 above, clear the large loaded batches now if needed\n        # del seismic_data_batch, velocity_map_batch\n        # gc.collect()\n\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        sample_info = self.samples[idx]\n        file_idx = sample_info['file_idx']\n        sample_idx_in_file = sample_info['sample_idx']\n\n        # Load data for the specific sample\n        # This might involve reloading the file or accessing pre-loaded data\n        try:\n            # --- Reload file (if not pre-loaded) ---\n            row = self.df.iloc[file_idx]\n            # Using mmap_mode='r' can help with large files if memory is an issue,\n            # but might be slower for random access depending on usage pattern.\n            seismic_data_full_batch = np.load(row['data_path']) # Potentially use mmap_mode='r'\n            velocity_map_full_batch = np.load(row['model_path']) # Potentially use mmap_mode='r'\n\n            seismic_sample = seismic_data_full_batch[sample_idx_in_file]\n            velocity_sample = velocity_map_full_batch[sample_idx_in_file]\n            # --- ---\n\n            # --- Access pre-loaded data (if using Option 2 in __init__) ---\n            # seismic_sample = sample_info['seismic']\n            # velocity_sample = sample_info['velocity']\n            # --- ---\n\n            # Apply transformations (preprocessing)\n            # The preprocessing functions expect batches, so add a batch dim temporarily\n            if self.transform_input:\n                seismic_tensor = self.transform_input(seismic_sample[np.newaxis, ...]) # Add batch dim\n                seismic_tensor = seismic_tensor.squeeze(0) # Remove batch dim\n            else:\n                seismic_tensor = torch.tensor(seismic_sample, dtype=torch.float32) # Basic tensor conversion\n\n            if self.transform_output:\n                velocity_tensor = self.transform_output(velocity_sample[np.newaxis, ...]) # Add batch dim\n                velocity_tensor = velocity_tensor.squeeze(0) # Remove batch dim\n            else:\n                velocity_tensor = torch.tensor(velocity_sample, dtype=torch.float32) # Basic tensor conversion\n\n\n            # Verify output tensor shape (should be ~ [C_out, H, W]) after processing\n            # print(f\"Sample {idx}: Input shape {seismic_tensor.shape}, Output shape {velocity_tensor.shape}\")\n            if velocity_tensor.shape[0] != N_OUTPUT_CHANNELS or velocity_tensor.shape[1] != IMG_HEIGHT or velocity_tensor.shape[2] != IMG_WIDTH:\n                 print(f\"Warning: Unexpected output tensor shape after processing: {velocity_tensor.shape}\")\n\n\n            return seismic_tensor, velocity_tensor\n\n        except Exception as e:\n             print(f\"Error getting item {idx} (file {file_idx}, sample {sample_idx_in_file}): {e}\")\n             # Return dummy data or raise error? For now, return None and handle in DataLoader.\n             # This needs robust error handling.\n             return None, None\n\n\n# --- Data Splitting ---\nif not train_df.empty:\n    train_indices, val_indices = train_test_split(\n        range(len(train_df)), # Split based on file indices\n        test_size=VALIDATION_SPLIT,\n        random_state=SEED\n    )\n\n    print(f\"\\nSplitting {len(train_df)} files into:\")\n    print(f\"Training files: {len(train_indices)}\")\n    print(f\"Validation files: {len(val_indices)}\")\n\n    # Create Datasets\n    train_dataset = WaveformDataset(train_df, train_indices, transform_input=preprocess_input, transform_output=preprocess_output)\n    val_dataset = WaveformDataset(train_df, val_indices, transform_input=preprocess_input, transform_output=preprocess_output)\n\n    # Create DataLoaders\n    # Handle potential None values returned by dataset __getitem__ due to errors\n    def collate_fn(batch):\n        batch = list(filter(lambda x: x[0] is not None and x[1] is not None, batch))\n        if not batch: return torch.Tensor(), torch.Tensor() # Return empty tensors if batch is empty\n        return torch.utils.data.dataloader.default_collate(batch)\n\n    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=os.cpu_count(), pin_memory=True, collate_fn=collate_fn)\n    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=os.cpu_count(), pin_memory=True, collate_fn=collate_fn)\n\n    print(f\"\\nDataLoaders created.\")\n    # Optional: Check a batch shape\n    # try:\n    #     inputs, targets = next(iter(train_loader))\n    #     print(f\"Sample batch input shape: {inputs.shape}\") # Should be [B, C_in, H_in, W_in]\n    #     print(f\"Sample batch target shape: {targets.shape}\") # Should be [B, C_out, H, W]\n    # except Exception as e:\n    #     print(f\"Could not load a sample batch: {e}\")\n\nelse:\n    print(\"\\nSkipping Dataset/DataLoader creation due to missing training files.\")\n    train_loader, val_loader = None, None","metadata":{"_uuid":"b8b02a19-1e0f-4cdc-9663-2c91297fd644","_cell_guid":"5725a1de-d5e8-444d-affd-fb7d0f78b01b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:23:45.203152Z","iopub.execute_input":"2025-04-10T05:23:45.203401Z","iopub.status.idle":"2025-04-10T05:23:45.224867Z","shell.execute_reply.started":"2025-04-10T05:23:45.203381Z","shell.execute_reply":"2025-04-10T05:23:45.224229Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Model Definition (U-Net)\n\nWe'll use a standard U-Net architecture. You might need to adjust `n_channels` (input features) based on your seismic data preprocessing. `n_classes` is 1 for the single velocity map output.","metadata":{"_uuid":"a1c24b22-b185-4d54-855c-7821cf4f7d5e","_cell_guid":"07fc982e-e01a-4e4f-9759-7f302ba03e5a","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# Basic U-Net implementation (taken from PyTorch Hub or common implementations)\n# (Ensure you have the necessary U-Net code defined here or imported)\n\nclass DoubleConv(nn.Module):\n    \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.double_conv(x)\n\nclass Down(nn.Module):\n    \"\"\"Downscaling with maxpool then double conv\"\"\"\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.maxpool_conv = nn.Sequential(\n            nn.MaxPool2d(2),\n            DoubleConv(in_channels, out_channels)\n        )\n\n    def forward(self, x):\n        return self.maxpool_conv(x)\n\nclass Up(nn.Module):\n    \"\"\"Upscaling then double conv\"\"\"\n    def __init__(self, in_channels, out_channels, bilinear=True):\n        super().__init__()\n        # if bilinear, use the normal convolutions to reduce the number of channels\n        if bilinear:\n            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n        else:\n            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n            self.conv = DoubleConv(in_channels, out_channels)\n\n    def forward(self, x1, x2):\n        x1 = self.up(x1)\n        # input is CHW\n        diffY = x2.size()[2] - x1.size()[2]\n        diffX = x2.size()[3] - x1.size()[3]\n        x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])\n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\nclass OutConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(OutConv, self).__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass UNet(nn.Module):\n    def __init__(self, n_channels, n_classes, bilinear=True):\n        super(UNet, self).__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n        self.bilinear = bilinear\n\n        self.inc = DoubleConv(n_channels, 64)\n        self.down1 = Down(64, 128)\n        self.down2 = Down(128, 256)\n        self.down3 = Down(256, 512)\n        factor = 2 if bilinear else 1\n        self.down4 = Down(512, 1024 // factor)\n        self.up1 = Up(1024, 512 // factor, bilinear)\n        self.up2 = Up(512, 256 // factor, bilinear)\n        self.up3 = Up(256, 128 // factor, bilinear)\n        self.up4 = Up(128, 64, bilinear)\n        self.outc = OutConv(64, n_classes)\n        # Add a final activation like Sigmoid if output is normalized to [0, 1]\n        self.final_activation = nn.Sigmoid() # Or nn.Identity() if not normalizing to [0,1]\n\n    def forward(self, x):\n        x1 = self.inc(x)\n        x2 = self.down1(x1)\n        x3 = self.down2(x2)\n        x4 = self.down3(x3)\n        x5 = self.down4(x4)\n        x = self.up1(x5, x4)\n        x = self.up2(x, x3)\n        x = self.up3(x, x2)\n        x = self.up4(x, x1)\n        logits = self.outc(x)\n        # Apply final activation\n        outputs = self.final_activation(logits)\n        return outputs\n\n# Instantiate the model\nmodel = UNet(n_channels=N_INPUT_CHANNELS, n_classes=N_OUTPUT_CHANNELS).to(DEVICE)\nprint(f\"U-Net model created with {N_INPUT_CHANNELS} input channels and {N_OUTPUT_CHANNELS} output classes.\")\n# Optional: Print model summary (requires torchinfo)\n# try:\n#     from torchinfo import summary\n#     # Input size needs to match preprocessed input tensor shape B,C,H,W\n#     # Use a dummy batch size, e.g., 1\n#     # WARNING: This assumes your preprocess_input correctly shapes data to (B, N_INPUT_CHANNELS, IMG_HEIGHT, IMG_WIDTH)\n#     # Adjust input_size if your preprocessing yields different H, W for the input tensor\n#     dummy_input_size = (BATCH_SIZE, N_INPUT_CHANNELS, IMG_HEIGHT, IMG_WIDTH)\n#     summary(model, input_size=dummy_input_size)\n# except ImportError:\n#     print(\"torchinfo not installed, skipping model summary.\")\n# except Exception as e:\n#      print(f\"Could not generate model summary. Check input size/preprocessing. Error: {e}\")","metadata":{"_uuid":"13dedae4-c84a-441c-96d6-ca2c1a9e1587","_cell_guid":"e2f5dc61-0452-48d1-9ebb-56df511b1f18","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:24:09.468692Z","iopub.execute_input":"2025-04-10T05:24:09.469334Z","iopub.status.idle":"2025-04-10T05:24:09.916767Z","shell.execute_reply.started":"2025-04-10T05:24:09.469307Z","shell.execute_reply":"2025-04-10T05:24:09.916127Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Training\n\n*   **Loss Function:** Mean Absolute Error (L1 Loss) as per competition metric.\n*   **Optimizer:** Adam is a common choice.\n*   **Training Loop:** Standard epoch iteration, batch processing, backpropagation.\n*   **Validation:** Calculate MAE on the validation set after each epoch to monitor progress and prevent overfitting. Save the best model based on validation MAE.","metadata":{"_uuid":"660f3120-c3ec-4d1f-bf62-99c6e980024b","_cell_guid":"244b0a20-73bd-4e84-ae88-fbd830aa54c2","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# --- Loss Function ---\n# MAE Loss\ncriterion = nn.L1Loss()\n\n# --- Optimizer ---\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n\n# --- Learning Rate Scheduler (Optional) ---\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.1, patience=3)\n\n# --- Training Loop ---\nbest_val_mae = float('inf')\nbest_model_path = \"best_model.pth\"\n\nif train_loader and val_loader: # Only train if data loaders are available\n    for epoch in range(EPOCHS):\n        model.train()\n        train_loss = 0.0\n        pbar_train = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS} [Train]\")\n\n        for inputs, targets in pbar_train:\n            if inputs.numel() == 0 or targets.numel() == 0: continue # Skip empty batches\n            inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)\n\n            optimizer.zero_grad()\n            outputs = model(inputs)\n\n            # Ensure shapes match for loss calculation\n            if outputs.shape != targets.shape:\n                 print(f\"Shape mismatch! Output: {outputs.shape}, Target: {targets.shape}. Skipping batch.\")\n                 # This often indicates an issue in preprocessing or model output layer\n                 continue\n\n            loss = criterion(outputs, targets)\n            loss.backward()\n            optimizer.step()\n\n            train_loss += loss.item() * inputs.size(0)\n            pbar_train.set_postfix(loss=loss.item())\n\n        train_loss /= len(train_loader.dataset) # Average loss over all samples\n\n        # --- Validation ---\n        model.eval()\n        val_mae = 0.0\n        pbar_val = tqdm(val_loader, desc=f\"Epoch {epoch+1}/{EPOCHS} [Val]\")\n        with torch.no_grad():\n            for inputs, targets in pbar_val:\n                if inputs.numel() == 0 or targets.numel() == 0: continue # Skip empty batches\n                inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)\n                outputs = model(inputs)\n\n                if outputs.shape != targets.shape:\n                     print(f\"Shape mismatch! Output: {outputs.shape}, Target: {targets.shape}. Skipping batch.\")\n                     continue\n\n                mae_batch = criterion(outputs, targets) # MAE loss\n                val_mae += mae_batch.item() * inputs.size(0)\n                pbar_val.set_postfix(mae=mae_batch.item())\n\n        val_mae /= len(val_loader.dataset) # Average MAE over all validation samples\n\n        print(f\"Epoch {epoch+1}/{EPOCHS} - Train Loss: {train_loss:.6f} - Val MAE: {val_mae:.6f}\")\n\n        # Optional: Update learning rate scheduler\n        # scheduler.step(val_mae)\n\n        # Save the best model\n        if val_mae < best_val_mae:\n            best_val_mae = val_mae\n            torch.save(model.state_dict(), best_model_path)\n            print(f\"✨ New best model saved with Val MAE: {best_val_mae:.6f} to {best_model_path}\")\n\n        # Clean up GPU memory\n        del inputs, targets, outputs\n        gc.collect()\n        if DEVICE == 'cuda':\n            torch.cuda.empty_cache()\n\n    print(f\"\\nTraining finished. Best Validation MAE: {best_val_mae:.6f}\")\n\nelse:\n     print(\"Skipping training as DataLoaders could not be created.\")","metadata":{"_uuid":"5d37e933-b0a3-4486-a1b3-4647de443a7b","_cell_guid":"7da7d235-d4dd-47c1-a26d-82e92b148d96","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:24:13.055141Z","iopub.execute_input":"2025-04-10T05:24:13.055435Z","iopub.status.idle":"2025-04-10T05:24:15.597097Z","shell.execute_reply.started":"2025-04-10T05:24:13.055415Z","shell.execute_reply":"2025-04-10T05:24:15.596436Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Prediction & Submission\n\n*   Load the best saved model.\n*   Iterate through the test files (`{oid}.npy`).\n*   Preprocess each test seismic data sample.\n*   Perform inference using the model.\n*   Postprocess the prediction (inverse transform normalization).\n*   Format the predictions according to the submission rules (odd columns only, stacked rows).","metadata":{"_uuid":"d7716f06-d4e5-4f91-923b-5949edccab90","_cell_guid":"01896cc5-82c0-4196-a8bc-2309311fae9d","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"class WaveformTestDataset(Dataset):\n    \"\"\"Dataset for loading test data (only seismic input).\"\"\"\n    def __init__(self, test_files, oids, transform_input=None):\n        self.test_files = test_files\n        self.oids = oids\n        self.transform_input = transform_input\n        # Assume test files also contain batches like training files\n        self.samples = []\n        print(f\"Mapping test samples...\")\n        for i, file_path in enumerate(tqdm(self.test_files)):\n            try:\n                # Get the number of samples in the file without loading everything (if possible)\n                # This is tricky with .npy without loading. Let's assume we load it to find out.\n                # Use mmap_mode potentially if files are huge.\n                test_data_batch = np.load(file_path) # Potentially use mmap_mode='r'\n                num_samples_in_file = test_data_batch.shape[0]\n                oid = self.oids[i]\n                for j in range(num_samples_in_file):\n                     # Store file path and index within file for each sample\n                     self.samples.append({'file_path': file_path, 'sample_idx': j, 'oid': oid})\n                del test_data_batch # Free memory\n                gc.collect()\n            except Exception as e:\n                 print(f\"Error processing test file {file_path}: {e}\")\n        print(f\"Total individual test samples found: {len(self.samples)}\")\n\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        sample_info = self.samples[idx]\n        file_path = sample_info['file_path']\n        sample_idx_in_file = sample_info['sample_idx']\n        oid = sample_info['oid']\n\n        try:\n            # Load the specific sample\n            # Again, potentially more efficient ways for huge files (mmap)\n            test_data_full_batch = np.load(file_path) # Potentially use mmap_mode='r'\n            seismic_sample = test_data_full_batch[sample_idx_in_file]\n            del test_data_full_batch # Free memory\n            gc.collect()\n\n            # Apply input transformation\n            if self.transform_input:\n                seismic_tensor = self.transform_input(seismic_sample[np.newaxis, ...]) # Add batch dim\n                seismic_tensor = seismic_tensor.squeeze(0) # Remove batch dim\n            else:\n                seismic_tensor = torch.tensor(seismic_sample, dtype=torch.float32)\n\n            # Return the processed input tensor and the original oid/index info\n            # OID might need careful handling if test files don't map 1:1 with submission OIDs\n            # Assuming each test file {oid}.npy corresponds to one oid in submission.\n            # If test files contain *batches* that need individual predictions, the OID logic needs adjustment.\n            # Let's assume each sample within a test file needs prediction, but they all share the root OID.\n            # The submission format implies one output map per OID.\n            # This part is ambiguous in the problem description if test files are batched.\n            # For now, assume one prediction needed per file {oid}.npy --> average predictions if batched?\n            # Let's refine: Assume one prediction per file {oid}.npy. We'll process the first sample only.\n            # *** THIS NEEDS CLARIFICATION FROM COMPETITION HOSTS or FORUM ***\n            # If each test file needs *one* map prediction, we might average predictions across samples inside, or just use one sample.\n            # Let's proceed assuming we predict only for the *first sample* in each test file for simplicity.\n            if sample_idx_in_file == 0:\n                 return seismic_tensor, oid\n            else:\n                 # Skip other samples within the same file for now\n                 return None, None # Need collate_fn to handle this\n\n        except Exception as e:\n             print(f\"Error getting test item {idx} (file {file_path}, sample {sample_idx_in_file}): {e}\")\n             return None, None\n\n\n# --- Create Test Dataset & DataLoader ---\n# Adjusting the test dataset assumption: Predict one map per file {oid}.npy\n# Modify the dataset to only yield the first sample from each file.\n\nclass WaveformTestDatasetPerFile(Dataset):\n    \"\"\"Dataset for loading test data (one sample per file).\"\"\"\n    def __init__(self, test_files, oids, transform_input=None):\n        self.test_files = test_files\n        self.oids = oids\n        self.transform_input = transform_input\n        print(f\"Found {len(test_files)} test files to process.\")\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        oid = self.oids[idx]\n\n        try:\n            test_data_batch = np.load(file_path) # Potentially use mmap_mode='r'\n            # Use only the first sample from the batch/file\n            seismic_sample = test_data_batch[0]\n            del test_data_batch\n            gc.collect()\n\n            # Apply input transformation\n            if self.transform_input:\n                seismic_tensor = self.transform_input(seismic_sample[np.newaxis, ...]) # Add batch dim\n                seismic_tensor = seismic_tensor.squeeze(0) # Remove batch dim\n            else:\n                seismic_tensor = torch.tensor(seismic_sample, dtype=torch.float32)\n\n            return seismic_tensor, oid\n\n        except Exception as e:\n             print(f\"Error getting test item {idx} (file {file_path}): {e}\")\n             return None, None # Handle potential errors\n\n\n# Helper collate function for test loader\ndef collate_fn_test(batch):\n    batch = list(filter(lambda x: x[0] is not None and x[1] is not None, batch))\n    if not batch: return torch.Tensor(), []\n    inputs = [item[0] for item in batch]\n    oids = [item[1] for item in batch]\n    inputs_collated = torch.utils.data.dataloader.default_collate(inputs)\n    return inputs_collated, oids\n\n\nif os.path.exists(best_model_path):\n    print(f\"Loading best model from {best_model_path}\")\n    # Ensure model architecture is defined before loading state_dict\n    model = UNet(n_channels=N_INPUT_CHANNELS, n_classes=N_OUTPUT_CHANNELS).to(DEVICE)\n    model.load_state_dict(torch.load(best_model_path, map_location=DEVICE))\n    model.eval()\n\n    test_dataset = WaveformTestDatasetPerFile(test_files, test_oids, transform_input=preprocess_input)\n    # Use batch size 1 for prediction if memory is tight or processing one file at a time\n    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=os.cpu_count(), collate_fn=collate_fn_test)\n\n    predictions = []\n    print(\"Starting prediction on test set...\")\n    with torch.no_grad():\n        for inputs, oids_batch in tqdm(test_loader):\n             if inputs.numel() == 0: continue # Skip empty batches\n\n             inputs = inputs.to(DEVICE)\n             outputs_norm = model(inputs) # Normalized output [B, 1, H, W]\n\n             # Postprocess (denormalize) and move to CPU\n             outputs_denorm = postprocess_output(outputs_norm) # Numpy array [B, H, W]\n\n             # Format for submission\n             for i, oid in enumerate(oids_batch):\n                 pred_map = outputs_denorm[i] # Shape (H, W), e.g., (70, 70)\n                 # Iterate through rows (ypos)\n                 for y_pos in range(pred_map.shape[0]): # Iterate through height (rows)\n                     row_data = {'oid_ypos': f\"{oid}_y_{y_pos}\"}\n                     # Extract odd columns (x_1, x_3, ..., x_69)\n                     # Assuming width is N_SUBMISSION_COLS (e.g., 70), so columns are 0 to 69\n                     # Odd indices are 1, 3, 5, ..., 69\n                     for x_pos in range(1, N_SUBMISSION_COLS, 2): # Step by 2\n                          if x_pos < pred_map.shape[1]: # Check bounds\n                              row_data[f'x_{x_pos}'] = pred_map[y_pos, x_pos]\n                          else:\n                              # Handle cases where prediction width might be smaller? Pad? Error?\n                              # For now, let's assume width matches or is larger.\n                              # If smaller, maybe fill with a default value (e.g., 0 or nan) or error out.\n                              # Setting to 0 for now if out of bounds.\n                              print(f\"Warning: x_pos {x_pos} out of bounds for prediction width {pred_map.shape[1]} for oid {oid}. Setting to 0.\")\n                              row_data[f'x_{x_pos}'] = 0.0\n                     predictions.append(row_data)\n\n    print(f\"Generated {len(predictions)} prediction rows.\")\n\n    # Create submission DataFrame\n    if predictions:\n        submission_df = pd.DataFrame(predictions)\n        # Ensure correct column order (oid_ypos, x_1, x_3, ...)\n        cols = ['oid_ypos'] + [f'x_{i}' for i in range(1, N_SUBMISSION_COLS, 2)]\n        submission_df = submission_df[cols]\n\n        submission_df.to_csv(\"submission.csv\", index=False)\n        print(\"submission.csv created successfully.\")\n        display(submission_df.head())\n    else:\n        print(\"No predictions were generated. Cannot create submission file.\")\n\nelif not train_loader or not val_loader:\n     print(\"Skipping prediction as model was not trained due to data loading issues.\")\n     # Create dummy submission based on sample if needed for platform checks\n     if os.path.exists(SAMPLE_SUB_PATH):\n          print(\"Creating dummy submission from sample file.\")\n          sample_sub = pd.read_csv(SAMPLE_SUB_PATH)\n          # Potentially populate with a constant value if required\n          # sample_sub.iloc[:, 1:] = 3000.0 # Example: Fill with constant velocity\n          sample_sub.to_csv(\"submission.csv\", index=False)\n     else:\n          print(\"Sample submission file not found, cannot create dummy submission.\")\n\nelse:\n     print(\"Best model file not found. Cannot run prediction.\")\n     # Create dummy submission if possible (as above)\n     if os.path.exists(SAMPLE_SUB_PATH):\n          print(\"Creating dummy submission from sample file.\")\n          sample_sub = pd.read_csv(SAMPLE_SUB_PATH)\n          sample_sub.to_csv(\"submission.csv\", index=False)\n     else:\n          print(\"Sample submission file not found, cannot create dummy submission.\")","metadata":{"_uuid":"f1a71bd1-b5e6-451d-9ef9-d08a1a4ceea7","_cell_guid":"8f8516c6-f617-4db3-a33f-9b790f0b07ff","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-04-10T05:24:20.622654Z","iopub.execute_input":"2025-04-10T05:24:20.623067Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Potential Improvements & Next Steps\n\n*   **Use Full OpenFWI Dataset:** Download and integrate the complete dataset for training.\n*   **Refine Input Preprocessing:** This is critical. Explore different ways to represent the 4D seismic data for the 2D U-Net (FFT, statistics over time/sources, using Conv1D/3D layers initially, attention mechanisms).\n*   **Hyperparameter Tuning:** Optimize learning rate, batch size, optimizer, U-Net depth/width.\n*   **Data Augmentation:** Apply relevant augmentations (if applicable to seismic data, e.g., noise injection, time shifts - requires domain expertise).\n*   **Advanced Architectures:** Explore alternatives like Vision Transformers (ViT), or networks specifically designed for sequence/spatial data.\n*   **Physics-Informed Neural Networks (PINNs):** Integrate the wave equation or other physical constraints directly into the loss function or network architecture. This is the core theme suggested by the competition organizers.\n*   **Post-processing:** Refine the output predictions (e.g., smoothing, enforcing physical bounds more strictly).\n*   **Ensembling:** Combine predictions from multiple models.\n*   **Error Analysis:** Analyze where the model performs poorly (e.g., specific geological structures, noisy inputs) to guide improvements.","metadata":{"_uuid":"fd713e0e-7370-446a-b732-b472f9952ea0","_cell_guid":"08edd6fd-6d7b-460f-a399-d41ae4e80244","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}}]}