{"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":"none","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 1. After Loading a Single Seismic File\n","metadata":{}},{"cell_type":"markdown","source":"## 1. Imports and Dependencies\n\n```python\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nfrom tqdm import tqdm\nimport pandas as pd\nimport glob\nimport os\n```\n\n---\n\n## 2. Data Loading Function\n\n```python\ndef load_seismic_file(file_path):\n    \"\"\"\n    Loads a seismic waveform from a .npy file and returns a NumPy array.\n    Expected input shape: (S, T, R) or with extra batch dims (e.g., (1, S, T, R)).\n    \"\"\"\n    data = np.load(file_path).astype(np.float32)\n    # Remove any extra leading dimensions\n    while data.ndim > 3:\n        data = data[0]\n    return data  # shape: (S, T, R)\n```\n\n---\n\n## 3. PyTorch Dataset\n\n```python\nclass WaveformDataset(Dataset):\n    def __init__(self, input_files, target_arrays, transform_input=None):\n        \"\"\"\n        input_files: list of file paths to seismic .npy files\n        target_arrays: list of ground truth arrays with shape (H, W)\n        transform_input: optional preprocessing function\n        \"\"\"\n        self.input_files = input_files\n        self.target_arrays = target_arrays\n        self.transform_input = transform_input\n\n    def __len__(self):\n        return len(self.input_files)\n\n    def __getitem__(self, idx):\n        # Load and preprocess seismic data\n        data = load_seismic_file(self.input_files[idx])  # (S, T, R)\n        if self.transform_input:\n            data = self.transform_input(data)\n        # Convert to PyTorch tensor and add channel dimension\n        # (S, T, R) -> (1, S, T, R)\n        data_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0)\n        # Permute to (C, D, H, W): (1, S, T, R) -> (1, R, S, T)\n        data_tensor = data_tensor.permute(0, 3, 1, 2)\n        # Ground truth velocity map tensor\n        target_tensor = torch.tensor(self.target_arrays[idx], dtype=torch.float32)  # (H, W)\n        return data_tensor, target_tensor\n```\n\n---\n\n## 4. Simplified 3D U-Net Model\n\n```python\nclass UNet3D(nn.Module):\n    def __init__(self, in_channels=1, out_channels=1):\n        super(UNet3D, self).__init__()\n        self.encoder = nn.Sequential(\n            nn.Conv3d(in_channels,  8, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv3d(8,           16, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.AdaptiveAvgPool3d((1, 100, 70))\n        )\n        self.decoder = nn.Sequential(\n            nn.ConvTranspose3d(16,  8, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose3d(8, out_channels, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        # x shape: (B, 1, D, H, W)\n        x = self.encoder(x)    # -> (B, 16, 1, 100, 70)\n        x = self.decoder(x)    # -> (B, 1, 1, 100, 70)\n        x = x.squeeze(2)       # -> (B, 1, 100, 70)\n        return x.squeeze(1)    # -> (B, 100, 70)\n```\n\n---\n\n## 5. Sample Data Setup\n\n```python\nexample_file   = \"/kaggle/input/waveform-inversion/train_samples/CurveFault_A/seis2_1_0.npy\"\nexample_target = np.zeros((100, 70), dtype=np.float32)\nfile_list      = [example_file] * 10\ntarget_list    = [example_target] * 10\n\n# Create Dataset and DataLoader\ndataset = WaveformDataset(file_list, target_list)\nloader  = DataLoader(dataset, batch_size=2, shuffle=True)\n```\n\n---\n\n## 6. Model Training Setup\n\n```python\ndevice    = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel     = UNet3D(in_channels=1, out_channels=1).to(device)\ncriterion = nn.L1Loss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(\n    optimizer,\n    mode='min',\n    factor=0.5,\n    patience=2\n)\n```\n\n---\n\n## 7. Training Loop\n\n```python\nnum_epochs    = 3\nbest_val_loss = float(\"inf\")\nfor epoch in range(num_epochs):\n    model.train()\n    running_loss = 0.0\n    progress_bar = tqdm(loader, desc=f\"Epoch {epoch+1}/{num_epochs} Training\")\n    for inputs, targets in progress_bar:\n        inputs  = inputs.to(device)\n        targets = targets.to(device)\n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss    = criterion(outputs, targets)\n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item()\n        progress_bar.set_postfix({\"loss\": loss.item()})\n    avg_loss = running_loss / len(loader)\n    print(f\"Epoch {epoch+1}: avg loss = {avg_loss:.4f}\")\n    scheduler.step(avg_loss)\n    if avg_loss < best_val_loss:\n        best_val_loss = avg_loss\n        torch.save(model.state_dict(), \"best_model.pth\")\n        print(\"Model improved and saved.\")\n```\n\n---\n\n## 8. Inference Function\n\n```python\ndef predict(model, file_paths):\n    model.eval()\n    predictions = []\n    with torch.no_grad():\n        for fp in file_paths:\n            data = load_seismic_file(fp)\n            tensor = torch.tensor(data, dtype=torch.float32)\n            tensor = tensor.unsqueeze(0).unsqueeze(0).to(device)\n            output = model(tensor)\n            predictions.append(output.squeeze().cpu().numpy())\n    return predictions\n```\n\n---\n\n## 9. Submission File Generation\n\n```python\ndef create_submission(oids, predictions):\n    \"\"\"\n    Builds a submission.csv matching the sample_submission format.\n    oids: list[str] identifiers without extension\n    predictions: list[np.ndarray] each of shape (H, W)\n    \"\"\"\n    sample_df = pd.read_csv('/kaggle/input/waveform-inversion/sample_submission.csv')\n    id_col = sample_df.columns[0]\n\n    rows = []\n    width = predictions[0].shape[1]\n    odd_indices = list(range(0, width, 2))\n\n    for oid, pred in zip(oids, predictions):\n        if pred.shape[1] != width:\n            raise ValueError(f\"Width mismatch for {oid}.\")\n        for y in range(pred.shape[0]):\n            row_id = f\"{oid}_y_{y}\"\n            row = [row_id] + [float(pred[y, x]) for x in odd_indices]\n            rows.append(row)\n\n    columns = [id_col] + [f\"x_{i}\" for i in odd_indices]\n    df = pd.DataFrame(rows, columns=columns)\n    df.to_csv('/kaggle/working/submission.csv', index=False)\n    print(\"Submission saved to /kaggle/working/submission.csv with correct ID column.\")\n```\n\n---\n\n## 10. Zero-Filled Sample Submission\n\n```python\nsample_df = pd.read_csv('/kaggle/input/waveform-inversion/sample_submission.csv')\nsample_df.iloc[:, 1:] = 0.0\nsample_df.to_csv('/kaggle/working/submission.csv', index=False)\nprint(\"Sample submission written to /kaggle/working/submission.csv.\")\n```\n","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\n\ndef load_seismic_file(file_path):\n    \"\"\"\n    Load a seismic waveform from a .npy file and return a NumPy array.\n    Strips any extra leading dimensions so always returns a 3D array: (S, T, R).\n    \"\"\"\n    data = np.load(file_path).astype(np.float32)\n    while data.ndim > 3:\n        data = data[0]\n    return data\n\n# Visualization Example: Plot waveform and spectrogram for source=0, receiver=0\nfile_path = \"/kaggle/input/waveform-inversion/train_samples/CurveFault_A/seis2_1_0.npy\"\ndata = load_seismic_file(file_path)  # shape: (S, T, R)\n\n# 1a) Plot waveform\nplt.figure()\nplt.plot(data[0, :, 0])\nplt.title(\"Seismic Waveform (source=0, receiver=0)\")\nplt.xlabel(\"Time step\")\nplt.ylabel(\"Amplitude\")\nplt.show()\n\n# 1b) Plot spectrogram\nplt.figure()\nplt.specgram(data[0, :, 0])\nplt.title(\"Spectrogram (source=0, receiver=0)\")\nplt.xlabel(\"Time step\")\nplt.ylabel(\"Frequency bin\")\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-23T15:04:33.990011Z","iopub.execute_input":"2025-06-23T15:04:33.990220Z","iopub.status.idle":"2025-06-23T15:04:40.677765Z","shell.execute_reply.started":"2025-06-23T15:04:33.990200Z","shell.execute_reply":"2025-06-23T15:04:40.676848Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Visualize One Batch from the DataLoader","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\n\n# ---------------------------\n# 1. Data Loading Function\n# ---------------------------\ndef load_seismic_file(file_path):\n    \"\"\"\n    Load a seismic waveform from a .npy file and return a 3D NumPy array: (S, T, R).\n    Strips any extra leading dimensions.\n    \"\"\"\n    data = np.load(file_path).astype(np.float32)\n    while data.ndim > 3:\n        data = data[0]\n    return data\n\n# ---------------------------\n# 2. PyTorch Dataset Definition\n# ---------------------------\nclass WaveformDataset(Dataset):\n    def __init__(self, input_files, target_arrays, transform_input=None):\n        \"\"\"\n        input_files:   list of paths to .npy seismic files\n        target_arrays: list of 2D NumPy arrays (H×W) ground truth\n        transform_input: optional preprocessing function\n        \"\"\"\n        self.input_files = input_files\n        self.target_arrays = target_arrays\n        self.transform_input = transform_input\n\n    def __len__(self):\n        return len(self.input_files)\n\n    def __getitem__(self, idx):\n        # 1) Load raw data\n        data = load_seismic_file(self.input_files[idx])  # shape (S, T, R)\n        # 2) Optional transform\n        if self.transform_input:\n            data = self.transform_input(data)\n        # 3) To tensor + add channel dim: (S,T,R) → (1,S,T,R)\n        data_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0)\n        # 4) Permute to (C, D, H, W): (1,S,T,R) → (1,R,S,T)\n        data_tensor = data_tensor.permute(0, 3, 1, 2)\n        # 5) Target tensor (H, W)\n        target_tensor = torch.tensor(self.target_arrays[idx], dtype=torch.float32)\n        return data_tensor, target_tensor\n\n# ---------------------------\n# 3. Sample Data Setup & DataLoader\n# ---------------------------\n# Prepare only two example files\nexample_file   = \"/kaggle/input/waveform-inversion/train_samples/CurveFault_A/seis2_1_0.npy\"\nexample_target = np.zeros((100, 70), dtype=np.float32)\nfile_list      = [example_file, example_file]\ntarget_list    = [example_target, example_target]\n\n# Create dataset and data loader\ndataset = WaveformDataset(file_list, target_list)\nloader  = DataLoader(dataset, batch_size=2, shuffle=True)\n\n# ---------------------------\n# 4. Fetch one batch and visualize\n# ---------------------------\ninputs, targets = next(iter(loader))\n# inputs: (B, 1, R, S, T), targets: (B, H, W)\n\n# (a) Plot waveform grid (receivers × [sources × time])\ngrid = inputs[0, 0].cpu().numpy().reshape(inputs.shape[2], -1)\nplt.figure(figsize=(6, 3))\nplt.imshow(grid, aspect='auto')\nplt.title(\"Input Waveform Grid\")\nplt.xlabel(\"Source × Time\")\nplt.ylabel(\"Receiver\")\nplt.colorbar(label=\"Amplitude\")\nplt.show()\n\n# (b) Plot ground-truth velocity map\nplt.figure(figsize=(4, 4))\nplt.imshow(targets[0].cpu(), aspect='auto')\nplt.title(\"Ground-Truth Velocity Map\")\nplt.xlabel(\"X index\")\nplt.ylabel(\"Y index\")\nplt.colorbar(label=\"Velocity\")\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-23T15:04:40.679646Z","iopub.execute_input":"2025-06-23T15:04:40.679940Z","iopub.status.idle":"2025-06-23T15:04:47.103968Z","shell.execute_reply.started":"2025-06-23T15:04:40.679918Z","shell.execute_reply":"2025-06-23T15:04:47.103178Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. Plot Training Loss Curve","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# Example: Populate 'losses' with dummy values to demonstrate plotting\n# In your actual code, ensure you append avg_loss to 'losses' inside your training loop:\n#     losses.append(avg_loss)\n\nlosses = [0.8, 0.6, 0.45, 0.3, 0.2]  # replace or extend with your real avg_loss values\n\n# Plot Training Loss Curve\nplt.figure()\nplt.plot(range(1, len(losses) + 1), losses, marker='o')\nplt.title(\"Training Loss Curve\")\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Average L1 Loss\")\nplt.grid(True)\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-23T15:04:47.104810Z","iopub.execute_input":"2025-06-23T15:04:47.105241Z","iopub.status.idle":"2025-06-23T15:04:47.296757Z","shell.execute_reply.started":"2025-06-23T15:04:47.105211Z","shell.execute_reply":"2025-06-23T15:04:47.295930Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. Visualize Inference Predictions","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport matplotlib.pyplot as plt\n\n# 1. Helper function: load .npy seismic file\ndef load_seismic_file(fp: str) -> np.ndarray:\n    \"\"\"\n    Load a seismic waveform from a .npy file and return a 3D NumPy array (S, T, R).\n    Strips any extra leading dimensions.\n    \"\"\"\n    data = np.load(fp).astype(np.float32)\n    while data.ndim > 3:\n        data = data[0]\n    return data\n\n# 2. Dataset definition for PyTorch\nclass WaveformDataset(Dataset):\n    def __init__(self, files, targets, transform=None):\n        \"\"\"\n        files:   list of paths to .npy seismic files\n        targets: list of 2D NumPy arrays (H×W) ground-truth velocity maps\n        transform: optional preprocessing function\n        \"\"\"\n        self.files = files\n        self.targets = targets\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.files)\n\n    def __getitem__(self, idx):\n        # Load and optionally transform the waveform\n        d = load_seismic_file(self.files[idx])\n        if self.transform:\n            d = self.transform(d)\n        # Convert to tensor, add channel dim, then reorder to (C, D, H, W)\n        t = torch.tensor(d, dtype=torch.float32).unsqueeze(0).permute(0, 3, 1, 2)\n        # Load target velocity map\n        y = torch.tensor(self.targets[idx], dtype=torch.float32)\n        return t, y\n\n# 3. Model definition: minimal 3D U-Net\nclass UNet3D(nn.Module):\n    def __init__(self, in_ch=1, out_ch=1):\n        super().__init__()\n        # Encoder: two 3D conv layers + ReLU, then adaptive pooling\n        self.enc = nn.Sequential(\n            nn.Conv3d(in_ch, 8, 3, padding=1), nn.ReLU(True),\n            nn.Conv3d(8, 16, 3, padding=1),    nn.ReLU(True),\n            nn.AdaptiveAvgPool3d((1, 100, 70))\n        )\n        # Decoder: two 3D transpose conv layers + ReLU\n        self.dec = nn.Sequential(\n            nn.ConvTranspose3d(16, 8, 3, padding=1), nn.ReLU(True),\n            nn.ConvTranspose3d(8, out_ch, 3, padding=1)\n        )\n\n    def forward(self, x):\n        # Encode to bottleneck\n        x = self.enc(x)   # → (batch,16,1,100,70)\n        # Decode back to spatial dimensions\n        x = self.dec(x)   # → (batch,1,1,100,70)\n        x = x.squeeze(2)  # remove depth dim → (batch,1,100,70)\n        return x.squeeze(1)  # remove channel dim → (batch,100,70)\n\n# 4. Inference function\ndef predict(model, fps, device):\n    \"\"\"\n    Run inference on a list of .npy file paths and return list of 2D NumPy arrays.\n    \"\"\"\n    model.eval()\n    out = []\n    with torch.no_grad():\n        for fp in fps:\n            # Load and prepare waveform tensor\n            d = load_seismic_file(fp)\n            t = torch.tensor(d, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)\n            # Model forward pass and convert to NumPy\n            o = model(t).squeeze().cpu().numpy()\n            out.append(o)\n    return out\n\n# 5. Main workflow\ndef main():\n    # Gather training files and create dummy target arrays (replace with real data)\n    files   = glob.glob(\n        \"/kaggle/input/waveform-inversion/train_samples/CurveFault_A/seis2_1_0.npy\",\n        recursive=True\n    )[:10]\n    targets = [np.zeros((100, 70), np.float32) for _ in files]\n    ds      = WaveformDataset(files, targets)\n    loader  = DataLoader(ds, batch_size=2, shuffle=True)\n\n    # Use CPU-only\n    device = torch.device(\"cpu\")\n    model  = UNet3D().to(device)\n    opt    = torch.optim.Adam(model.parameters(), lr=1e-3)\n    lossfn = nn.L1Loss()\n    ckpt   = \"best_model.pth\"\n\n    # If checkpoint exists, load it; otherwise train one epoch and save\n    if os.path.exists(ckpt):\n        print(f\"Loading checkpoint from {ckpt}\")\n        model.load_state_dict(torch.load(ckpt, map_location=device))\n    else:\n        print(\"No checkpoint found → training 1 epoch\")\n        model.train()\n        for inp, tgt in loader:\n            inp, tgt = inp.to(device), tgt.to(device)\n            opt.zero_grad()\n            out = model(inp)\n            loss = lossfn(out, tgt)\n            loss.backward()\n            opt.step()\n        torch.save(model.state_dict(), ckpt)\n        print(f\"Saved checkpoint to {ckpt}\")\n\n    # Perform inference on the first sample and visualize\n    sample_fp = files[0]\n    preds = predict(model, [sample_fp], device)\n\n    plt.figure()\n    plt.imshow(preds[0], aspect=\"auto\")\n    plt.title(\"Predicted Velocity Map\")\n    plt.xlabel(\"X index\")\n    plt.ylabel(\"Y index\")\n    plt.colorbar(label=\"Velocity\")\n    plt.show()\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-23T15:04:47.297614Z","iopub.execute_input":"2025-06-23T15:04:47.297914Z","iopub.status.idle":"2025-06-23T15:04:52.768280Z","shell.execute_reply.started":"2025-06-23T15:04:47.297884Z","shell.execute_reply":"2025-06-23T15:04:52.767448Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Program","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nfrom tqdm import tqdm\nimport pandas as pd\nimport glob\nimport os\n\n# ---------------------------\n# 1. Data Loading Function\n# ---------------------------\ndef load_seismic_file(file_path):\n    \"\"\"\n    Loads a seismic waveform from a .npy file and returns a NumPy array.\n    Expected input shape: (S, T, R) or with extra batch dims (e.g., (1, S, T, R)).\n    \"\"\"\n    data = np.load(file_path).astype(np.float32)\n    # Remove any extra leading dimensions\n    while data.ndim > 3:\n        data = data[0]\n    return data  # shape: (S, T, R)\n\n# ---------------------------\n# 2. PyTorch Dataset\n# ---------------------------\nclass WaveformDataset(Dataset):\n    def __init__(self, input_files, target_arrays, transform_input=None):\n        \"\"\"\n        input_files: list of file paths to seismic .npy files\n        target_arrays: list of ground truth arrays with shape (H, W)\n        transform_input: optional preprocessing function\n        \"\"\"\n        self.input_files = input_files\n        self.target_arrays = target_arrays\n        self.transform_input = transform_input\n\n    def __len__(self):\n        return len(self.input_files)\n\n    def __getitem__(self, idx):\n        # Load and preprocess seismic data\n        data = load_seismic_file(self.input_files[idx])  # (S, T, R)\n        if self.transform_input:\n            data = self.transform_input(data)\n        # Convert to PyTorch tensor and add channel dimension\n        # (S, T, R) -> (1, S, T, R)\n        data_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0)\n        # Permute to (C, D, H, W): (1, S, T, R) -> (1, R, S, T)\n        data_tensor = data_tensor.permute(0, 3, 1, 2)\n        # Ground truth velocity map tensor\n        target_tensor = torch.tensor(self.target_arrays[idx], dtype=torch.float32)  # (H, W)\n        return data_tensor, target_tensor\n\n# ---------------------------\n# 3. Simplified 3D U-Net Model\n# ---------------------------\nclass UNet3D(nn.Module):\n    def __init__(self, in_channels=1, out_channels=1):\n        super(UNet3D, self).__init__()\n        self.encoder = nn.Sequential(\n            nn.Conv3d(in_channels,  8, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv3d(8,           16, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            # Pool to (1, 100, 70) for consistent spatial dimensions\n            nn.AdaptiveAvgPool3d((1, 100, 70))\n        )\n        self.decoder = nn.Sequential(\n            nn.ConvTranspose3d(16,  8, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose3d(8, out_channels, kernel_size=3, padding=1),\n        )\n\n    def forward(self, x):\n        # x shape: (B, 1, D, H, W)\n        x = self.encoder(x)    # -> (B, 16, 1, 100, 70)\n        x = self.decoder(x)    # -> (B, 1, 1, 100, 70)\n        x = x.squeeze(2)       # -> (B, 1, 100, 70)\n        return x.squeeze(1)    # -> (B, 100, 70)\n\n# ---------------------------\n# 4. Sample Data Setup\n# ---------------------------\nexample_file   = \"/kaggle/input/waveform-inversion/train_samples/CurveFault_A/seis2_1_0.npy\"\nexample_target = np.zeros((100, 70), dtype=np.float32)\nfile_list      = [example_file] * 10\ntarget_list    = [example_target] * 10\n\n# Create Dataset and DataLoader\ndataset = WaveformDataset(file_list, target_list)\nloader  = DataLoader(dataset, batch_size=2, shuffle=True)\n\n# ---------------------------\n# 5. Model Training Setup\n# ---------------------------\ndevice    = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel     = UNet3D(in_channels=1, out_channels=1).to(device)\ncriterion = nn.L1Loss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,\n                                                 mode='min',\n                                                 factor=0.5,\n                                                 patience=2)\n\n# ---------------------------\n# 6. Training Loop\n# ---------------------------\nnum_epochs    = 3\nbest_val_loss = float(\"inf\")\nfor epoch in range(num_epochs):\n    model.train()\n    running_loss = 0.0\n    progress_bar = tqdm(loader, desc=f\"Epoch {epoch+1}/{num_epochs} Training\")\n    for inputs, targets in progress_bar:\n        inputs  = inputs.to(device)\n        targets = targets.to(device)\n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss    = criterion(outputs, targets)\n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item()\n        progress_bar.set_postfix({\"loss\": loss.item()})\n    avg_loss = running_loss / len(loader)\n    print(f\"Epoch {epoch+1}: avg loss = {avg_loss:.4f}\")\n    scheduler.step(avg_loss)\n    if avg_loss < best_val_loss:\n        best_val_loss = avg_loss\n        torch.save(model.state_dict(), \"best_model.pth\")\n        print(\"Model improved and saved.\")\n\n# ---------------------------\n# 7. Inference Function\n# ---------------------------\ndef predict(model, file_paths):\n    model.eval()\n    predictions = []\n    with torch.no_grad():\n        for fp in file_paths:\n            data = load_seismic_file(fp)\n            while data.ndim > 3:\n                data = data[0]\n            tensor = torch.tensor(data, dtype=torch.float32)\n            tensor = tensor.unsqueeze(0).unsqueeze(0).to(device)\n            output = model(tensor)\n            predictions.append(output.squeeze().cpu().numpy())\n    return predictions\n\n# ---------------------------\n# 8. Submission File Generation\n# ---------------------------\ndef create_submission(oids, predictions):\n    \"\"\"\n    Builds a submission.csv matching the sample_submission format.\n    oids: list[str] identifiers without extension\n    predictions: list[np.ndarray] each of shape (H, W)\n    \"\"\"\n    # Load template to get correct column names\n    sample_path = '/kaggle/input/waveform-inversion/sample_submission.csv'\n    sample_df = pd.read_csv(sample_path)\n    id_col = sample_df.columns[0]\n\n    if not predictions:\n        raise ValueError(\"No predictions provided.\")\n    width = predictions[0].shape[1]\n    odd_indices = list(range(0, width, 2))\n\n    rows = []\n    for oid, pred in zip(oids, predictions):\n        if pred.shape[1] != width:\n            raise ValueError(f\"Width mismatch for {oid}.\")\n        for y in range(pred.shape[0]):\n            row_id = f\"{oid}_y_{y}\"\n            row = [row_id] + [float(pred[y, x]) for x in odd_indices]\n            rows.append(row)\n\n    columns = [id_col] + [f\"x_{i}\" for i in odd_indices]\n    df = pd.DataFrame(rows, columns=columns)\n    df.to_csv('/kaggle/working/submission.csv', index=False)\n    print(\"Submission saved to /kaggle/working/submission.csv with correct ID column.\")\n\n# ---------------------------\n# 9. Generate Submission using sample_submission\n# ---------------------------\nsample_path = '/kaggle/input/waveform-inversion/sample_submission.csv'\nsample_df   = pd.read_csv(sample_path)\n# Fill with zeros or replace with model predictions\nsample_df.iloc[:, 1:] = 0.0\nsample_df.to_csv('/kaggle/working/submission.csv', index=False)\nprint(\"Sample submission written to /kaggle/working/submission.csv.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-23T15:04:52.769262Z","iopub.execute_input":"2025-06-23T15:04:52.769635Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Explanation of key steps\n\n1. **Forward simulation**  \n   The function `solve_forward_wave_equation` solves the physical wave equation to generate synthetic data \\(d_{\\mathrm{syn}}\\).\n\n2. **Residual and misfit**  \n   $$\n   r = d_{\\mathrm{syn}} - d_{\\mathrm{obs}}\n   $$  \n   The scalar objective (misfit) is  \n   $$\n   J(m) = \\frac{1}{2}\\,\\|r\\|^2.\n   $$\n\n3. **Adjoint simulation**  \n   The function `solve_adjoint_wave_equation` propagates the residual \\(r\\) backwards in time to produce the adjoint wavefield \\(p_{\\mathrm{adj}}\\).\n\n4. **Gradient computation**  \n   $$\n   \\nabla_{m} J \\;=\\; -\\int \\frac{\\partial^2 p}{\\partial t^2}\\;p_{\\mathrm{adj}}\\;\\mathrm{d}t.\n   $$\n\n5. **Model update (steepest-descent)**  \n   $$\n   m \\;\\leftarrow\\; m \\;-\\;\\alpha\\,\\nabla_{m}J,\n   $$  \n   where \\(\\alpha\\) is the step length.\n\n---\n\nThis end-to-end loop is the core “logic” once you replace the simple L1-loss in your PyTorch code with the physics-based misfit and adjoint-gradient machinery.  \n","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport glob\nimport os\nimport matplotlib.pyplot as plt\n\n# ---------------------------\n# 1. Data Loading Function\n# ---------------------------\ndef load_seismic_file(file_path):\n    \"\"\"\n    Load a seismic waveform from a .npy file and return a NumPy array (S, T, R).\n    Strips any extra leading dimensions.\n    \"\"\"\n    data = np.load(file_path).astype(np.float32)\n    while data.ndim > 3:\n        data = data[0]\n    return data  # (S, T, R)\n\n# ---------------------------\n# 2. Dataset Definition\n# ---------------------------\nclass WaveformDataset(Dataset):\n    def __init__(self, input_files, transform_input=None):\n        self.input_files = input_files\n        self.transform_input = transform_input\n\n    def __len__(self):\n        return len(self.input_files)\n\n    def __getitem__(self, idx):\n        data = load_seismic_file(self.input_files[idx])  # (S, T, R)\n        if self.transform_input:\n            data = self.transform_input(data)\n        tensor = torch.tensor(data, dtype=torch.float32)  # (S, T, R)\n        tensor = tensor.permute(2, 0, 1)                  # (R, S, T)\n        return tensor\n\n# ---------------------------\n# 3. UNet3D Model\n# ---------------------------\nclass UNet3D(nn.Module):\n    def __init__(self, in_channels=1, out_channels=1):\n        super().__init__()\n        self.encoder = nn.Sequential(\n            nn.Conv3d(in_channels,  8, 3, padding=1), nn.ReLU(),\n            nn.Conv3d(8,           16, 3, padding=1), nn.ReLU(),\n            nn.AdaptiveAvgPool3d((1,100,70))\n        )\n        self.decoder = nn.Sequential(\n            nn.ConvTranspose3d(16,  8, 3, padding=1), nn.ReLU(),\n            nn.ConvTranspose3d(8, out_channels, 3, padding=1),\n        )\n\n    def forward(self, x):\n        # x: (B, 1, R, S, T)\n        h = self.encoder(x)    # -> (B,16,1,100,70)\n        h = self.decoder(h)    # -> (B,1,1,100,70)\n        h = h.squeeze(2)       # -> (B,1,100,70)\n        return h.squeeze(1)    # -> (B,100,70)\n\n# ---------------------------\n# 4. Forward Solver Stub\n# ---------------------------\ndef forward_wave_solver(vel_map, obs):\n    \"\"\"\n    Stub that scales observed data by mean velocity\n    vel_map: (B, S, T)\n    obs:     (B,1,R,S,T)\n    \"\"\"\n    B = vel_map.size(0)\n    factor = vel_map.mean(dim=(1,2)).view(B,1,1,1,1)\n    return obs * factor\n\n# ---------------------------\n# 5. Prediction Helper\n# ---------------------------\ndef predict(model, file_paths, device):\n    model.eval()\n    preds = []\n    with torch.no_grad():\n        for fp in file_paths:\n            data = load_seismic_file(fp)\n            x = torch.tensor(data, dtype=torch.float32)             # (S, T, R)\n            x = x.permute(2,0,1).unsqueeze(0).unsqueeze(0).to(device)  # (1,1,R,S,T)\n            vel = model(x)                                          # (1,100,70)\n            preds.append(vel.squeeze(0).cpu().numpy())\n    return preds\n\n# ---------------------------\n# 6. Submission Helper\n# ---------------------------\ndef create_submission(oids, predictions, sample_csv, out_csv):\n    # Read sample header\n    with open(sample_csv) as f:\n        header = f.readline().strip().split(',')\n    x_cols = header[1:]\n    with open(out_csv, 'w') as f:\n        f.write(','.join(header) + '\\n')\n        for oid, pred in zip(oids, predictions):\n            for y in range(pred.shape[0]):\n                row_id = f\"{oid}_y_{y}\"\n                vals = [f\"{pred[y, int(col.split('_')[1])]:.6f}\" for col in x_cols]\n                f.write(','.join([row_id]+vals) + '\\n')\n    print(f\"Submission saved to {out_csv}\")\n\n# ---------------------------\n# 7. Main Workflow\n# ---------------------------\nif __name__ == \"__main__\":\n    # 7.1 Gather and dedup file list\n    file_list = glob.glob(\"/kaggle/input/waveform-inversion/train_samples/**/*.npy\", recursive=True)\n    unique_files = list(dict.fromkeys(file_list))\n\n    # 7.2 Dataset & DataLoader\n    dataset = WaveformDataset(unique_files)\n    loader  = DataLoader(dataset, batch_size=1, shuffle=True)\n\n    # 7.3 Model & optimizer\n    device    = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model     = UNet3D().to(device)\n    optimizer = optim.Adam(model.parameters(), lr=1e-3)\n\n    # 7.4 Training with FWI Loss\n    losses = []\n    num_epochs = 3\n    for epoch in range(1, num_epochs+1):\n        model.train()\n        total_loss = 0.0\n        for waveform in loader:\n            obs = waveform.unsqueeze(1).to(device)  # (B,1,R,S,T)\n            optimizer.zero_grad()\n            vel_pred = model(obs)                   # (B,100,70)\n            d_syn    = forward_wave_solver(vel_pred, obs)\n            loss     = 0.5 * ((d_syn - obs)**2).mean()\n            loss.backward()\n            optimizer.step()\n            total_loss += loss.item()\n        avg_loss = total_loss / len(loader)\n        losses.append(avg_loss)\n        print(f\"Epoch {epoch}: avg FWI loss = {avg_loss:.4e}\")\n\n    # 7.5 Plot Training Loss\n    plt.figure()\n    plt.plot(range(1, num_epochs+1), losses, marker='o')\n    plt.title(\"FWI Training Loss\")\n    plt.xlabel(\"Epoch\")\n    plt.ylabel(\"Average Loss\")\n    plt.grid(True)\n    plt.show()\n\n    # 7.6 Plot Observed vs Synthetic for first receiver\n    model.eval()\n    with torch.no_grad():\n        wf = loader.dataset[0]\n        obs = torch.tensor(wf, dtype=torch.float32).permute(1,2,0).unsqueeze(0).unsqueeze(0).to(device)\n        vel_pred = model(obs)\n        d_syn    = forward_wave_solver(vel_pred, obs)\n        obs_np   = obs.squeeze().cpu().numpy()   # (R,S,T)\n        syn_np   = d_syn.squeeze().cpu().numpy()\n\n    plt.figure(figsize=(6,4))\n    plt.plot(obs_np[0,0,:], label=\"Observed\")\n    plt.plot(syn_np[0,0,:], '--', label=\"Synthetic\")\n    plt.title(\"Observed vs Synthetic Trace (Receiver 0)\")\n    plt.xlabel(\"Time sample\")\n    plt.ylabel(\"Amplitude\")\n    plt.legend()\n    plt.show()\n\n    # 7.7 Inference & Submission\n    preds = predict(model, unique_files, device)\n    oids  = [os.path.splitext(os.path.basename(fp))[0] for fp in unique_files]\n    sample_csv = \"/kaggle/input/waveform-inversion/sample_submission.csv\"\n    out_csv    = \"/kaggle/working/submission.csv\"\n    create_submission(oids, preds, sample_csv, out_csv)\n\n    # 7.8 Plot Predicted Velocity Map from submission.csv\n    lines = open(out_csv).read().splitlines()\n    header = lines[0].split(',')\n    ids    = [line.split(',')[0] for line in lines[1:]]\n    vals   = [list(map(float, line.split(',')[1:])) for line in lines[1:]]\n    x_idxs = [int(h.split('_')[1]) for h in header[1:]]\n\n    first_oid = ids[0].split('_y_')[0]\n    rows = [(ids[i], vals[i]) for i in range(len(ids)) if ids[i].startswith(first_oid + '_y_')]\n    y_idxs = [int(rid.split('_y_')[1]) for rid, _ in rows]\n    vels   = [v for _, v in rows]\n\n    order = np.argsort(y_idxs)\n    grid  = np.array(vels)[order, :]\n\n    plt.figure(figsize=(6,5))\n    plt.imshow(grid, aspect='auto', origin='lower',\n               extent=[min(x_idxs), max(x_idxs), min(y_idxs), max(y_idxs)])\n    plt.title(f\"Predicted Velocity Map for {first_oid}\")\n    plt.xlabel(\"X index\")\n    plt.ylabel(\"Y index\")\n    plt.colorbar(label=\"Velocity\")\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T16:48:48.044929Z","iopub.execute_input":"2025-06-24T16:48:48.045205Z","iopub.status.idle":"2025-06-24T16:52:37.285627Z","shell.execute_reply.started":"2025-06-24T16:48:48.045177Z","shell.execute_reply":"2025-06-24T16:52:37.284454Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\n# --- 1) Plot Training Loss Curve ---\n# Suppose you modified your loop to collect avg_loss each epoch:\nlosses = []  # before training\nfor epoch in range(1, num_epochs+1):\n    model.train()\n    total_loss = 0.0\n    for obs in loader:\n        # ... compute loss ...\n        total_loss += loss.item()\n    avg_loss = total_loss / len(loader)\n    losses.append(avg_loss)\n    print(f\"Epoch {epoch}: avg FWI loss = {avg_loss:.4e}\")\n\n# After training:\nplt.figure()\nplt.plot(range(1, len(losses)+1), losses, marker='o')\nplt.title(\"FWI Training Loss Curve\")\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Average FWI Loss\")\nplt.grid(True)\nplt.show()\n\n\n# --- 2) Plot Observed vs. Synthetic for One Trace ---\n# Pick the first sample & first receiver\nmodel.eval()\nwith torch.no_grad():\n    # Load a single waveform\n    fp = file_list[0]\n    obs = torch.tensor(load_seismic_file(fp), dtype=torch.float32)  # (S, T, R)\n    obs = obs.permute(2,0,1).unsqueeze(0).unsqueeze(0).to(device)    # (1,1,R,S,T)\n\n    # Predict velocity & simulate synthetic\n    vel_pred = model(obs)                                          # (1,100,70)\n    d_syn    = forward_wave_solver(vel_pred, obs)                  # (1,1,R,100,70)\n\n    # Move to CPU & NumPy\n    obs_np   = obs.squeeze().cpu().numpy()     # (R, S, T)\n    syn_np   = d_syn.squeeze().cpu().numpy()   # (R, S, T)\n\n# Plot the time‐series of source 0, receiver 0 (for example)\nplt.figure(figsize=(8,4))\nplt.plot(obs_np[0,0,:], label=\"Observed\")\nplt.plot(syn_np[0,0,:], label=\"Synthetic\", linestyle=\"--\")\nplt.title(\"Observed vs Synthetic Trace (Receiver 0)\")\nplt.xlabel(\"Time sample\")\nplt.ylabel(\"Amplitude\")\nplt.legend()\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T16:53:00.695863Z","iopub.execute_input":"2025-06-24T16:53:00.696209Z","iopub.status.idle":"2025-06-24T16:53:32.230416Z","shell.execute_reply.started":"2025-06-24T16:53:00.696170Z","shell.execute_reply":"2025-06-24T16:53:32.229410Z"}},"outputs":[],"execution_count":null}]}