Upload data_loader.py with huggingface_hub

data_loader.py

@@ -0,0 +1,412 @@
+"""
+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.")