{"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"},{"sourceId":11367935,"sourceType":"datasetVersion","datasetId":7116013},{"sourceId":11368083,"sourceType":"datasetVersion","datasetId":7116134},{"sourceId":11368169,"sourceType":"datasetVersion","datasetId":7116196},{"sourceId":11368268,"sourceType":"datasetVersion","datasetId":7116272},{"sourceId":11368499,"sourceType":"datasetVersion","datasetId":7116445},{"sourceId":11368545,"sourceType":"datasetVersion","datasetId":7116479},{"sourceId":11368547,"sourceType":"datasetVersion","datasetId":7116481},{"sourceId":11376433,"sourceType":"datasetVersion","datasetId":7122462},{"sourceId":11376448,"sourceType":"datasetVersion","datasetId":7122476},{"sourceId":11376464,"sourceType":"datasetVersion","datasetId":7122489},{"sourceId":11376742,"sourceType":"datasetVersion","datasetId":7122712},{"sourceId":11376868,"sourceType":"datasetVersion","datasetId":7122812},{"sourceId":11376871,"sourceType":"datasetVersion","datasetId":7122814},{"sourceId":11376872,"sourceType":"datasetVersion","datasetId":7122815},{"sourceId":11376935,"sourceType":"datasetVersion","datasetId":7122866},{"sourceId":11377083,"sourceType":"datasetVersion","datasetId":7122981},{"sourceId":11377231,"sourceType":"datasetVersion","datasetId":7123090},{"sourceId":11377291,"sourceType":"datasetVersion","datasetId":7123138},{"sourceId":11377325,"sourceType":"datasetVersion","datasetId":7123163},{"sourceId":11377334,"sourceType":"datasetVersion","datasetId":7123172},{"sourceId":11377594,"sourceType":"datasetVersion","datasetId":7123380},{"sourceId":11377614,"sourceType":"datasetVersion","datasetId":7123394},{"sourceId":11377741,"sourceType":"datasetVersion","datasetId":7123490},{"sourceId":11377752,"sourceType":"datasetVersion","datasetId":7123499},{"sourceId":11377756,"sourceType":"datasetVersion","datasetId":7123503},{"sourceId":11377935,"sourceType":"datasetVersion","datasetId":7123649},{"sourceId":11377970,"sourceType":"datasetVersion","datasetId":7123675},{"sourceId":11378141,"sourceType":"datasetVersion","datasetId":7123813},{"sourceId":11378162,"sourceType":"datasetVersion","datasetId":7123831},{"sourceId":11378178,"sourceType":"datasetVersion","datasetId":7123842}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Initialize","metadata":{}},{"cell_type":"code","source":"!pip install faiss-cpu\nimport numpy as np\nimport pandas as pd\nimport os\nfrom pathlib import Path\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch import nn\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import mean_squared_error\nfrom sklearn.model_selection import train_test_split\nimport seaborn as sns\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nimport matplotlib.colors as mcolors\nimport gc\nimport torchvision.models as models\nimport torchvision.transforms as transforms\nfrom torchvision.models import resnet18, ResNet18_Weights\nimport faiss\nfrom tqdm import tqdm\n\n# Verify GPU availability\nprint(f\"PyTorch Version: {torch.__version__}\")\nprint(f\"CUDA Available: {torch.cuda.is_available()}\")\n\nif torch.cuda.is_available(): \n    device = torch.device(\"cuda:0\") \n    print(f\"GPU Name: {torch.cuda.get_device_name(0)}\") \n    print(f\"CUDA Version: {torch.version.cuda}\")\n    print(f\"Number of GPUs: {torch.cuda.device_count()}\")\nelse: \n    device = torch.device(\"cpu\") \n    print(f\"Using device: {device}\")\n    \n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T23:44:55.577421Z","iopub.execute_input":"2025-05-07T23:44:55.578322Z","iopub.status.idle":"2025-05-07T23:45:04.014624Z","shell.execute_reply.started":"2025-05-07T23:44:55.578296Z","shell.execute_reply":"2025-05-07T23:45:04.013725Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Find all sample directories\n### Use Kaggle Original and Open FWI database","metadata":{}},{"cell_type":"code","source":"# New updated \n# Collect Original Kaggle input files\ninput_dir = Path('/kaggle/input/waveform-inversion/train_samples')\n\ndata_files = []\nmodel_files = []\n\n# Add files from Vel and Style families\n\ndata_files = sorted(input_dir.rglob('data*.npy'))\nmodel_files = sorted(input_dir.rglob('model*.npy'))\n\n# Add Fault family\nfault_data_files = sorted(input_dir.rglob('seis*.npy'))\nfault_model_files = sorted(input_dir.rglob('vel*.npy'))\ndata_files += fault_data_files\nmodel_files += fault_model_files\n\n#LOAD the OPEN FWI training files\nroot_dir = Path(\"/kaggle/input\")\n\nall_data_files = []\nall_model_files = []\nseis_files = []\nvel_files = []\nseis2_files = []\nvel2_files = []\n\n# Loop over waveform-inversion-1 to waveform-inversion-30\nfor i in range(1,31):\n    subroot = root_dir / f\"waveform-inversion-{i}\"\n    seis_files = sorted(subroot.rglob(\"seis*.npy\"))\n    vel_files = sorted(subroot.rglob(\"vel*.npy\"))\n    all_data_files.extend(seis_files)\n    all_model_files.extend(vel_files)\n\n    seis2_files = sorted(subroot.rglob(\"data*.npy\"))\n    vel2_files = sorted(subroot.rglob(\"model*.npy\"))\n    all_data_files.extend(seis2_files)\n    all_model_files.extend(vel2_files)\n    seis_files = []\n    vel_files = [] \n    seis2_files = []\n    vel2_files = []\n  \n## File check\n#for i in range(1, 20):\n#    print(f\"index: {i}\" )\n#    print(\"data\", all_data_files[i])\n#    print(\"model\", all_model_files[i])   \n\n# Randomly select some Open FWI consistent pairs\n\nprint(\"Number of Open FWI seismic training files found:\", len(all_data_files))\n\nimport random\n# Set a fixed seed for reproducibility\nrandom.seed(42)\ndata_FWIfiles = []\nmodel_FWIfiles = []\n\n#Max number tested is 80, larger may crash memory\nselect_FWI  = 80\nprint(\"Number of Open FWI to add to original set:\", select_FWI)\nindices = random.sample(range(len(all_data_files)), select_FWI)\ndata_FWIfiles = [all_data_files[i] for i in indices]\nmodel_FWIfiles = [all_model_files[i] for i in indices]\n\n# Combine Original Kaggle and select Open FWI files\ndata_files += data_FWIfiles\nmodel_files += model_FWIfiles\n\nprint(\"Number of combined origianl/Open FWI training files found:\", len(data_files))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load and preprocess sample files ","metadata":{}},{"cell_type":"code","source":"\nX_list = []\ny_list = []\n\nfor count, (data_path, model_path) in enumerate(zip(data_files, model_files), start=1):\n    try:\n        #print(f\"Loading file {count} of {len(data_files)}\")\n        #print(\"File\", data_path)\n        \n        # Memory-mapped loading\n        X = np.load(data_path, mmap_mode='r')  # (500, 5, 980, 70)\n        y = np.load(model_path, mmap_mode='r') # (500, 1, 70, 70)\n\n        # Preprocessing\n        X = X[:, :, 20:1000, :]        # -> (500, 5, 980, 70)\n        X = X[:, :, ::14, :]           # -> (500, 5, 70, 70)\n        X = X.transpose(0, 2, 1, 3)    # -> (500, 70, 5, 70)\n        X = X.reshape(500, 70, 350)    # -> (500, 70, 350)\n        X = X[:, :, ::5]               # -> (500, 70, 70)\n\n        if y.ndim == 4:\n            y = np.squeeze(y, axis=1)  # -> (500, 70, 70)\n\n        y = (y - 1500.0) / (4500.0 - 1500.0)\n\n        X_list.append(X)\n        y_list.append(y)\n\n    except Exception as e:\n        print(f\"Error loading {data_path} or {model_path}: {str(e)}\")\n        continue\n\n# Final conversion (only once, fast)\nX_data = np.concatenate(X_list, axis=0)\ny_data = np.concatenate(y_list, axis=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T23:50:12.710331Z","iopub.execute_input":"2025-05-07T23:50:12.711159Z","iopub.status.idle":"2025-05-07T23:53:11.254450Z","shell.execute_reply.started":"2025-05-07T23:50:12.711125Z","shell.execute_reply":"2025-05-07T23:53:11.253154Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train/Val Split (80/20)\nX_train, X_val, y_train, y_val = train_test_split(\n    X_data, y_data, test_size=0.2, random_state=42\n)\n\nprint(\"X_train shape:\", X_train.shape)\nprint(\"X_val shape:\", X_val.shape)\nprint(\"y_train shape:\", y_train.shape)\nprint(\"y_val shape:\", y_val.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T23:58:24.520266Z","iopub.execute_input":"2025-05-07T23:58:24.520754Z","iopub.status.idle":"2025-05-07T23:58:24.678272Z","shell.execute_reply.started":"2025-05-07T23:58:24.520719Z","shell.execute_reply":"2025-05-07T23:58:24.677360Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualize training samples","metadata":{}},{"cell_type":"code","source":"# Loop through all samples\nfor sample_idx in range(0, len(X_train), 5000):\n    seismic = X_train[sample_idx]      # shape (70, 70)\n    velocity = y_train[sample_idx]     # shape (70, 70)\n\n    # Plot them\n    plt.figure(figsize=(12, 5))\n\n    # Seismic Input\n    plt.subplot(1, 2, 1)\n    plt.imshow(seismic, aspect='auto', cmap='seismic', origin='lower')\n    plt.title(f\"Seismic Input (Sample {sample_idx})\")\n    plt.xlabel(\"Fused Shots Offset\")\n    plt.ylabel(\"Time\")\n    plt.colorbar(label=\"Amplitude\")\n\n    # Velocity Model (Target)\n    plt.subplot(1, 2, 2)\n    plt.imshow(velocity, aspect='equal', cmap='viridis', origin='lower', vmin=0, vmax=1)\n    plt.title(\"Velocity Model (Target)\")\n    plt.xlabel(\"X Position\")\n    plt.ylabel(\"Depth\")\n    plt.colorbar(label=\"Velocity\")\n\n    plt.tight_layout()\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T23:58:27.436029Z","iopub.execute_input":"2025-05-07T23:58:27.436791Z","iopub.status.idle":"2025-05-07T23:58:29.009655Z","shell.execute_reply.started":"2025-05-07T23:58:27.436763Z","shell.execute_reply":"2025-05-07T23:58:29.008550Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## ResNet","metadata":{}},{"cell_type":"code","source":"\n# Define the feature extractor\nclass ResNetFeatureExtractor(nn.Module):\n    def __init__(self):\n        super(ResNetFeatureExtractor, self).__init__()\n        base_model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)\n        base_model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n        self.features = nn.Sequential(*list(base_model.children())[:-1])  # Remove the final FC layer\n\n    def forward(self, x):\n        x = self.features(x)  # Output shape: (batch_size, 512, 1, 1)\n        return x.view(x.size(0), -1)  # Flatten to (batch_size, 512)\n\n# Instantiate the model\nfeature_extractor = ResNetFeatureExtractor().to(device)\nfeature_extractor.eval()\n\n\nX_train_tensor = torch.tensor(X_train, dtype=torch.float32).unsqueeze(1)  # Shape: \n\n# Generate embeddings\ntrain_embeddings = []\nwith torch.no_grad():\n    for i in tqdm(range(0, len(X_train_tensor), 512)):\n        batch = X_train_tensor[i:i+512].to(device)\n        embeddings = feature_extractor(batch).cpu().numpy()\n        train_embeddings.append(embeddings)\n\ntrain_embeddings = np.vstack(train_embeddings)  # Shape: \n\n# Build the FAISS index\nembedding_dim = train_embeddings.shape[1]\nfaiss_index = faiss.IndexFlatL2(embedding_dim)\nfaiss_index.add(train_embeddings)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T00:48:28.127677Z","iopub.execute_input":"2025-05-08T00:48:28.128952Z","iopub.status.idle":"2025-05-08T00:49:30.415272Z","shell.execute_reply.started":"2025-05-08T00:48:28.128914Z","shell.execute_reply":"2025-05-08T00:49:30.414023Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prediction vs Validation Plots\n","metadata":{}},{"cell_type":"code","source":"num_samples = 8\nsample_indices = random.sample(range(len(X_val)), num_samples)\nfeature_extractor.eval()\n\nplt.figure(figsize=(12, 4 * num_samples))\n\nfor i, idx in enumerate(sample_indices):\n    X = torch.tensor(X_val[idx], dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)\n    y_true = y_val[idx] * (4500 - 1500) + 1500\n\n    # Extract embedding and search\n    with torch.no_grad():\n        embedding = feature_extractor(X).cpu().numpy()\n    _, indices = faiss_index.search(embedding, 1)\n    pred = y_train[indices[0][0]] * (4500 - 1500) + 1500\n\n    # Difference plot\n    diff = pred - y_true\n\n    # Plot ground truth\n    plt.subplot(num_samples, 3, 3*i + 1)\n    plt.imshow(y_true, cmap='viridis', vmin=1500, vmax=4500, origin='lower')\n    plt.title(f\"Ground Truth #{idx}\")\n    plt.colorbar()\n\n    # Plot prediction\n    plt.subplot(num_samples, 3, 3*i + 2)\n    plt.imshow(pred, cmap='viridis', vmin=1500, vmax=4500, origin='lower')\n    plt.title(f\"Prediction #{idx}\")\n    plt.colorbar()\n\n    # Plot difference\n    plt.subplot(num_samples, 3, 3*i + 3)\n    plt.imshow(diff, cmap='bwr', vmin=-500, vmax=500, origin='lower')\n    plt.title(f\"Diff (Pred - GT) #{idx}\")\n    plt.colorbar()\n\nplt.tight_layout()\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T00:53:45.731976Z","iopub.execute_input":"2025-05-08T00:53:45.732412Z","iopub.status.idle":"2025-05-08T00:53:52.391432Z","shell.execute_reply.started":"2025-05-08T00:53:45.732382Z","shell.execute_reply":"2025-05-08T00:53:52.390006Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prepare Submission File","metadata":{}},{"cell_type":"code","source":"# Path to test files\ntest_dir = Path('/kaggle/input/waveform-inversion/test')\n\nsubmission_rows = []\ncount = 0\ntest_examples = []\n\nfor file_path in test_dir.rglob('*.npy'):\n    # Load and preprocess test file\n    X = np.load(file_path)  # Shape: (5, 1000, 70)\n    X = X[:, 20:1000, :]\n    X = X[:, ::14, :]\n    X = X.transpose(1, 0, 2).reshape(70, -1)  # Shape: (70, 350)\n    X = X[:, ::5]  # Shape: (70, 70)\n    \n    # Convert to tensor\n    X_tensor = torch.tensor(X, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # Shape: (1, 1, 70, 70)\n\n    # Extract feature vector\n    with torch.no_grad():\n        embedding = feature_extractor(X_tensor).cpu().numpy()  # Shape: (1, 512)\n\n    # Search FAISS index\n    _, indices = faiss_index.search(embedding, 1)\n    nearest_idx = indices[0][0]\n\n    # Retrieve predicted velocity map\n    pred = y_train[nearest_idx]  # Shape: (70, 70)\n    pred = pred * (4500.0 - 1500.0) + 1500.0  # Unnormalize\n\n    # Format for submission\n    file_stem = file_path.stem\n    for i in range(pred.shape[0]):\n        oid_ypos = f\"{file_stem}_y_{i}\"\n        x_values = pred[i, 1::2]  # x_1, x_3, ..., x_69\n        row = [oid_ypos] + x_values.tolist()\n        submission_rows.append(row)\n\n    count += 1\n    if count % 5000 == 0:\n        print(f\"Processed {count} files.\")\n        test_examples.append(file_path)\n        \n# Create submission DataFrame\ncolumns = ['oid_ypos'] + [f'x_{i}' for i in range(1, 70, 2)]\nsubmission_df = pd.DataFrame(submission_rows, columns=columns)\n\n# Save to CSV\nsubmission_df.to_csv('submission.csv', index=False)\nprint(\"✅ Submission file saved!\")\nsubmission_df.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T00:57:18.071371Z","iopub.execute_input":"2025-05-08T00:57:18.071846Z","iopub.status.idle":"2025-05-08T01:43:56.906839Z","shell.execute_reply.started":"2025-05-08T00:57:18.071819Z","shell.execute_reply":"2025-05-08T01:43:56.903921Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualize a test sample","metadata":{}},{"cell_type":"code","source":"\nfor i in range (len(test_examples)):\n    X = np.load(test_examples[i])  # shape (5, 1000, 70)\n    print(test_examples[i])\n    X = X[:, 20:1000, :]\n    X = X[:, ::14, :]\n    X = X.transpose(1, 0, 2).reshape(70, -1)  # (70, 350)\n    X = X[:, ::5]  # (70, 70)\n\n    X_tensor = torch.tensor(X, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # (1, 1, 70, 70)\n\n    # Extract embedding and search\n    with torch.no_grad():\n        embedding = feature_extractor(X_tensor).cpu().numpy()\n    _, indices = faiss_index.search(embedding, 1)\n    pred = y_train[indices[0][0]] * (4500 - 1500) + 1500\n\n    # Plot Prediction\n    plt.imshow(pred, cmap='viridis', aspect='equal', origin='lower', vmin=1500, vmax=4500)\n    plt.colorbar()\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-08T01:51:46.828526Z","iopub.execute_input":"2025-05-08T01:51:46.828958Z","iopub.status.idle":"2025-05-08T01:51:51.067602Z","shell.execute_reply.started":"2025-05-08T01:51:46.828931Z","shell.execute_reply":"2025-05-08T01:51:51.066447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}