🌊 VARUNA
Spatiotemporal Multi-Task AI Engine for Extreme Weather Nowcasting
A highly optimized Hybrid TransUNet architecture designed to simultaneously predict the probability of Cloudbursts, Severe Thunderstorms, and Flash Floods across highly complex topographical regions (the Himalayas) with a 2–6 hour actionable lead time.
🔬 System Architecture
VARUNA utilizes a Hybrid TransUNet (17.6M parameters), bridging 3D Convolutional temporal compression with 2D U-Net spatial feature extraction, stabilized by a Multi-Head Self-Attention (MHSA) bottleneck (64×64 spatial resolution, 4096 tokens).
Modalities & Input Tensors
The network accepts three synchronized, multi-modal data branches corresponding to physical meteorology:
- IMDAA Thermodynamics
(B, 30, 6, 256, 256): 30 atmospheric parameters (TMP, UGRD, VGRD, RH, PRMSL, CAPE, etc.) measured across multiple atmospheric pressure levels (1000hPa to 300hPa), fed as 6 temporal frames. - INSAT-3D Satellite
(B, 4, 6, 256, 256): 4 engineered multi-spectral channels:- Water Vapor (WV)
- Cloud Top Temperature (CTT)
- Hydro-Estimator Method (HEM) Rainfall Rate
- Temporal CTT Drop-Rate (Calculated dynamically to track rapid vertical storm development)
- CartoDEM Terrain
(B, 2, 256, 256): Static high-resolution (30m) elevation and slope masks for the target topography, injected via late-stage decoder fusion to dictate flash flood trajectories.
📊 Data Scale & Curating
The model was trained on a meticulously synchronized 190 GB meteorological corpus encompassing the Himalayan states of Uttarakhand and Himachal Pradesh for the severe monsoon month of August 2019.
Raw Corpus:
- IMDAA: 7,440 NetCDF4 reanalysis files (Hourly resolution).
- INSAT-3D: 435 HDF5 L1B (Imager), 443 HDF5 L2B-CTP, and 426 HDF5 L2B-HEM files (Half-hourly resolution).
- CartoDEM: High-resolution GeoTIFF topography.
Temporal Windowing: The raw data was temporally aggregated and aligned into 55 discrete spatiotemporal windows. Each window consists of 6 contiguous historical hours (represented by 6 frames) pointing to a target label 2–6 hours in the future. The dataset was split chronologically: 44 windows for training, and 11 strictly unseen windows (Aug 23–25, 2019) for holdout validation.
📈 Training Progression & Breakthroughs
Training a heavily imbalanced, multi-modal meteorological model required iterative architectural evolution. Below is the technical history of model development across 7 distinct training cycles:
| Iteration | Key Modifications | CB (F1) | TS (F1) | FF (F1) |
|---|---|---|---|---|
| v1.0 (Runs 1-3) | Baseline TransUNet. Encountered critical timestamp misalignment between IMDAA and INSAT inputs. | 0.08% | ~40% | 0.00% |
| v2.0 (Runs 4-5) | Implemented strict 30-min temporal binning. Switched to DiceFocalLoss due to extreme class imbalance (1:460 positive pixel ratio). |
0.08% | ~72% | 0.00% |
| v2.1 (Run 6) | Applied nearest-neighbor label resizing. Introduced per-task gradient loss amplification (4×CB, 1×TS, 6×FF). |
10.60% | 75.82% | 1.52% |
| v3.0 (Final) | Discarded coarse 25km GPI labels in favor of 1km HEM labels (yielding a 461× increase in Cloudburst signal density). Expanded INSAT branch to Conv3d(4, ...) to ingest dynamic CTT Drop-Rates. Flash Flood labels dynamically constrained by DEM slope mask (>10°). |
30.00% | 85.97% | 14.94% |
Loss Dynamics
To prevent the high-frequency Thunderstorm (TS) class from dominating the U-Net decoders and extinguishing the gradients for the ultra-rare Cloudburst (CB) and Flash Flood (FF) classes, we utilized a custom amplified multi-task loss function:
Total Loss = (4.0 * Loss_CB) + (1.0 * Loss_TS) + (6.0 * Loss_FF)
🎯 Production Performance & Applicability
Evaluated strictly on the Aug 23–25 temporal holdout, VARUNA v3.0 achieved the following metrics:
| Hazard | Precision | Recall | F1-Score | IoU |
|---|---|---|---|---|
| 🌧️ Cloudburst | 18.29% | 83.36% | 30.00% | 17.65% |
| ⛈️ Thunderstorm | 80.56% | 92.17% | 85.97% | 75.40% |
| 🌊 Flash Flood | 8.24% | 80.58% | 14.94% | 8.08% |
Use-Case Viability
For disaster management and early warning systems, Recall is the strictly dominant metric. VARUNA operates as an ultra-sensitive triggering engine. It successfully captures the physical atmospheric signatures of >80% of all extreme events (CB/FF). While precision is comparatively low (resulting in regional false-positive triggers), this is standard and highly acceptable in operational meteorology. In life-critical contexts, over-warning is survivable; missing a catastrophic cloudburst is not.
🔭 Future Scope
- Long-Term Temporal Dependencies: Upgrading the temporal compression module from standard
Conv3dto a TimeSformer or VideoMAE backbone to capture 24-48 hour atmospheric momentum. - Additional Spectral Channels: Integrating Short-Wave Infrared (SWIR) and Middle-Infrared (MIR) INSAT channels to better isolate specific cloud microphysics.
- Data Expansion: Expanding the dataset from 55 windows (1 month) to a multi-year climatological baseline (2015-2023) to allow the network to learn seasonal variance, which will drastically improve Precision by teaching the model when not to trigger.
⚡ Quickstart: Real-World Inference Demo
Since the full meteorological dataset is 190 GB, we have extracted a single 6-hour physical window from the strict unseen holdout set (August 23-25, 2019) and hosted it directly on Hugging Face.
You can instantly download the AI Engine and watch it predict real cloudbursts and thunderstorms by running our standalone demo script:
# 1. Download the demo script
wget https://huggingface.co/JayF14/varuna-HybridTransU-net/resolve/main/demo_inference.py
# 2. Run it (The script automatically downloads the model weights and the 52MB sample tensor)
python demo_inference.py
(The script calculates predicted activated pixels versus ground truth pixels to physically demonstrate the model's accuracy).
💻 Custom Loading & Inference
VARUNA requires three tensors of exact dimensionality. The weights (17.6M parameters) can be downloaded seamlessly in Python:
import torch
from huggingface_hub import hf_hub_download
from model import SpatiotemporalMultiTaskModel
# 1. Download weights (auto-cached locally)
ckpt_path = hf_hub_download(
repo_id="JayF14/varuna-HybridTransU-net",
filename="best_model.pth"
)
# 2. Initialize Architecture
model = SpatiotemporalMultiTaskModel()
# 3. Load state dict
ckpt = torch.load(ckpt_path, map_location="cpu")
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
# 4. Predict
# imdaa: (B, 30, 6, 256, 256)
# insat: (B, 4, 6, 256, 256)
# terrain: (B, 2, 256, 256)
with torch.no_grad():
logits = model(imdaa, insat, terrain)
probs = torch.sigmoid(logits) # → (B, 3, 256, 256) Output Maps