imvetri/afno-v4-pm25

AFNO V4 - PM2.5 Air Quality Forecasting

A neural-operator model that forecasts PM2.5 over India for the next 16 hours from the previous 10 hours of gridded meteorology and pollution data (140x124 grid, ~25 km resolution, hourly). Built for the ANRF AISEHack Phase 2, Theme 2 pollution-forecasting competition (IIT Delhi). ~7M parameters, trained from scratch.

Results

• Private score 0.8833 on the official metric (public 0.8959) - places 2nd on the final competition leaderboard (1st = 0.8893).
• Validation (leak-free holdout): 0.9110 / 0.9114 for the two seeds.
• Baselines: prior v3 0.8609; persistence NormGlobalSMAPE 0.851.

Metric: w1*NormGlobalSMAPE + w2*NormEpisodeCorr + w3*NormEpisodeSMAPE (weights undisclosed by the organizers).

Architecture

Input tensor: (batch, 20 channels, 10 time steps, 140, 124).

1. Per-frame 2D U-Net encoder (20->64->128->256, MaxPool x2).
2. Multiscale AFNO (AdaptFNO-style): 1 spectral block at the 35x31 encoder scale.
3. Bottleneck AFNO: 6 spectral blocks at 18x16 (global 2D-FFT token mixing).
4. SimVP temporal translator (Incep1D, 4 layers): 10 -> 16 frames.
5. Lead-time embedding (per-forecast-hour bias) + 2 post-translator AFNO blocks.
6. U-Net decoder (256->128->64->32) with last+mean skip fusion.
7. Persistence-residual head: output = last_observed_frame + delta.

Input channels (20)

• observed (9): PM2.5, q2, t2, u10, v10, pblh, psfc, swdown, rain
• engineered (5): wind speed, ventilation coefficient (PBLH x wind), advection tendency, PM tendency, diurnal climatology anomaly
• time (2): hour sin, hour cos
• static (4): lat, lon, PM2.5 climatology, episode sigma

All channels are derived from the competition data - no external or pretrained inputs.

Files

• s42_phaseA_best.pt, s1337_phaseA_best.pt - EMA (eval) weights for the two seeds, each {'model': state_dict, 'metrics': {...}} for AFNOv4.
• horizon_scales.npy - 16 per-horizon multiplicative calibration factors.
• stats.npz - per-pixel normalization stats.
• afno_v4.py - full model, data, and inference code.

Usage

import ast, torch, torch.nn as nn, torch.nn.functional as F, numpy as np, os
from huggingface_hub import hf_hub_download

repo = "imvetri/afno-v4-pm25"
src = open(hf_hub_download(repo, "afno_v4.py")).read()
ns = {"os": os, "torch": torch, "nn": nn, "F": F, "np": np}
blocks = "".join(ast.get_source_segment(src, n) + "\n" for n in ast.parse(src).body
                 if isinstance(n, ast.ClassDef) and n.name in
                 {"C","AFNO2D","AFNOBlock","ConvBlock","Incep1D","Translator","AFNOv4"})
exec(blocks, ns)

model = ns["AFNOv4"]().eval()
ck = torch.load(hf_hub_download(repo, "s42_phaseA_best.pt"), map_location="cpu", weights_only=False)
model.load_state_dict(ck["model"])

# x: (B, 20, 10, 140, 124) normalized inputs (build with stats.npz as in PM25Test)
# last_log: (B, 140, 124) log1p of the last observed PM2.5 frame
# with torch.no_grad(): pred_log = model(x, last_log)   # -> (B, 16, 140, 124)

Inference used to reproduce the competition result: average the two seeds' predictions and multiply by horizon_scales per forecast hour.

How it was built

The biggest gain came from fixing a validation data leak (random splits over overlapping windows) that had been masking real progress. On trustworthy leak-free validation, the model then adds atmospheric-physics features, multiscale spectral mixing, a 2-seed ensemble, and per-horizon calibration.

Scope and limitations

This model is trained on 2016 WRF-Chem simulation data with competition-specific inputs. It is a research / benchmark artifact, not an operational real-world forecaster - it cannot ingest live observations as-is. A real-world version would require retraining on reanalysis inputs (ERA5 / MERRA-2 / CAMS) with ground-truth validation (OpenAQ / CPCB).

Citation / references

• Adaptive Fourier Neural Operator (Guibas et al., ICLR 2022)
• FourCastNet (Pathak et al., 2022)
• SimVP (Gao et al., CVPR 2022)
• STL (Cleveland et al., 1990)

Code: https://github.com/Vetri-78640/AISEHack-2026 (models/model-4)