{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ===============================\n# Yale/UNC-CH - FWI Competition Starter (PyTorch)\n# Physics-Guided ML for Velocity Map Prediction\n# ===============================\n\nimport torch\nimport torch.nn as nn\nimport pandas as pd\nimport numpy as np\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\n# ===============================\n# Custom Dataset for Seismic Data\n# ===============================\nclass SeismicDataset(Dataset):\n    def __init__(self, waveforms, labels=None):\n        self.waveforms = waveforms  # shape: (N, 1, H, W)\n        self.labels = labels        # shape: (N, H, W) or None\n\n    def __len__(self):\n        return len(self.waveforms)\n\n    def __getitem__(self, idx):\n        x = self.waveforms[idx]\n        if self.labels is not None:\n            y = self.labels[idx]\n            return x, y\n        return x\n\n# ===============================\n# A Simple U-Net Style CNN\n# ===============================\nclass SimpleUNet(nn.Module):\n    def __init__(self):\n        super(SimpleUNet, self).__init__()\n        self.enc1 = nn.Sequential(nn.Conv2d(1, 32, 3, padding=1), nn.ReLU())\n        self.enc2 = nn.Sequential(nn.Conv2d(32, 64, 3, padding=1), nn.ReLU())\n        self.dec1 = nn.Sequential(nn.ConvTranspose2d(64, 32, 3, padding=1), nn.ReLU())\n        self.out = nn.Conv2d(32, 1, 3, padding=1)  # Final output layer\n\n    def forward(self, x):\n        x1 = self.enc1(x)\n        x2 = self.enc2(x1)\n        x3 = self.dec1(x2)\n        out = self.out(x3)\n        return out.squeeze(1)  # Remove channel dimension\n\n# ===============================\n# Training Function\n# ===============================\ndef train_model(model, train_loader, val_loader, num_epochs=10, lr=1e-3):\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model.to(device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    criterion = nn.L1Loss()  # MAE loss\n\n    for epoch in range(num_epochs):\n        model.train()\n        total_loss = 0\n        for x, y in tqdm(train_loader, desc=f\"Epoch {epoch+1}\"):\n            x, y = x.to(device), y.to(device)\n            optimizer.zero_grad()\n            y_pred = model(x)\n            loss = criterion(y_pred, y)\n            loss.backward()\n            optimizer.step()\n            total_loss += loss.item()\n        print(f\"Epoch {epoch+1}, Train Loss: {total_loss / len(train_loader):.4f}\")\n\n        # Validation loop\n        model.eval()\n        with torch.no_grad():\n            val_loss = 0\n            for x, y in val_loader:\n                x, y = x.to(device), y.to(device)\n                y_pred = model(x)\n                val_loss += criterion(y_pred, y).item()\n            print(f\"Val Loss: {val_loss / len(val_loader):.4f}\")\n\n# ===============================\n# Create Submission CSV\n# ===============================\ndef make_submission(preds, oids, filename=\"submission.csv\"):\n    # preds: np.array (N, H, W), oids: list of strings\n    rows = []\n    for i, pred in enumerate(preds):\n        for y in range(pred.shape[0]):\n            values = pred[y, 1::2]  # Only odd x columns: x_1, x_3, ..., x_69\n            row_id = f\"{oids[i]}_y_{y}\"\n            row = [row_id] + values.tolist()\n            rows.append(row)\n    columns = [\"oid_ypos\"] + [f\"x_{i}\" for i in range(1, 70, 2)]\n    df = pd.DataFrame(rows, columns=columns)\n    df.to_csv(filename, index=False)\n    print(f\"Saved submission to {filename}\")\n\n# ===============================\n# Main Driver Code (Simulated Example)\n# ===============================\nif __name__ == \"__main__\":\n    # Simulated data (replace with real .npy/.csv data)\n    N, H, W = 10, 100, 70  # 10 samples, 100 rows, 70 columns\n    waveforms = np.random.randn(N, 1, H, W).astype(np.float32)\n    velocities = np.random.uniform(2500, 3500, size=(N, H, W)).astype(np.float32)\n\n    # Split data into training and validation\n    train_x, val_x, train_y, val_y = train_test_split(waveforms, velocities, test_size=0.2)\n\n    # Create Dataloaders\n    train_dataset = SeismicDataset(train_x, train_y)\n    val_dataset = SeismicDataset(val_x, val_y)\n    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=2)\n\n    # Initialize model and train\n    model = SimpleUNet()\n    train_model(model, train_loader, val_loader, num_epochs=5)\n\n    # Predict on new test data (simulated)\n    test_waveforms = np.random.randn(2, 1, H, W).astype(np.float32)\n    test_dataset = SeismicDataset(test_waveforms)\n    test_loader = DataLoader(test_dataset, batch_size=1)\n\n    model.eval()\n    preds = []\n    with torch.no_grad():\n        for x in test_loader:\n            pred = model(x)\n            preds.append(pred.cpu().numpy())\n\n    preds = np.concatenate(preds, axis=0)\n\n    # Dummy OIDs for submission example\n    make_submission(preds, oids=[\"00001\", \"00002\"])","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}