code_corrdiff

Introduction

This repository accompanies a manuscript on precipitation downscaling over the United Kingdom using a lightweight CorrDiff-style framework (RainUNet++ regression + PrecipDiff residual diffusion). The citation will be added once the paper is published.

The objectives of this project are:

• Provide the model architecture used in the paper (RainUNet++ and PrecipDiff);
• Evaluate forecasts against NIMROD observations and UKV baselines (SAL, RAPSD, neighbourhood reliability, CRPS, WD, FSS);
• Produce paper figures for accumulated rainfall, intensity distributions, and selected UK events.

Scripts and data

1_model

File	Role
unet_plus_plus.py	RainUNet++ backbone (depth=3, base_channels=48)
lightweight_corrdiff.py	Two-stage wrapper, EDM preconditioning, t-EDM loss & sampler
advanced_loss_functions.py	Regression loss (L_{\mathrm{reg}}=\lambda_H L_{WH}+\lambda_z L_{\mathrm{latent}}+\lambda_l L_{\mathrm{lowres}})
dataset.py	HDF5 dataset loader

Paper defaults: (\lambda_H=1.0), (\lambda_z=0.1), (\lambda_l=0.3), Huber (\delta=1.0); sampler (\sigma_{\max}\approx 80), (\sigma_{\min}\approx 0.002), (\rho=7), t-EDM (\nu=5).

2_evaluation

Num.	Metric	Folder	Script
2.1	SAL	sal	calculate_sal.py
2.2	RAPSD; neighbourhood exceedance reliability	rapsd_reliability	run_rapsd.py, run_reliability_neighbourhood.py
2.3	CRPS	crps	calculate_crps.py
2.4	Wasserstein distance (WD)	wasserstein	calculate_wasserstein.py
2.5	FSS	fss	calculate_fss.py

cd 2_evaluation/sal
python calculate_sal.py \
  --checkpoint /path/to/best_diffusion.pth \
  --regression_checkpoint /path/to/best_regression.pth \
  --data_dir /path/to/qpe_dataset

3_visualization

Num.	Subject	Folder	Script
3.1	Spatial distribution of total accumulated rainfall	composite_rainfall	calculate_composite_lightweight.py, generate_paper_figures.py
3.2	Rainfall intensity distribution	intensity_histogram	calculate_histogram.py, plot_histogram_grouped_bars.py
3.3	Selected UK event rainfall maps	event_comparison	visualize_diffusion_stochastic.py

cd 3_visualization/event_comparison
python visualize_diffusion_stochastic.py \
  --diffusion_checkpoint /path/to/best_diffusion.pth \
  --regression_checkpoint /path/to/best_regression.pth \
  --data_dir /path/to/qpe_dataset \
  --timestamp 20240526_1200 \
  --output_dir ./panels

Requirements

• Python 3 with PyTorch (GPU recommended)
• numpy, scipy, matplotlib, h5py, netCDF4, pyproj, and Cartopy as needed for map figures

Acknowledgments

• This work used computing facilities at The University of Manchester.
• The authors declare no conflict of interest.
• This work was supported by the Natural Environment Research Council [grant number UKRI1294].