3DTCR: Vortex-following generative reconstruction improves tropical cyclone intensity prediction from global AI forecasts

Abstract

3DTCR Description 3DTCR is a generative framework for high resolution reconstruction of three dimensional tropical cyclone (TC) structure from global AI weather forecasts, with vortex following and adaptive regional reconstruction. The project addresses the scale gap between global atmospheric prediction and localized high resolution representation of TC structure. Global forecasting systems can efficiently characterize the large scale environmental conditions in which TCs evolve, but further improvement in intensity prediction also depends on resolving the fine three dimensional vortex structures associated with intensity change, including sharp inner core gradients, asymmetric circulation, and organized vertical circulation. These features are more effectively represented at higher spatial resolution, motivating a framework that can recover localized structural detail from coarse resolution global forecasts. The framework is developed using a dedicated seven year dataset of 3 km Weather Research and Forecasting (WRF) simulations constructed around the TC center. Training proceeds in two stages. The first stage learns the mapping from coarse resolution ERA5 fields to high resolution WRF structure under spatially aligned conditions. The second stage adapts this mapping to FuXi forecasts using alignment around the TC center, forecast lead time conditioning, and maximum mean discrepancy (MMD) regularization in latent space. This strategy is designed to reduce the mismatch between reanalysis and forecast domains and to mitigate structural smoothing associated with spatial pattern shifts. Unlike high resolution reconstruction defined over a prescribed geographical region, 3DTCR organizes reconstruction around the evolving TC and allows the target region to adapt as the vortex moves. In this way, the mobility of the weather system is incorporated directly into the reconstruction process, rather than relying solely on generalization across fixed geographical domains. By combining the continuous large scale coverage of global AI forecasts with localized three dimensional reconstruction, 3DTCR provides an efficient pathway from global atmospheric prediction to high resolution representation of TC vortex structure. 3DTCR Inference 3DTCR is a vortex-following, region-adaptive generative framework that reconstructs high-resolution three-dimensional tropical cyclone structures from global AI forecasts to improve TC intensity representation and prediction. Structure 3DTCR/ ├── configs/inference.yaml # Data, model, and inference settings ├── inference.py # End-to-end inference ├── run_inference.sh # Run entry point ├── data/ │ ├── fuxi_forecasts.zarr # FuXi forecast fields │ ├── fuxi_statistics/ # Included FuXi mean.nc and std.nc │ ├── wrf_3km_statistics/ # Included WRF 3-km mean.nc and std.nc │ └── fuxi2_finetune_test.csv # Inference sample metadata ├── model/ # SafeTensors model checkpoint └── tests/ # Unit tests Setup Python 3.11 is recommended. Create and activate a virtual environment: python3.11 -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip python -m pip install -r requirements-runtime.txt \ --extra-index-url https://download.pytorch.org/whl/cu124 Use requirements-runtime.txt for normal inference. requirements.txt is only a full reference-environment export. Model Weights The checkpoint uses the SafeTensors format. Place diffusion_pytorch_model.safetensors with its config.json, or a standalone model.safetensors, under model/. model/ ├── config.json └── diffusion_pytorch_model.safetensors Data Update the paths in configs/inference.yaml: data: fuxi_path: data/fuxi_forecasts.zarr fuxi_statistics_path: data/fuxi_statistics sample_csv: data/fuxi2_finetune_test.csv target_statistics_path: data/wrf_3km_statistics The files have separate purposes: fuxi_path contains the raw FuXi forecast field named data. fuxi_statistics_path/mean.nc and std.nc are the FuXi statistics used during training. They normalize the 21-channel model input. target_statistics_path/mean.nc and std.nc are the WRF 3-km target statistics used during training. They restore generated normalized fields to physical units. sample_csv must contain init_time, valid_time, lead, id, Latitude, and Longitude. Both sets of statistics are included under data/ and should not be interchanged: FuXi statistics are for input normalization, while WRF statistics are for output denormalization. The 21-channel order is: z200, z500, z700, z850, t200, t500, t700, t850, u200, u500, u700, u850, v200, v500, v700, v850, t2m, u10m, v10m, msl, ws10m This order must match the checkpoint and both statistics files. Missing or infinite FuXi values are replaced with the corresponding channel mean before spatial interpolation and normalization. Run ./run_inference.sh Before running, make sure configs/inference.yaml points to the FuXi data, sample CSV, target statistics, and checkpoint directory. The equivalent direct command is: python inference.py --config configs/inference.yaml The default uses 10 Euler sampling steps and writes NetCDF files to outputs/inference/netcdf/. Sampling This project uses Diffusers for the UNet2DModel architecture and checkpoint loading, but it does not use a DDPM or DDIM scheduler. The checkpoint is trained with Rectified Flow, so inference directly integrates the learned velocity field with Euler steps: Normalize the FuXi condition and resize it to 256 × 256. Initialize the 21-channel target state with Gaussian noise. Concatenate the FuXi condition and current state into 42 channels. Let the U-Net predict the flow velocity at time t. Update the state with state = state + velocity / steps. Repeat for the configured number of steps, then resize and denormalize with WRF statistics. model(...).sample is the Diffusers name for the U-Net output tensor; it does not call a separate sampling pipeline. The Euler loop in reconstruct() performs the actual sampling. Test python -m unittest discover -s tests -v