""" SIH26077 — Spatiotemporal Data Loader ====================================== Loads and aligns multi-modal atmospheric data from three sources: 1. IMDAA Reanalysis (.nc) → (30, 6, 256, 256) thermodynamic + kinematic fields 2. INSAT-3DR Satellite (.h5) → (3, 6, 256, 256) WV, CTT, HEM 3. CartoDEM (.tif) → (2, 256, 256) elevation + slope All modalities are spatially aligned to a unified 256×256 grid and channel-wise z-score normalized for stable training. """ import os import json import torch import numpy as np import xarray as xr import h5py import rasterio from torch.utils.data import Dataset, DataLoader import torch.nn.functional as F import warnings from config import ( INDEX_PATH, DATASET_ROOT, DEM_PATH, GRID_SIZE, DEFAULT_LEAD_TIME, IMDAA_MEAN, IMDAA_STD, INSAT_MEAN, INSAT_STD, TERRAIN_MEAN, TERRAIN_STD, ) warnings.filterwarnings('ignore') def resize_tensor(tensor, size): """Resizes a 2D or 3D/4D spatial tensor to the target (H, W) grid. Uses bilinear interpolation. Handles tensors of shape: - (H, W) → unsqueeze to (1, 1, H, W) → resize → squeeze back - (C, H, W) → unsqueeze to (1, C, H, W) → resize → squeeze back - (N, C, H, W) → resize directly """ original_shape = tensor.shape if len(tensor.shape) == 2: tensor = tensor.unsqueeze(0).unsqueeze(0) elif len(tensor.shape) == 3: tensor = tensor.unsqueeze(0) tensor = tensor.float() resized = F.interpolate(tensor, size=size, mode='bilinear', align_corners=False) if len(original_shape) == 2: return resized.squeeze(0).squeeze(0) elif len(original_shape) == 3: return resized.squeeze(0) return resized def resize_binary_target(tensor, size): """Resize a binary (0/1) label map using nearest-neighbor interpolation. WHY NOT bilinear for binary masks: Bilinear on a sparse binary mask (e.g. CB at 2816×2805 → 256×256) smears each '1' pixel into a tiny float blur (~0.0001). After downscaling 11× the signal is effectively zero — the model sees no positive pixels and the loss gradient collapses. Nearest-neighbor keeps every '1' as a hard '1' and every '0' as '0', preserving the true label distribution at the new resolution. """ original_shape = tensor.shape if len(tensor.shape) == 2: tensor = tensor.unsqueeze(0).unsqueeze(0) elif len(tensor.shape) == 3: tensor = tensor.unsqueeze(0) tensor = tensor.float() resized = F.interpolate(tensor, size=size, mode='nearest') # Hard-threshold to guarantee strict binary output (no float residuals) resized = (resized > 0.5).float() if len(original_shape) == 2: return resized.squeeze(0).squeeze(0) elif len(original_shape) == 3: return resized.squeeze(0) return resized class SpatiotemporalDataset(Dataset): """Multi-modal spatiotemporal dataset for severe weather nowcasting. Each sample is a 6-timestep window (~18 hours) containing: - IMDAA reanalysis: 30 atmospheric channels across 6 timesteps - INSAT satellite: 3 observational channels across 6 timesteps - CartoDEM terrain: 2 static channels (elevation + slope) - Targets: 3 binary risk maps (cloudburst, thunderstorm, flash flood) Args: index_path: Path to window_index.json root_dir: Path to dataset_root/ lead_time: Which lead time to predict ('2', '3', '4', '5', or '6' hours) grid_size: Unified spatial grid (H, W) for all modalities normalize: Whether to apply z-score normalization (default: True) """ def __init__(self, index_path=INDEX_PATH, root_dir=DATASET_ROOT, lead_time=DEFAULT_LEAD_TIME, grid_size=GRID_SIZE, normalize=True): with open(index_path, 'r') as f: self.windows = json.load(f) self.root_dir = root_dir self.lead_time = str(lead_time) self.grid_size = grid_size self.normalize = normalize # Load static DEM once (same for all windows) self.dem_tensor = self._load_dem() def __len__(self): return len(self.windows) def _load_dem(self): """Load and preprocess the Digital Elevation Model.""" print("Loading and downsampling static DEM...") with rasterio.open(DEM_PATH) as src: # Downsample immediately to save memory factor = max(src.width // 1000, 1) elevation = src.read( 1, out_shape=(src.height // factor, src.width // factor) ).astype(np.float32) # Compute slope from elevation gradients dy, dx = np.gradient(elevation) slope = np.sqrt(dx**2 + dy**2) # Stack into (C=2, H, W): [elevation, slope] terrain = np.stack([elevation, slope], axis=0) terrain_tensor = torch.from_numpy(terrain) terrain_tensor = torch.nan_to_num(terrain_tensor, nan=0.0) # Resize to unified grid terrain_tensor = resize_tensor(terrain_tensor, self.grid_size) # Normalize if self.normalize: terrain_tensor = (terrain_tensor - TERRAIN_MEAN) / (TERRAIN_STD + 1e-8) return terrain_tensor def _load_imdaa(self, paths): """Load IMDAA reanalysis data into (Channels=30, Time=6, H, W). CRITICAL FIX — group by TIMESTAMP, not by alphabetical sort: The 180 paths span 6 timestamps × 30 channels (5 vars × 6 levels). When sorted alphabetically and chunked by 30, each 'timestep' chunk contains only ONE variable repeated (e.g. all HGT-1000mb through HGT-925mb at 6 different times) — NOT all variables at one time. The 3D temporal conv was therefore operating on var-grouped slices with no actual temporal meaning. Correct grouping: extract the 10-digit timestamp from each filename (e.g. '2019081600'), group files sharing the same timestamp, then sort groups chronologically. Each group is a genuine snapshot of all 5 atmospheric variables at one moment in time. """ import re from collections import defaultdict # Group by 10-digit timestamp embedded in filename (YYYYMMDDHH) ts_groups = defaultdict(list) for p in paths: fname = os.path.basename(p.replace('\\', '/')) m = re.search(r'_(\d{10})_', fname) if m: ts_groups[m.group(1)].append(p) if not ts_groups: # Fallback to old alphabetical chunking if regex fails paths_sorted = sorted(paths) chunks = [paths_sorted[i:i+30] for i in range(0, len(paths_sorted), 30)] else: # Sort chronologically; within each timestamp, sort alphabetically # (alphabetical within-timestamp gives: HGT, RH, TMP, UGRD, VGRD × levels) chunks = [sorted(ts_groups[ts]) for ts in sorted(ts_groups.keys())] time_steps = [] for chunk in chunks: channels = [] for path in chunk: try: path = path.replace('\\', '/') ds = xr.open_dataset(path) var_name = list(ds.data_vars)[0] data = ds[var_name].squeeze().values data = np.nan_to_num(data, nan=0.0) channels.append(torch.from_numpy(data)) ds.close() except Exception as e: print(f"Error loading {path}: {e}") channels.append(torch.zeros((501, 751))) # (30, H_raw, W_raw) — one full atmospheric state snapshot timestep_tensor = torch.stack(channels, dim=0) time_steps.append(timestep_tensor) # (Time=6, C=30, H, W) → (C=30, Time=6, H, W) imdaa_tensor = torch.stack(time_steps, dim=0).permute(1, 0, 2, 3) imdaa_tensor = resize_tensor(imdaa_tensor, self.grid_size) # Normalize channel-wise (broadcast across Time, H, W) if self.normalize: # IMDAA_MEAN/STD shape: (30, 1, 1, 1) — broadcasts over (30, 6, 256, 256) imdaa_tensor = (imdaa_tensor - IMDAA_MEAN) / (IMDAA_STD + 1e-8) return imdaa_tensor def _load_insat(self, l1b_paths, ctp_paths, hem_paths): """Load INSAT satellite data into (Channels=4, Time=6, H, W). Four channels per timestep: - Ch 0: WV — Water Vapor brightness temperature (L1B IMG_WV key) - Ch 1: CTT — Cloud Top Temperature (L2B CTP 'CTT' key, Kelvin) - Ch 2: HEM — Hydro-Estimator precipitation rate (L2B HEM 'HEM' key, mm/30min) - Ch 3: CTT_RATE — Frame-to-frame CTT change (K per 3h step, NEGATIVE = cooling = storm building) CTT_RATE (ps.md explicit requirement): "Rapid cooling of cloud tops (CTT Drop Rate) provides real-time validation of explosive vertical updrafts within the system." Without CTT_RATE, the model sees static snapshots and cannot detect convective intensification. A -15K/step drop in CTT is the clearest single-variable cloudburst precursor available from satellite. """ time_steps = [] num_steps = max(len(l1b_paths), len(ctp_paths), len(hem_paths), 6) # --- Load raw CTT values for all timesteps first --- # Needed to compute temporal differences (drop rate) raw_ctts = [] for i in range(num_steps): if i < len(ctp_paths): try: with h5py.File(ctp_paths[i].replace('\\', '/'), 'r') as f: ctt = np.squeeze(f['CTT'][:]).astype(np.float32) ctt = np.ma.filled(np.ma.masked_where(ctt < 0, ctt), 0.0) raw_ctts.append(ctt) except Exception: raw_ctts.append(np.zeros((313, 312), dtype=np.float32)) else: raw_ctts.append(np.zeros((313, 312), dtype=np.float32)) # CTT drop rate: diff between consecutive frames (negative = cooling) # At t=0 there is no prior frame — use zero (no rate info). ctt_rates = [np.zeros_like(raw_ctts[0])] for i in range(1, len(raw_ctts)): ctt_rates.append(raw_ctts[i] - raw_ctts[i - 1]) # neg = cooling # --- Build per-timestep channel stacks --- for i in range(num_steps): channels = [] # Ch 0: WV (L1B IMG_WV) if i < len(l1b_paths): try: with h5py.File(l1b_paths[i].replace('\\', '/'), 'r') as f: wv = np.squeeze(f['IMG_WV'][:]).astype(np.float32) wv = wv[::4, ::4] # subsample: 1408×1402 → 352×351 channels.append(torch.from_numpy(np.nan_to_num(wv, nan=0.0))) except Exception: channels.append(torch.zeros((352, 351))) else: channels.append(torch.zeros((352, 351))) # Ch 1: CTT (L2B CTP) channels.append(torch.from_numpy(raw_ctts[i])) # Ch 2: HEM precipitation rate (L2B HEM) if i < len(hem_paths): try: with h5py.File(hem_paths[i].replace('\\', '/'), 'r') as f: hem = np.squeeze(f['HEM'][:]).astype(np.float32) hem[hem < 0] = 0.0 hem[hem > 500] = 0.0 hem = hem[::8, ::8] # subsample: 2816×2805 → 352×351 channels.append(torch.from_numpy(hem)) except Exception: channels.append(torch.zeros((352, 351))) else: channels.append(torch.zeros((352, 351))) # Ch 3: CTT drop rate (K/step, negative = explosive cooling) channels.append(torch.from_numpy(ctt_rates[i])) # Resize each channel to grid_size (different raw sizes per channel) resized = [resize_tensor(c, self.grid_size) for c in channels] time_steps.append(torch.stack(resized, dim=0)) # (Time=6, C=4, H, W) → (C=4, Time=6, H, W) insat_tensor = torch.stack(time_steps, dim=0).permute(1, 0, 2, 3) # Normalize channel-wise if self.normalize: insat_tensor = (insat_tensor - INSAT_MEAN) / (INSAT_STD + 1e-8) return insat_tensor def _load_targets(self, target_dict): """Load 3 binary target maps into (Channels=3, H, W). Channel order: [Cloudburst, Thunderstorm, FlashFlood] IMPORTANT — Flash Flood label fix: flash_flood.npy may contain raw QPE values (mm/3hr) rather than a pre-thresholded binary mask. Thunderstorm was stored pre-binarized (CTT < 208.15K → 1) but FF was not. We apply the threshold here. If max value > 10.0 → treat as continuous mm values → threshold at 50mm. If max value ≤ 2.0 → already binary, load as-is. """ cb = np.load(target_dict['cloudburst'].replace('\\', '/')) ts = np.load(target_dict['thunderstorm'].replace('\\', '/')) ff = np.load(target_dict['flash_flood'].replace('\\', '/')) cb_np = cb.squeeze().astype(np.float32) ts_np = ts.squeeze().astype(np.float32) ff_np = ff.squeeze().astype(np.float32) # ── Flash Flood threshold fix ──────────────────────────────────── # If values are continuous QPE (mm), binarize at 50mm (cloudburst threshold) # AND require slope > 12° (encoded as a multiplier below — slope mask # is applied during the DEM-overlay step in the model output, not here, # but the QPE threshold alone creates valid FF labels) if ff_np.max() > 10.0: ff_np = (ff_np >= 50.0).astype(np.float32) # Same check for CB in case it also stores raw QPE if cb_np.max() > 10.0: cb_np = (cb_np >= 50.0).astype(np.float32) cb = torch.from_numpy(cb_np).float() ts = torch.from_numpy(ts_np).float() ff = torch.from_numpy(ff_np).float() # Ensure 2D (take first channel if 3D) if len(cb.shape) > 2: cb = cb[0] if len(ts.shape) > 2: ts = ts[0] if len(ff.shape) > 2: ff = ff[0] # Use nearest-neighbor for all binary targets — preserves sparse 0/1 signals. # Bilinear would smear a 25-pixel CB mask over a 2816×2805 grid into # near-zero floats after downscaling to 256×256 (label destruction). cb = resize_binary_target(cb, self.grid_size) ts = resize_binary_target(ts, self.grid_size) ff = resize_binary_target(ff, self.grid_size) # Stack: [Cloudburst, Thunderstorm, FlashFlood] return torch.stack([cb, ts, ff], dim=0) def __getitem__(self, idx): window = self.windows[idx] # 1. IMDAA (30, 6, H, W) imdaa = self._load_imdaa(window['imdaa_paths']) # 2. INSAT (3, 6, H, W) insat = self._load_insat( window.get('insat_l1b', []), window.get('insat_l2b_ctp', []), window.get('insat_l2b_hem', []), ) # 3. Terrain (2, H, W) — pre-loaded and shared terrain = self.dem_tensor # 4. Targets (3, H, W) target_dict = window['targets_by_lead'][self.lead_time] targets = self._load_targets(target_dict) return { 'imdaa': imdaa, # (30, 6, 256, 256) 'insat': insat, # (3, 6, 256, 256) 'terrain': terrain, # (2, 256, 256) 'targets': targets, # (3, 256, 256) } # ============================================================ # Self-Test # ============================================================ if __name__ == "__main__": print("Testing SpatiotemporalDataset...") dataset = SpatiotemporalDataset() print(f"Dataset length: {len(dataset)}") sample = dataset[0] print("\nSample Shapes:") print(f"IMDAA: {sample['imdaa'].shape} (Channels, Time, H, W)") print(f"INSAT: {sample['insat'].shape} (Channels, Time, H, W)") print(f"Terrain: {sample['terrain'].shape} (Channels, H, W)") print(f"Targets: {sample['targets'].shape} (Channels, H, W)") print("\nNormalization Check (should be near mean=0, std=1):") imdaa = sample['imdaa'] print(f" IMDAA ch0 mean: {imdaa[0].mean():.2f}, std: {imdaa[0].std():.2f}") print(f" IMDAA ch15 mean: {imdaa[15].mean():.2f}, std: {imdaa[15].std():.2f}") # Check for NaNs for k, v in sample.items(): nan_count = torch.isnan(v).sum().item() print(f"{k.capitalize()} NaNs: {nan_count}") print("\n[OK] DataLoader test passed.")