{"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":11368499,"sourceType":"datasetVersion","datasetId":7116445},{"sourceId":11368545,"sourceType":"datasetVersion","datasetId":7116479},{"sourceId":11448436,"sourceType":"datasetVersion","datasetId":7155401},{"sourceId":235617897,"sourceType":"kernelVersion"}],"dockerImageVersionId":31012,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Faiss\n\nThis is just for a bit of fun. I've trained a classifier [here](https://www.kaggle.com/code/johnnyhyland/waveform-classifier) which determines the folder the waveform would belong in. The idea is to speed up the search by only looking in the relevant folders.\n\nThen I use faiss, necessary packages are [here](https://www.kaggle.com/code/johnnyhyland/faiss-cpu). The aim is to find the most similar waveform and return the velocity map corresponding to it.\n\nI've only included some of the extra data due to processing time. If you find a way to get more data in and improve on this, please share!","metadata":{}},{"cell_type":"code","source":"!pip install /kaggle/input/faiss-cpu/faiss_cpu-1.10.0-cp311-cp311-manylinux_2_28_x86_64.whl","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport numpy as np\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nimport csv\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nfrom torchvision.models import efficientnet_b0\nfrom collections import defaultdict\n\n# --- Classifier setup ---\nclass_map = {\n    \"CurveFault_A\": 0, \"CurveFault_B\": 1, \"CurveVel_A\": 2, \"CurveVel_B\": 3,\n    \"FlatFault_A\": 4, \"FlatFault_B\": 5, \"FlatVel_A\": 6, \"FlatVel_B\": 7,\n    \"Style_A\": 8, \"Style_B\": 9,\n}\ninv_class_map = {v: k for k, v in class_map.items()}\n\nclass SeismicEfficientNetClassifier(nn.Module):\n    def __init__(self, num_classes):\n        super().__init__()\n        self.backbone = efficientnet_b0(weights=None)\n        if self.backbone.features[0][0].in_channels != 1:\n            self.backbone.features[0][0] = nn.Conv2d(\n                1, 32, kernel_size=3, stride=2, padding=1, bias=False\n            )\n        in_features = self.backbone.classifier[1].in_features\n        self.backbone.classifier[1] = nn.Linear(in_features, num_classes)\n    def forward(self, x):\n        return self.backbone(x)\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nclassifier = SeismicEfficientNetClassifier(num_classes=10)\nclassifier.load_state_dict(torch.load('/kaggle/input/waveform-inversion-models/classifier.pth', map_location=device))\nclassifier.to(device)\nclassifier.eval()\n\n# --- Submission setup ---\ntest_files = list(Path('/kaggle/input/waveform-inversion/test').glob('*.npy'))\nx_cols = [f'x_{i}' for i in range(1, 70, 2)]\nfieldnames = ['oid_ypos'] + x_cols\n\ndef get_class_train_files(category):\n    # Define all data sources to search\n    data_sources = [\n        # Original data source\n        # {'base_path': '/kaggle/input/waveform-inversion/train_samples', 'subdir': True},\n        # Additional data sources - add more as needed\n        {'base_path': '/kaggle/input/waveform-inversion-1', 'subdir': False},\n        {'base_path': '/kaggle/input/waveform-inversion-2', 'subdir': False},\n        # {'base_path': '/kaggle/input/waveform-inversion-3', 'subdir': False},\n        # Add more sources as needed\n    ]\n    \n    all_input_files = []\n    all_output_files = []\n    \n    for source in data_sources:\n        base_path = source['base_path']\n        if source['subdir']:\n            # Original structure with train_samples subdirectory\n            folder = Path(f\"{base_path}/{category}\")\n        else:\n            # Direct category structure without train_samples\n            folder = Path(f\"{base_path}/{category}\")\n        \n        # Skip if folder doesn't exist\n        if not folder.exists():\n            continue\n            \n        # Find input files\n        files = [f for f in folder.rglob('*.npy') if ('seis' in f.stem) or ('data' in f.stem)]\n        # Generate corresponding output files\n        out_files = [Path(str(f).replace('seis', 'vel').replace('data', 'model')) for f in files]\n        # Only keep pairs where both input and output exist\n        valid_pairs = [(f, o) for f, o in zip(files, out_files) if o.exists()]\n        \n        if valid_pairs:\n            files, out_files = zip(*valid_pairs)\n            all_input_files.extend(files)\n            all_output_files.extend(out_files)\n    \n    print(f\"Found {len(all_input_files)} training samples for category {category}\")\n    return all_input_files, all_output_files\n\ndef load_train_batch(inputs_files, outputs_files, start_idx, batch_size):\n    \"\"\"Load a batch of training samples starting from start_idx\"\"\"\n    end_idx = min(start_idx + batch_size, len(inputs_files))\n    x_batch = []\n    y_batch = []\n    for i in range(start_idx, end_idx):\n        x = np.load(inputs_files[i])  # Shape: [500, 5, 1000, 70]\n        y = np.load(outputs_files[i])\n        x_batch.append(x)\n        y_batch.append(y)\n    return torch.tensor(np.concatenate(x_batch), dtype=torch.float32), np.concatenate(y_batch)\n\n# Step 1: Classify all test files efficiently\nprint(\"Classifying all test files...\")\nclass TestClassificationDataset(Dataset):\n    def __init__(self, test_files):\n        self.test_files = test_files\n\n\n    def __len__(self):\n        return len(self.test_files)\n\n\n    def __getitem__(self, i):\n        test_file = self.test_files[i]\n\n        return np.load(test_file), test_file.stem\n\n# Create dataset and dataloader for classification\nclassification_dataset = TestClassificationDataset(test_files)\nclassification_loader = DataLoader(\n    classification_dataset, \n    batch_size=128,  # Adjust based on GPU memory\n    num_workers=6,\n    pin_memory=True\n)\n\n# Simple dictionary to store file → category mapping\ntest_file_to_category = {}\n\n# Classify all test files\nfor batch_tensors, batch_files in tqdm(classification_loader, desc=\"Classifying test files\"):\n    batch_tensors = batch_tensors.to(device)\n    \n    with torch.no_grad():\n        x_for_class = batch_tensors.mean(dim=1, keepdim=True)\n        class_logits = classifier(x_for_class)\n        pred_classes = class_logits.argmax(dim=1).cpu().numpy()\n    \n    # Store results in the dictionary\n    for file_path, pred_class in zip(batch_files, pred_classes):\n        test_file_to_category[file_path] = inv_class_map[pred_class]\n\n# Group files by category for efficient processing\ntest_files_by_category = defaultdict(list)\nfor test_file, category in test_file_to_category.items():\n    test_files_by_category[category].append(test_file)\n\n# Print summary of classification\nfor category, files in test_files_by_category.items():\n    print(f\"Category {category}: {len(files)} test files\")\n\n# Free classification resources\ndel classifier\ntorch.cuda.empty_cache()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-23T09:52:21.053684Z","iopub.execute_input":"2025-04-23T09:52:21.054033Z","iopub.status.idle":"2025-04-23T10:02:27.100237Z","shell.execute_reply.started":"2025-04-23T09:52:21.054002Z","shell.execute_reply":"2025-04-23T10:02:27.096339Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport csv\nfrom tqdm import tqdm\nimport faiss\nimport gc\nimport torch\n\nwith open('submission.csv', 'wt', newline='') as csvfile:\n    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)\n    writer.writeheader()\n    \n    # Process each category one by one\n    for category, category_test_files in test_files_by_category.items():\n        print(f\"Processing category: {category} with {len(category_test_files)} files\")\n        \n        # Load training data for this category only once\n        train_inputs, train_outputs = get_class_train_files(category)\n        \n        # Preprocess all training data for this category\n        print(f\"Preprocessing training data for {category}...\")\n        all_train_data = []\n        file_positions = []  # To keep track of which file and position each sample is from\n        \n        for file_idx, input_file in enumerate(tqdm(train_inputs, desc=\"Building index\")):\n            # Load training file\n            train_data = np.load(input_file)\n            num_samples = train_data.shape[0]\n            \n            # Flatten the dimensions (keeping batch dimension)\n            flattened_data = train_data.reshape(num_samples, -1).astype(np.float32)\n            \n            # Add to our dataset\n            all_train_data.append(flattened_data)\n            \n            # Keep track of file index and position within file\n            for pos in range(num_samples):\n                file_positions.append((file_idx, pos))\n        \n        # Concatenate all data\n        all_train_data = np.vstack(all_train_data)\n\n        # Use more efficient index types\n        dimension = all_train_data.shape[1]\n        \n        # For large datasets (>1M vectors), use IVF index\n        nlist = min(4096, max(int(len(file_positions)/50), 256))  # Rule of thumb: nlist ≈ sqrt(N)\n        quantizer = faiss.IndexFlat(dimension)\n        quantizer.metric_type = faiss.METRIC_L1\n        index = faiss.IndexIVFFlat(quantizer, dimension, nlist)\n        index.train(all_train_data)\n        index.add(all_train_data)\n        index.nprobe = min(nlist//4, 256)\n        \n        # For even better speed, consider:\n        # index = faiss.IndexHNSWFlat(dimension, 32, faiss.METRIC_L1)  # HNSW with 32 neighbors\n        \n        # Use GPU if available and dataset fits\n        if torch.cuda.is_available() and all_train_data.shape[0] < 10_000_000:\n            gpu_resource = faiss.StandardGpuResources()\n            index = faiss.index_cpu_to_gpu(gpu_resource, 0, index)\n        \n        # # Create FAISS index with L1 distance (MAE) metric\n        # print(\"Creating FAISS index...\")\n        \n        # # Use METRIC_L1 (value=1) for L1/Manhattan distance\n        # index = faiss.IndexFlat(dimension, faiss.METRIC_L1)\n        \n        # # Add all training vectors to the index\n        # index.add(all_train_data)\n        \n        # Free memory\n        del all_train_data\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n        \n        # Process each test file in this category\n        for test_file in tqdm(category_test_files, desc=f\"Processing {category} test files\"):\n            # Load the test sample\n            test_sample = np.load(f'/kaggle/input/waveform-inversion/test/{test_file}.npy')\n            \n            # Reshape to match the indexing dimensions\n            test_flat = test_sample.reshape(1, -1).astype(np.float32)\n            \n            # Search for nearest neighbor\n            distances, indices = index.search(test_flat, 1)  # find 1 nearest neighbor\n            \n            # Get the original file index and position\n            file_idx, pos_in_file = file_positions[indices[0][0]]\n            \n            # Load the best matching velocity map\n            vel_file = train_outputs[file_idx]\n            all_velocities = np.load(vel_file)\n            best_vel = all_velocities[pos_in_file]\n                        \n            # Write to CSV - fixed to ensure one value per column\n            test_id = test_file\n            for y_pos in range(best_vel.shape[1]):\n                # Create a dictionary with 'oid_ypos' and individual x_col values\n                row = {'oid_ypos': f\"{test_id}_y_{y_pos}\"}\n                \n                # Add each x column value individually - use item() to extract scalar\n                for i, x_pos in enumerate(range(1, 70, 2)):\n                    col_name = f'x_{x_pos}'\n                    # Use .item() to convert NumPy array element to Python scalar\n                    row[col_name] = best_vel[0, y_pos, x_pos].item()\n                \n                writer.writerow(row)\n        \n        # After processing all files for this category, free memory\n        del index, file_positions, train_inputs, train_outputs\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\nprint(\"Submission complete!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-23T10:12:20.805890Z","iopub.execute_input":"2025-04-23T10:12:20.806742Z","execution_failed":"2025-04-23T10:19:04.372Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}