{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":39763,"databundleVersionId":11756775,"sourceType":"competition"}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Cell 1: System Resource Check and Configuration Validation\nimport os\nimport sys\nimport torch\nimport numpy as np\nfrom pathlib import Path\n\n# Check Python version and critical libraries\nprint(f\"Python version: {sys.version}\")\nprint(f\"PyTorch version: {torch.__version__}\")\nprint(f\"NumPy version: {np.__version__}\")\n\n# Check system resources\nprint(\"\\n=== System Resources ===\")\n\n# Check CPU\ntry:\n import psutil\n cpu_count = psutil.cpu_count(logical=True)\n cpu_physical = psutil.cpu_count(logical=False)\n memory_gb = psutil.virtual_memory().total / (1024**3)\n available_memory_gb = psutil.virtual_memory().available / (1024**3)\n print(f\"CPU: {cpu_physical} physical cores, {cpu_count} logical cores\")\n print(f\"Memory: {memory_gb:.1f} GB total, {available_memory_gb:.1f} GB available\")\n \n # Calculate optimal worker count for DataLoader\n optimal_workers = min(16, max(4, cpu_count - 2)) # Cap at 16, but at least 4\n print(f\"Recommended num_workers for DataLoader: {optimal_workers}\")\nexcept ImportError:\n cpu_count = os.cpu_count()\n print(f\"CPU: {cpu_count} cores detected\")\n optimal_workers = min(16, max(4, cpu_count - 2))\n print(f\"Recommended num_workers for DataLoader: {optimal_workers}\")\n\n# Check GPU\nif torch.cuda.is_available():\n device_count = torch.cuda.device_count()\n print(f\"GPU: {device_count} device(s) available\")\n for i in range(device_count):\n gpu_name = torch.cuda.get_device_name(i)\n gpu_memory = torch.cuda.get_device_properties(i).total_memory / (1024**3)\n print(f\" - GPU {i}: {gpu_name}, {gpu_memory:.1f} GB memory\")\n \n # Test GPU speed (optional)\n print(\"\\nGPU quick benchmark...\")\n x = torch.randn(1000, 1000, device=\"cuda\")\n start = torch.cuda.Event(enable_timing=True)\n end = torch.cuda.Event(enable_timing=True)\n \n start.record()\n result = torch.matmul(x, x)\n end.record()\n torch.cuda.synchronize()\n print(f\"Matrix multiplication time: {start.elapsed_time(end):.2f} ms\")\n \n # Clear memory\n del x, result, start, end\n torch.cuda.empty_cache()\n \n # GPU is available, so provide optimal settings\n print(\"\\n=== Recommended GPU Settings ===\")\n print(\"- mixed_precision: True\")\n print(\"- pin_memory: True\")\n print(\"- persistent_workers: True\")\n print(f\"- prefetch_factor: {min(4, optimal_workers // 2)}\")\nelse:\n print(\"GPU: Not available, using CPU only\")\n print(\"\\nWARNING: Training will be much slower without GPU acceleration\")\n\n# Check where we are\nprint(\"\\n=== Current Directory ===\")\nprint(f\"Current path: {os.getcwd()}\")\nprint(f\"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}\")\n\nprint(\"\\n=== Configuration Checklist ===\")\nchecks = [\n \"Set num_workers to optimal value\",\n \"Enable mixed_precision=True for GPU training\",\n \"Use 200 epochs for full training\",\n \"Increase early_stopping patience for long runs\",\n \"Update post_process_binary function for dimension handling\", \n \"Enable pin_memory and persistent_workers in DataLoader\",\n]\n\nfor i, check in enumerate(checks, 1):\n print(f\"{i}. [ ] {check}\")\n\nprint(\"\\nReady to begin 200-epoch training. Make sure all checkboxes are addressed.\")\nprint(\"Estimated training time: ~6-7 hours with GPU acceleration and optimized settings.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:11.659341Z","iopub.execute_input":"2025-05-12T17:58:11.659597Z","iopub.status.idle":"2025-05-12T17:58:16.747299Z","shell.execute_reply.started":"2025-05-12T17:58:11.659574Z","shell.execute_reply":"2025-05-12T17:58:16.746540Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 2: Import Libraries and Major Configuration Block\nimport os\nimport time\nimport glob\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom datetime import datetime\nfrom sklearn.model_selection import train_test_split\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\n\n# Set random seeds for reproducibility\nSEED = 42\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n torch.cuda.manual_seed_all(SEED)\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n\n# ==== MASTER CONFIGURATION BLOCK ====\nCONFIG = {\n # Runtime configuration\n 'approach': 'physics_guided', # Options: 'thresholding', 'physics_guided', 'feature_detection', 'geo-aware'\n 'device': torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n 'mixed_precision': True, # Use mixed precision training for GPU acceleration\n 'save_models': True, # Save models during training\n 'compare_approaches': False, # Run multiple approaches and compare\n 'generate_submission': False, # Toggle submission file generation\n 'run_analysis_tools': True,\n \n 'in_channels': 5,\n 'out_channels': 1,\n 'hidden_dim': 64,\n \n # Data configuration \n 'data_dir': Path(\"/kaggle/input/waveform-inversion/train_samples\"),\n 'test_dir': Path(\"/kaggle/input/waveform-inversion/test\"),\n 'output_dir': Path(\"./outputs\"),\n \n # Dataset parameters\n 'val_size': 0.15, # Validation set size\n 'batch_size': 4, # GPU-friendly batch size\n 'num_workers': 4, # DataLoader workers\n \n # Training parameters\n 'num_epochs': 50,\n 'learning_rate': 5e-5, # Careful with numeric instability\n 'weight_decay': 1e-5,\n 'early_stopping': 10,\n 'scheduler_patience': 5,\n 'prefetch_factor': 2, # Controls batch prefetching per worker\n \n 'grad_clip_value': 1.0, # NaN issues earlier\n 'warmup_epochs': 5, # NaN issues earlier\n\n \n # Model parameters\n 'in_channels': 5,\n 'out_channels': 1,\n 'hidden_dim': 64, # Base dimensionality for models\n \n # Physics-guided parameters\n 'wave_eq_weight': 0.15,\n 'slowness_weight': 0.1,\n 'layering_weight': 0.05,\n 'contrast_weight': 0.2,\n \n # Feature detection parameters; will revist as these may be too small\n 'salt_weight': 0.2,\n 'fault_weight': 0.1,\n 'layer_weight': 0.15,\n 'geological_constraint_weight': 0.3,\n \n # Thresholding parameters\n 'threshold_method': 'otsu', # Options: 'otsu', 'mean', 'adaptive'\n 'edge_enhancement': 1.5,\n 'use_morphology': True,\n \n # Submission parameters\n 'ensemble_submission': False, # Use ensemble of multiple models\n 'post_process': True, # Apply post-processing to predictions\n 'submission_path': \"submission.csv\",\n}\n\n# Create output directory\nos.makedirs(CONFIG['output_dir'], exist_ok=True)\nos.makedirs(CONFIG['output_dir'] / 'models', exist_ok=True)\nos.makedirs(CONFIG['output_dir'] / 'visualizations', exist_ok=True)\n\n# Print configuration summary\nprint(f\"Running with approach: {CONFIG['approach']}\")\nprint(f\"Device: {CONFIG['device']}\")\nif torch.cuda.is_available():\n print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n print(f\"Mixed precision: {CONFIG['mixed_precision']}\")\n\n# Initialize experiment name with timestamp\ntimestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\nexperiment_name = f\"{CONFIG['approach']}_{timestamp}\"\nprint(f\"Experiment name: {experiment_name}\")\n\n# Save configuration\nconfig_path = CONFIG['output_dir'] / f\"config_{experiment_name}.txt\"\nwith open(config_path, 'w') as f:\n for key, value in CONFIG.items():\n f.write(f\"{key}: {value}\\n\")\n\nCONFIG['confirmed'] = True\nCONFIG['experiment_name'] = experiment_name\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:16.748448Z","iopub.execute_input":"2025-05-12T17:58:16.748833Z","iopub.status.idle":"2025-05-12T17:58:17.727918Z","shell.execute_reply.started":"2025-05-12T17:58:16.748803Z","shell.execute_reply":"2025-05-12T17:58:17.727100Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 3: Data Loading and Dataset Classes\nclass SeismicDataset(Dataset):\n \"\"\"Dataset class for seismic data with flexible input/output handling\"\"\"\n def __init__(self, input_files, output_files=None, transform=None, \n normalize=True, gain=True, augment=False, \n binary_threshold=None, feature_mode=False):\n \"\"\"\n Args:\n input_files: List of seismic data file paths\n output_files: List of velocity map file paths (None for test set)\n transform: Optional transforms to apply\n normalize: Whether to normalize data\n gain: Whether to apply time-dependent gain\n augment: Whether to apply data augmentation\n binary_threshold: If not None, convert velocity maps to binary using this threshold\n feature_mode: If True, extract geological features from velocity maps\n \"\"\"\n self.input_files = input_files\n self.output_files = output_files\n self.transform = transform\n self.normalize = normalize\n self.gain = gain\n self.augment = augment\n self.binary_threshold = binary_threshold\n self.feature_mode = feature_mode\n \n # Build index map for efficient access\n self.index_map = []\n for i, f in enumerate(self.input_files):\n try:\n data = np.load(f, mmap_mode='r')\n for j in range(data.shape[0]):\n self.index_map.append((i, j))\n except Exception as e:\n print(f\"Error loading {f}: {e}\")\n \n # Calculate statistics if normalizing\n if normalize and output_files:\n self.calc_stats()\n \n def calc_stats(self, max_files=10, max_samples=100):\n \"\"\"Calculate dataset statistics for normalization\"\"\"\n print(\"Calculating dataset statistics...\")\n # Input statistics\n in_samples = []\n for i in range(min(max_files, len(self.input_files))):\n try:\n data = np.load(self.input_files[i], mmap_mode='r')\n idx = np.random.choice(data.shape[0], min(max_samples, data.shape[0]), replace=False)\n in_samples.append(data[idx])\n except Exception as e:\n print(f\"Error in stats calculation for {self.input_files[i]}: {e}\")\n \n if in_samples:\n in_array = np.concatenate(in_samples, axis=0)\n self.in_mean = float(np.mean(in_array))\n self.in_std = float(np.std(in_array))\n else:\n self.in_mean, self.in_std = 0.0, 1.0\n \n # Output statistics (if available)\n if self.output_files:\n out_samples = []\n for i in range(min(max_files, len(self.output_files))):\n try:\n data = np.load(self.output_files[i], mmap_mode='r')\n idx = np.random.choice(data.shape[0], min(max_samples, data.shape[0]), replace=False)\n \n # Handle different output shapes\n if len(data.shape) == 4: # [batch, channel, height, width]\n out_samples.append(data[idx])\n elif len(data.shape) == 3: # [batch, height, width]\n out_samples.append(data[idx, np.newaxis])\n except Exception as e:\n print(f\"Error in stats calculation for {self.output_files[i]}: {e}\")\n \n if out_samples:\n out_array = np.concatenate(out_samples, axis=0)\n self.out_mean = float(np.mean(out_array))\n self.out_std = float(np.std(out_array))\n \n # For binary threshold detection\n if self.binary_threshold is None:\n try:\n from skimage.filters import threshold_otsu\n self.auto_threshold = threshold_otsu(out_array.flatten())\n except:\n self.auto_threshold = self.out_mean\n else:\n self.auto_threshold = self.binary_threshold\n else:\n self.out_mean, self.out_std = 0.0, 1.0\n self.auto_threshold = 0.5\n \n print(f\"Input stats: mean={self.in_mean:.4f}, std={self.in_std:.4f}\")\n if self.output_files:\n print(f\"Output stats: mean={self.out_mean:.4f}, std={self.out_std:.4f}\")\n print(f\"Auto threshold: {self.auto_threshold:.4f}\")\n \n def __len__(self):\n return len(self.index_map)\n \n def apply_gain(self, x):\n \"\"\"Apply time-dependent gain to seismic data\"\"\"\n # Time-dependent gain increases amplitude with time\n time_steps = x.shape[1]\n time = np.linspace(0, 1, time_steps)\n gain = (time ** 2)[:, np.newaxis] # Square gain with time\n \n # Apply to each channel\n gained_x = x.copy()\n for c in range(x.shape[0]):\n gained_x[c] = x[c] * gain\n \n return gained_x\n \n def extract_geological_features(self, y):\n \"\"\"Extract geological features from velocity map\"\"\"\n # 1. Salt body detection (high velocity regions)\n salt_mask = (y > self.auto_threshold + 0.2 * self.out_std).astype(np.float32)\n \n # 2. Fault detection (using edge detection)\n from scipy import ndimage\n sobel_x = ndimage.sobel(y, axis=1)\n sobel_y = ndimage.sobel(y, axis=0)\n grad_mag = np.sqrt(sobel_x**2 + sobel_y**2)\n grad_mag = grad_mag / (np.max(grad_mag) + 1e-8)\n fault_mask = (grad_mag > 0.3).astype(np.float32)\n \n # 3. Layer detection (horizontal structures)\n # Use horizontal gradient\n layer_edges = np.abs(sobel_y)\n layer_edges = layer_edges / (np.max(layer_edges) + 1e-8)\n layer_mask = (layer_edges > 0.2).astype(np.float32)\n \n # Stack into feature channels [salt, fault, layer]\n features = np.stack([salt_mask, fault_mask, layer_mask], axis=0)\n return features\n \n def apply_augmentation(self, x, y=None):\n \"\"\"Apply data augmentation\"\"\"\n # Flip horizontally with 50% probability\n if np.random.random() > 0.5:\n x = np.flip(x, axis=2).copy() # Flip along receiver dimension\n if y is not None:\n y = np.flip(y, axis=-1).copy() # Flip along width\n \n # Add small random noise with 30% probability\n if np.random.random() > 0.7:\n noise_level = np.random.uniform(0, 0.05)\n x = x + np.random.normal(0, noise_level, size=x.shape)\n \n return x, y\n \n def __getitem__(self, idx):\n \"\"\"Get a single sample from the dataset\"\"\"\n file_idx, sample_idx = self.index_map[idx]\n \n # Load input data\n x = np.load(self.input_files[file_idx], mmap_mode='r')[sample_idx].astype(np.float32)\n \n # Apply gain if requested\n if self.gain:\n x = self.apply_gain(x)\n \n # Load output data if available (for training)\n if self.output_files:\n y_data = np.load(self.output_files[file_idx], mmap_mode='r')[sample_idx]\n \n # Handle different output shapes\n if len(y_data.shape) == 2: # (height, width)\n y = y_data\n elif len(y_data.shape) == 3 and y_data.shape[0] == 1: # (1, height, width)\n y = y_data[0]\n elif len(y_data.shape) > 2: # Multi-dimensional\n y = y_data.reshape(-1, y_data.shape[-2], y_data.shape[-1])[0]\n else:\n raise ValueError(f\"Unexpected velocity map shape: {y_data.shape}\")\n \n y = y.astype(np.float32)\n \n # Convert to binary if threshold is provided\n if self.binary_threshold is not None:\n threshold = self.binary_threshold if self.binary_threshold > 0 else self.auto_threshold\n y = (y > threshold).astype(np.float32)\n \n # Extract geological features if in feature mode\n if self.feature_mode:\n y = self.extract_geological_features(y)\n else:\n # For test set, create a dummy y\n y = np.zeros((1, x.shape[-1], x.shape[-1]), dtype=np.float32)\n \n # Apply data augmentation if enabled\n if self.augment:\n x, y = self.apply_augmentation(x, y)\n \n # Normalize if enabled\n if self.normalize:\n x = (x - self.in_mean) / (self.in_std + 1e-6)\n if self.output_files and not self.feature_mode and self.binary_threshold is None:\n y = (y - self.out_mean) / (self.out_std + 1e-6)\n \n # Convert to PyTorch tensors\n x_tensor = torch.from_numpy(x).float()\n \n # Ensure y has proper dimensions for loss functions\n if not self.feature_mode and len(y.shape) == 2:\n y = y[np.newaxis, ...] # Add channel dimension if needed\n y_tensor = torch.from_numpy(y).float()\n \n # Apply transforms if provided\n if self.transform:\n x_tensor, y_tensor = self.transform(x_tensor, y_tensor)\n \n return x_tensor, y_tensor\n\n\ndef load_and_prepare_data(config):\n \"\"\"Prepare datasets and dataloaders based on configuration\"\"\"\n data_dir = config['data_dir']\n \n # Find input and output files\n print(\"Loading competition dataset...\")\n input_files = sorted([f for f in data_dir.rglob(\"*.npy\") \n if 'seis' in f.name or 'data' in f.name])\n output_files = [Path(str(f).replace(\"seis\", \"vel\").replace(\"data\", \"model\")) \n for f in input_files]\n \n print(f\"Found {len(input_files)} input files and {len(output_files)} output files\")\n \n # Verify files exist\n input_files = [f for f in input_files if f.exists()]\n output_files = [f for f in output_files if f.exists()]\n if len(input_files) != len(output_files):\n print(f\"Warning: Mismatched file counts - {len(input_files)} inputs, {len(output_files)} outputs\")\n # Keep only matching pairs\n input_basenames = [f.stem for f in input_files]\n output_basenames = [f.stem.replace(\"vel\", \"seis\").replace(\"model\", \"data\") for f in output_files]\n common_basenames = set(input_basenames) & set(output_basenames)\n input_files = [f for f in input_files if f.stem in common_basenames]\n output_files = [f for f in output_files if f.stem.replace(\"vel\", \"seis\").replace(\"model\", \"data\") in common_basenames]\n print(f\"Kept {len(input_files)} matching file pairs\")\n \n # Split into train and validation sets\n train_in, val_in, train_out, val_out = train_test_split(\n input_files, output_files, test_size=config['val_size'], random_state=SEED\n )\n \n print(f\"Training set: {len(train_in)} files\")\n print(f\"Validation set: {len(val_in)} files\")\n \n # Determine binary threshold and feature mode based on approach\n binary_threshold = None\n feature_mode = False\n \n if config['approach'] == 'thresholding':\n binary_threshold = 0 # Use auto-detected threshold\n elif config['approach'] == 'feature_detection':\n feature_mode = True\n \n # Create datasets\n train_ds = SeismicDataset(\n train_in, train_out,\n normalize=True,\n gain=True,\n augment=True,\n binary_threshold=binary_threshold,\n feature_mode=feature_mode\n )\n \n val_ds = SeismicDataset(\n val_in, val_out,\n normalize=True,\n gain=True,\n augment=False,\n binary_threshold=binary_threshold,\n feature_mode=feature_mode\n )\n \n# Optimized DataLoader creation\n train_loader = DataLoader(\n train_ds,\n batch_size=config['batch_size'],\n shuffle=True,\n num_workers=config['num_workers'],\n pin_memory=torch.cuda.is_available(),\n persistent_workers=config['num_workers'] > 0,\n prefetch_factor=config.get('prefetch_factor', 2)\n )\n \n val_loader = DataLoader(\n val_ds,\n batch_size=config['batch_size'],\n shuffle=False,\n num_workers=config['num_workers'],\n pin_memory=torch.cuda.is_available(),\n persistent_workers=config['num_workers'] > 0,\n prefetch_factor=config.get('prefetch_factor', 2)\n )\n \n # Get dataset statistics for normalization and thresholding\n stats = {\n 'input_mean': train_ds.in_mean,\n 'input_std': train_ds.in_std,\n 'output_mean': train_ds.out_mean if hasattr(train_ds, 'out_mean') else 0.0,\n 'output_std': train_ds.out_std if hasattr(train_ds, 'out_std') else 1.0,\n 'auto_threshold': train_ds.auto_threshold if hasattr(train_ds, 'auto_threshold') else 0.5\n }\n \n return train_loader, val_loader, stats","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.728901Z","iopub.execute_input":"2025-05-12T17:58:17.729327Z","iopub.status.idle":"2025-05-12T17:58:17.756912Z","shell.execute_reply.started":"2025-05-12T17:58:17.729300Z","shell.execute_reply":"2025-05-12T17:58:17.756211Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 4: Model Architectures\n# 1. Physics-Guided Model for Velocity Prediction\nclass PhysicsGuidedUNet(nn.Module):\n \"\"\"U-Net with physics-guided components for velocity prediction\"\"\"\n def __init__(self, in_channels=5, out_channels=1, hidden_dim=64):\n super().__init__()\n \n # Encoder blocks\n self.enc1 = self._make_encoder_block(in_channels, hidden_dim)\n self.enc2 = self._make_encoder_block(hidden_dim, hidden_dim*2)\n self.enc3 = self._make_encoder_block(hidden_dim*2, hidden_dim*4)\n self.enc4 = self._make_encoder_block(hidden_dim*4, hidden_dim*8)\n \n # Bottleneck\n self.bottleneck = nn.Sequential(\n nn.Conv2d(hidden_dim*8, hidden_dim*16, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*16),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(hidden_dim*16, hidden_dim*16, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*16),\n nn.LeakyReLU(0.2, inplace=True)\n )\n \n # Decoder blocks with skip connections\n self.dec4 = self._make_decoder_block(hidden_dim*16 + hidden_dim*8, hidden_dim*8)\n self.dec3 = self._make_decoder_block(hidden_dim*8 + hidden_dim*4, hidden_dim*4)\n self.dec2 = self._make_decoder_block(hidden_dim*4 + hidden_dim*2, hidden_dim*2)\n self.dec1 = self._make_decoder_block(hidden_dim*2 + hidden_dim, hidden_dim)\n \n # Final output layer\n self.final = nn.Sequential(\n nn.Conv2d(hidden_dim, out_channels, kernel_size=1),\n nn.Sigmoid() # Output in [0,1] range\n )\n \n # Pooling layer for downsampling\n self.pool = nn.MaxPool2d(2)\n \n # Physics-guided layers\n # Edge detection for geological boundaries\n self.edge_detector = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)\n self.edge_detector.weight.data = torch.tensor([\n [[-1, -1, -1],\n [2, 2, 2],\n [-1, -1, -1]]\n ], dtype=torch.float32).view(1, 1, 3, 3).repeat(out_channels, out_channels, 1, 1)\n self.edge_detector.weight.requires_grad = False\n \n def _make_encoder_block(self, in_channels, out_channels):\n return nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.BatchNorm2d(out_channels),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n nn.BatchNorm2d(out_channels),\n nn.LeakyReLU(0.2, inplace=True)\n )\n \n def _make_decoder_block(self, in_channels, out_channels):\n return nn.Sequential(\n nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n nn.BatchNorm2d(out_channels),\n nn.LeakyReLU(0.2, inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n nn.BatchNorm2d(out_channels),\n nn.LeakyReLU(0.2, inplace=True)\n )\n \n def enhance_contrast(self, x):\n \"\"\"Enhance contrast in the prediction to match high-contrast ground truth\"\"\"\n # Apply tanh with scaling to enhance contrast\n enhanced = torch.tanh(5 * (x - 0.5)) * 0.5 + 0.5\n return enhanced\n \n def enhance_edges(self, x):\n \"\"\"Enhance geological edges in the prediction\"\"\"\n edge_features = self.edge_detector(x)\n enhanced = x + 0.1 * edge_features\n return torch.clamp(enhanced, 0, 1)\n \n def forward(self, x):\n # Input format adjustment - seismic data comes as [B, C, T, R]\n # We'll permute to [B, C, R, T] for 2D convolutions\n x = x.permute(0, 1, 3, 2)\n \n # Encoder path with skip connections\n e1 = self.enc1(x)\n p1 = self.pool(e1)\n \n e2 = self.enc2(p1)\n p2 = self.pool(e2)\n \n e3 = self.enc3(p2)\n p3 = self.pool(e3)\n \n e4 = self.enc4(p3)\n p4 = self.pool(e4)\n \n # Bottleneck\n b = self.bottleneck(p4)\n \n # Decoder path with skip connections\n d4 = F.interpolate(b, size=e4.shape[2:], mode='bilinear', align_corners=False)\n d4 = self.dec4(torch.cat([d4, e4], dim=1))\n \n d3 = F.interpolate(d4, size=e3.shape[2:], mode='bilinear', align_corners=False)\n d3 = self.dec3(torch.cat([d3, e3], dim=1))\n \n d2 = F.interpolate(d3, size=e2.shape[2:], mode='bilinear', align_corners=False)\n d2 = self.dec2(torch.cat([d2, e2], dim=1))\n \n d1 = F.interpolate(d2, size=e1.shape[2:], mode='bilinear', align_corners=False)\n d1 = self.dec1(torch.cat([d1, e1], dim=1))\n \n # Final output and enhancement\n out = self.final(d1)\n out = self.enhance_contrast(out)\n out = self.enhance_edges(out)\n \n # Ensure output is 70x70 as expected\n out = F.interpolate(out, size=(70, 70), mode='bilinear', align_corners=False)\n \n return out\n\n\n# 2. Geological Feature Detection Model\nclass GeologicalFeatureNet(nn.Module):\n \"\"\"Neural network for detecting geological features from seismic data\"\"\"\n def __init__(self, in_channels=5, hidden_dim=64):\n super().__init__()\n \n # Encoder (shared for all features)\n self.encoder = nn.Sequential(\n nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim),\n nn.ReLU(),\n nn.MaxPool2d(2),\n \n nn.Conv2d(hidden_dim, hidden_dim*2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*2),\n nn.ReLU(),\n nn.MaxPool2d(2),\n \n nn.Conv2d(hidden_dim*2, hidden_dim*4, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*4),\n nn.ReLU(),\n nn.MaxPool2d(2),\n \n nn.Conv2d(hidden_dim*4, hidden_dim*8, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*8),\n nn.ReLU(),\n nn.MaxPool2d(2)\n )\n \n # Specialized decoders for different geological features\n \n # 1. Salt body detector - MODIFIED to remove sigmoid for mixed precision compatibility\n self.salt_decoder = nn.Sequential(\n nn.Conv2d(hidden_dim*8, hidden_dim*4, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*4),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim*4, hidden_dim*2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*2),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim*2, hidden_dim, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim, hidden_dim//2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim//2),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim//2, 1, kernel_size=1)\n # Sigmoid removed for BCEWithLogitsLoss compatibility\n )\n \n # 2. Fault detector with edge awareness - MODIFIED to remove sigmoid\n self.fault_decoder = nn.Sequential(\n nn.Conv2d(hidden_dim*8, hidden_dim*4, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*4),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim*4, hidden_dim*2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*2),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim*2, hidden_dim, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim, hidden_dim//2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim//2),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim//2, 1, kernel_size=1)\n # Sigmoid removed for BCEWithLogitsLoss compatibility\n )\n \n # 3. Layer boundary detector with horizontal bias - MODIFIED to remove sigmoid\n self.layer_decoder = nn.Sequential(\n nn.Conv2d(hidden_dim*8, hidden_dim*4, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*4),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim*4, hidden_dim*2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim*2),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim*2, hidden_dim, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim, hidden_dim//2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim//2),\n nn.ReLU(),\n nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),\n \n nn.Conv2d(hidden_dim//2, 1, kernel_size=1)\n # Sigmoid removed for BCEWithLogitsLoss compatibility\n )\n \n # Optional velocity reconstruction from geological features - Keep sigmoid here\n self.velocity_reconstruction = nn.Sequential(\n nn.Conv2d(3, hidden_dim//2, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim//2),\n nn.ReLU(),\n \n nn.Conv2d(hidden_dim//2, hidden_dim//4, kernel_size=3, padding=1),\n nn.BatchNorm2d(hidden_dim//4),\n nn.ReLU(),\n \n nn.Conv2d(hidden_dim//4, 1, kernel_size=1),\n nn.Sigmoid() # Keep sigmoid here as this is not used with BCE loss\n )\n \n def forward(self, x):\n # Input format adjustment\n x = x.permute(0, 1, 3, 2) # [B, C, T, R] to [B, C, R, T]\n \n # Encode features\n encoded = self.encoder(x)\n \n # Decode for each geological feature\n salt_mask = self.salt_decoder(encoded)\n fault_lines = self.fault_decoder(encoded)\n layer_boundaries = self.layer_decoder(encoded)\n \n # Ensure all outputs are 70x70\n salt_mask = F.interpolate(salt_mask, size=(70, 70), mode='bilinear', align_corners=False)\n fault_lines = F.interpolate(fault_lines, size=(70, 70), mode='bilinear', align_corners=False)\n layer_boundaries = F.interpolate(layer_boundaries, size=(70, 70), mode='bilinear', align_corners=False)\n \n # Apply sigmoid for the velocity reconstruction input\n salt_prob = torch.sigmoid(salt_mask)\n fault_prob = torch.sigmoid(fault_lines)\n layer_prob = torch.sigmoid(layer_boundaries)\n \n # Combine features for velocity reconstruction \n features = torch.cat([salt_prob, fault_prob, layer_prob], dim=1)\n velocity = self.velocity_reconstruction(features)\n \n return {\n 'salt': salt_mask, # Now returns logits\n 'faults': fault_lines, # Now returns logits\n 'layers': layer_boundaries, # Now returns logits\n 'velocity': velocity # Still returns probabilities\n }\n\n\ndef create_model(config):\n \"\"\"Create model based on configuration\"\"\"\n approach = config['approach']\n device = config['device']\n\n if approach == 'physics_guided':\n model = PhysicsGuidedUNet(\n in_channels=config['in_channels'],\n out_channels=config['out_channels'],\n hidden_dim=config['hidden_dim']\n )\n print(\"Created Physics-Guided U-Net model\")\n\n elif approach == 'feature_detection':\n model = GeologicalFeatureNet(\n in_channels=config['in_channels'],\n hidden_dim=config['hidden_dim']\n )\n print(\"Created Geological Feature Detection model\")\n\n elif approach == 'geo_aware':\n model = GeoAwareConvNet(\n in_channels=config['in_channels'],\n out_channels=config['out_channels'],\n hidden_dim=config['hidden_dim']\n )\n print(\"Created Geo-Aware ConvNet model\")\n\n else:\n model = PhysicsGuidedUNet(\n in_channels=config['in_channels'],\n out_channels=config['out_channels'],\n hidden_dim=config['hidden_dim']\n )\n print(f\"Unknown approach '{approach}', defaulting to Physics-Guided model\")\n\n model = model.to(device)\n num_params = sum(p.numel() for p in model.parameters())\n trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n print(f\"Model has {num_params:,} parameters ({trainable_params:,} trainable)\")\n \n return model\n\n \n # Move model to device\n model = model.to(device)\n \n # Print model summary\n num_params = sum(p.numel() for p in model.parameters())\n trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n print(f\"Model has {num_params:,} parameters ({trainable_params:,} trainable)\")\n \n return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.757795Z","iopub.execute_input":"2025-05-12T17:58:17.758077Z","iopub.status.idle":"2025-05-12T17:58:17.787767Z","shell.execute_reply.started":"2025-05-12T17:58:17.758053Z","shell.execute_reply":"2025-05-12T17:58:17.787042Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 5: GeoAwareConvNet Architecture (2D-Aware)\n# {\"_collapsed\": false}\n\nclass GeoAwareConvBlock(nn.Module):\n def __init__(self, in_channels, out_channels):\n super().__init__()\n self.conv_standard = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)\n self.conv_dilated = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=(0, 2), dilation=(1, 2))\n self.bn = nn.BatchNorm2d(out_channels)\n self.act = nn.ReLU()\n\n def forward(self, x):\n x = self.conv_standard(x)\n x = self.conv_dilated(x)\n x = self.bn(x)\n return self.act(x)\n\n\nclass GeoAwareConvNet(nn.Module):\n def __init__(self, in_channels, out_channels, hidden_dim=64):\n super().__init__()\n self.encoder = nn.Sequential(\n GeoAwareConvBlock(in_channels, hidden_dim),\n GeoAwareConvBlock(hidden_dim, hidden_dim),\n )\n self.decoder = nn.Sequential(\n nn.Conv2d(hidden_dim, hidden_dim // 2, kernel_size=3, padding=1),\n nn.ReLU(),\n nn.Conv2d(hidden_dim // 2, out_channels, kernel_size=1)\n )\n\n def forward(self, x):\n x = self.encoder(x)\n x = self.decoder(x)\n # Resize output to match label dimensions\n x = F.interpolate(x, size=(70, 70), mode='bilinear', align_corners=False)\n # Optionally clamp if you're outputting probabilities\n # x = torch.clamp(x, 0.0, 1.0)\n return x\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.790018Z","iopub.execute_input":"2025-05-12T17:58:17.790237Z","iopub.status.idle":"2025-05-12T17:58:17.811593Z","shell.execute_reply.started":"2025-05-12T17:58:17.790210Z","shell.execute_reply":"2025-05-12T17:58:17.811048Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 5: Geological Discriminator Network\n# {\"_collapsed\": true}\n\nclass GeoDiscriminator(nn.Module):\n \"\"\"Small CNN to assess geological plausibility of velocity models\"\"\"\n def __init__(self, in_channels=1, hidden_dim=32):\n super().__init__()\n self.net = nn.Sequential(\n nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1),\n nn.BatchNorm2d(hidden_dim),\n nn.LeakyReLU(0.2),\n\n nn.Conv2d(hidden_dim, hidden_dim*2, kernel_size=3, stride=2, padding=1),\n nn.BatchNorm2d(hidden_dim*2),\n nn.LeakyReLU(0.2),\n\n nn.Conv2d(hidden_dim*2, hidden_dim*4, kernel_size=3, stride=2, padding=1),\n nn.BatchNorm2d(hidden_dim*4),\n nn.LeakyReLU(0.2),\n\n nn.AdaptiveAvgPool2d(1),\n nn.Flatten(),\n nn.Linear(hidden_dim*4, 1),\n nn.Sigmoid()\n )\n\n def forward(self, x):\n return self.net(x)\n\n\ndef train_discriminator(discriminator, dataloader, device, epochs=5):\n disc = discriminator.to(device)\n optimizer = torch.optim.Adam(disc.parameters(), lr=1e-4)\n loss_fn = nn.BCELoss()\n\n disc.train()\n for epoch in range(epochs):\n total_loss = 0\n for _, targets in dataloader:\n targets = targets.to(device)\n preds = disc(targets)\n labels = torch.ones_like(preds)\n loss = loss_fn(preds, labels)\n\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n total_loss += loss.item()\n\n print(f\"Epoch {epoch+1} | Discriminator loss: {total_loss/len(dataloader):.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.812264Z","iopub.execute_input":"2025-05-12T17:58:17.812471Z","iopub.status.idle":"2025-05-12T17:58:17.835861Z","shell.execute_reply.started":"2025-05-12T17:58:17.812457Z","shell.execute_reply":"2025-05-12T17:58:17.835153Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 6: Loss Functions (Fixed for all tedious Mixed Precision Issues)\nclass PhysicsGuidedLoss(nn.Module):\n \"\"\"Loss function for physics-guided velocity prediction with improved numerical stability\"\"\"\n def __init__(self, wave_eq_weight=0.05, slowness_weight=0.1, \n layering_weight=0.05, contrast_weight=0.2):\n super().__init__()\n self.wave_eq_weight = wave_eq_weight\n self.slowness_weight = slowness_weight\n self.layering_weight = layering_weight\n self.contrast_weight = contrast_weight\n \n # Edge detection kernels\n self.sobel_x = torch.tensor([[1,0,-1],[2,0,-2],[1,0,-1]], dtype=torch.float32).view(1, 1, 3, 3)\n self.sobel_y = torch.tensor([[1,2,1],[0,0,0],[-1,-2,-1]], dtype=torch.float32).view(1, 1, 3, 3)\n \n def _prepare_kernels(self, device):\n \"\"\"Move kernels to the correct device and data type\"\"\"\n dtype = torch.float32 # Always use float32 for kernels\n if self.sobel_x.device != device:\n self.sobel_x = self.sobel_x.to(device=device, dtype=dtype)\n self.sobel_y = self.sobel_y.to(device=device, dtype=dtype)\n \n def data_fidelity_loss(self, pred, target):\n \"\"\"Standard L1 loss for overall accuracy\"\"\"\n return F.l1_loss(pred, target)\n \n def wave_equation_constraint(self, pred):\n \"\"\"Simplified wave equation physics constraint - FIXED\"\"\"\n device = pred.device\n dtype = pred.dtype # Get the current dtype of pred\n self._prepare_kernels(device)\n \n # Create kernels with matching dtype\n kernel_x = torch.tensor([[[[1.0, -2.0, 1.0]]]], device=device, dtype=dtype)\n kernel_y = torch.tensor([[[[1.0], [-2.0], [1.0]]]], device=device, dtype=dtype)\n \n # Second-order spatial derivatives (Laplacian)\n laplacian_x = F.conv2d(\n F.pad(pred, (1, 1, 0, 0), mode='replicate'),\n kernel_x,\n padding=0\n )\n \n laplacian_y = F.conv2d(\n F.pad(pred, (0, 0, 1, 1), mode='replicate'),\n kernel_y,\n padding=0\n )\n \n laplacian = laplacian_x + laplacian_y\n \n # Clip to prevent extreme values and NaN\n laplacian_abs = torch.clamp(torch.abs(laplacian), 0.0, 100.0)\n \n # Prevent NaN by taking mean with eps\n return torch.mean(laplacian_abs)\n \n def slowness_gradient_constraint(self, pred):\n \"\"\"Constraint to encourage smoothness in slowness (1/v) - FIXED\"\"\"\n device = pred.device\n self._prepare_kernels(device)\n \n # Add small epsilon and clamp to prevent division by zero\n eps = 1e-6\n pred_safe = torch.clamp(pred, min=eps)\n \n # Convert velocity to slowness (1/v) with safe division\n slowness = 1.0 / pred_safe\n \n # Clip slowness to prevent extreme values\n slowness = torch.clamp(slowness, max=100.0)\n \n # Calculate gradients of slowness\n grad_x = F.conv2d(slowness, self.sobel_x, padding=1)\n grad_y = F.conv2d(slowness, self.sobel_y, padding=1)\n \n # Penalize large gradients in slowness (encourage smoothness)\n # Clip gradients to avoid NaN\n grad_x_abs = torch.clamp(torch.abs(grad_x), 0.0, 100.0)\n grad_y_abs = torch.clamp(torch.abs(grad_y), 0.0, 100.0)\n \n return torch.mean(grad_x_abs) + torch.mean(grad_y_abs)\n \n def geological_layering_constraint(self, pred):\n \"\"\"Encourage horizontal layering typical in geological structures - FIXED\"\"\"\n device = pred.device\n self._prepare_kernels(device)\n \n # Calculate horizontal and vertical gradients\n grad_x = F.conv2d(pred, self.sobel_x, padding=1)\n grad_y = F.conv2d(pred, self.sobel_y, padding=1)\n \n # Add epsilon to denominators\n eps = 1e-6\n grad_x_abs = torch.abs(grad_x) + eps\n grad_y_abs = torch.abs(grad_y) + eps\n \n # Ratio of vertical to horizontal gradients (safe division)\n ratio = grad_y_abs / grad_x_abs\n \n # Clip ratio to prevent extreme values\n ratio = torch.clamp(ratio, 0.0, 100.0)\n \n # Use a safer exponential function with clipping\n exp_term = torch.clamp(-ratio, -20.0, 20.0) # Clamp to avoid overflow\n return torch.mean(torch.exp(exp_term))\n \n def contrast_enhancement_loss(self, pred, target):\n \"\"\"Loss to encourage high contrast similar to ground truth - FIXED\"\"\"\n # Calculate histograms safely by clipping inputs\n pred_safe = torch.clamp(pred, 0.0, 1.0)\n target_safe = torch.clamp(target, 0.0, 1.0)\n \n # Calculate histograms\n try:\n pred_hist = torch.histc(pred_safe, bins=10, min=0, max=1)\n target_hist = torch.histc(target_safe, bins=10, min=0, max=1)\n \n # Normalize histograms (with epsilon for safe division)\n eps = 1e-6\n pred_hist = pred_hist / (pred_hist.sum() + eps)\n target_hist = target_hist / (target_hist.sum() + eps)\n \n # Calculate difference between histograms (EMD approximation)\n hist_diff = F.l1_loss(\n torch.cumsum(pred_hist, dim=0),\n torch.cumsum(target_hist, dim=0)\n )\n \n # Penalize low contrast (encourage bi-modal distribution)\n # Calculate variance safely\n contrast = torch.var(pred_safe)\n contrast_target = torch.var(target_safe)\n \n return hist_diff + torch.abs(contrast - contrast_target)\n except RuntimeError:\n # Fallback if histc fails (should be rare)\n return F.l1_loss(pred, target)\n \n def forward(self, pred, target):\n try:\n # Data fidelity (supervised loss)\n l1_loss = self.data_fidelity_loss(pred, target)\n \n # Physics-guided constraints with safety\n wave_eq_loss = torch.clamp(self.wave_equation_constraint(pred), 0.0, 100.0)\n slowness_loss = torch.clamp(self.slowness_gradient_constraint(pred), 0.0, 100.0)\n layering_loss = torch.clamp(self.geological_layering_constraint(pred), 0.0, 100.0)\n contrast_loss = torch.clamp(self.contrast_enhancement_loss(pred, target), 0.0, 100.0)\n \n # Debug prints to catch any potential issues\n # print(f\"L1: {l1_loss.item():.4f}, Wave: {wave_eq_loss.item():.4f}, Slow: {slowness_loss.item():.4f}, Layer: {layering_loss.item():.4f}\")\n \n # Combine losses with weights\n total_loss = (\n l1_loss +\n self.wave_eq_weight * wave_eq_loss +\n self.slowness_weight * slowness_loss +\n self.layering_weight * layering_loss +\n self.contrast_weight * contrast_loss\n )\n \n # Final safety check\n if torch.isnan(total_loss) or torch.isinf(total_loss):\n print(\"Warning: NaN or Inf detected in loss calculation. Falling back to L1 loss.\")\n return l1_loss\n \n return total_loss\n \n except Exception as e:\n # If any error occurs, fall back to L1 loss\n print(f\"Error in loss calculation: {e}. Falling back to L1 loss.\")\n return self.data_fidelity_loss(pred, target)\n\n\nclass DiceLoss(nn.Module):\n \"\"\"Dice loss for better boundary detection - FIXED for tensor compatibility\"\"\"\n def __init__(self, smooth=1e-6):\n super().__init__()\n self.smooth = smooth\n \n def forward(self, pred, target):\n # Ensure tensors have the same shape and use reshape instead of view\n # This fixes the \"view size is not compatible\" error\n pred_flat = pred.reshape(-1)\n target_flat = target.reshape(-1)\n \n # Calculate Dice coefficient\n intersection = torch.sum(pred_flat * target_flat)\n pred_sum = torch.sum(pred_flat * pred_flat)\n target_sum = torch.sum(target_flat * target_flat)\n \n dice = (2.0 * intersection + self.smooth) / (pred_sum + target_sum + self.smooth)\n \n return 1.0 - dice\n\n\nclass GeologicalFeatureLoss(nn.Module):\n \"\"\"Loss function for geological feature detection - FIXED for mixed precision\"\"\"\n def __init__(self, salt_weight=1.0, fault_weight=1.0, layer_weight=1.0, \n constraint_weight=0.3):\n super().__init__()\n self.salt_weight = salt_weight\n self.fault_weight = fault_weight\n self.layer_weight = layer_weight\n self.constraint_weight = constraint_weight\n \n # Base loss functions - FIXED: using BCEWithLogitsLoss for mixed precision compatibility\n self.bce_loss = nn.BCEWithLogitsLoss()\n self.dice_loss = DiceLoss()\n \n def forward(self, predictions, targets):\n \"\"\"\n Args:\n predictions: dict with keys 'salt', 'faults', 'layers', 'velocity'\n targets: tensor of shape [B, 3, H, W] with channels [salt, faults, layers]\n \"\"\"\n # Separate target channels\n salt_target = targets[:, 0:1]\n fault_target = targets[:, 1:2]\n layer_target = targets[:, 2:3]\n \n # Feature detection losses (combine BCE and Dice for better boundary detection)\n # Apply sigmoid for Dice loss since we're now using BCEWithLogitsLoss\n salt_loss = self.bce_loss(predictions['salt'], salt_target) + \\\n self.dice_loss(torch.sigmoid(predictions['salt']), salt_target)\n \n fault_loss = self.bce_loss(predictions['faults'], fault_target) + \\\n self.dice_loss(torch.sigmoid(predictions['faults']), fault_target)\n \n layer_loss = self.bce_loss(predictions['layers'], layer_target) + \\\n self.dice_loss(torch.sigmoid(predictions['layers']), layer_target)\n \n # Apply sigmoid to convert logits to probabilities for constraint functions\n salt_prob = torch.sigmoid(predictions['salt'])\n fault_prob = torch.sigmoid(predictions['faults'])\n layer_prob = torch.sigmoid(predictions['layers'])\n \n # Geological constraints\n # 1. Salt bodies should be continuous (salt body coherence)\n salt_coherence = self.continuity_constraint(salt_prob)\n \n # 2. Faults should be thin, continuous lines\n fault_thin = self.thinness_constraint(fault_prob)\n \n # 3. Layers should be predominantly horizontal\n layer_horizontal = self.horizontal_bias_constraint(layer_prob)\n \n # 4. Geological compatibility between features\n compatibility = self.feature_compatibility_constraint(\n salt_prob, fault_prob, layer_prob\n )\n \n # Combine all losses with weights\n total_loss = (\n self.salt_weight * salt_loss +\n self.fault_weight * fault_loss +\n self.layer_weight * layer_loss +\n self.constraint_weight * (\n salt_coherence + \n fault_thin + \n layer_horizontal + \n compatibility\n )\n )\n \n return total_loss\n \n def continuity_constraint(self, pred):\n \"\"\"Penalize fragmented structures\"\"\"\n # Calculate gradient magnitude\n grad_x = torch.abs(pred[:, :, :, 1:] - pred[:, :, :, :-1])\n grad_y = torch.abs(pred[:, :, 1:, :] - pred[:, :, :-1, :])\n \n # Only penalize gradients within the predicted region\n mask = (pred > 0.5).float()\n masked_grad_x = grad_x * mask[:, :, :, :-1]\n masked_grad_y = grad_y * mask[:, :, :-1, :]\n \n return torch.mean(masked_grad_x) + torch.mean(masked_grad_y)\n \n def thinness_constraint(self, pred):\n \"\"\"Encourage thin fault lines\"\"\"\n # Dilated prediction\n dilated = F.max_pool2d(pred, kernel_size=3, stride=1, padding=1)\n \n # Eroded prediction\n kernel = torch.ones(1, 1, 3, 3).to(pred.device)\n eroded = 1.0 - F.max_pool2d(1.0 - pred, kernel_size=3, stride=1, padding=1)\n \n # Difference between dilated and eroded should be large for thin structures\n thinness = torch.mean(dilated - eroded)\n \n return -thinness # Negative to encourage thinness\n \n def horizontal_bias_constraint(self, pred):\n \"\"\"Encourage predominantly horizontal layers\"\"\"\n # Vertical vs horizontal gradient ratio\n grad_x = torch.abs(pred[:, :, :, 1:] - pred[:, :, :, :-1])\n grad_y = torch.abs(pred[:, :, 1:, :] - pred[:, :, :-1, :])\n \n # Mean gradients\n mean_grad_x = torch.mean(grad_x)\n mean_grad_y = torch.mean(grad_y)\n \n # Horizontal layers have larger vertical gradients\n return mean_grad_x - mean_grad_y # Minimize for horizontal bias\n \n def feature_compatibility_constraint(self, salt, faults, layers):\n \"\"\"Enforce compatibility between different geological features - FIXED for mixed precision\"\"\"\n # 1. Salt bodies should have clear boundaries that align with layer boundaries\n salt_boundary = torch.abs(\n F.avg_pool2d(salt, kernel_size=3, stride=1, padding=1) - salt\n )\n \n # This eliminates the autocast error while maintaining the same functionality\n boundary_alignment = F.mse_loss(salt_boundary, layers)\n \n # Alternative approach using BCE:\n # with torch.cuda.amp.autocast(enabled=False):\n # boundary_alignment = F.binary_cross_entropy(salt_boundary, layers)\n \n # 2. Faults should often terminate layers\n fault_layer_interaction = torch.mean(faults * layers)\n \n # 3. Faults rarely cut through salt bodies\n fault_salt_exclusion = torch.mean(faults * salt)\n \n return boundary_alignment + fault_layer_interaction - fault_salt_exclusion\n\n\ndef create_loss_function(config):\n \"\"\"Create loss function based on configuration\"\"\"\n approach = config['approach']\n device = config['device']\n \n if approach == 'physics_guided':\n criterion = PhysicsGuidedLoss(\n wave_eq_weight=config['wave_eq_weight'],\n slowness_weight=config['slowness_weight'],\n layering_weight=config['layering_weight'],\n contrast_weight=config['contrast_weight']\n )\n print(\"Created Physics-Guided Loss\")\n elif approach == 'feature_detection':\n criterion = GeologicalFeatureLoss(\n salt_weight=config['salt_weight'],\n fault_weight=config['fault_weight'],\n layer_weight=config['layer_weight'],\n constraint_weight=config['geological_constraint_weight']\n )\n print(\"Created Geological Feature Loss\")\n elif approach == 'geo_aware':\n # You could use a combined loss function for GeoAware approach\n criterion = nn.L1Loss() # Base loss\n print(\"Created Geo-Aware Loss (L1)\")\n else:\n # Default to L1 loss\n criterion = nn.L1Loss()\n print(f\"Unknown approach '{approach}', defaulting to L1 Loss\")\n \n return criterion.to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.836696Z","iopub.execute_input":"2025-05-12T17:58:17.836913Z","iopub.status.idle":"2025-05-12T17:58:17.865289Z","shell.execute_reply.started":"2025-05-12T17:58:17.836888Z","shell.execute_reply":"2025-05-12T17:58:17.864770Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 7: Training and Evaluation Functions with Enhanced Visualizations\n\ndef train_model(model, train_loader, val_loader, criterion, config):\n \"\"\"Train model with validation and early stopping - FIXED for PyTorch 2.0+ compatibility\"\"\"\n device = config['device']\n num_epochs = config['num_epochs']\n learning_rate = config['learning_rate']\n weight_decay = config['weight_decay']\n early_stopping = config['early_stopping']\n output_dir = config['output_dir']\n grad_clip_value = config.get('grad_clip_value', 1.0) # Default to 1.0 if not specified\n \n # Create experiment subdirectory\n experiment_path = output_dir / 'models' / config['experiment_name']\n os.makedirs(experiment_path, exist_ok=True)\n \n # Initialize optimizer\n optimizer = torch.optim.AdamW(\n model.parameters(),\n lr=learning_rate,\n weight_decay=weight_decay\n )\n \n # Create warmup scheduler if requested\n if config.get('warmup_epochs', 0) > 0:\n warmup_scheduler = torch.optim.lr_scheduler.LinearLR(\n optimizer, \n start_factor=0.1, \n end_factor=1.0, \n total_iters=config['warmup_epochs'] * len(train_loader)\n )\n \n # Learning rate scheduler\n scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n optimizer, \n mode='min', \n factor=0.5, \n patience=config['scheduler_patience']\n )\n \n # Mixed precision training (updated for PyTorch 2.0+)\n if config['mixed_precision'] and torch.cuda.is_available():\n # Use proper 'cuda' device parameter\n scaler = torch.amp.GradScaler(device_type='cuda')\n else:\n scaler = None\n \n # Training metrics\n best_val_loss = float('inf')\n early_stop_counter = 0\n train_losses = []\n val_losses = []\n \n # Start training\n print(f\"Starting training for {num_epochs} epochs...\")\n start_time = time.time()\n \n for epoch in range(1, num_epochs + 1):\n epoch_start = time.time()\n \n # Training phase\n model.train()\n epoch_train_loss = 0\n \n # Progress tracking\n progress = tqdm(train_loader, desc=f\"Epoch {epoch}/{num_epochs} [Train]\")\n \n for batch_idx, (x, y) in enumerate(progress):\n x, y = x.to(device), y.to(device)\n \n # Mixed precision forward pass\n if scaler:\n with torch.amp.autocast(device_type='cuda'):\n if config['approach'] == 'feature_detection':\n pred = model(x)\n loss = criterion(pred, y)\n else:\n pred = model(x)\n loss = criterion(pred, y)\n \n # Mixed precision backward pass with gradient clipping\n optimizer.zero_grad()\n scaler.scale(loss).backward()\n \n # Unscale gradients for clipping\n scaler.unscale_(optimizer)\n torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip_value)\n \n scaler.step(optimizer)\n scaler.update()\n else:\n # Standard forward pass\n if config['approach'] == 'feature_detection':\n pred = model(x)\n loss = criterion(pred, y)\n else:\n pred = model(x)\n loss = criterion(pred, y)\n \n # Standard backward pass with gradient clipping\n optimizer.zero_grad()\n loss.backward()\n torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip_value)\n optimizer.step()\n \n # Apply warmup scheduler if in warmup period\n if config.get('warmup_epochs', 0) > 0 and epoch <= config['warmup_epochs']:\n warmup_scheduler.step()\n \n # Skip NaN losses in the average calculation\n if not torch.isnan(loss) and not torch.isinf(loss):\n epoch_train_loss += loss.item()\n else:\n print(f\"Warning: NaN/Inf loss detected at batch {batch_idx}. Skipping in average calculation.\")\n \n # Update progress bar with non-NaN loss\n if not torch.isnan(loss) and not torch.isinf(loss):\n progress.set_postfix({'loss': loss.item()})\n else:\n progress.set_postfix({'loss': 'NaN/Inf'})\n \n # Calculate average training loss (safely handling NaN batches)\n if len(train_loader) > 0:\n avg_train_loss = epoch_train_loss / len(train_loader)\n else:\n avg_train_loss = float('nan')\n \n train_losses.append(avg_train_loss)\n \n # Validation phase\n model.eval()\n epoch_val_loss = 0\n \n # Additional metrics for physics-guided model\n l1_error = 0\n ssim_values = 0\n \n progress = tqdm(val_loader, desc=f\"Epoch {epoch}/{num_epochs} [Valid]\")\n \n with torch.no_grad():\n for x, y in progress:\n x, y = x.to(device), y.to(device)\n \n # Forward pass\n if config['approach'] == 'feature_detection':\n pred = model(x)\n loss = criterion(pred, y)\n \n # Use reconstructed velocity for metrics\n pred_velocity = pred['velocity']\n else:\n pred = model(x)\n loss = criterion(pred, y)\n pred_velocity = pred\n \n # Update metrics\n epoch_val_loss += loss.item()\n \n # Calculate additional metrics\n l1_error += F.l1_loss(pred_velocity, y[:, 0:1] if y.shape[1] > 1 else y).item()\n \n # Normalized for SSIM calculation\n pred_norm = (pred_velocity - pred_velocity.min()) / (pred_velocity.max() - pred_velocity.min() + 1e-8)\n target_norm = (y[:, 0:1] if y.shape[1] > 1 else y - y.min()) / (y.max() - y.min() + 1e-8)\n \n # Simple SSIM approximation\n c1, c2 = 0.01**2, 0.03**2\n mu_x = F.avg_pool2d(pred_norm, kernel_size=11, stride=1, padding=5)\n mu_y = F.avg_pool2d(target_norm, kernel_size=11, stride=1, padding=5)\n \n sigma_x = F.avg_pool2d(pred_norm**2, kernel_size=11, stride=1, padding=5) - mu_x**2\n sigma_y = F.avg_pool2d(target_norm**2, kernel_size=11, stride=1, padding=5) - mu_y**2\n sigma_xy = F.avg_pool2d(pred_norm * target_norm, kernel_size=11, stride=1, padding=5) - mu_x * mu_y\n \n ssim = ((2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)) / \\\n ((mu_x**2 + mu_y**2 + c1) * (sigma_x + sigma_y + c2))\n \n ssim_values += torch.mean(ssim).item()\n \n # Calculate average validation metrics\n avg_val_loss = epoch_val_loss / len(val_loader)\n avg_l1_error = l1_error / len(val_loader)\n avg_ssim = ssim_values / len(val_loader)\n val_losses.append(avg_val_loss)\n \n # Update learning rate\n scheduler.step(avg_val_loss)\n \n # Print epoch results\n epoch_time = time.time() - epoch_start\n print(f\"Epoch {epoch}/{num_epochs} completed in {epoch_time:.2f}s | \"\n f\"Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f} | \"\n f\"L1 Error: {avg_l1_error:.6f} | SSIM: {avg_ssim:.6f}\")\n \n # Visualize predictions - Enhanced with display\n visualize_frequency = config.get('visualize_every', 5)\n if epoch % visualize_frequency == 0 or epoch == 1 or epoch == num_epochs:\n visualize_predictions(model, val_loader, device, epoch, config)\n \n # Check for improvement\n if avg_val_loss < best_val_loss:\n improvement = (best_val_loss - avg_val_loss) / best_val_loss * 100\n best_val_loss = avg_val_loss\n early_stop_counter = 0\n \n # Save best model\n if config['save_models']:\n best_model_path = experiment_path / f\"best_model_epoch_{epoch}.pt\"\n torch.save({\n 'epoch': epoch,\n 'model_state_dict': model.state_dict(),\n 'optimizer_state_dict': optimizer.state_dict(),\n 'loss': best_val_loss,\n 'train_losses': train_losses,\n 'val_losses': val_losses\n }, best_model_path)\n \n # Also save with standard name\n torch.save(model.state_dict(), experiment_path / \"best_model.pt\")\n \n print(f\"✅ New best model saved with {improvement:.2f}% improvement\")\n else:\n early_stop_counter += 1\n \n # Early stopping check\n if early_stop_counter >= early_stopping:\n print(f\"Early stopping triggered after {epoch} epochs\")\n break\n \n # Save checkpoint every 10 epochs\n if epoch % 10 == 0 and config['save_models']:\n checkpoint_path = experiment_path / f\"checkpoint_epoch_{epoch}.pt\"\n torch.save({\n 'epoch': epoch,\n 'model_state_dict': model.state_dict(),\n 'optimizer_state_dict': optimizer.state_dict(),\n 'loss': avg_val_loss,\n 'train_losses': train_losses,\n 'val_losses': val_losses\n }, checkpoint_path)\n print(f\"Checkpoint saved at epoch {epoch}\")\n \n # Save final model\n if config['save_models']:\n final_path = experiment_path / \"final_model.pt\"\n torch.save({\n 'model_state_dict': model.state_dict(),\n 'loss': avg_val_loss,\n 'train_losses': train_losses,\n 'val_losses': val_losses\n }, final_path)\n print(f\"✅ Final model saved at: {final_path}\")\n \n # Plot loss curves\n plt.figure(figsize=(10, 5))\n plt.plot(train_losses, label='Training Loss')\n plt.plot(val_losses, label='Validation Loss')\n plt.xlabel('Epoch')\n plt.ylabel('Loss')\n plt.title(f'{config[\"approach\"].capitalize()} Model - Loss Curves')\n plt.legend()\n plt.grid(True)\n loss_curve_path = output_dir / 'visualizations' / f\"{experiment_name}_loss_curves.png\"\n plt.savefig(loss_curve_path)\n \n # Display loss curves\n try:\n from IPython.display import display, Image\n plt.close()\n display(Image(str(loss_curve_path)))\n except:\n plt.show()\n plt.close()\n \n # Report training summary\n total_time = time.time() - start_time\n print(f\"Training completed in {total_time/60:.2f} minutes\")\n print(f\"Best validation loss: {best_val_loss:.6f}\")\n \n return model, train_losses, val_losses\n\ndef visualize_predictions(model, loader, device, epoch, config):\n \"\"\"Generate visualizations based on the selected approach with enhanced display\"\"\"\n model.eval()\n \n # Get a batch of data\n x, y = next(iter(loader))\n x, y = x.to(device), y.to(device)\n \n # Make predictions\n with torch.no_grad():\n if config['approach'] == 'feature_detection':\n predictions = model(x)\n else:\n predictions = model(x)\n \n # Create visualization based on approach\n if config['approach'] == 'feature_detection':\n viz_path = visualize_geological_features(x, y, predictions, epoch, config)\n else:\n viz_path = visualize_velocity_prediction(x, y, predictions, epoch, config)\n \n # Display the visualization\n try:\n from IPython.display import display, Image\n print(f\"Visualizations for epoch {epoch}:\")\n display(Image(str(viz_path)))\n except Exception as e:\n print(f\"Could not display image directly: {e}\")\n print(f\"Visualization saved at: {viz_path}\")\n\ndef visualize_velocity_prediction(x, y, pred, epoch, config):\n \"\"\"Visualize velocity model predictions with enhanced display\"\"\"\n # Convert tensors to numpy\n x_np = x.cpu().numpy()\n y_np = y.cpu().numpy()\n pred_np = pred.cpu().numpy()\n \n # Create figure\n n_samples = min(4, x.size(0))\n fig, axs = plt.subplots(n_samples, 4, figsize=(16, 4 * n_samples))\n \n # Handle single sample case\n if n_samples == 1:\n axs = np.array([axs])\n \n for i in range(n_samples):\n # Input seismic data (first channel)\n axs[i, 0].imshow(x_np[i, 0], cmap='seismic', aspect='auto')\n axs[i, 0].set_title(f\"Input Seismic #{i+1}\")\n \n # Ground truth velocity\n if y_np.shape[1] > 1:\n y_plot = y_np[i, 0] # For feature mode, use first channel (salt)\n else:\n y_plot = y_np[i, 0]\n \n axs[i, 1].imshow(y_plot, cmap='magma', origin='lower')\n axs[i, 1].set_title(f\"Ground Truth #{i+1}\")\n \n # Predicted velocity\n axs[i, 2].imshow(pred_np[i, 0], cmap='magma', origin='lower')\n axs[i, 2].set_title(f\"Prediction #{i+1}\")\n \n # Absolute error\n abs_error = np.abs(pred_np[i, 0] - y_plot)\n axs[i, 3].imshow(abs_error, cmap='inferno', origin='lower')\n axs[i, 3].set_title(f\"Abs Error #{i+1}\")\n \n # Remove axis ticks\n for j in range(4):\n axs[i, j].set_xticks([])\n axs[i, j].set_yticks([])\n \n plt.tight_layout()\n plt.suptitle(f\"Physics-Guided Velocity Prediction - Epoch {epoch}\", y=1.02, fontsize=16)\n \n # Save figure\n viz_path = config['output_dir'] / 'visualizations' / f\"velocity_pred_epoch_{epoch}.png\"\n plt.savefig(viz_path, bbox_inches='tight')\n plt.close()\n \n return viz_path\n \ndef visualize_geological_features(x, y, predictions, epoch, config):\n \"\"\"Visualize geological feature predictions with enhanced display\"\"\"\n # Convert tensors to numpy\n x_np = x.cpu().numpy()\n y_np = y.cpu().numpy()\n \n # Apply sigmoid to convert logits to probabilities for visualization\n salt_pred = torch.sigmoid(predictions['salt']).cpu().numpy()\n fault_pred = torch.sigmoid(predictions['faults']).cpu().numpy()\n layer_pred = torch.sigmoid(predictions['layers']).cpu().numpy()\n velocity_pred = predictions['velocity'].cpu().numpy()\n \n # Create figure\n n_samples = min(3, x.size(0))\n fig, axs = plt.subplots(n_samples, 5, figsize=(20, 5 * n_samples))\n \n # Handle single sample case\n if n_samples == 1:\n axs = np.array([axs])\n \n for i in range(n_samples):\n # Input seismic data (first channel)\n axs[i, 0].imshow(x_np[i, 0], cmap='seismic', aspect='auto')\n axs[i, 0].set_title(f\"Input Seismic #{i+1}\")\n \n # Salt body detection\n salt_true = y_np[i, 0]\n axs[i, 1].imshow(salt_pred[i, 0], cmap='viridis')\n axs[i, 1].contour(salt_true, colors='r', linewidths=0.5, levels=[0.5])\n axs[i, 1].set_title(f\"Salt Detection #{i+1}\")\n \n # Fault detection\n fault_true = y_np[i, 1]\n axs[i, 2].imshow(fault_pred[i, 0], cmap='viridis')\n axs[i, 2].contour(fault_true, colors='r', linewidths=0.5, levels=[0.5])\n axs[i, 2].set_title(f\"Fault Detection #{i+1}\")\n \n # Layer detection\n layer_true = y_np[i, 2]\n axs[i, 3].imshow(layer_pred[i, 0], cmap='viridis')\n axs[i, 3].contour(layer_true, colors='r', linewidths=0.5, levels=[0.5])\n axs[i, 3].set_title(f\"Layer Detection #{i+1}\")\n \n # Reconstructed velocity\n axs[i, 4].imshow(velocity_pred[i, 0], cmap='magma', origin='lower')\n axs[i, 4].set_title(f\"Reconstructed Velocity #{i+1}\")\n \n # Remove axis ticks\n for j in range(5):\n axs[i, j].set_xticks([])\n axs[i, j].set_yticks([])\n \n plt.tight_layout()\n plt.suptitle(f\"Geological Feature Detection - Epoch {epoch}\", y=1.02, fontsize=16)\n \n # Save figure\n viz_path = config['output_dir'] / 'visualizations' / f\"features_epoch_{epoch}.png\" \n plt.savefig(viz_path, bbox_inches='tight')\n plt.close()\n \n return viz_path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.865956Z","iopub.execute_input":"2025-05-12T17:58:17.866175Z","iopub.status.idle":"2025-05-12T17:58:17.901208Z","shell.execute_reply.started":"2025-05-12T17:58:17.866153Z","shell.execute_reply":"2025-05-12T17:58:17.900719Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 8: Discriminator Loss Hook (Modular Helper)\n# {\"_collapsed\": true}\n\n\n\ndef add_geodisc_loss(config, pred_velocity, discriminator):\n \"\"\"Apply geological discriminator loss if enabled in config.\"\"\"\n if config.get('use_geodiscriminator', False):\n with torch.no_grad():\n realism_score = discriminator(pred_velocity).squeeze()\n realism_loss = (1 - realism_score).mean()\n return config['geodisc_weight'] * realism_loss\n return 0.\n\n\ndef train_model(model, train_loader, val_loader, criterion, config, discriminator=None):\n device = config['device']\n num_epochs = config['num_epochs']\n learning_rate = config['learning_rate']\n weight_decay = config['weight_decay']\n early_stopping = config['early_stopping']\n output_dir = config['output_dir']\n\n experiment_path = output_dir / 'models' / config['experiment_name']\n os.makedirs(experiment_path, exist_ok=True)\n\n optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)\n scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=config['scheduler_patience'])\n scaler = torch.cuda.amp.GradScaler() if config['mixed_precision'] and torch.cuda.is_available() else None\n\n best_val_loss = float('inf')\n early_stop_counter = 0\n train_losses, val_losses = [], []\n\n print(f\"Starting training for {num_epochs} epochs...\")\n start_time = time.time()\n\n for epoch in range(1, num_epochs + 1):\n model.train()\n epoch_train_loss = 0\n\n for x, y in tqdm(train_loader, desc=f\"Epoch {epoch}/{num_epochs} [Train]\"):\n x, y = x.to(device), y.to(device)\n\n if scaler:\n with torch.amp.autocast('cuda'):\n pred = model(x)\n loss = criterion(pred, y)\n optimizer.zero_grad()\n scaler.scale(loss).backward()\n scaler.step(optimizer)\n scaler.update()\n else:\n pred = model(x)\n loss = criterion(pred, y)\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n\n if isinstance(pred, dict):\n pred_velocity = pred.get('velocity', pred)\n else:\n pred_velocity = pred\n\n loss += add_geodisc_loss(config, pred_velocity, discriminator)\n epoch_train_loss += loss.item()\n\n avg_train_loss = epoch_train_loss / len(train_loader)\n train_losses.append(avg_train_loss)\n\n model.eval()\n epoch_val_loss = 0\n with torch.no_grad():\n for x, y in tqdm(val_loader, desc=f\"Epoch {epoch}/{num_epochs} [Valid]\"):\n x, y = x.to(device), y.to(device)\n pred = model(x)\n val_loss = criterion(pred, y)\n\n if isinstance(pred, dict):\n pred_velocity = pred.get('velocity', pred)\n else:\n pred_velocity = pred\n\n val_loss += add_geodisc_loss(config, pred_velocity, discriminator)\n epoch_val_loss += val_loss.item()\n\n avg_val_loss = epoch_val_loss / len(val_loader)\n val_losses.append(avg_val_loss)\n scheduler.step(avg_val_loss)\n\n print(f\"Epoch {epoch}: Train Loss = {avg_train_loss:.4f} | Val Loss = {avg_val_loss:.4f}\")\n\n if avg_val_loss < best_val_loss:\n best_val_loss = avg_val_loss\n early_stop_counter = 0\n if config['save_models']:\n torch.save(model.state_dict(), experiment_path / \"best_model.pt\")\n else:\n early_stop_counter += 1\n\n if early_stop_counter >= early_stopping:\n print(f\"Early stopping at epoch {epoch}\")\n break\n\n total_time = time.time() - start_time\n print(f\"Training completed in {total_time / 60:.2f} minutes\")\n return model, train_losses, val_losses\n\n\ndef load_model_for_submission(config):\n from models import GeoAwareConvNet, GeologicalFeatureNet, PhysicsGuidedUNet\n\n model_name = config.get('approach', 'physics_guided')\n if model_name == 'geo_aware':\n model = GeoAwareConvNet(\n in_channels=config['in_channels'],\n out_channels=config['out_channels'],\n hidden_dim=config['hidden_dim']\n )\n elif model_name == 'feature_detection':\n model = GeologicalFeatureNet(\n in_channels=config['in_channels'],\n hidden_dim=config['hidden_dim']\n )\n else:\n model = PhysicsGuidedUNet(\n in_channels=config['in_channels'],\n out_channels=config['out_channels'],\n hidden_dim=config['hidden_dim']\n )\n return model.to(config['device'])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.901805Z","iopub.execute_input":"2025-05-12T17:58:17.901962Z","iopub.status.idle":"2025-05-12T17:58:17.929309Z","shell.execute_reply.started":"2025-05-12T17:58:17.901950Z","shell.execute_reply":"2025-05-12T17:58:17.928828Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 9: Threshold-Based Approach\n\ndef apply_thresholding(data, method='otsu', threshold=None, stats=None):\n \"\"\"Apply thresholding to create binary velocity model\"\"\"\n from skimage import filters\n \n # Determine threshold\n if threshold is not None:\n # Use provided threshold\n thresh = threshold\n elif method == 'otsu':\n # Otsu's adaptive thresholding\n try:\n thresh = filters.threshold_otsu(data)\n except:\n # Fallback to mean if Otsu fails\n thresh = np.mean(data)\n elif method == 'mean':\n # Mean-based thresholding\n thresh = np.mean(data)\n elif method == 'adaptive':\n # Local adaptive thresholding\n thresh = filters.threshold_local(data, block_size=15)\n elif stats is not None:\n # Use pre-calculated statistics\n thresh = stats['auto_threshold']\n else:\n # Default to mean\n thresh = np.mean(data)\n \n # Apply threshold\n binary = (data > thresh).astype(np.float32)\n \n return binary\n\n\ndef post_process_binary(binary, config):\n \"\"\"Apply post-processing to binary velocity model - FIXED for dimension mismatch\"\"\"\n from scipy import ndimage\n \n # Copy to avoid modifying original\n result = binary.copy()\n \n # Check input dimensions\n input_dims = len(result.shape)\n \n # Apply morphological operations if enabled\n if config.get('use_morphology', True):\n # Create structure element with matching dimensions\n if input_dims == 2:\n # 2D input\n structure = np.ones((2, 2))\n structure_close = np.ones((3, 3))\n elif input_dims == 3:\n # 3D input (batch, height, width)\n if result.shape[0] == 1:\n # Single channel, treat as 2D after squeezing\n result = result.squeeze(0)\n structure = np.ones((2, 2))\n structure_close = np.ones((3, 3))\n # Flag that we need to unsqueeze later\n needs_unsqueeze = True\n else:\n # Multi-batch, use 3D structure\n structure = np.ones((1, 2, 2))\n structure_close = np.ones((1, 3, 3))\n needs_unsqueeze = False\n elif input_dims == 4:\n # 4D input (batch, channels, height, width)\n structure = np.ones((1, 1, 2, 2))\n structure_close = np.ones((1, 1, 3, 3))\n needs_unsqueeze = False\n else:\n # Unexpected dimension, skip morphology\n print(f\"Warning: Skipping morphology for unusual dimensions: {result.shape}\")\n structure = None\n \n # Apply morphology if we have a valid structure\n if structure is not None:\n try:\n # Clean up small artifacts\n result = ndimage.binary_opening(result, structure=structure)\n result = ndimage.binary_closing(result, structure=structure_close)\n \n # Restore original dimensions if needed\n if input_dims == 3 and 'needs_unsqueeze' in locals() and needs_unsqueeze:\n result = result[np.newaxis, ...]\n except Exception as e:\n print(f\"Warning: Morphology operation failed: {e}\")\n print(f\"Input shape: {binary.shape}, Structure shape: {structure.shape}\")\n \n # Apply edge enhancement if specified\n if config.get('edge_enhancement', 0) > 0:\n # Calculate edges\n edge_x = ndimage.sobel(binary, axis=-2)\n edge_y = ndimage.sobel(binary, axis=-1)\n edges = np.sqrt(edge_x**2 + edge_y**2)\n \n # Normalize edge strength\n edge_strength = edges / (np.max(edges) + 1e-8)\n \n # Enhance edges\n enhanced = result.copy()\n enhanced[edge_strength > 0.3] = 1.0\n \n # Blend with original based on enhancement strength\n alpha = config['edge_enhancement']\n result = (1 - alpha) * result + alpha * enhanced\n \n return result\n\n\ndef process_with_thresholding(input_data, config, stats=None):\n \"\"\"Process seismic data using thresholding approach\"\"\"\n # Apply gain to seismic data\n time_steps = input_data.shape[1]\n time = np.linspace(0, 1, time_steps)\n gain = (time ** 2)[:, np.newaxis]\n \n gained_data = input_data.copy()\n for c in range(input_data.shape[0]):\n gained_data[c] = input_data[c] * gain\n \n # Take mean across channels and time\n # We need a 2D map that we can threshold\n mean_data = np.mean(gained_data, axis=(0, 1))\n \n # Apply thresholding\n binary = apply_thresholding(\n mean_data, \n method=config['threshold_method'],\n stats=stats\n )\n \n # Apply post-processing\n processed = post_process_binary(binary, config)\n \n return processed\n\n\ndef generate_threshold_submission(test_dir, output_path, config, stats=None):\n \"\"\"Generate submission using thresholding approach\"\"\"\n # Get test files\n test_files = sorted(glob.glob(os.path.join(test_dir, \"*.npy\")))\n print(f\"Found {len(test_files)} test files\")\n \n rows = []\n \n # Process files with progress tracking\n progress = tqdm(test_files, desc=\"Generating threshold-based submission\")\n \n for filepath in progress:\n # Get file ID\n oid = os.path.splitext(os.path.basename(filepath))[0]\n \n # Load data\n data = np.load(filepath)\n \n # Check shape and extract if needed\n if len(data.shape) == 4: # (batch, channels, time, receivers)\n data = data[0] # Use first sample if batched\n \n # Process with thresholding\n binary_pred = process_with_thresholding(data, config, stats)\n \n # Ensure correct shape (70x70)\n if binary_pred.shape != (70, 70):\n from skimage.transform import resize\n binary_pred = resize(binary_pred, (70, 70), order=0, preserve_range=True)\n \n # Format for submission (all rows, odd columns)\n for y in range(70):\n row_id = f\"{oid}_y_{y}\"\n row_data = binary_pred[y, 1:70:2] # Extract odd-indexed columns\n rows.append([row_id] + row_data.tolist())\n \n # Create submission DataFrame\n columns = [\"ID\"] + [f\"x_{i}\" for i in range(1, 70, 2)]\n submission_df = pd.DataFrame(rows, columns=columns)\n \n # Save submission\n submission_df.to_csv(output_path, index=False)\n print(f\"✅ Threshold-based submission saved at: {output_path}\")\n \n return submission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.929942Z","iopub.execute_input":"2025-05-12T17:58:17.930126Z","iopub.status.idle":"2025-05-12T17:58:17.953973Z","shell.execute_reply.started":"2025-05-12T17:58:17.930112Z","shell.execute_reply":"2025-05-12T17:58:17.953437Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 10: Implement edge-based velocity modeling as a post-processing step\n\nimport numpy as np\nfrom scipy import ndimage\nfrom skimage import feature, segmentation, filters, morphology, measure, color\nimport matplotlib.pyplot as plt\n\ndef extract_edges(feature_map, method='sobel', threshold=0.2):\n \"\"\"Extract edges from a feature probability map\"\"\"\n if method == 'sobel':\n edges_x = ndimage.sobel(feature_map, axis=1)\n edges_y = ndimage.sobel(feature_map, axis=0)\n edges = np.sqrt(edges_x**2 + edges_y**2)\n elif method == 'canny':\n edges = feature.canny(feature_map, sigma=1.0, low_threshold=0.1, high_threshold=0.3)\n else:\n edges = filters.scharr(feature_map)\n \n # Normalize and threshold\n if edges.max() > 0:\n edges = edges / edges.max()\n edges = (edges > threshold).astype(np.float32)\n \n return edges\n\ndef multi_scale_edge_detection(seismic_data, scales=[1, 2, 3]):\n \"\"\"Apply edge detection at multiple scales for robust edge extraction\"\"\"\n # If multi-channel, convert to single channel\n if seismic_data.ndim > 2:\n # Use mean across channels or first channel\n if seismic_data.shape[0] > 1:\n data = np.mean(seismic_data, axis=0)\n else:\n data = seismic_data[0]\n else:\n data = seismic_data\n \n # Apply edge detection at multiple scales\n edges_multi_scale = np.zeros_like(data)\n \n for scale in scales:\n # Smooth data at this scale\n smoothed = ndimage.gaussian_filter(data, sigma=scale)\n \n # Detect edges\n edges_x = ndimage.sobel(smoothed, axis=1)\n edges_y = ndimage.sobel(smoothed, axis=0)\n edges = np.sqrt(edges_x**2 + edges_y**2)\n \n # Normalize\n if edges.max() > 0:\n edges = edges / edges.max()\n \n # Accumulate\n edges_multi_scale = np.maximum(edges_multi_scale, edges)\n \n return edges_multi_scale\n\ndef combine_edges(edge_map, feature_edges, weights={'main': 0.5, 'salt': 0.3, 'faults': 0.2, 'layers': 0.1}):\n \"\"\"Combine multiple edge maps with weights\"\"\"\n # Start with main edge map\n combined = weights['main'] * edge_map\n \n # Add feature-specific edges\n for feature_name, edge in feature_edges.items():\n if feature_name in weights:\n combined += weights[feature_name] * edge\n \n # Normalize\n if combined.max() > 0:\n combined = combined / combined.max()\n \n return combined\n\ndef close_gaps(edge_map, threshold=0.3, min_size=3):\n \"\"\"Close small gaps in edge map to create closed contours\"\"\"\n # Threshold edge map\n binary_edges = (edge_map > threshold).astype(np.uint8)\n \n # Apply morphological closing to bridge small gaps\n closed_edges = morphology.binary_closing(binary_edges, morphology.disk(min_size//2))\n \n # Clean up small isolated edge segments\n cleaned_edges = morphology.remove_small_objects(closed_edges, min_size=min_size)\n \n return cleaned_edges.astype(np.float32)\n\ndef segment_image(edge_map, threshold=0.3, compactness=30):\n \"\"\"Segment image based on edge map using watershed\"\"\"\n # Invert edge map to get ridge map (watershed works on valleys)\n ridge_map = 1 - edge_map\n \n # Use watershed algorithm to segment\n # Create markers for watershed (can be improved with feature info)\n binary_edges = (edge_map > threshold)\n distance = ndimage.distance_transform_edt(~binary_edges)\n \n # Create watershed markers\n local_max = feature.peak_local_max(distance, footprint=np.ones((3, 3)), labels=~binary_edges)\n markers = np.zeros(distance.shape, dtype=np.int32)\n markers[tuple(local_max.T)] = np.arange(1, len(local_max) + 1)\n \n # Apply watershed\n segmented = segmentation.watershed(ridge_map, markers, mask=~binary_edges)\n \n # Count regions\n num_regions = len(np.unique(segmented)) - 1 # Exclude 0 (background)\n \n return segmented, num_regions\n\ndef classify_region(region_mask, feature_maps, thresholds={'salt': 0.5, 'faults': 0.3, 'layers': 0.4}):\n \"\"\"Classify region based on feature map probabilities\"\"\"\n # For each feature type, calculate average probability within the region\n region_probs = {}\n for feature_name, feature_map in feature_maps.items():\n avg_prob = np.mean(feature_map[region_mask])\n region_probs[feature_name] = avg_prob\n \n # Classify based on highest probability\n if region_probs['salt'] > thresholds['salt']:\n return 'salt'\n elif region_probs['layers'] > thresholds['layers']:\n # Further classify layers by depth\n y_indices = np.where(region_mask)[0]\n avg_depth = np.mean(y_indices) / region_mask.shape[0]\n layer_number = int(avg_depth * 10) # 0-9 based on depth\n return f'layer_{layer_number}'\n elif region_probs['faults'] > thresholds['faults']:\n return 'fault'\n else:\n # Background - classify by depth for realistic layering\n y_indices = np.where(region_mask)[0]\n avg_depth = np.mean(y_indices) / region_mask.shape[0]\n layer_number = int(avg_depth * 5) # 0-4 based on depth\n return f'background_{layer_number}'\n\ndef get_velocity_for_region_type(region_type):\n \"\"\"Assign velocity based on region type\"\"\"\n # Typical velocity values (m/s) for different geological structures\n if region_type == 'salt':\n # Salt bodies have high velocity\n return 4500.0\n elif region_type.startswith('layer_'):\n # Layers get increasing velocity with depth\n layer_number = int(region_type.split('_')[1])\n return 2000.0 + layer_number * 200.0\n elif region_type == 'fault':\n # Fault zones often have slightly lower velocity\n return 2300.0\n elif region_type.startswith('background_'):\n # Background also increases with depth\n layer_number = int(region_type.split('_')[1])\n return 2200.0 + layer_number * 150.0\n else:\n # Default velocity\n return 2500.0\n\ndef apply_geological_constraints(velocity_model, edge_map, edge_weight=0.2):\n \"\"\"Apply geological constraints to ensure realistic velocity model\"\"\"\n # 1. Ensure velocity increases with depth (general trend in geology)\n rows, cols = velocity_model.shape\n depth_trend = np.linspace(0.9, 1.1, rows).reshape(-1, 1)\n velocity_model = velocity_model * depth_trend\n \n # 2. Enhance velocity contrast at edges\n edge_enhanced = velocity_model.copy()\n edge_gradient = ndimage.gaussian_gradient_magnitude(velocity_model, sigma=1.0)\n \n # Only enhance where we have detected edges\n where_edges = (edge_map > 0.2)\n edge_enhanced[where_edges] = velocity_model[where_edges] * (1.0 + edge_weight * edge_gradient[where_edges])\n \n # 3. Smooth very small regions while preserving edges\n # Create edge-preserving smoothing mask\n smooth_mask = 1.0 - (edge_map > 0.2).astype(np.float32)\n \n # Apply selective smoothing\n smoothed = ndimage.gaussian_filter(edge_enhanced, sigma=1.0)\n final_velocity = edge_enhanced * (1.0 - smooth_mask) + smoothed * smooth_mask\n \n return final_velocity\n\ndef edge_based_velocity_modeling(feature_maps, seismic_data=None):\n \"\"\"Generate velocity model using edge detection and region filling\"\"\"\n # 1. Get feature maps and convert to numpy if needed\n if isinstance(feature_maps, dict) and 'velocity' in feature_maps:\n # If we're using the feature detection model output, it contains velocity\n if isinstance(feature_maps['velocity'], torch.Tensor):\n initial_velocity = feature_maps['velocity'].cpu().numpy().squeeze()\n else:\n initial_velocity = feature_maps['velocity'].squeeze()\n else:\n # Fallback if no velocity is provided\n initial_velocity = np.ones((70, 70)) * 2500.0\n \n salt_map = feature_maps['salt'] if isinstance(feature_maps, dict) else feature_maps\n if isinstance(salt_map, torch.Tensor):\n salt_map = torch.sigmoid(salt_map).cpu().numpy().squeeze()\n else:\n salt_map = salt_map.squeeze()\n \n # Convert other feature maps if available\n if isinstance(feature_maps, dict):\n processed_maps = {}\n for key, value in feature_maps.items():\n if key in ['salt', 'faults', 'layers']:\n if isinstance(value, torch.Tensor):\n processed_maps[key] = torch.sigmoid(value).cpu().numpy().squeeze()\n else:\n processed_maps[key] = value.squeeze()\n else:\n # If not a dictionary, create dummy feature maps\n processed_maps = {\n 'salt': salt_map,\n 'faults': np.zeros_like(salt_map),\n 'layers': np.zeros_like(salt_map)\n }\n \n # 2. Extract edges from each feature map\n feature_edges = {\n 'salt': extract_edges(processed_maps['salt']),\n 'faults': extract_edges(processed_maps['faults']),\n 'layers': extract_edges(processed_maps['layers'])\n }\n \n # 3. Get edge map from seismic data if available\n if seismic_data is not None and not isinstance(seismic_data, type(None)):\n # Direct edge detection from seismic\n seismic_edges = multi_scale_edge_detection(seismic_data)\n else:\n # Fallback to using feature maps for edge detection\n combined_feature = (processed_maps['salt'] + \n processed_maps.get('faults', np.zeros_like(salt_map)) + \n processed_maps.get('layers', np.zeros_like(salt_map))) / 3.0\n seismic_edges = extract_edges(combined_feature)\n \n # 4. Combine edges\n combined_edges = combine_edges(seismic_edges, feature_edges)\n enhanced_edges = close_gaps(combined_edges)\n \n # 5. Segment image into regions\n segmented, num_regions = segment_image(enhanced_edges)\n \n # 6. Assign velocity values to each region\n velocity_model = np.zeros_like(segmented, dtype=float)\n region_types = {}\n \n for region_id in range(1, num_regions + 1):\n # Create mask for this region\n region_mask = (segmented == region_id)\n \n # Skip very small regions\n if np.sum(region_mask) < 5:\n continue\n \n # Determine region characteristics\n region_type = classify_region(region_mask, processed_maps)\n region_types[region_id] = region_type\n \n # Assign appropriate velocity\n velocity_value = get_velocity_for_region_type(region_type)\n \n # Fill the region\n velocity_model[region_mask] = velocity_value\n \n # Handle any unfilled areas\n if np.any(velocity_model == 0):\n velocity_model[velocity_model == 0] = 2500.0 # Default velocity\n \n # 7. Apply geological constraints\n final_velocity = apply_geological_constraints(velocity_model, enhanced_edges)\n \n # 8. If initial velocity is provided, blend with edge-based model\n alpha = 0.7 # Weight for edge-based model\n final_velocity = alpha * final_velocity + (1 - alpha) * initial_velocity\n \n # 9. Ensure output range matches expected range\n # Assuming output should be normalized to [0,1]\n if np.max(final_velocity) > 0:\n normalized = (final_velocity - np.min(final_velocity)) / (np.max(final_velocity) - np.min(final_velocity))\n else:\n normalized = final_velocity\n \n # For debugging: visualize the edge detection and segmentation\n # visualize_edge_segmentation(processed_maps, enhanced_edges, segmented, normalized)\n \n return normalized\n\ndef bimodality_coefficient(data):\n \"\"\"Calculate the bimodality coefficient of a distribution.\n \n A value greater than 0.555 suggests bimodality.\n Formula: (skewness²+1)/(kurtosis+3*(n-1)²/((n-2)*(n-3)))\n \"\"\"\n n = len(data)\n if n < 4:\n return 0.0\n \n # Calculate moments\n mean = np.mean(data)\n var = np.var(data)\n std = np.sqrt(var)\n \n # Skewness\n skewness = np.mean(((data - mean) / std) ** 3) if std > 0 else 0\n \n # Kurtosis (using Fisher's definition: normal = 0)\n kurtosis = np.mean(((data - mean) / std) ** 4) - 3 if std > 0 else 0\n \n # Bimodality coefficient\n numerator = skewness**2 + 1\n denominator = kurtosis + (3 * (n-1)**2)/((n-2)*(n-3))\n \n return numerator / denominator\n\ndef visualize_edge_segmentation(feature_maps, edges, segmented, velocity):\n \"\"\"Visualize the edge-based segmentation process\"\"\"\n plt.figure(figsize=(16, 12))\n \n plt.subplot(231)\n plt.imshow(feature_maps['salt'], cmap='viridis')\n plt.title('Salt Feature Map')\n plt.colorbar()\n \n plt.subplot(232)\n plt.imshow(edges, cmap='gray')\n plt.title('Enhanced Edges')\n plt.colorbar()\n \n plt.subplot(233)\n plt.imshow(segmented, cmap='nipy_spectral')\n plt.title(f'Segmentation ({len(np.unique(segmented))-1} regions)')\n plt.colorbar()\n \n plt.subplot(234)\n if 'faults' in feature_maps:\n plt.imshow(feature_maps['faults'], cmap='plasma')\n plt.title('Faults Feature Map')\n else:\n plt.imshow(np.zeros_like(feature_maps['salt']), cmap='plasma')\n plt.title('No Fault Features')\n plt.colorbar()\n \n plt.subplot(235)\n if 'layers' in feature_maps:\n plt.imshow(feature_maps['layers'], cmap='cividis')\n plt.title('Layers Feature Map')\n else:\n plt.imshow(np.zeros_like(feature_maps['salt']), cmap='cividis')\n plt.title('No Layer Features')\n plt.colorbar()\n \n plt.subplot(236)\n plt.imshow(velocity, cmap='magma')\n plt.title('Final Velocity Model')\n plt.colorbar()\n \n plt.tight_layout()\n plt.savefig('edge_based_segmentation.png')\n plt.close()\n\ndef post_process_prediction_with_edges(pred_np, config, stats=None, seismic_data=None):\n \"\"\"Apply edge-based post-processing to model predictions\"\"\"\n # For feature detection, use the edge-based velocity modeling\n if config['approach'] == 'feature_detection':\n # Apply edge-based velocity modeling\n edge_based_velocity = edge_based_velocity_modeling(pred_np, seismic_data)\n \n if config['post_process']:\n # Apply binary thresholding to get clean velocity model\n binary = (edge_based_velocity > 0.5).astype(np.float32)\n processed = post_process_binary(binary, config)\n return processed\n else:\n # Return edge-based velocity directly\n return edge_based_velocity\n else:\n # For other approaches, use the standard post-processing instead of recursively calling\n # Original implementation (copied here to avoid recursion)\n if config['approach'] == 'thresholding':\n # For thresholding, apply edge enhancement\n if config['post_process']:\n return post_process_binary(pred_np, config)\n else:\n return pred_np\n else: # physics_guided or geo_aware\n # For physics-guided approach\n if config['post_process']:\n # Enhance contrast for clearer boundaries\n from scipy import ndimage\n \n # Smooth with edge-preserving bilateral filter if available\n try:\n from skimage.restoration import denoise_bilateral\n smoothed = denoise_bilateral(\n pred_np.squeeze(), \n sigma_color=0.1, \n sigma_spatial=1,\n multichannel=False\n )\n except:\n # Fallback to gaussian filter\n smoothed = ndimage.gaussian_filter(pred_np.squeeze(), sigma=0.5)\n \n # Apply contrast adjustment\n p2, p98 = np.percentile(smoothed, (2, 98))\n enhanced = np.clip((smoothed - p2) / (p98 - p2), 0, 1)\n \n # Convert to binary if very high contrast\n hist, bins = np.histogram(enhanced, bins=50)\n if np.max(hist) > 0.4 * np.sum(hist):\n # Likely bimodal distribution, apply thresholding\n thresh = apply_thresholding(enhanced, method=config['threshold_method'], stats=stats)\n processed = post_process_binary(thresh, config)\n return processed\n else:\n # Not clearly bimodal, return enhanced version\n return enhanced\n else:\n # Return raw prediction\n return pred_np.squeeze()\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:17.954718Z","iopub.execute_input":"2025-05-12T17:58:17.954893Z","iopub.status.idle":"2025-05-12T17:58:18.163576Z","shell.execute_reply.started":"2025-05-12T17:58:17.954879Z","shell.execute_reply":"2025-05-12T17:58:18.162868Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 12: Inference and Submission Functions\n\ndef load_best_model(config):\n \"\"\"Load the best model for inference\"\"\"\n # Path to best model\n model_path = config['output_dir'] / 'models' / experiment_name / \"best_model.pt\"\n standard_path = config['output_dir'] / 'models' / experiment_name / \"best_model.pt\"\n \n # Create model\n model = create_model(config)\n \n # Try to load model weights\n try:\n # First try the full checkpoint format\n checkpoint = torch.load(model_path, map_location=config['device'])\n if 'model_state_dict' in checkpoint:\n model.load_state_dict(checkpoint['model_state_dict'])\n print(f\"Loaded model from checkpoint at epoch {checkpoint.get('epoch', 'unknown')}\")\n else:\n # If not a checkpoint, try direct state dict\n model.load_state_dict(checkpoint)\n print(f\"Loaded model state dict from {model_path}\")\n except FileNotFoundError:\n try:\n # Try standard path as fallback\n model.load_state_dict(torch.load(standard_path, map_location=config['device']))\n print(f\"Loaded model from standard path: {standard_path}\")\n except FileNotFoundError:\n print(f\"Warning: No model found at {model_path} or {standard_path}\")\n print(\"Using untrained model\")\n \n # Set to evaluation mode\n model.eval()\n \n return model\n\n\ndef post_process_prediction(pred_np, config, stats=None):\n \"\"\"Apply post-processing to model predictions\"\"\"\n # Determine approach-specific processing\n if config['approach'] == 'feature_detection':\n # For feature detection, use the reconstructed velocity\n velocity = pred_np['velocity']\n if config['post_process']:\n # Apply binary thresholding to get clean velocity model\n binary = (velocity > 0.5).astype(np.float32)\n processed = post_process_binary(binary, config)\n return processed\n else:\n # Return raw velocity reconstruction\n return velocity \n \n elif config['approach'] == 'thresholding':\n # For thresholding, apply edge enhancement\n if config['post_process']:\n return post_process_binary(pred_np, config)\n else:\n return pred_np\n \n else: # physics_guided\n # For physics-guided approach\n if config['post_process']:\n # Enhance contrast for clearer boundaries\n from scipy import ndimage\n \n # Smooth with edge-preserving bilateral filter if available\n try:\n from skimage.restoration import denoise_bilateral\n smoothed = denoise_bilateral(\n pred_np.squeeze(), \n sigma_color=0.1, \n sigma_spatial=1,\n multichannel=False\n )\n except:\n # Fallback to gaussian filter\n smoothed = ndimage.gaussian_filter(pred_np.squeeze(), sigma=0.5)\n \n # Apply contrast adjustment\n p2, p98 = np.percentile(smoothed, (2, 98))\n enhanced = np.clip((smoothed - p2) / (p98 - p2), 0, 1)\n \n # Convert to binary if very high contrast\n hist, bins = np.histogram(enhanced, bins=50)\n if np.max(hist) > 0.4 * np.sum(hist):\n # Likely bimodal distribution, apply thresholding\n thresh = apply_thresholding(enhanced, method=config['threshold_method'], stats=stats)\n processed = post_process_binary(thresh, config)\n return processed\n else:\n # Not clearly bimodal, return enhanced version\n return enhanced\n else:\n # Return raw prediction\n return pred_np.squeeze()\n\n\ndef generate_model_submission(model, test_dir, output_path, config, stats=None):\n \"\"\"Generate submission using trained model\"\"\"\n # Get test files\n test_files = sorted(glob.glob(os.path.join(test_dir, \"*.npy\")))\n print(f\"Found {len(test_files)} test files\")\n \n rows = []\n device = config['device']\n \n # Process files with progress tracking\n progress = tqdm(test_files, desc=f\"Generating {config['approach']} submission\")\n \n with torch.no_grad():\n for filepath in progress:\n # Get file ID\n oid = os.path.splitext(os.path.basename(filepath))[0]\n \n # Load data\n data = np.load(filepath)\n \n # Check shape and extract if needed\n if len(data.shape) == 4: # (batch, channels, time, receivers)\n data = data[0] # Use first sample if batched\n \n # Apply gain\n time_steps = data.shape[1]\n time = np.linspace(0, 1, time_steps)\n gain = (time ** 2)[:, np.newaxis]\n \n gained_data = data.copy()\n for c in range(data.shape[0]):\n gained_data[c] = data[c] * gain\n \n # Normalize if statistics are available\n if stats:\n gained_data = (gained_data - stats['input_mean']) / (stats['input_std'] + 1e-6)\n \n # Convert to tensor and add batch dimension\n x = torch.from_numpy(gained_data).float().unsqueeze(0).to(device)\n \n # Predict with model\n if config['approach'] == 'feature_detection':\n pred = model(x)\n # Convert dict of tensors to dict of numpy arrays\n pred_np = {k: v.cpu().numpy() for k, v in pred.items()}\n else:\n pred = model(x)\n pred_np = pred.cpu().numpy()\n \n # Apply post-processing\n processed = post_process_prediction(pred_np, config, stats)\n \n # Ensure correct shape (70x70)\n if processed.shape != (70, 70):\n processed = processed.squeeze()\n if processed.shape != (70, 70):\n from skimage.transform import resize\n processed = resize(processed, (70, 70), order=1, preserve_range=True)\n \n # Format for submission (all rows, odd columns)\n for y in range(70):\n row_id = f\"{oid}_y_{y}\"\n row_data = processed[y, 1:70:2] # Extract odd-indexed columns\n rows.append([row_id] + row_data.tolist())\n \n # Create submission DataFrame\n columns = [\"ID\"] + [f\"x_{i}\" for i in range(1, 70, 2)]\n submission_df = pd.DataFrame(rows, columns=columns)\n \n # Save submission\n submission_df.to_csv(output_path, index=False)\n print(f\"✅ Model-based submission saved at: {output_path}\")\n \n return submission_df\n\n\ndef ensemble_predictions(models, x, config):\n \"\"\"Generate ensemble prediction from multiple models\"\"\"\n device = config['device']\n ensemble_preds = []\n \n with torch.no_grad():\n # Basic predictions\n for model in models:\n if config['approach'] == 'feature_detection':\n pred = model(x)\n # Use velocity reconstruction\n ensemble_preds.append(pred['velocity'])\n else:\n pred = model(x)\n ensemble_preds.append(pred)\n \n # Add test-time augmentation if enabled\n if config.get('use_tta', False):\n # Horizontal flip\n x_flip = torch.flip(x, dims=[3])\n if config['approach'] == 'feature_detection':\n pred_flip = model(x_flip)\n # Flip back the prediction\n velocity_flip = torch.flip(pred_flip['velocity'], dims=[3])\n ensemble_preds.append(velocity_flip)\n else:\n pred_flip = model(x_flip)\n pred_flip = torch.flip(pred_flip, dims=[3])\n ensemble_preds.append(pred_flip)\n \n # Average all predictions\n ensemble_pred = torch.mean(torch.stack(ensemble_preds), dim=0)\n \n return ensemble_pred\n\n\ndef generate_ensemble_submission(models, test_dir, output_path, config, stats=None):\n \"\"\"Generate submission using ensemble of models\"\"\"\n # Get test files\n test_files = sorted(glob.glob(os.path.join(test_dir, \"*.npy\")))\n print(f\"Found {len(test_files)} test files\")\n \n rows = []\n device = config['device']\n \n # Process files with progress tracking\n progress = tqdm(test_files, desc=\"Generating ensemble submission\")\n \n with torch.no_grad():\n for filepath in progress:\n # Get file ID\n oid = os.path.splitext(os.path.basename(filepath))[0]\n \n # Load data\n data = np.load(filepath)\n \n # Check shape and extract if needed\n if len(data.shape) == 4: # (batch, channels, time, receivers)\n data = data[0] # Use first sample if batched\n \n # Apply gain\n time_steps = data.shape[1]\n time = np.linspace(0, 1, time_steps)\n gain = (time ** 2)[:, np.newaxis]\n \n gained_data = data.copy()\n for c in range(data.shape[0]):\n gained_data[c] = data[c] * gain\n \n # Normalize if statistics are available\n if stats:\n gained_data = (gained_data - stats['input_mean']) / (stats['input_std'] + 1e-6)\n \n # Convert to tensor and add batch dimension\n x = torch.from_numpy(gained_data).float().unsqueeze(0).to(device)\n \n # Get ensemble prediction\n ensemble_pred = ensemble_predictions(models, x, config)\n pred_np = ensemble_pred.cpu().numpy()\n \n # Apply post-processing\n processed = post_process_prediction(pred_np, config, stats)\n \n # Ensure correct shape (70x70)\n if processed.shape != (70, 70):\n processed = processed.squeeze()\n if processed.shape != (70, 70):\n from skimage.transform import resize\n processed = resize(processed, (70, 70), order=1, preserve_range=True)\n \n # Format for submission (all rows, odd columns)\n for y in range(70):\n row_id = f\"{oid}_y_{y}\"\n row_data = processed[y, 1:70:2] # Extract odd-indexed columns\n rows.append([row_id] + row_data.tolist())\n \n # Create submission DataFrame\n columns = [\"ID\"] + [f\"x_{i}\" for i in range(1, 70, 2)]\n submission_df = pd.DataFrame(rows, columns=columns)\n \n # Save submission\n submission_df.to_csv(output_path, index=False)\n print(f\"✅ Ensemble submission saved at: {output_path}\")\n \n return submission_df\n\n\ndef generate_submission(config, stats=None):\n \"\"\"Generate submission based on selected approach\"\"\"\n approach = config['approach']\n output_path = config['submission_path']\n \n if approach == 'thresholding':\n # Use direct thresholding approach (no training)\n return generate_threshold_submission(\n config['test_dir'], \n output_path, \n config, \n stats\n )\n elif config['ensemble_submission']:\n # Load multiple models for ensemble\n model_dir = config['output_dir'] / 'models' / experiment_name\n model_files = list(model_dir.glob(\"*model_epoch_*.pt\"))\n \n if len(model_files) > 1:\n print(f\"Found {len(model_files)} models for ensemble\")\n models = []\n \n for model_file in model_files:\n model = create_model(config)\n # Load checkpoint\n checkpoint = torch.load(model_file, map_location=config['device'])\n if 'model_state_dict' in checkpoint:\n model.load_state_dict(checkpoint['model_state_dict'])\n else:\n model.load_state_dict(checkpoint)\n \n model.eval()\n models.append(model)\n \n return generate_ensemble_submission(\n models,\n config['test_dir'],\n output_path,\n config,\n stats\n )\n else:\n # Fall back to single model if not enough models found\n print(\"Not enough models for ensemble, using best model\")\n model = load_best_model(config)\n return generate_model_submission(\n model, \n config['test_dir'], \n output_path, \n config, \n stats\n )\n else:\n # Use single best model\n model = load_best_model(config)\n return generate_model_submission(\n model, \n config['test_dir'], \n output_path, \n config, \n stats\n )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:18.164331Z","iopub.execute_input":"2025-05-12T17:58:18.164667Z","iopub.status.idle":"2025-05-12T17:58:18.290034Z","shell.execute_reply.started":"2025-05-12T17:58:18.164642Z","shell.execute_reply":"2025-05-12T17:58:18.289268Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 13: Run Main Pipeline\n\ndef run_pipeline(config):\n \"\"\"Run the complete pipeline based on configuration\"\"\"\n # 1. Prepare datasets\n train_loader, val_loader, stats = load_and_prepare_data(config)\n \n # 2. Create model and loss function\n if config['approach'] != 'thresholding':\n model = create_model(config)\n criterion = create_loss_function(config)\n \n # 3. Train model\n model, train_losses, val_losses = train_model(\n model, train_loader, val_loader, criterion, config\n )\n else:\n print(\"Using thresholding approach, skipping model training\")\n model = None\n \n # 4. Generate submission\n submission_df = generate_submission(config, stats)\n \n print(f\"Pipeline completed for {config['approach']} approach\")\n print(f\"Submission saved at: {config['submission_path']}\")\n \n return submission_df\n\n\n# Execute the pipeline\nif __name__ == \"__main__\":\n # Run the selected approach\n submission_df = run_pipeline(CONFIG)\n \n # If comparison mode is enabled, run multiple approaches\n if CONFIG['compare_approaches']:\n approaches = ['thresholding', 'physics_guided', 'feature_detection']\n results = {}\n \n for approach in approaches:\n if approach != CONFIG['approach']: # Skip already run approach\n print(f\"\\nRunning comparison pipeline for {approach} approach\")\n \n # Create copy of config with new approach\n comp_config = CONFIG.copy()\n comp_config['approach'] = approach\n comp_config['submission_path'] = f\"submission_{approach}.csv\"\n \n # Run pipeline with this approach\n results[approach] = run_pipeline(comp_config)\n \n # Show comparison summary\n print(\"\\nComparison of approaches completed:\")\n for approach in approaches:\n path = f\"submission_{approach}.csv\" if approach != CONFIG['approach'] else CONFIG['submission_path']\n print(f\"- {approach.capitalize()}: {path}\")\n ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-12T17:58:18.290855Z","iopub.execute_input":"2025-05-12T17:58:18.291524Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 14: Edge-Based Post-Processing Using Existing Model\n# {\\\"_collapsed\\\": false}\\n\n\n# Create configuration for edge-based processing\nCONFIG_EDGE = CONFIG.copy()\nCONFIG_EDGE.update({\n 'post_process': True,\n 'use_edge_based': True,\n 'submission_path': \"submission_edge_based.csv\",\n})\n\n# Dynamically resolve saved model path\nbest_model_path = CONFIG['output_dir'] / 'models' / CONFIG['experiment_name'] / 'best_model.pt'\nprint(f\"Using existing trained model: {best_model_path}\")\n\n# Extract model approach from experiment name\nif CONFIG['experiment_name'].startswith('geo_aware'):\n original_approach = 'geo_aware'\nelif CONFIG['experiment_name'].startswith('physics_guided'):\n original_approach = 'physics_guided'\nelif CONFIG['experiment_name'].startswith('feature_detection'):\n original_approach = 'feature_detection'\nelif CONFIG['experiment_name'].startswith('thresholding'):\n original_approach = 'thresholding'\nelse:\n # Try to infer from directory structure\n dir_name = str(best_model_path.parent)\n if 'geo_aware' in dir_name:\n original_approach = 'geo_aware'\n elif 'physics' in dir_name:\n original_approach = 'physics_guided'\n elif 'feature' in dir_name:\n original_approach = 'feature_detection'\n elif 'threshold' in dir_name:\n original_approach = 'thresholding'\n else:\n print(\"WARNING: Could not determine model approach from experiment name.\")\n print(\"Attempting to load model to determine architecture...\")\n # Default to current approach and we'll try to adapt later\n original_approach = CONFIG['approach']\n\nprint(f\"Detected original model approach: {original_approach}\")\nCONFIG_EDGE['approach'] = original_approach\n\n# Load the model using the detected approach\nmodel = create_model(CONFIG_EDGE)\n\nif best_model_path.exists():\n try:\n model.load_state_dict(torch.load(best_model_path, map_location=CONFIG_EDGE['device']))\n model.eval()\n print(\"✅ Model loaded successfully\")\n except Exception as e:\n print(f\"Error loading model with detected approach: {e}\")\n print(\"Attempting to try other model architectures...\")\n \n # Try all approaches if first attempt fails\n for approach in ['geo_aware', 'physics_guided', 'feature_detection', 'thresholding']:\n if approach == original_approach:\n continue # Skip already tried approach\n \n print(f\"Trying {approach} architecture...\")\n CONFIG_EDGE['approach'] = approach\n model = create_model(CONFIG_EDGE)\n \n try:\n model.load_state_dict(torch.load(best_model_path, map_location=CONFIG_EDGE['device']))\n model.eval()\n print(f\"✅ Model loaded successfully using {approach} architecture\")\n break\n except Exception:\n print(f\"Failed to load using {approach} architecture\")\n \n if model is None:\n raise FileNotFoundError(f\"❌ Could not load model with any architecture. Please check the model path.\")\nelse:\n raise FileNotFoundError(f\"❌ No model found at {best_model_path}. Run training first.\")\n\n# Get dataset statistics\n_, _, stats = load_and_prepare_data(CONFIG_EDGE)\n\n# Generate submission with edge-based post-processing\ndef generate_edge_based_submission():\n \"\"\"Generate submission using edge-based post-processing on existing model\"\"\"\n original_post_process = globals()['post_process_prediction']\n globals()['post_process_prediction'] = post_process_prediction_with_edges\n\n print(\"Generating submission with edge-based post-processing...\")\n submission_df = generate_model_submission(\n model,\n CONFIG_EDGE['test_dir'],\n CONFIG_EDGE['submission_path'],\n CONFIG_EDGE,\n stats\n )\n\n globals()['post_process_prediction'] = original_post_process\n print(f\"✅ Edge-based submission saved to: {CONFIG_EDGE['submission_path']}\")\n return submission_df\n\n\n\n# Run it\nsubmission_df = generate_edge_based_submission()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 15: Additional Analysis and Visualization Tools (Optional)\n# \n# {\\\"_collapsed\\\": true}\\n\nif CONFIG.get('run_analysis_tools', True):\n\n import numpy as np\n import matplotlib.pyplot as plt\n from scipy import stats as scipy_stats\n from scipy import ndimage\n from skimage import filters\n from scipy import fft\n import os\n\n def analyze_velocity_models(output_files, max_samples=10):\n \"\"\"Analyze velocity models to understand their characteristics\"\"\"\n sample_files = output_files[:min(len(output_files), max_samples)]\n samples = []\n\n for f in sample_files:\n try:\n data = np.load(f, mmap_mode='r')\n if len(data.shape) == 4:\n samples.append(data[:5])\n elif len(data.shape) == 3:\n samples.append(data[:5, np.newaxis])\n except Exception as e:\n print(f\"Error loading {f}: {e}\")\n\n if not samples:\n print(\"No samples loaded for analysis\")\n return\n\n all_samples = np.concatenate(samples, axis=0).squeeze()\n print(f\"Analyzing {all_samples.shape[0]} velocity models\")\n\n mean_vel = np.mean(all_samples)\n std_vel = np.std(all_samples)\n min_vel = np.min(all_samples)\n max_vel = np.max(all_samples)\n\n print(f\"Velocity statistics: mean={mean_vel:.2f}, std={std_vel:.2f}, min={min_vel:.2f}, max={max_vel:.2f}\")\n\n plt.figure(figsize=(15, 5))\n plt.subplot(131)\n plt.hist(all_samples.flatten(), bins=50)\n plt.title(\"Velocity Distribution\")\n plt.xlabel(\"Velocity Value\")\n plt.ylabel(\"Frequency\")\n\n kde = scipy_stats.gaussian_kde(all_samples.flatten())\n x = np.linspace(min_vel, max_vel, 1000)\n plt.plot(x, kde(x) * len(all_samples.flatten()) * (max_vel - min_vel) / 50, 'r-', linewidth=2)\n\n grad_y, grad_x = np.gradient(all_samples[0])\n grad_mag = np.sqrt(grad_x**2 + grad_y**2)\n\n plt.subplot(132)\n plt.hist(grad_mag.flatten(), bins=50)\n plt.title(\"Gradient Magnitude Distribution\")\n plt.xlabel(\"Gradient Magnitude\")\n plt.ylabel(\"Frequency\")\n\n plt.subplot(133)\n layer_profile = np.mean(all_samples, axis=(0, 2))\n plt.plot(layer_profile, np.arange(len(layer_profile)))\n plt.title(\"Average Horizontal Layer Profile\")\n plt.xlabel(\"Average Velocity\")\n plt.ylabel(\"Depth\")\n plt.gca().invert_yaxis()\n\n os.makedirs(CONFIG['output_dir'] / 'visualizations', exist_ok=True)\n plt.tight_layout()\n plt.savefig(CONFIG['output_dir'] / 'visualizations' / \"velocity_analysis.png\")\n plt.show()\n\n plt.figure(figsize=(15, 5 * min(3, len(all_samples))))\n for i in range(min(3, len(all_samples))):\n plt.subplot(min(3, len(all_samples)), 3, i*3 + 1)\n plt.imshow(all_samples[i], cmap='magma', origin='lower')\n plt.title(f\"Velocity Model #{i+1}\")\n plt.colorbar()\n\n grad_y, grad_x = np.gradient(all_samples[i])\n edges = np.sqrt(grad_x**2 + grad_y**2)\n\n plt.subplot(min(3, len(all_samples)), 3, i*3 + 2)\n plt.imshow(edges, cmap='viridis', origin='lower')\n plt.title(f\"Edge Detection #{i+1}\")\n plt.colorbar()\n\n try:\n thresh = filters.threshold_otsu(all_samples[i])\n except:\n thresh = np.mean(all_samples[i])\n binary = (all_samples[i] > thresh).astype(float)\n\n plt.subplot(min(3, len(all_samples)), 3, i*3 + 3)\n plt.imshow(binary, cmap='gray', origin='lower')\n plt.title(f\"Binary Threshold (t={thresh:.2f}) #{i+1}\")\n\n plt.tight_layout()\n plt.savefig(CONFIG['output_dir'] / 'visualizations' / \"velocity_examples.png\")\n plt.show()\n\n \n \n return {\n 'mean': mean_vel,\n 'std': std_vel,\n 'min': min_vel,\n 'max': max_vel,\n 'bimodal': bimodality_coefficient(all_samples.flatten()) > 0.555\n }\n \n def visualize_seismic_data(input_files, max_samples=5):\n \"\"\"Visualize seismic data to understand input characteristics\"\"\"\n sample_files = input_files[:min(len(input_files), max_samples)]\n samples = []\n\n for f in sample_files:\n try:\n data = np.load(f, mmap_mode='r')\n samples.append(data[:1])\n except Exception as e:\n print(f\"Error loading {f}: {e}\")\n\n if not samples:\n print(\"No samples loaded for visualization\")\n return\n\n all_samples = np.concatenate(samples, axis=0)\n print(f\"Visualizing {all_samples.shape[0]} seismic data samples\")\n\n plt.figure(figsize=(15, 4 * min(3, len(all_samples))))\n\n for i in range(min(3, len(all_samples))):\n plt.subplot(min(3, len(all_samples)), 3, i*3 + 1)\n plt.imshow(all_samples[i, 0], cmap='seismic', aspect='auto')\n plt.title(f\"Raw Seismic (Ch 1) #{i+1}\")\n plt.colorbar()\n\n time_steps = all_samples.shape[2]\n time = np.linspace(0, 1, time_steps)\n gain = (time ** 2)[:, np.newaxis]\n gained = all_samples[i, 0] * gain\n\n plt.subplot(min(3, len(all_samples)), 3, i*3 + 2)\n plt.imshow(gained, cmap='seismic', aspect='auto')\n plt.title(f\"With Time Gain #{i+1}\")\n plt.colorbar()\n\n freq = np.abs(fft.fft2(all_samples[i, 0]))\n freq = np.fft.fftshift(freq)\n\n plt.subplot(min(3, len(all_samples)), 3, i*3 + 3)\n plt.imshow(np.log(freq + 1), cmap='viridis', aspect='auto')\n plt.title(f\"Frequency Spectrum #{i+1}\")\n plt.colorbar()\n\n os.makedirs(CONFIG['output_dir'] / 'visualizations', exist_ok=True)\n plt.tight_layout()\n plt.savefig(CONFIG['output_dir'] / 'visualizations' / \"seismic_visualization.png\")\n plt.show()\n \n # Actually call the analysis functions with data from the configuration\n print(\"Running analysis and visualization tools...\")\n \n # Find input and output files using the same pattern as in load_and_prepare_data\n input_files = sorted([f for f in CONFIG['data_dir'].rglob(\"*.npy\") \n if 'seis' in f.name or 'data' in f.name])\n output_files = [Path(str(f).replace(\"seis\", \"vel\").replace(\"data\", \"model\")) \n for f in input_files]\n \n # Verify files exist (same as in load_and_prepare_data function)\n input_files = [f for f in input_files if f.exists()]\n output_files = [f for f in output_files if f.exists()]\n \n if len(input_files) > 0 and len(output_files) > 0:\n print(f\"Found {len(input_files)} input files and {len(output_files)} output files for analysis\")\n \n # Run the visualization functions\n visualize_seismic_data(input_files)\n velocity_stats = analyze_velocity_models(output_files)\n \n # Display the statistics summary\n if velocity_stats:\n print(\"\\nVelocity Model Statistics Summary:\")\n for key, value in velocity_stats.items():\n print(f\"- {key}: {value}\")\n else:\n print(\"No files found for analysis. Please check data paths in configuration.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell: Visualizing Ground Truth vs Predictions for Best Model\n\ndef visualize_model_predictions(model, data_loader, config, num_samples=5):\n \"\"\"\n Generate visualizations comparing ground truth to model predictions\n \n Args:\n model: Trained PyTorch model\n data_loader: DataLoader containing validation or test data\n config: Configuration dictionary\n num_samples: Number of samples to visualize\n \"\"\"\n device = config['device']\n model.eval()\n \n # Get samples from the data loader\n samples = []\n targets = []\n \n with torch.no_grad():\n for x, y in data_loader:\n samples.append(x)\n targets.append(y)\n if len(samples) >= num_samples:\n break\n \n # Concatenate samples if we got multiple batches\n if len(samples) > 1:\n x_samples = torch.cat(samples, dim=0)[:num_samples]\n y_samples = torch.cat(targets, dim=0)[:num_samples]\n else:\n x_samples = samples[0][:num_samples]\n y_samples = targets[0][:num_samples]\n \n # Move to device\n x_samples = x_samples.to(device)\n y_samples = y_samples.to(device)\n \n # Generate predictions\n predictions = []\n with torch.no_grad():\n for i in range(num_samples):\n x = x_samples[i:i+1] # Add batch dimension\n pred = model(x)\n predictions.append(pred.cpu())\n \n # Visualize comparisons\n fig, axes = plt.subplots(num_samples, 4, figsize=(20, 5*num_samples))\n \n # Handle single sample case\n if num_samples == 1:\n axes = axes.reshape(1, -1)\n \n for i in range(num_samples):\n # Extract single sample, ground truth and prediction\n seismic = x_samples[i].cpu().numpy()\n truth = y_samples[i].cpu().numpy()\n pred = predictions[i].numpy()\n \n # Ensure we're working with the right dimensions\n if seismic.ndim > 3:\n seismic = seismic[0] # Take first channel if multi-channel\n \n if truth.ndim > 2:\n if truth.shape[0] == 1: # Single channel\n truth = truth[0]\n else: # Multi-channel (for feature detection)\n truth = truth[0] # Just take first channel for visualization\n \n if pred.ndim > 2:\n if pred.shape[0] == 1: # Single channel\n pred = pred[0]\n elif isinstance(pred, dict): # Feature detection output\n pred = pred['velocity'][0]\n else:\n pred = pred[0] # Just take first channel\n \n # Calculate error\n error = np.abs(pred - truth)\n \n # Plot seismic data (first channel)\n axes[i, 0].imshow(seismic[0], cmap='seismic', aspect='auto')\n axes[i, 0].set_title(f\"Input Seismic #{i+1}\")\n axes[i, 0].axis('off')\n \n # Plot ground truth\n axes[i, 1].imshow(truth, cmap='magma', aspect='auto')\n axes[i, 1].set_title(f\"Ground Truth #{i+1}\")\n axes[i, 1].axis('off')\n \n # Plot prediction\n axes[i, 2].imshow(pred, cmap='magma', aspect='auto')\n axes[i, 2].set_title(f\"Prediction #{i+1}\")\n axes[i, 2].axis('off')\n \n # Plot error\n im = axes[i, 3].imshow(error, cmap='inferno', aspect='auto')\n axes[i, 3].set_title(f\"Absolute Error #{i+1}\")\n axes[i, 3].axis('off')\n \n # Add colorbar for error\n plt.colorbar(im, ax=axes[i, 3])\n \n plt.tight_layout()\n \n # Save the figure\n os.makedirs(config['output_dir'] / 'visualizations', exist_ok=True)\n fig_path = config['output_dir'] / 'visualizations' / 'ground_truth_vs_predictions.png'\n plt.savefig(fig_path)\n \n print(f\"Visualizations saved to {fig_path}\")\n \n # Display the figure if in notebook environment\n plt.show()\n \n return fig_path\n\n# Now call the function with your best model and validation loader\nprint(\"Generating ground truth vs. prediction visualizations...\")\n\n# Get the best model\nbest_model = load_best_model(CONFIG)\n\n# Get a data loader for visualization\n# Re-use validation data loader from training\n_, val_loader, _ = load_and_prepare_data(CONFIG)\n\n# Generate visualizations\nvis_path = visualize_model_predictions(best_model, val_loader, CONFIG, num_samples=5)\n\nprint(f\"✅ Visualization completed and saved to: {vis_path}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}