NRFormer+

A Physics-Guided Spatio-Temporal Transformer for Nationwide Nuclear Radiation Forecasting.

This repository hosts the official PyTorch implementation of NRFormer+, the journal extension of our KDD 2025 paper NRFormer. NRFormer+ couples non-stationary temporal attention and density-adaptive spatial attention with a new physics-informed atmospheric diffusion module, and is trained on two nationwide benchmarks (Japan-4H and Japan-1D) covering 3,627 radiation and 228 meteorological stations over four years.

Nuclear radiation, the energy released during atomic decay, poses persistent risks to public health and the environment, and concerns have only grown since the Fukushima accident and the recent commencement of treated-water discharge. Modern monitoring networks now record radiation levels and accompanying weather conditions at thousands of stations, opening the door to nationwide forecasting that can inform emergency response, agricultural advisories, and routine public-safety decisions.

Turning this abundance of monitoring data into reliable forecasts is difficult for three reasons. (i) The series at each station are highly non-stationary, shaped by radioactive decay, weather variability, and irregular human interventions. (ii) Monitoring stations are severely unevenly distributed in space — roughly 78% of Japan's stations sit in less than 6% of the country, clustered near Fukushima, which breaks the assumptions of standard graph-based models. (iii) Radiation co-evolves with heterogeneous meteorological context (wind, temperature, humidity) through atmospheric transport processes that purely data-driven models struggle to capture from observations alone.

We introduce NRFormer+, a spatio-temporal Transformer that couples non-stationary temporal attention and density-adaptive spatial attention with a new atmospheric diffusion module that estimates how meteorology drives radiation dispersion and injects this physical signal into the network as an architectural prior. NRFormer+ delivers state-of-the-art accuracy on both benchmarks against 13 baselines, reducing sudden-change MAE by up to 19.1% over the strongest baseline at comparable inference latency.

NRFormer+ is organised around four components:

1. Physics-Informed Atmospheric Diffusion — encodes the governing equation $\partial C / \partial t = D,\nabla^2 C$ with a station-wise diffusion coefficient $D_i$ estimated from local meteorology.
2. Enhanced Meteorological Encoder — separately models wind dynamics (advection) and thermal stability (mixing) through physically-motivated architectural designs.
3. Context-Aware Propagation Prompting — a deep location encoder + a day-of-year seasonal embedding + radiation–location cross features.
4. Physics-Enriched Spatial Attention — radiation features directly guide the query–key routing so spatial dependencies are grounded in radiation patterns rather than purely temporal embeddings.

The non-stationary temporal attention and imbalance-aware spatial attention from the KDD 2025 NRFormer are retained.

NRFormer_Plus/
├── src/
│   ├── DataProcessing.py        # Radiation + NOAA loader, scalers, adjacency
│   ├── trainer.py               # Training loop with early stopping
│   ├── utils.py                 # Metrics, graph helpers, logging
│   └── model/
│       ├── NRFormer.py          # Conference-version baseline (KDD'25)
│       └── NRFormer_Plus.py     # Proposed PGRT2 (NRFormer+)
├── model_settings/
│   ├── NRFormer.yaml            # Hyper-parameters for the baseline
│   └── NRFormer_Plus.yaml       # Hyper-parameters for NRFormer+
├── train.py                     # Training entry point
├── test.py                      # Generic test entry point
├── test_sudden_change.py        # Sudden-change benchmark probe
├── assets/                      # README images
├── requirements.txt
├── LICENSE                      # MIT
└── README.md

git clone https://github.com/tfeilyu/NRFormer_Plus.git
cd NRFormer_Plus
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

Tested with Python 3.10 and PyTorch 2.0+ on a single NVIDIA A100. One CUDA-capable GPU is sufficient for the Japan-1D experiments.

NRFormer+ is trained on Japan-1D (3,627 radiation stations paired with 228 NOAA ISD meteorological stations). The raw data are public:

• Radiation — Japan Nuclear Regulation Authority (NRA), 10-minute readings from March 2021 to May 2025
  📡 https://radioactivity.nra.go.jp/
• Meteorology — NOAA Integrated Surface Dataset (ISD)
  🌤 https://www.ncei.noaa.gov/products/land-based-station/integrated-surface-database

The data path is configured by the DATA_PATH key in model_settings/NRFormer_Plus.yaml. Point it at a directory laid out as

<DATA_PATH>/
└── 1D-data/
    ├── radiation_*.csv          # one file per station
    ├── locations.csv            # station coordinates
    └── noaa/                    # interpolated meteorological series

(See src/DataProcessing.py for the exact filenames the loader expects.)

Single-GPU training with default hyper-parameters:

python train.py \
  --dataset 1D-data \
  --model_name NRFormer_Plus \
  --batch_size 8 \
  --gpu_ids 0

Common flags:

Run python train.py --help for the full list.

General test run:

python test.py --model_name NRFormer_Plus --dataset 1D-data --gpu_ids 0

Sudden-change benchmark (Table V in the paper):

python test_sudden_change.py --model_name NRFormer_Plus --dataset 1D-data --gpu_ids 0

train.py logs metrics to Weights & Biases. Set your API key before training:

export WANDB_API_KEY=<your-key>

To run offline or disable the dashboard entirely:

export WANDB_MODE=offline    # or 'disabled'

If you find this work useful, please cite the conference version:

@inproceedings{lyu2025nrformer,
  title={NRFormer: Nationwide Nuclear Radiation Forecasting with Spatio-Temporal Transformer},
  author={Lyu, Tengfei and Han, Jindong and Liu, Hao},
  booktitle={Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 2},
  pages={4705--4716},
  year={2025}
}

Released under the MIT License.