""" SIH26077 — Hybrid TransUNet: Spatiotemporal Multi-Task Model ============================================================== A multi-task deep learning model for simultaneous prediction of: 1. Cloudburst probability 2. Severe Thunderstorm probability 3. Flash Flood probability Architecture (Hybrid CNN-Transformer): - Temporal Compression: 3D convolutions collapse 6 timesteps -> 1 feature map - Early Fusion: IMDAA (64ch) + INSAT (32ch) + Terrain (16ch) = 112 channels - Shared Encoder: 2-level U-Net backbone (128 -> 256 -> 512 bottleneck) - Transformer Bottleneck: Multi-Head Self-Attention (8 heads) with positional encoding at the bottleneck (64x64 spatial resolution). This provides global receptive field awareness, allowing the model to cross-reference distant atmospheric patterns (e.g., moisture from Arabian Sea triggering storms in the Himalayas) — satisfying the MoES requirement for "attention mechanisms". - Multi-Task Decoders: 3 independent U-Net decoders with skip connections Input shapes: - IMDAA: (B, 30, 6, 256, 256) — 30 atmospheric channels, 6 timesteps - INSAT: (B, 4, 6, 256, 256) — 4 satellite channels (WV, CTT, HEM, CTT_RATE), 6 timesteps - Terrain: (B, 2, 256, 256) — elevation + slope (static) Output: - (B, 3, 256, 256) — logits for [Cloudburst, Thunderstorm, FlashFlood] """ import torch import torch.nn as nn import torch.nn.functional as F import math class ConvBlock2D(nn.Module): """Double convolution block: Conv -> BN -> ReLU -> Conv -> BN -> ReLU. Args: dropout: If > 0, adds Dropout2d after the second ReLU. Use in decoder heads (not encoder) to prevent overfitting on small datasets (< 100 training samples). """ def __init__(self, in_c, out_c, dropout=0.0): super().__init__() layers = [ nn.Conv2d(in_c, out_c, kernel_size=3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True), nn.Conv2d(out_c, out_c, kernel_size=3, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(inplace=True), ] if dropout > 0.0: layers.append(nn.Dropout2d(p=dropout)) self.conv = nn.Sequential(*layers) def forward(self, x): return self.conv(x) class SelfAttention(nn.Module): """Multi-Head Self-Attention block for the TransUNet bottleneck. Flattens spatial dimensions (H*W) into a sequence of tokens, applies standard Multi-Head Attention with learned positional encoding, then reshapes back to 2D spatial feature maps. This enables the model to capture long-range spatial dependencies (global receptive field) that pure convolutions cannot achieve, while keeping compute manageable at 64x64 resolution (4096 tokens). Args: embed_dim: Number of channels in the feature map (default: 512) num_heads: Number of attention heads (default: 8) spatial_size: Spatial resolution at bottleneck (default: 64) dropout: Dropout rate for attention weights (default: 0.1) """ def __init__(self, embed_dim=512, num_heads=8, spatial_size=64, dropout=0.1): super().__init__() self.embed_dim = embed_dim self.spatial_size = spatial_size # Layer normalization (pre-norm architecture for stability) self.norm = nn.LayerNorm(embed_dim) # Standard Multi-Head Attention self.attn = nn.MultiheadAttention( embed_dim=embed_dim, num_heads=num_heads, dropout=dropout, batch_first=True, ) # Learned 2D positional encoding num_tokens = spatial_size * spatial_size # 64*64 = 4096 self.pos_embed = nn.Parameter( torch.randn(1, num_tokens, embed_dim) * 0.02 ) # Feed-forward network after attention (standard Transformer block) self.ffn = nn.Sequential( nn.LayerNorm(embed_dim), nn.Linear(embed_dim, embed_dim * 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(embed_dim * 2, embed_dim), nn.Dropout(dropout), ) def forward(self, x): """ Args: x: (B, C, H, W) — feature map from the CNN encoder Returns: (B, C, H, W) — attention-enhanced feature map (same shape) """ B, C, H, W = x.shape # Flatten spatial dims: (B, C, H, W) -> (B, H*W, C) x_flat = x.flatten(2).transpose(1, 2) # (B, N, C) where N = H*W # Add positional encoding x_pos = x_flat + self.pos_embed[:, :H*W, :] # Pre-norm + Multi-Head Self Attention + Residual x_norm = self.norm(x_pos) attn_out, _ = self.attn(x_norm, x_norm, x_norm) x_flat = x_flat + attn_out # Residual connection # Feed-forward + Residual x_flat = x_flat + self.ffn(x_flat) # Reshape back to 2D: (B, N, C) -> (B, C, H, W) out = x_flat.transpose(1, 2).reshape(B, C, H, W) return out class SpatiotemporalMultiTaskModel(nn.Module): """Hybrid TransUNet: Multi-task spatiotemporal model for severe weather nowcasting. Combines a CNN encoder/decoder with a Transformer attention mechanism at the bottleneck for global atmospheric pattern recognition. """ def __init__(self): super().__init__() # -- 1. Temporal Compression (3D -> 2D) ---------------------- # IMDAA: (B, 30, T=6, H, W) -> (B, 64, 1, H, W) -> squeeze -> (B, 64, H, W) self.imdaa_time_compress = nn.Sequential( nn.Conv3d(30, 32, kernel_size=(3, 3, 3), padding=(0, 1, 1)), # T: 6->4 nn.BatchNorm3d(32), nn.ReLU(inplace=True), nn.Conv3d(32, 64, kernel_size=(4, 3, 3), padding=(0, 1, 1)), # T: 4->1 nn.BatchNorm3d(64), nn.ReLU(inplace=True), ) # INSAT: (B, 4, T=6, H, W) -> (B, 32, 1, H, W) -> squeeze -> (B, 32, H, W) # 4 channels: WV, CTT, HEM, CTT_RATE (drop rate, ps.md requirement) self.insat_time_compress = nn.Sequential( nn.Conv3d(4, 16, kernel_size=(3, 3, 3), padding=(0, 1, 1)), # T: 6->4 nn.BatchNorm3d(16), nn.ReLU(inplace=True), nn.Conv3d(16, 32, kernel_size=(4, 3, 3), padding=(0, 1, 1)), # T: 4->1 nn.BatchNorm3d(32), nn.ReLU(inplace=True), ) # Terrain embedding: (B, 2, H, W) -> (B, 16, H, W) self.terrain_embed = nn.Sequential( nn.Conv2d(2, 16, kernel_size=3, padding=1), nn.BatchNorm2d(16), nn.ReLU(inplace=True), ) # -- 2. Shared Encoder ---------------------------------------- # Fused channels: 64 (IMDAA) + 32 (INSAT) + 16 (Terrain) = 112 self.enc1 = ConvBlock2D(112, 128) # -> (B, 128, 256, 256) self.pool1 = nn.MaxPool2d(2) # -> (B, 128, 128, 128) self.enc2 = ConvBlock2D(128, 256) # -> (B, 256, 128, 128) self.pool2 = nn.MaxPool2d(2) # -> (B, 256, 64, 64) # -- 3. Hybrid Bottleneck (CNN + Transformer) ----------------- self.bottleneck = ConvBlock2D(256, 512) # -> (B, 512, 64, 64) self.attention = SelfAttention( # Transformer attention embed_dim=512, num_heads=8, spatial_size=64, dropout=0.1, ) # -- 4. Multi-Task Decoder Heads ------------------------------ # Each head has its own upsampling path for task-specific specialization # Head 1: Thunderstorm self.ts_up1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.ts_dec1 = ConvBlock2D(256 + 256, 256, dropout=0.3) # skip from enc2 self.ts_up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.ts_dec2 = ConvBlock2D(128 + 128, 128, dropout=0.3) # skip from enc1 self.ts_out = nn.Conv2d(128, 1, kernel_size=1) # Head 2: Cloudburst self.cb_up1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.cb_dec1 = ConvBlock2D(256 + 256, 256, dropout=0.3) self.cb_up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.cb_dec2 = ConvBlock2D(128 + 128, 128, dropout=0.3) self.cb_out = nn.Conv2d(128, 1, kernel_size=1) # Head 3: Flash Flood self.ff_up1 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) self.ff_dec1 = ConvBlock2D(256 + 256, 256, dropout=0.3) self.ff_up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) self.ff_dec2 = ConvBlock2D(128 + 128, 128, dropout=0.3) self.ff_out = nn.Conv2d(128, 1, kernel_size=1) def forward(self, imdaa, insat, terrain): """ Args: imdaa: (B, 30, 6, 256, 256) -- IMDAA reanalysis insat: (B, 3, 6, 256, 256) -- INSAT satellite terrain: (B, 2, 256, 256) -- DEM terrain Returns: (B, 3, 256, 256) -- logits for [Cloudburst, Thunderstorm, FlashFlood] """ # 1. Temporal compression x_imdaa = self.imdaa_time_compress(imdaa).squeeze(2) # (B, 64, H, W) x_insat = self.insat_time_compress(insat).squeeze(2) # (B, 32, H, W) x_terrain = self.terrain_embed(terrain) # (B, 16, H, W) # 2. Early fusion (concatenation) x = torch.cat([x_imdaa, x_insat, x_terrain], dim=1) # (B, 112, H, W) # 3. Shared encoder e1 = self.enc1(x) # (B, 128, 256, 256) x_pool1 = self.pool1(e1) # (B, 128, 128, 128) e2 = self.enc2(x_pool1) # (B, 256, 128, 128) x_pool2 = self.pool2(e2) # (B, 256, 64, 64) # 4. Hybrid bottleneck: CNN + Transformer Self-Attention b = self.bottleneck(x_pool2) # (B, 512, 64, 64) — CNN features b = self.attention(b) # (B, 512, 64, 64) — attention-enhanced # 5. Multi-task decoders # Thunderstorm head ts_u1 = self.ts_up1(b) ts_d1 = self.ts_dec1(torch.cat([ts_u1, e2], dim=1)) ts_u2 = self.ts_up2(ts_d1) ts_d2 = self.ts_dec2(torch.cat([ts_u2, e1], dim=1)) ts_pred = self.ts_out(ts_d2) # (B, 1, 256, 256) # Cloudburst head cb_u1 = self.cb_up1(b) cb_d1 = self.cb_dec1(torch.cat([cb_u1, e2], dim=1)) cb_u2 = self.cb_up2(cb_d1) cb_d2 = self.cb_dec2(torch.cat([cb_u2, e1], dim=1)) cb_pred = self.cb_out(cb_d2) # (B, 1, 256, 256) # Flash Flood head ff_u1 = self.ff_up1(b) ff_d1 = self.ff_dec1(torch.cat([ff_u1, e2], dim=1)) ff_u2 = self.ff_up2(ff_d1) ff_d2 = self.ff_dec2(torch.cat([ff_u2, e1], dim=1)) ff_pred = self.ff_out(ff_d2) # (B, 1, 256, 256) # Output: [Cloudburst, Thunderstorm, FlashFlood] — matches target order return torch.cat([cb_pred, ts_pred, ff_pred], dim=1) # ============================================================ # Self-Test # ============================================================ if __name__ == "__main__": print("Testing Hybrid TransUNet (SpatiotemporalMultiTaskModel)...") # Dummy tensors simulating batch size 2 imdaa_dummy = torch.randn(2, 30, 6, 256, 256) insat_dummy = torch.randn(2, 3, 6, 256, 256) terrain_dummy = torch.randn(2, 2, 256, 256) model = SpatiotemporalMultiTaskModel() # Count parameters total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"Total parameters: {total_params:,}") print(f"Trainable parameters: {trainable_params:,}") # Count attention-specific parameters attn_params = sum(p.numel() for p in model.attention.parameters()) print(f"Attention parameters: {attn_params:,} ({100*attn_params/total_params:.1f}% of total)") # Test forward pass preds = model(imdaa_dummy, insat_dummy, terrain_dummy) print(f"Model output shape: {preds.shape} (Expected: 2, 3, 256, 256)") print("[OK] Hybrid TransUNet model definition passed.")