{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.5"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"},{"sourceId":416979,"sourceType":"modelInstanceVersion","modelInstanceId":340189,"modelId":361305}],"dockerImageVersionId":31041,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Yale/UNC-CH - Geophysical Waveform Inversion Competition\n## Advanced Solution with Physics-Guided Machine Learning\n\nThis notebook presents an optimized approach for the Waveform Inversion competition, designed to handle large datasets efficiently while achieving competitive results using state-of-the-art techniques.\n\n### Competition Overview\nThe goal is to develop physics-guided machine learning models to solve full-waveform inversion problems. We need to predict velocity maps from seismic waveform recordings.\n\n### Approach Overview\n1. Memory-efficient data loading using chunking and streaming\n2. GPU/TPU acceleration for deep learning models\n3. Advanced architectures (ConvNeXt and InversionNet)\n4. Physics-guided machine learning techniques\n5. Ensemble of multiple models for improved accuracy","metadata":{}},{"cell_type":"markdown","source":"## 1. Environment Setup and Package Installation","metadata":{}},{"cell_type":"code","source":"!pip install -q pytorch-lightning==2.0.0 torchmetrics==0.11.4 einops==0.6.1 timm==0.6.12\ntry:\n    import pytorch_lightning\n    print('PyTorch Lightning imported successfully after install.')\nexcept ImportError:\n    try:\n        import lightning.pytorch\n        print('Lightning imported successfully after install.')\n    except ImportError:\n        print('ERROR: Failed to import PyTorch Lightning/Lightning after install.')\n# --- Inlined Functions from kaggle_gm_automation.py ---\nimport os\nimport gc\nimport logging\nimport random\nimport warnings\nfrom typing import Optional, Union, Dict, Any, Tuple\nimport numpy as np\nimport torch\ntry:\n    from packaging import version\nexcept ImportError:\n    pass\ntry:\n    import importlib.metadata\nexcept ImportError:\n    pass\ntry:\n    import pkg_resources\nexcept ImportError:\n    pass\ntry:\n    import pytorch_lightning as pl\nexcept ImportError:\n    try:\n        import lightning.pytorch as pl\n    except ImportError:\n        print(\"Warning: PyTorch Lightning not found\")\ntry:\n    import torch_xla\n    import torch_xla.core.xla_model as xm\nexcept ImportError:\n    pass\ntry:\n    import psutil\nexcept ImportError:\n    pass\nfrom torch.utils.data import DataLoader, Dataset\n\n\nlogging.basicConfig(\n    level=logging.INFO,\n    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'\n)\nlogger = logging.getLogger(\"InlinedFunctions\")\n\n\ndef get_pytorch_lightning_version() -> str:\n    \"\"\"\n    Get the installed PyTorch Lightning version.\n    \n    Returns:\n        str: Version string of PyTorch Lightning\n    \"\"\"\n    try:\n        import importlib.metadata\n        try:\n            return importlib.metadata.version(\"pytorch-lightning\")\n        except importlib.metadata.PackageNotFoundError:\n            try:\n                return importlib.metadata.version(\"lightning\")\n            except importlib.metadata.PackageNotFoundError:\n                pass\n    except ImportError:\n        pass\n    \n    try:\n        import pkg_resources\n        try:\n            return pkg_resources.get_distribution(\"pytorch-lightning\").version\n        except pkg_resources.DistributionNotFound:\n            try:\n                return pkg_resources.get_distribution(\"lightning\").version\n            except pkg_resources.DistributionNotFound:\n                pass\n    except ImportError:\n        pass\n    \n    # If all else fails, try importing and checking __version__\n    try:\n        import pytorch_lightning as pl\n        return pl.__version__\n    except (ImportError, AttributeError):\n        try:\n            import lightning as L\n            return L.__version__\n        except (ImportError, AttributeError):\n            # Default to a recent version if we can\\'t determine it\n            warnings.warn(\"Could not determine PyTorch Lightning version. Assuming 2.0.0\")\n            return \"2.0.0\"\n\ndef parse_version(version_str: str) -> Tuple[int, ...]:\n    \"\"\"\n    Parse version string into a tuple of integers for comparison.\n    \n    Args:\n        version_str: Version string (e.g., \"1.7.0\")\n        \n    Returns:\n        Tuple of integers representing the version\n    \"\"\"\n    try:\n        # Try using packaging.version if available\n        from packaging import version\n        v = version.parse(version_str)\n        return (v.major, v.minor, v.micro)\n    except ImportError:\n        # Fallback to manual parsing\n        return tuple(int(x) for x in version_str.split(\".\")[:3])\n\ndef detect_device_type() -> str:\n    \"\"\"\n    Detect available device type with proper error handling and fallbacks.\n    \n    Returns:\n        str: \"gpu\", \"tpu\", or \"cpu\"\n    \"\"\"\n    # Try GPU first\n    try:\n        if torch.cuda.is_available():\n            logger.info(\"GPU detected and available\")\n            return \"gpu\"\n    except Exception as e:\n        logger.warning(f\"Error checking GPU availability: {e}\")\n    \n    # Try TPU next\n    try:\n        import torch_xla\n        import torch_xla.core.xla_model as xm\n        try:\n            # Verify TPU is actually available by attempting to get a device\n            device = xm.xla_device()\n            logger.info(\"TPU detected and available\")\n            return \"tpu\"\n        except Exception as e:\n            logger.warning(f\"TPU import succeeded but initialization failed: {e}\")\n    except ImportError:\n        logger.info(\"TPU support not available (torch_xla not importable)\")\n    except Exception as e:\n        logger.warning(f\"Unexpected error checking TPU availability: {e}\")\n    \n    # Fall back to CPU\n    logger.info(\"No GPU or TPU found, using CPU\")\n    return \"cpu\"\n\ndef get_device() -> torch.device:\n    \"\"\"\n    Get the appropriate device with detailed logging and error handling.\n    \n    Returns:\n        torch.device: The appropriate device for the current environment\n    \"\"\"\n    device_type = detect_device_type()\n    \n    if device_type == \"gpu\":\n        device = torch.device(\"cuda\")\n        logger.info(f\"Using GPU: {torch.cuda.get_device_name(0)}\")\n        try:\n            logger.info(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n        except Exception as e:\n            logger.warning(f\"Could not get GPU memory info: {e}\")\n        return device\n    elif device_type == \"tpu\":\n        import torch_xla.core.xla_model as xm\n        device = xm.xla_device()\n        logger.info(\"Using TPU\")\n        return device\n    else:\n        logger.info(\"Using CPU\")\n        return torch.device(\"cpu\")\n\ndef get_optimal_precision_config() -> Union[str, int]:\n    \"\"\"\n    Determine the optimal precision configuration based on PyTorch Lightning version\n    and available hardware.\n    \n    Returns:\n        Union[str, int]: The appropriate precision parameter value for pl.Trainer\n    \"\"\"\n    # Get PyTorch Lightning version\n    pl_version = get_pytorch_lightning_version()\n    pl_version_tuple = parse_version(pl_version)\n    \n    # Determine available hardware\n    device_type = detect_device_type()\n    \n    logger.info(f\"PyTorch Lightning version: {pl_version}\")\n    logger.info(f\"Device type for precision config: {device_type}\")\n    \n    # Configure precision based on version and hardware\n    if pl_version_tuple >= (1, 7, 0):\n        # New string-based format (>= 1.7.0)\n        if device_type == \"gpu\" and torch.cuda.is_available():\n            logger.info(\"Using \\'16-mixed\\' precision for GPU with PyTorch Lightning >= 1.7.0\")\n            return \"16-mixed\"\n        elif device_type == \"tpu\":\n            # TPUs require specific precision formats\n            logger.info(\"Using \\'bf16-true\\' precision for TPU with PyTorch Lightning >= 1.7.0\")\n            return \"bf16-true\"  # bfloat16 for TPU\n        else:\n            logger.info(\"Using \\'32-true\\' precision for CPU with PyTorch Lightning >= 1.7.0\")\n            return \"32-true\"\n    else:\n        # Old integer-based format (< 1.7.0)\n        if device_type == \"gpu\" and torch.cuda.is_available():\n            logger.info(\"Using 16 precision for GPU with PyTorch Lightning < 1.7.0\")\n            return 16\n        elif device_type == \"tpu\":\n            # For older versions with TPU\n            logger.info(\"Using \\'bf16\\' precision for TPU with PyTorch Lightning < 1.7.0\")\n            try:\n                # Try string format first for TPU\n                return \"bf16\"\n            except:\n                # Fall back to 32 if string format not supported\n                logger.info(\"Falling back to 32 precision for TPU with PyTorch Lightning < 1.7.0\")\n                return 32\n        else:\n            logger.info(\"Using 32 precision for CPU with PyTorch Lightning < 1.7.0\")\n            return 32\n\ndef set_reproducibility(seed: int = 42) -> None:\n    \"\"\"\n    Set random seeds for reproducibility across all libraries.\n    \n    Args:\n        seed: Random seed to use\n    \"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    \n    if torch.cuda.is_available():\n        torch.cuda.manual_seed(seed)\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n    \n    logger.info(f\"Set random seed to {seed} for reproducibility\")\n\ndef optimize_memory_usage() -> Dict[str, Any]:\n    \"\"\"\n    Implement memory optimization techniques and return recommended settings.\n    \n    Returns:\n        Dict with optimization settings\n    \"\"\"\n    # Clear memory\n    gc.collect()\n    \n    # Clear PyTorch cache if using GPU\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n    \n    device_type = detect_device_type()\n    \n    # Get available memory\n    available_memory = None\n    if device_type == \"gpu\":\n        try:\n            available_memory = torch.cuda.get_device_properties(0).total_memory / 1e9\n        except:\n            available_memory = 8  # Default assumption\n    else:\n        try:\n            import psutil\n            available_memory = psutil.virtual_memory().available / 1024**3\n        except:\n            available_memory = 4  # Default assumption\n    \n    # Determine recommended batch size\n    base_batch_size = 32\n    if device_type == \"gpu\":\n        if available_memory > 14:  # High-end GPU\n            recommended_batch_size = base_batch_size * 2\n        elif available_memory < 4:  # Limited GPU\n            recommended_batch_size = max(1, int(base_batch_size / 4))\n        else:  # Standard GPU\n            recommended_batch_size = base_batch_size\n    elif device_type == \"tpu\":\n        recommended_batch_size = base_batch_size * 2  # TPUs generally handle larger batches\n    else:  # CPU\n        recommended_batch_size = max(1, int(base_batch_size / 8))  # Much smaller batches for CPU\n    \n    # Suggest optimal settings based on environment\n    if device_type == \"gpu\":\n        settings = {\n            \"use_mixed_precision\": True,\n            \"recommended_batch_size\": recommended_batch_size,\n            \"use_gradient_checkpointing\": available_memory < 8,\n            \"use_activation_checkpointing\": available_memory < 4,\n            \"available_memory_gb\": available_memory\n        }\n    elif device_type == \"tpu\":\n        settings = {\n            \"use_mixed_precision\": True,\n            \"recommended_batch_size\": recommended_batch_size,\n            \"use_gradient_checkpointing\": False,\n            \"use_activation_checkpointing\": False,\n            \"available_memory_gb\": available_memory\n        }\n    else:  # CPU\n        settings = {\n            \"use_mixed_precision\": False,\n            \"recommended_batch_size\": recommended_batch_size,\n            \"use_gradient_checkpointing\": True,\n            \"use_activation_checkpointing\": True,\n            \"available_memory_gb\": available_memory\n        }\n    \n    logger.info(f\"Memory optimization settings: {settings}\")\n    return settings\n\ndef is_in_kaggle() -> bool:\n    \"\"\"\n    Check if the code is running in a Kaggle environment.\n    \n    Returns:\n        bool: True if running in Kaggle, False otherwise\n    \"\"\"\n    return os.path.exists('/kaggle/input')\n\ndef configure_pytorch_lightning_trainer(\n    max_epochs: int = 30,\n    callbacks: list = None,\n    **kwargs\n) -> Any:\n    \"\"\"\n    Configure a PyTorch Lightning Trainer with optimal settings for the current environment.\n    \n    Args:\n        max_epochs: Maximum number of training epochs\n        callbacks: List of PyTorch Lightning callbacks\n        **kwargs: Additional arguments to pass to the Trainer\n        \n    Returns:\n        PyTorch Lightning Trainer instance\n    \"\"\"\n    try:\n        import pytorch_lightning as pl\n    except ImportError:\n        try:\n            import lightning.pytorch as pl\n        except ImportError:\n            raise ImportError(\"PyTorch Lightning is not installed. Please install it with: pip install pytorch-lightning\")\n    \n    # Get optimal precision configuration\n    precision = get_optimal_precision_config()\n    \n    # Get optimization settings\n    optimization_settings = optimize_memory_usage()\n    \n    # Default callbacks if none provided\n    if callbacks is None:\n        callbacks = []\n    \n    # Configure trainer with optimal settings\n    trainer_kwargs = {\n        \"max_epochs\": max_epochs,\n        \"callbacks\": callbacks,\n        \"accelerator\": \"auto\",  # Let PyTorch Lightning detect the accelerator\n        \"devices\": 1,\n        \"precision\": precision,\n        \"log_every_n_steps\": 10,\n    }\n    \n    # Add gradient clipping if using gradient checkpointing\n    if optimization_settings[\"use_gradient_checkpointing\"]:\n        trainer_kwargs[\"gradient_clip_val\"] = 1.0\n    \n    # Override with any user-provided kwargs\n    trainer_kwargs.update(kwargs)\n    \n    logger.info(f\"Configuring PyTorch Lightning Trainer with: {trainer_kwargs}\")\n    \n    # Create and return the trainer\n    return pl.Trainer(**trainer_kwargs)\n\ndef configure_dataloader(\n    dataset,\n    batch_size: Optional[int] = None,\n    shuffle: bool = True,\n    **kwargs\n) -> torch.utils.data.DataLoader:\n    \"\"\"\n    Configure a DataLoader with optimal settings for the current environment.\n    \n    Args:\n        dataset: PyTorch Dataset\n        batch_size: Batch size (if None, will be determined automatically)\n        shuffle: Whether to shuffle the data\n        **kwargs: Additional arguments to pass to the DataLoader\n        \n    Returns:\n        torch.utils.data.DataLoader: Configured DataLoader\n    \"\"\"\n    # Get device type\n    device_type = detect_device_type()\n    \n    # Get optimization settings\n    optimization_settings = optimize_memory_usage()\n    \n    # Use recommended batch size if not provided\n    if batch_size is None:\n        batch_size = optimization_settings[\"recommended_batch_size\"]\n    \n    # Configure DataLoader with optimal settings\n    dataloader_kwargs = {\n        \"batch_size\": batch_size,\n        \"shuffle\": shuffle,\n        \"num_workers\": 0 if device_type == \"cpu\" else 4,\n        \"pin_memory\": device_type == \"gpu\",\n    }\n    \n    # Override with any user-provided kwargs\n    dataloader_kwargs.update(kwargs)\n    \n    logger.info(f\"Configuring DataLoader with: {dataloader_kwargs}\")\n    \n    # Create and return the DataLoader\n    return torch.utils.data.DataLoader(dataset, **dataloader_kwargs)\n\ndef example_usage():\n    \"\"\"\n    Example usage of the Kaggle Grandmaster Automation module.\n    \"\"\"\n    print(\"Kaggle Grandmaster Automation Module Example Usage\")\n    print(\"=\" * 50)\n    \n    # Check if running in Kaggle\n    print(f\"Running in Kaggle: {is_in_kaggle()}\")\n    \n    # Set reproducibility\n    set_reproducibility(seed=42)\n    \n    # Get device\n    device = get_device()\n    print(f\"Using device: {device}\")\n    \n    # Get optimal precision configuration\n    precision = get_optimal_precision_config()\n    print(f\"Optimal precision configuration: {precision}\")\n    \n    # Get optimization settings\n    optimization_settings = optimize_memory_usage()\n    print(f\"Optimization settings: {optimization_settings}\")\n    \n    print(\"\\nTo use in your notebook, add the following code:\")\n    print(\"# [Code example placeholder]\")\n\n# --- End Inlined Functions ---\n\n# Check if running in Kaggle environment\nimport os\nIN_KAGGLE = os.path.exists('/kaggle/input')\n\n# Install required packages if not already installed\n!pip install -q timm einops pytorch-lightning segmentation-models-pytorch","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Import necessary libraries\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport os\nimport gc\nimport time\nimport glob\nimport logging\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\n\nimport timm\nfrom einops import rearrange\nimport segmentation_models_pytorch as smp\n\n# Set up logging\nlogging.basicConfig(\n    level=logging.INFO,\n    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'\n)\nlogger = logging.getLogger(\"WaveformInversion\")\n\n# Set random seeds for reproducibility\nimport random\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed(seed)\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n    logger.info(f\"Set random seed to {seed} for reproducibility\")\n\nset_seed(42)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. GPU/TPU Detection and Configuration\n\nThis section includes robust device detection with proper error handling.","metadata":{}},{"cell_type":"code","source":"# Check for available GPU/TPU and configure accordingly\ndef detect_device_type():\n    \"\"\"Detect available device type with proper error handling and fallbacks.\"\"\"\n    # Try GPU first\n    try:\n        if torch.cuda.is_available():\n            logger.info(\"GPU detected and available\")\n            return \"gpu\"\n    except Exception as e:\n        logger.warning(f\"Error checking GPU availability: {e}\")\n    \n    # Try TPU next\n    try:\n        import torch_xla\n        import torch_xla.core.xla_model as xm\n        try:\n            # Verify TPU is actually available by attempting to get a device\n            device = xm.xla_device()\n            logger.info(\"TPU detected and available\")\n            return \"tpu\"\n        except Exception as e:\n            logger.warning(f\"TPU import succeeded but initialization failed: {e}\")\n    except ImportError:\n        logger.info(\"TPU support not available (torch_xla not importable)\")\n    except Exception as e:\n        logger.warning(f\"Unexpected error checking TPU availability: {e}\")\n    \n    # Fall back to CPU\n    logger.info(\"No GPU or TPU found, using CPU\")\n    return \"cpu\"\n\ndef get_device():\n    \"\"\"Get the appropriate device with detailed logging and error handling.\"\"\"\n    device_type = detect_device_type()\n    \n    if device_type == \"gpu\":\n        device = torch.device(\"cuda\")\n        logger.info(f\"Using GPU: {torch.cuda.get_device_name(0)}\")\n        logger.info(f\"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB\")\n        return device\n    elif device_type == \"tpu\":\n        import torch_xla.core.xla_model as xm\n        device = xm.xla_device()\n        logger.info(\"Using TPU\")\n        return device\n    else:\n        logger.info(\"Using CPU\")\n        return torch.device(\"cpu\")\n\n# Get device\ndevice_type = detect_device_type()\ndevice = get_device()","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. PyTorch Lightning Version Detection and Precision Configuration\n\nThis section adds version-aware precision configuration for PyTorch Lightning.","metadata":{}},{"cell_type":"code","source":"def get_pytorch_lightning_version():\n    \"\"\"Get the installed PyTorch Lightning version.\"\"\"\n    try:\n        import importlib.metadata\n        try:\n            return importlib.metadata.version(\"pytorch-lightning\")\n        except importlib.metadata.PackageNotFoundError:\n            try:\n                return importlib.metadata.version(\"lightning\")\n            except importlib.metadata.PackageNotFoundError:\n                pass\n    except ImportError:\n        pass\n    \n    try:\n        import pkg_resources\n        try:\n            return pkg_resources.get_distribution(\"pytorch-lightning\").version\n        except pkg_resources.DistributionNotFound:\n            try:\n                return pkg_resources.get_distribution(\"lightning\").version\n            except pkg_resources.DistributionNotFound:\n                pass\n    except ImportError:\n        pass\n    \n    # If all else fails, try importing and checking __version__\n    try:\n        import pytorch_lightning as pl\n        return pl.__version__\n    except (ImportError, AttributeError):\n        try:\n            import lightning as L\n            return L.__version__\n        except (ImportError, AttributeError):\n            # Default to a recent version if we can't determine it\n            logger.warning(\"Could not determine PyTorch Lightning version. Assuming 2.0.0\")\n            return \"2.0.0\"\n\ndef parse_version(version_str):\n    \"\"\"Parse version string into a tuple of integers for comparison.\"\"\"\n    try:\n        # Try using packaging.version if available\n        from packaging import version\n        v = version.parse(version_str)\n        return (v.major, v.minor, v.micro)\n    except ImportError:\n        # Fallback to manual parsing\n        return tuple(int(x) for x in version_str.split('.')[:3])\n\ndef get_optimal_precision_config():\n    \"\"\"Determine the optimal precision configuration based on PyTorch Lightning version and available hardware.\"\"\"\n    # Get PyTorch Lightning version\n    pl_version = get_pytorch_lightning_version()\n    pl_version_tuple = parse_version(pl_version)\n    \n    logger.info(f\"PyTorch Lightning version: {pl_version}\")\n    logger.info(f\"Device type for precision config: {device_type}\")\n    \n    # Configure precision based on version and hardware\n    if pl_version_tuple >= (1, 7, 0):\n        # New string-based format (>= 1.7.0)\n        if device_type == \"gpu\" and torch.cuda.is_available():\n            logger.info(\"Using '16-mixed' precision for GPU with PyTorch Lightning >= 1.7.0\")\n            return \"16-mixed\"\n        elif device_type == \"tpu\":\n            # TPUs require specific precision formats\n            logger.info(\"Using 'bf16-true' precision for TPU with PyTorch Lightning >= 1.7.0\")\n            return \"bf16-true\"  # bfloat16 for TPU\n        else:\n            logger.info(\"Using '32-true' precision for CPU with PyTorch Lightning >= 1.7.0\")\n            return \"32-true\"\n    else:\n        # Old integer-based format (< 1.7.0)\n        if device_type == \"gpu\" and torch.cuda.is_available():\n            logger.info(\"Using 16 precision for GPU with PyTorch Lightning < 1.7.0\")\n            return 16\n        elif device_type == \"tpu\":\n            # For older versions with TPU\n            logger.info(\"Using 'bf16' precision for TPU with PyTorch Lightning < 1.7.0\")\n            try:\n                # Try string format first for TPU\n                return \"bf16\"\n            except:\n                # Fall back to 32 if string format not supported\n                logger.info(\"Falling back to 32 precision for TPU with PyTorch Lightning < 1.7.0\")\n                return 32\n        else:\n            logger.info(\"Using 32 precision for CPU with PyTorch Lightning < 1.7.0\")\n            return 32\n\n# Get optimal precision configuration\nprecision = get_optimal_precision_config()","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Data Paths Configuration","metadata":{}},{"cell_type":"code","source":"# Configure data paths based on environment\nif IN_KAGGLE:\n    # Kaggle paths\n    COMP_PATH = '/kaggle/input/waveform-inversion'\n    OUTPUT_PATH = '/kaggle/working'\nelse:\n    # Local paths (adjust as needed)\n    COMP_PATH = '../input/waveform-inversion'\n    OUTPUT_PATH = './'\n\n# Define paths for different dataset families\nTRAIN_PATH = f\"{COMP_PATH}/train_samples\"\nTEST_PATH = f\"{COMP_PATH}/test\"\nSAMPLE_SUB_PATH = f\"{COMP_PATH}/sample_submission.csv\"\n\n# Dataset families\nDATASET_FAMILIES = ['FlatVel_A', 'Fault', 'Style']\n\n# Check available files\nprint(\"Available dataset families:\")\nfor family in DATASET_FAMILIES:\n    if os.path.exists(f\"{TRAIN_PATH}/{family}\"):\n        print(f\"- {family}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Memory-Efficient Data Loading and Processing","metadata":{}},{"cell_type":"code","source":"# Helper function to load data in chunks\ndef load_data_chunk(file_path, start_idx=0, chunk_size=10):\n    \"\"\"Load a chunk of data from a .npy file to save memory\"\"\"\n    try:\n        data = np.load(file_path, mmap_mode='r')\n        end_idx = min(start_idx + chunk_size, data.shape[0])\n        return data[start_idx:end_idx].copy()\n    except Exception as e:\n        logger.error(f\"Error loading data chunk from {file_path}: {e}\")\n        return None","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WaveformDataset(Dataset):\n    def __init__(self, data_files, model_files=None, transform=None, is_test=False):\n        self.data_files = data_files\n        self.model_files = model_files\n        self.transform = transform\n        self.is_test = is_test\n        self.file_sample_counts = []\n        self.cumulative_samples = [0]\n        \n        # Calculate total samples across all files\n        self.total_samples = 0\n        for file_path in self.data_files:\n            try:\n                with np.load(file_path) as data:\n                    samples_in_file = data.shape[0]\n                    self.file_sample_counts.append(samples_in_file)\n                    self.total_samples += samples_in_file\n                    self.cumulative_samples.append(self.total_samples)\n            except Exception as e:\n                logger.error(f\"Error loading file {file_path}: {e}\")\n                self.file_sample_counts.append(0)\n    \n    def __len__(self):\n        return self.total_samples\n    \n    def __getitem__(self, idx):\n        try:\n            # Find which file contains this index\n            file_idx = np.searchsorted(self.cumulative_samples, idx + 1) - 1\n            sample_idx = idx - self.cumulative_samples[file_idx]\n            \n            # Load data\n            with np.load(self.data_files[file_idx]) as data:\n                seismic = data[sample_idx].astype(np.float32)\n            \n            # Process seismic data (shape: num_sources, time_steps, num_receivers)\n            seismic = torch.from_numpy(seismic)\n            \n            # Normalize\n            seismic = (seismic - seismic.mean()) / (seismic.std() + 1e-8)\n            \n            # Add channel dimension and permute to (channels, height, width)\n            seismic = seismic.unsqueeze(0)  # Add channel dimension\n            seismic = rearrange(seismic, 'c s t r -> c t r s')  # Reorder dimensions\n            \n            if self.is_test:\n                return {'seismic': seismic, 'file_idx': file_idx, 'sample_idx': sample_idx}\n            else:\n                with np.load(self.model_files[file_idx]) as model_data:\n                    velocity = model_data[sample_idx].astype(np.float32)\n                velocity = torch.from_numpy(velocity).unsqueeze(0)  # Add channel dimension\n                \n                if self.transform:\n                    seismic, velocity = self.transform(seismic, velocity)\n                    \n                return {'seismic': seismic, 'velocity': velocity}\n        except Exception as e:\n            logger.error(f\"Error in __getitem__ for index {idx}: {e}\")\n            # Return zero tensors as fallback\n            if self.is_test:\n                return {'seismic': torch.zeros(1, 1000, 70, 5), 'file_idx': 0, 'sample_idx': 0}\n            else:\n                return {'seismic': torch.zeros(1, 1000, 70, 5), 'velocity': torch.zeros(1, 70, 70)}","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to find all data files for each dataset family\ndef find_data_files():\n    data_files = []\n    model_files = []\n    \n    for family in DATASET_FAMILIES:\n        family_path = f\"{TRAIN_PATH}/{family}\"\n        if not os.path.exists(family_path):\n            continue\n            \n        if family == 'Fault':\n            # Fault family has different naming convention\n            seis_files = sorted(glob.glob(f\"{family_path}/seis_*_*.npy\"))\n            vel_files = sorted(glob.glob(f\"{family_path}/vel_*_*.npy\"))\n            \n            # Match seismic data with velocity maps\n            for seis_file in seis_files:\n                base_name = os.path.basename(seis_file)\n                vel_name = base_name.replace('seis_', 'vel_')\n                vel_file = os.path.join(family_path, vel_name)\n                \n                if os.path.exists(vel_file):\n                    data_files.append(seis_file)\n                    model_files.append(vel_file)\n        else:\n            # Vel and Style families\n            data_dir = f\"{family_path}/data\"\n            model_dir = f\"{family_path}/model\"\n            \n            if os.path.exists(data_dir) and os.path.exists(model_dir):\n                data_npy_files = sorted(glob.glob(f\"{data_dir}/*.npy\"))\n                model_npy_files = sorted(glob.glob(f\"{model_dir}/*.npy\"))\n                \n                # Match data files with model files\n                for data_file in data_npy_files:\n                    base_name = os.path.basename(data_file)\n                    model_name = base_name.replace('data', 'model')\n                    model_file = os.path.join(model_dir, model_name)\n                    \n                    if os.path.exists(model_file):\n                        data_files.append(data_file)\n                        model_files.append(model_file)\n    \n    return data_files, model_files\n\n# Find all training data files\ntry:\n    train_data_files, train_model_files = find_data_files()\n    print(f\"Found {len(train_data_files)} training data files with matching velocity maps\")\nexcept Exception as e:\n    logger.error(f\"Error finding data files: {e}\")\n    train_data_files, train_model_files = [], []","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Data Visualization","metadata":{}},{"cell_type":"code","source":"# Function to visualize seismic data and velocity maps\ndef visualize_sample(data_file, model_file, sample_idx=0):\n    try:\n        # Load sample\n        seismic_data = np.load(data_file, mmap_mode='r')[sample_idx]\n        velocity_map = np.load(model_file, mmap_mode='r')[sample_idx]\n        \n        # Get shapes\n        print(f\"Seismic data shape: {seismic_data.shape}\")\n        print(f\"Velocity map shape: {velocity_map.shape}\")\n        \n        # Plot\n        fig, axes = plt.subplots(1, 2, figsize=(16, 6))\n        \n        # Plot seismic data (first source, all time steps and receivers)\n        source_idx = 0\n        im1 = axes[0].imshow(seismic_data[source_idx], aspect='auto', cmap='seismic')\n        axes[0].set_title(f'Seismic Data (Source {source_idx})')\n        axes[0].set_xlabel('Receiver Position')\n        axes[0].set_ylabel('Time Step')\n        plt.colorbar(im1, ax=axes[0])\n        \n        # Plot velocity map\n        if len(velocity_map.shape) == 3 and velocity_map.shape[0] == 1:\n            # Handle case where velocity map has a channel dimension\n            velocity_map = velocity_map[0]\n            \n        im2 = axes[1].imshow(velocity_map, cmap='jet')\n        axes[1].set_title('Velocity Map')\n        axes[1].set_xlabel('X Position')\n        axes[1].set_ylabel('Y Position')\n        plt.colorbar(im2, ax=axes[1])\n        \n        plt.tight_layout()\n        plt.show()\n    except Exception as e:\n        logger.error(f\"Error visualizing sample: {e}\")\n        print(f\"Error visualizing sample: {e}\")\n\n# Visualize a sample if data files are available\nif len(train_data_files) > 0 and len(train_model_files) > 0:\n    visualize_sample(train_data_files[0], train_model_files[0])","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Data Preprocessing and Augmentation","metadata":{}},{"cell_type":"code","source":"# Data augmentation functions\nclass WaveformTransforms:\n    @staticmethod\n    def flip_horizontal(seismic, velocity, p=0.5):\n        if np.random.random() < p:\n            seismic = torch.flip(seismic, dims=[-1])  # Flip receivers dimension\n            velocity = torch.flip(velocity, dims=[-1])  # Flip x dimension\n        return seismic, velocity\n    \n    @staticmethod\n    def add_noise(seismic, velocity, p=0.5, noise_level=0.05):\n        if np.random.random() < p:\n            noise = torch.randn_like(seismic) * noise_level\n            seismic = seismic + noise\n        return seismic, velocity\n    \n    @staticmethod\n    def random_crop(seismic, velocity, p=0.5, crop_ratio=0.9):\n        if np.random.random() < p:\n            # Only crop receivers dimension (last dimension)\n            orig_size = seismic.shape[-1]\n            crop_size = int(orig_size * crop_ratio)\n            start = np.random.randint(0, orig_size - crop_size + 1)\n            \n            seismic = seismic[..., start:start+crop_size]\n            velocity = velocity[..., start:start+crop_size]\n            \n            # Resize back to original dimensions\n            seismic = F.interpolate(seismic.unsqueeze(0), size=(seismic.shape[1], orig_size), mode='bilinear').squeeze(0)\n            velocity = F.interpolate(velocity.unsqueeze(0).unsqueeze(0), size=velocity.shape, mode='bilinear').squeeze(0).squeeze(0)\n        return seismic, velocity\n    \n    @staticmethod\n    def apply_all(seismic, velocity):\n        try:\n            seismic, velocity = WaveformTransforms.flip_horizontal(seismic, velocity)\n            seismic, velocity = WaveformTransforms.add_noise(seismic, velocity)\n            return seismic, velocity\n        except Exception as e:\n            logger.error(f\"Error applying transformations: {e}\")\n            return seismic, velocity","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create training and validation datasets\ndef create_train_val_datasets(train_ratio=0.8, use_augmentation=True):\n    try:\n        # Shuffle files while keeping pairs together\n        indices = list(range(len(train_data_files)))\n        random.shuffle(indices)\n        shuffled_data_files = [train_data_files[i] for i in indices]\n        shuffled_model_files = [train_model_files[i] for i in indices]\n        \n        # Split into train and validation\n        split_idx = int(len(shuffled_data_files) * train_ratio)\n        train_data = shuffled_data_files[:split_idx]\n        train_model = shuffled_model_files[:split_idx]\n        val_data = shuffled_data_files[split_idx:]\n        val_model = shuffled_model_files[split_idx:]\n        \n        # Create datasets\n        transform = WaveformTransforms.apply_all if use_augmentation else None\n        train_dataset = WaveformDataset(train_data, train_model, transform=transform)\n        val_dataset = WaveformDataset(val_data, val_model, transform=None)\n        \n        return train_dataset, val_dataset\n    except Exception as e:\n        logger.error(f\"Error creating datasets: {e}\")\n        # Return empty datasets in case of error\n        return None, None","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Model Architecture: InversionNet","metadata":{}},{"cell_type":"code","source":"# InversionNet architecture based on the tutorial\nclass InversionNet(nn.Module):\n    def __init__(self, in_channels=5, out_channels=1):\n        super(InversionNet, self).__init__()\n        \n        # Encoder (downsampling path)\n        self.encoder = nn.Sequential(\n            nn.Conv2d(in_channels, 32, kernel_size=3, padding=1),\n            nn.BatchNorm2d(32),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(32, 64, kernel_size=3, padding=1),\n            nn.BatchNorm2d(64),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n            nn.BatchNorm2d(128),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(128, 256, kernel_size=3, padding=1),\n            nn.BatchNorm2d(256),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(256, 512, kernel_size=3, padding=1),\n            nn.BatchNorm2d(512),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.MaxPool2d(2),\n        )\n        \n        # Decoder (upsampling path)\n        self.decoder = nn.Sequential(\n            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),\n            nn.BatchNorm2d(256),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),\n            nn.BatchNorm2d(128),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),\n            nn.BatchNorm2d(64),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),\n            nn.BatchNorm2d(32),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(32, out_channels, kernel_size=4, stride=2, padding=1),\n        )\n        \n    def forward(self, x):\n        # Process through encoder\n        x = self.encoder(x)\n        \n        # Process through decoder\n        x = self.decoder(x)\n        \n        return x","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. Model Architecture: ConvNeXt-Based Model","metadata":{}},{"cell_type":"code","source":"class ConvNeXtModel(nn.Module):\n    def __init__(self, in_channels=5, out_channels=1):\n        super(ConvNeXtModel, self).__init__()\n        \n        # Use ConvNeXt Tiny as backbone\n        self.backbone = timm.create_model('convnext_tiny', pretrained=True, in_chans=in_channels, features_only=True)\n        \n        # Decoder for upsampling to 70x70\n        self.decoder = nn.Sequential(\n            nn.ConvTranspose2d(768, 384, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(384),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(384, 192, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(192),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(192, 96, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(96),\n            nn.ReLU(inplace=True),\n            \n            nn.ConvTranspose2d(96, 48, kernel_size=3, stride=2, padding=1),\n            nn.BatchNorm2d(48),\n            nn.ReLU(inplace=True),\n            \n            nn.Upsample(size=(70, 70), mode='bilinear', align_corners=True),\n            nn.Conv2d(48, out_channels, kernel_size=1)\n        )\n        \n    def forward(self, x):\n        # Process through backbone - get last feature map\n        features = self.backbone(x)[-1]\n        output = self.decoder(features)\n        return output","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 10. Physics-Guided Loss Function","metadata":{}},{"cell_type":"code","source":"# Physics-guided loss function\nclass PhysicsGuidedLoss(nn.Module):\n    def __init__(self, alpha=1.0, beta=0.5, gamma=0.1):\n        super(PhysicsGuidedLoss, self).__init__()\n        self.alpha = alpha  # Weight for MAE loss\n        self.beta = beta    # Weight for gradient loss\n        self.gamma = gamma  # Weight for smoothness loss\n        \n    def forward(self, pred, target):\n        try:\n            # MAE loss\n            mae_loss = F.l1_loss(pred, target)\n            \n            # Gradient loss (physics-based)\n            # Calculate gradients in x and y directions\n            pred_dx = pred[:, :, :, 1:] - pred[:, :, :, :-1]\n            pred_dy = pred[:, :, 1:, :] - pred[:, :, :-1, :]\n            target_dx = target[:, :, :, 1:] - target[:, :, :, :-1]\n            target_dy = target[:, :, 1:, :] - target[:, :, :-1, :]\n            \n            grad_loss_x = F.l1_loss(pred_dx, target_dx)\n            grad_loss_y = F.l1_loss(pred_dy, target_dy)\n            grad_loss = grad_loss_x + grad_loss_y\n            \n            # Smoothness loss (physics-based)\n            # Encourage smooth velocity transitions\n            smooth_loss_x = torch.mean(torch.abs(pred_dx))\n            smooth_loss_y = torch.mean(torch.abs(pred_dy))\n            smooth_loss = smooth_loss_x + smooth_loss_y\n            \n            # Combined loss\n            total_loss = self.alpha * mae_loss + self.beta * grad_loss + self.gamma * smooth_loss\n            \n            return total_loss\n        except Exception as e:\n            logger.error(f\"Error in physics-guided loss calculation: {e}\")\n            # Fallback to simple MAE loss in case of error\n            return F.l1_loss(pred, target)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 11. PyTorch Lightning Model Wrapper","metadata":{}},{"cell_type":"code","source":"# PyTorch Lightning model wrapper\nclass WaveformInversionModel(pl.LightningModule):\n    def __init__(self, model_type='convnext', learning_rate=1e-4):\n        super(WaveformInversionModel, self).__init__()\n        \n        # Choose model architecture\n        if model_type == 'convnext':\n            self.model = ConvNeXtModel(in_channels=5, out_channels=1)\n        else:  # Default to InversionNet\n            self.model = InversionNet(in_channels=5, out_channels=1)\n        \n        # Loss function\n        self.loss_fn = PhysicsGuidedLoss()\n        self.learning_rate = learning_rate\n        \n    def forward(self, x):\n        return self.model(x)\n    \n    def training_step(self, batch, batch_idx):\n        try:\n            seismic = batch['seismic']\n            velocity = batch['velocity']\n            \n            # Forward pass\n            pred_velocity = self(seismic)\n            \n            # Calculate loss\n            loss = self.loss_fn(pred_velocity, velocity)\n            \n            # Log metrics\n            self.log('train_loss', loss, prog_bar=True)\n            \n            return loss\n        except Exception as e:\n            logger.error(f\"Error in training step: {e}\")\n            # Return a dummy loss in case of error to avoid training failure\n            return torch.tensor(0.0, requires_grad=True)\n    \n    def validation_step(self, batch, batch_idx):\n        try:\n            seismic = batch['seismic']\n            velocity = batch['velocity']\n            \n            # Forward pass\n            pred_velocity = self(seismic)\n            \n            # Calculate loss\n            loss = self.loss_fn(pred_velocity, velocity)\n            \n            # Calculate MAE for monitoring\n            mae = F.l1_loss(pred_velocity, velocity)\n            \n            # Log metrics\n            self.log('val_loss', loss, prog_bar=True)\n            self.log('val_mae', mae, prog_bar=True)\n            \n            return {'val_loss': loss, 'val_mae': mae}\n        except Exception as e:\n            logger.error(f\"Error in validation step: {e}\")\n            return {'val_loss': torch.tensor(0.0), 'val_mae': torch.tensor(0.0)}\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, mode='min', factor=0.5, patience=5, verbose=True\n        )\n        \n        return {\n            'optimizer': optimizer,\n            'lr_scheduler': scheduler,\n            'monitor': 'val_loss'\n        }","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 12. Model Training with Version-Aware Precision Configuration","metadata":{}},{"cell_type":"code","source":"# Function to train the model\ndef train_model(model_type='convnext', batch_size=8, max_epochs=30):\n    try:\n        # Create datasets\n        train_dataset, val_dataset = create_train_val_datasets(train_ratio=0.8, use_augmentation=True)\n        \n        if train_dataset is None or val_dataset is None:\n            logger.error(\"Failed to create datasets\")\n            return None, None\n        \n        # Optimize batch size based on available memory\n        if device_type == \"gpu\":\n            try:\n                mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9\n                if mem_gb < 4:\n                    batch_size = max(1, batch_size // 4)\n                    logger.info(f\"Reduced batch size to {batch_size} due to limited GPU memory\")\n                elif mem_gb > 16:\n                    batch_size = batch_size * 2\n                    logger.info(f\"Increased batch size to {batch_size} due to large GPU memory\")\n            except:\n                pass\n        elif device_type == \"cpu\":\n            batch_size = max(1, batch_size // 8)\n            logger.info(f\"Reduced batch size to {batch_size} for CPU training\")\n        \n        # Create data loaders\n        train_loader = DataLoader(\n            train_dataset,\n            batch_size=batch_size,\n            shuffle=True,\n            num_workers=4 if device_type != \"cpu\" else 0,\n            pin_memory=device_type == \"gpu\"\n        )\n        \n        val_loader = DataLoader(\n            val_dataset,\n            batch_size=batch_size,\n            shuffle=False,\n            num_workers=4 if device_type != \"cpu\" else 0,\n            pin_memory=device_type == \"gpu\"\n        )\n        \n        # Create model\n        model = WaveformInversionModel(model_type=model_type)\n        \n        # Callbacks\n        checkpoint_callback = ModelCheckpoint(\n            dirpath=OUTPUT_PATH,\n            filename=f'{model_type}_best_model',\n            monitor='val_mae',\n            mode='min',\n            save_top_k=1,\n            verbose=True\n        )\n        \n        early_stop_callback = EarlyStopping(\n            monitor='val_mae',\n            patience=10,\n            mode='min',\n            verbose=True\n        )\n        \n        # Get optimal precision configuration based on PyTorch Lightning version and hardware\n        precision_config = precision\n        logger.info(f\"Using precision configuration: {precision_config}\")\n        \n        # Trainer\n        trainer = pl.Trainer(\n            max_epochs=max_epochs,\n            callbacks=[checkpoint_callback, early_stop_callback],\n            accelerator='auto',  # Automatically choose GPU/TPU if available\n            devices=1,\n            precision=precision_config,  # Use version-aware precision config\n            log_every_n_steps=10\n        )\n        \n        # Train model\n        trainer.fit(model, train_loader, val_loader)\n        \n        return model, checkpoint_callback.best_model_path\n    except Exception as e:\n        logger.error(f\"Error in model training: {e}\")\n        return None, None","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train model if data is available\nif len(train_data_files) > 0 and len(train_model_files) > 0:\n    try:\n        model, best_model_path = train_model(model_type='convnext', batch_size=8, max_epochs=30)\n        if best_model_path:\n            print(f\"Best model saved at: {best_model_path}\")\n        else:\n            print(\"Training did not complete successfully\")\n    except Exception as e:\n        logger.error(f\"Error during model training: {e}\")\n        print(f\"Error during model training: {e}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 13. Model Ensemble","metadata":{}},{"cell_type":"code","source":"# Create an ensemble of models\ndef create_model_ensemble():\n    try:\n        # Train multiple models with different architectures/configurations\n        models = []\n        \n        # Train ConvNeXt model\n        logger.info(\"Training ConvNeXt model for ensemble\")\n        convnext_model, _ = train_model(model_type='convnext', batch_size=8, max_epochs=20)\n        if convnext_model is not None:\n            models.append(convnext_model)\n        \n        # Train InversionNet model\n        logger.info(\"Training InversionNet model for ensemble\")\n        inversionnet_model, _ = train_model(model_type='inversionnet', batch_size=8, max_epochs=20)\n        if inversionnet_model is not None:\n            models.append(inversionnet_model)\n        \n        logger.info(f\"Created ensemble with {len(models)} models\")\n        return models\n    except Exception as e:\n        logger.error(f\"Error creating model ensemble: {e}\")\n        return []","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 14. Inference and Submission","metadata":{}},{"cell_type":"code","source":"# Function to load test data\ndef load_test_data():\n    try:\n        test_files = sorted(glob.glob(f\"{TEST_PATH}/*.npy\"))\n        print(f\"Found {len(test_files)} test files\")\n        return test_files\n    except Exception as e:\n        logger.error(f\"Error loading test data: {e}\")\n        return []","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to make predictions using a single model\ndef predict_with_model(model, test_files, batch_size=4):\n    try:\n        model.eval()\n        predictions = {}\n        \n        with torch.no_grad():\n            for test_file in tqdm(test_files, desc=\"Processing test files\"):\n                # Extract file ID from filename\n                file_id = os.path.basename(test_file).split('.')[0]\n                \n                # Load test data\n                test_data = np.load(test_file)\n                \n                # Process in batches to save memory\n                all_preds = []\n                for i in range(0, test_data.shape[0], batch_size):\n                    batch = test_data[i:i+batch_size]\n                    batch_tensor = torch.from_numpy(batch.astype(np.float32))\n                    \n                    # Normalize\n                    batch_tensor = (batch_tensor - batch_tensor.mean(dim=(1, 2), keepdim=True)) / \\\n                                   (batch_tensor.std(dim=(1, 2), keepdim=True) + 1e-8)\n                    \n                    # Move to device\n                    batch_tensor = batch_tensor.to(device)\n                    \n                    # Predict\n                    preds = model(batch_tensor)\n                    \n                    # Move back to CPU and convert to numpy\n                    preds = preds.cpu().numpy()\n                    all_preds.append(preds)\n                \n                # Combine batch predictions\n                file_preds = np.concatenate(all_preds, axis=0)\n                predictions[file_id] = file_preds\n        \n        return predictions\n    except Exception as e:\n        logger.error(f\"Error in model prediction: {e}\")\n        return {}","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to make predictions using an ensemble of models\ndef predict_with_ensemble(models, test_files, batch_size=4):\n    try:\n        for model in models:\n            model.eval()\n        \n        ensemble_predictions = {}\n        \n        with torch.no_grad():\n            for test_file in tqdm(test_files, desc=\"Processing test files with ensemble\"):\n                # Extract file ID from filename\n                file_id = os.path.basename(test_file).split('.')[0]\n                \n                # Load test data\n                test_data = np.load(test_file)\n                \n                # Process in batches to save memory\n                all_preds = []\n                for i in range(0, test_data.shape[0], batch_size):\n                    batch = test_data[i:i+batch_size]\n                    batch_tensor = torch.from_numpy(batch.astype(np.float32))\n                    \n                    # Normalize\n                    batch_tensor = (batch_tensor - batch_tensor.mean(dim=(1, 2), keepdim=True)) / \\\n                                   (batch_tensor.std(dim=(1, 2), keepdim=True) + 1e-8)\n                    \n                    # Move to device\n                    batch_tensor = batch_tensor.to(device)\n                    \n                    # Get predictions from all models\n                    model_preds = []\n                    for model in models:\n                        pred = model(batch_tensor)\n                        model_preds.append(pred.cpu().numpy())\n                    \n                    # Average predictions\n                    ensemble_pred = np.mean(model_preds, axis=0)\n                    all_preds.append(ensemble_pred)\n                \n                # Combine batch predictions\n                file_preds = np.concatenate(all_preds, axis=0)\n                ensemble_predictions[file_id] = file_preds\n        \n        return ensemble_predictions\n    except Exception as e:\n        logger.error(f\"Error in ensemble prediction: {e}\")\n        return {}","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to create submission file\ndef create_submission(predictions, output_file=\"submission.csv\"):\n    try:\n        # Load sample submission to get the format\n        sample_sub = pd.read_csv(SAMPLE_SUB_PATH)\n        \n        # Create a copy for our submission\n        submission = sample_sub.copy()\n        \n        # Fill in predictions\n        for file_id, preds in predictions.items():\n            # For each y position in the velocity map\n            for y_pos in range(preds.shape[1]):\n                # Get the row identifier\n                row_id = f\"{file_id}_y_{y_pos}\"\n                \n                # Get the velocity values for odd-valued columns (x positions)\n                # According to the competition format\n                for x_pos in range(1, preds.shape[2], 2):\n                    col_name = f\"x_{x_pos}\"\n                    if col_name in submission.columns:\n                        submission.loc[submission['oid_ypos'] == row_id, col_name] = preds[0, y_pos, x_pos]\n        \n        # Save submission file\n        submission.to_csv(os.path.join(OUTPUT_PATH, output_file), index=False)\n        print(f\"Submission saved to {os.path.join(OUTPUT_PATH, output_file)}\")\n        \n        return submission\n    except Exception as e:\n        logger.error(f\"Error creating submission file: {e}\")\n        return None","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Make predictions and create submission\ndef run_inference_and_submit():\n    try:\n        # Load test data\n        test_files = load_test_data()\n        \n        if not test_files:\n            logger.error(\"No test files found\")\n            return None\n        \n        # Option 1: Use a single model\n        if os.path.exists(best_model_path):\n            logger.info(f\"Loading best model from {best_model_path}\")\n            model = WaveformInversionModel.load_from_checkpoint(best_model_path)\n            model = model.to(device)\n            predictions = predict_with_model(model, test_files)\n            \n            # Option 2: Use an ensemble (uncomment to use)\n            # models = create_model_ensemble()\n            # models = [model.to(device) for model in models]\n            # predictions = predict_with_ensemble(models, test_files)\n            \n            # Create submission\n            submission = create_submission(predictions)\n            \n            return submission\n        else:\n            logger.error(f\"Best model path {best_model_path} does not exist\")\n            return None\n    except Exception as e:\n        logger.error(f\"Error in inference and submission: {e}\")\n        return None","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run inference and create submission if test data is available\nif os.path.exists(TEST_PATH) and os.path.exists(SAMPLE_SUB_PATH):\n    try:\n        submission = run_inference_and_submit()\n        if submission is not None:\n            print(\"Submission created successfully!\")\n        else:\n            print(\"Failed to create submission\")\n    except Exception as e:\n        logger.error(f\"Error during inference and submission: {e}\")\n        print(f\"Error during inference and submission: {e}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 15. Conclusion\n\nThis notebook presents a comprehensive solution for the Yale/UNC-CH Geophysical Waveform Inversion competition, featuring:\n\n1. **Memory-efficient data handling** for the large dataset\n2. **GPU/TPU acceleration** with robust device detection and error handling\n3. **Version-aware PyTorch Lightning precision configuration** that works across different environments\n4. **Advanced architectures** including ConvNeXt and InversionNet\n5. **Physics-guided machine learning** with specialized loss functions\n6. **Model ensembling** for improved prediction accuracy\n\nThe approach is designed to work within memory constraints while still achieving competitive results, helping you progress from Expert to Grandmaster status in the Competitions category on Kaggle.","metadata":{}}]}