varuna-HybridTransU-net

<div align="center">

🌊 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.

</div>

---

🔬 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:

1. 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.
2. 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)
3. 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:

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:

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

1. Long-Term Temporal Dependencies: Upgrading the temporal compression module from standard Conv3d to a TimeSformer or VideoMAE backbone to capture 24-48 hour atmospheric momentum.
2. Additional Spectral Channels: Integrating Short-Wave Infrared (SWIR) and Middle-Infrared (MIR) INSAT channels to better isolate specific cloud microphysics.
3. 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