Fix q unit: kg/kg → g/kg (scaled ×1000 from ERA5); clarify input/output normalization space

README.md

@@ -1,286 +1,132 @@
----
-license: cc-by-4.0
-metrics:
-- mse
-- csi
-pipeline_tag: image-to-image
-language:
-- en
-library_name: pytorch
-tags:
-- weather
-- weather-forecasting
-- climate
-- machine-learning-weather-prediction
-- fuxi
-- transformer
----
-
-# FuXi-2.1
-
-[![License](https://img.shields.io/badge/license-CC--BY--4.0-blue.svg)](https://creativecommons.org/licenses/by/4.0/)
-[![Task](https://img.shields.io/badge/task-Weather%20Forecasting-green)]()
-[![Model](https://img.shields.io/badge/model-Transformer-purple)]()
-
-**FuXi-2.1** is a global, deterministic machine-learning weather forecasting model
-developed by **Fudan University** & **[SAIS](https://www.sais.com.cn/)**.
-It produces global forecasts at **0.25°** resolution, on **6-hourly** steps, out to
-**10 days**.
-
-FuXi-2.1 targets the defining failure mode of data-driven weather prediction: forecasts
-that blur into a smooth spatial average as lead time grows, erasing the small-scale
-structure that matters most for extremes. FuXi-2.1 produces markedly **sharper** fields
-whose spatial power spectra track observations across the full wavenumber range, while
-keeping deterministic skill (RMSE) **comparable to** FuXi-1.0 — and substantially
-improving extreme-event detection for heavy precipitation and strong wind.
-
-This model is released as part of the **FuXi Single** collection.
-
-## Table of contents
-
-- [What's new in 2.1](#whats-new-in-21)
-- [Quickstart](#quickstart)
-- [Model overview](#model-overview)
-- [Data details](#data-details)
-- [Evaluation](#evaluation)
-- [Known limitations](#known-limitations)
-- [Citation](#citation)
-
----
-
-## What's new in 2.1
-
-Relative to FuXi-1.0, FuXi-2.1 introduces:
-
-- **A flat Transformer backbone** replacing FuXi-1.0's U-Transformer (ResNet
-  downsample → Swin Transformer → upsample). FuXi-2.1 drops the U-shaped
-  up/down-sampling in favour of a single full-resolution Transformer trunk.
-- **Rotary position embeddings (RoPE)** inside the Swin windowed attention, replacing
-  the learned relative-position bias used in FuXi-1.0.
-- **adaLN time conditioning** that injects time-period information — forecast lead step,
-  time-of-day and day-of-year phase — into every block, inspired by diffusion
-  transformers in image generation.
-- **A variable-aware multi-head decoder** that gives pressure-level, surface and derived
-  variables their own specialised output heads.
-
-The combined effect is sharper, spectrally faithful forecasts with no penalty on
-mean-error skill.
-
----
-
-## Quickstart
-
-This repository ships the exported model (`fuxi-2.1.pt2`), normalization statistics
-(`mean.nc`, `std.nc`), a sample pre-normalized input (`input.nc`), and minimal inference
-code.
+# FuXi 2.1 — Global Weather Forecasting Model
+
+FuXi 2.1 is a 0.25-degree global weather forecasting model producing 85-channel predictions at 6-hour intervals. This repository provides minimal inference code for autoregressive rollout using the PyTorch PT2 (torch.export) backend.
+
+## Quick Start
 
 ```bash
-# 1. Install dependencies
+# Install dependencies
 pip install -r requirements.txt
 
-# 2. End-to-end demo (inference + plots)
-bash run.sh --model_dir . --input input.nc --steps 5
+# Run end-to-end (inference + plots)
+bash run.sh --model_dir /path/to/model --input /path/to/input.nc --steps 5
 
-# Or run inference directly (40 steps = 10-day forecast):
+# Or run directly:
 python inference.py \
-    --model_dir . \
-    --input input.nc \
+    --model_dir /path/to/model \
+    --input /path/to/input.nc \
     --output_dir ./output \
     --steps 40 \
     --forecast_time 2024092900
 
-# 3. Plot selected channels
+# Plot results
 python plot.py --output_dir ./output --channels t2m z500 tp --discrete
 ```
 
-**Input** — a NetCDF with a variable `input` of shape `(time=2, channel=85, lat=721,
-lon=1440)`, z-score normalized, coordinates `lat` 90→−90 and `lon` 0→359.75. The provided
-`input.nc` is a sample for 2024-09-29 00Z.
-
-**Output** — each step saved as `{output_dir}/{step:03d}.nc`, shape `(channel=85, lat=721,
-lon=1440)` in physical units (denormalized), with a `valid_time` attribute. Steps are
-1-based: `001.nc` = +6 h, … `040.nc` = +240 h (10 days).
-
-> **GPU:** the device is baked into the exported graph; load on CUDA. ~8 GB GPU memory is
-> enough (model ~4 GB + recurrent state ~1.4 GB + working memory). Tested on A100, V100,
-> RTX 3090/4090. See `variables.py` for the full ordered channel list.
-
----
-
-## Model overview
-
-### Model description
-
-FuXi-2.1 is a single Transformer. The global atmospheric state is split into patches and
-embedded into tokens, processed by a stack of windowed-attention blocks, and read out by
-a variable-aware multi-head decoder. The model is **deterministic** — one forward pass
-per step, with no adversarial or diffusion sampling at inference — and is rolled out
-**autoregressively** at 6-hourly steps.
-
-- **Developed by:** Fudan University & [SAIS](https://www.sais.com.cn/)
-- **Model type:** Transformer (patch-embed → Swin attention with RoPE + adaLN → multi-head decoder)
-- **Forecast type:** Global, deterministic, autoregressive
-- **License:** CC BY 4.0
-- **Predecessor:** [FuXi-1.0](https://github.com/tpys/FuXi)
-
-<div align="center">
-  <img src="assets/arch.png" alt="FuXi-2.1 architecture" style="width: 95%;"/>
-</div>
-
-### Architecture details
-
-| Component | Specification |
-|---|---|
-| Backbone | Single Transformer trunk (no U-Net up/down-sampling) |
-| Attention | Swin windowed attention |
-| Position encoding | Rotary (RoPE, 1-D) |
-| Normalisation / conditioning | adaLN, conditioned on lead step, time-of-day, day-of-year |
-| Feed-forward | SwiGLU |
-| Decoder | Variable-aware multi-head (pressure / surface / derived) |
-| Input frames | 2 (states at t−6h and t₀) |
-| Output | State at t+6h, rolled out autoregressively |
-
-### Model resolution
-
-| Model | Horizontal resolution | Vertical resolution [pressure levels] (hPa) |
-|:---|:---:|:---|
-| FuXi-2.1 | 0.25° (721×1440) | 13: 50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000 |
-
----
-
-## Data details
-
-### Training data
-
-FuXi-2.1 is trained and evaluated on **ERA5** reanalysis at 0.25° resolution, 6-hourly.
-
-- **Training period:** 2002–2023
-- **Test period:** 2024 (held out)
-
-### Data parameters
-
-FuXi-2.1 operates on **85 channels** per time step: **65 pressure-level** channels
-(5 variables × 13 levels) and **20 surface** channels, plus static forcings supplied as
-constant inputs. Most channels are **prognostic** — the same channels are input and
-output and fed back during roll-out. Radiation fluxes and total precipitation are
-**diagnostic** outputs produced through a dedicated decoder head (they are predicted but
-not fed back as inputs).
-
-Channel order (exact): the 65 pressure-level channels first (z, then t, u, v, q, each
-over the 13 levels 50→1000 hPa), followed by the 20 surface channels:
-`msl, t2m, d2m, sst, ws10m, ws100m, u10m, v10m, u100m, v100m, lcc, mcc, hcc, tcc, ssr,
-ssrd, fdir, ttr, tcw, tp`.
-
-#### Pressure-level parameters (13 levels: 50–1000 hPa)
-
-| Short name | Name | Units | Input/Output |
-|:---:|:---|:---:|:---:|
-| z | Geopotential | m²·s⁻² | Both (prognostic) |
-| t | Temperature | K | Both (prognostic) |
-| u | Eastward wind | m·s⁻¹ | Both (prognostic) |
-| v | Northward wind | m·s⁻¹ | Both (prognostic) |
-| q | Specific humidity | kg·kg⁻¹ | Both (prognostic) |
-
-#### Surface parameters (20)
-
-| Short name | Name | Units | Input/Output |
-|:---:|:---|:---:|:---:|
-| msl | Mean sea-level pressure | Pa | Both (prognostic) |
-| t2m | 2 m temperature | K | Both (prognostic) |
-| d2m | 2 m dewpoint temperature | K | Both (prognostic) |
-| sst | Sea-surface temperature | K | Both (prognostic) |
-| ws10m | 10 m wind speed | m·s⁻¹ | Both (prognostic) |
-| ws100m | 100 m wind speed | m·s⁻¹ | Both (prognostic) |
-| u10m | 10 m eastward wind | m·s⁻¹ | Both (prognostic) |
-| v10m | 10 m northward wind | m·s⁻¹ | Both (prognostic) |
-| u100m | 100 m eastward wind | m·s⁻¹ | Both (prognostic) |
-| v100m | 100 m northward wind | m·s⁻¹ | Both (prognostic) |
-| lcc | Low cloud cover | 0–1 | Both (prognostic) |
-| mcc | Medium cloud cover | 0–1 | Both (prognostic) |
-| hcc | High cloud cover | 0–1 | Both (prognostic) |
-| tcc | Total cloud cover | 0–1 | Both (prognostic) |
-| tcw | Total column water | kg·m⁻² | Both (prognostic) |
-| ssr | Surface net solar radiation | J·m⁻² | Output (diagnostic) |
-| ssrd | Surface solar radiation downwards | J·m⁻² | Output (diagnostic) |
-| fdir | Total-sky direct solar radiation at surface | J·m⁻² | Output (diagnostic) |
-| ttr | Top net thermal radiation | J·m⁻² | Output (diagnostic) |
-| tp | Total precipitation | mm | Output (diagnostic) |
-
-| Field | Level type | Input/Output |
-|---|---|---|
-| Land-sea mask, orography/geopotential, latitude/longitude encodings, time-of-day / day-of-year | Surface / static | Input (forcings) |
-
----
-
-## Evaluation
-
-We compare FuXi-2.1 against FuXi-1.0 under an identical protocol: forecasts initialised
-from ERA5 and rolled out to 240 h in 6-hour steps. CSI is computed over **land only,
-globally**. These numbers come from a **limited set of sample cases**, not a full-year
-evaluation — they are indicative, and broader scorecards will follow.
-
-**Headline:** RMSE stays comparable to FuXi-1.0 across variables, while structural and
-extreme-event scores improve substantially.
-
-<div align="center">
-  <img src="assets/chart_tp_csi.png" alt="Precipitation CSI" style="width: 49%;"/>
-  <img src="assets/chart_ws10m_csi.png" alt="Wind-speed CSI" style="width: 49%;"/>
-</div>
-
-### Precipitation — Critical Success Index (CSI)
-
-| Threshold | FuXi-1.0 | FuXi-2.1 | Δ |
-|:---|---:|---:|---:|
-| ≥ 5 mm | 0.265 | 0.284 | +7.3% |
-| ≥ 20 mm | 0.131 | 0.146 | +11.4% |
-| ≥ 50 mm | 0.074 | 0.084 | +13.4% |
-| ≥ 100 mm | 0.014 | 0.024 | **+68.3%** |
-
-### 10 m wind speed — Critical Success Index (CSI)
-
-| Threshold | FuXi-1.0 | FuXi-2.1 | Δ |
-|:---|---:|---:|---:|
-| ≥ 10.8 m·s⁻¹ | 0.544 | 0.571 | +4.8% |
-| ≥ 24.5 m·s⁻¹ | 0.165 | 0.198 | +20.3% |
-| ≥ 28.5 m·s⁻¹ | 0.000 | 0.044 | newly resolved |
-
-The relative gain grows with event intensity, peaking at the extreme tail. At the
-28.5 m·s⁻¹ wind threshold FuXi-1.0 scores zero — it never predicts such winds — whereas
-FuXi-2.1 attains a non-zero CSI. Spatial power spectra of FuXi-2.1 track the observed
-spectra across the full wavenumber range, in contrast to FuXi-1.0's high-wavenumber
-energy deficit.
-
----
-
-## Known limitations
-
-- FuXi-2.1 is a deterministic model; it does not provide a calibrated ensemble spread.
-- The CSI numbers reported here are computed on land only, over a limited set of sample
-  cases rather than a full-year evaluation; treat them as indicative. Comprehensive
-  global scorecards will be added.
-- As with all ERA5-trained models, skill depends on the quality and resolution of the
-  initial conditions.
-
----
+## Model Details
 
-## Citation
+| Property | Value |
+|----------|-------|
+| Resolution | 0.25° (721 × 1440 grid) |
+| Channels | 85 (65 pressure-level + 20 surface) |
+| Time step | 6 hours |
+| Input frames | 2 (t-6h, t) |
+| Architecture | FuXi HR |
+| Format | torch.export (.pth) |
+| Size | ~3.7 GB |
+
+## Input Format
+
+The input NetCDF file must contain a variable named `input` with:
+- **Shape**: `(time=2, channel=85, lat=721, lon=1440)`
+- **Normalization**: z-score normalized (already in normalized space)
+- **Coordinates**: `time`, `channel` (C85 names), `lat` (90 to -90), `lon` (0 to 359.75)
+
+The provided `input.nc` is a sample input for 2024-09-29 00Z. Both `input.nc` and the model's internal weights operate in normalized space.
+
+## Output Format
+
+Each forecast step is saved as `{output_dir}/{step:03d}.nc`:
+- **Shape**: `(channel=85, lat=721, lon=1440)`
+- **Units**: Physical units (denormalized via `output = output * std + mean`)
+- **Coordinates**: `channel`, `lat`, `lon`
+- **Attribute**: `valid_time` — the forecast valid time for this step
+
+Step numbering is 1-based: `001.nc` = +6h, `002.nc` = +12h, ..., `040.nc` = +240h (10 days).
+
+> **Note:** The `tp` (total precipitation) channel is log1p-transformed during training. Denormalization reverses this with `expm1` and clips to ≥ 0.
+
+## Channel Table (C85)
+
+**Pressure-level variables** (5 vars × 13 levels = 65 channels):
+- z (geopotential), t (temperature), u (u-wind), v (v-wind), q (specific humidity, **g/kg** — scaled ×1000 from ERA5's kg/kg)
+- Levels: 50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000 hPa
+
+**Surface variables** (20 channels):
+| Channel | Variable | Units |
+|---------|----------|-------|
+| msl | Mean sea level pressure | Pa |
+| t2m | 2m temperature | K |
+| d2m | 2m dewpoint temperature | K |
+| sst | Sea surface temperature | K |
+| ws10m | 10m wind speed | m/s |
+| ws100m | 100m wind speed | m/s |
+| u10m / v10m | 10m wind components | m/s |
+| u100m / v100m | 100m wind components | m/s |
+| lcc / mcc / hcc / tcc | Cloud cover | 0-1 |
+| ssr / ssrd / fdir / ttr | Radiation fluxes | J/m² |
+| tcw | Total column water | kg/m² |
+| tp | Total precipitation | m (log1p-transformed, reversed on output) |
+
+See `variables.py` for the full ordered list.
 
-If you use FuXi-2.1, please cite the FuXi series:
+## GPU Requirements
+
+- **Minimum GPU memory**: ~8 GB (model ~4 GB + state ~1.4 GB + working memory)
+- The model device is baked into the exported graph. It must be loaded on a CUDA device.
+- Tested on: A100, V100, RTX 3090/4090
+
+## File Structure
+
+```
+fuxi-2.1/
+├── fuxi-2.1.pt2         # Model weights (torch.export)
+├── mean.nc              # Channel means for denormalization
+├── std.nc               # Channel stds for denormalization
+├── input.nc             # Sample input (pre-normalized)
+├── inference.py         # Rollout engine
+├── data_util.py         # Data loading + postprocessing
+├── variables.py         # C85 channel definitions
+├── plot.py              # Visualization (uses fuxi_viz)
+├── run.sh               # End-to-end demo
+├── requirements.txt     # Python dependencies
+└── README.md            # This file
+```
+
+## How It Works
+
+1. **Load** pre-normalized input (2 frames at t-6h and t)
+2. **Rollout** autoregressively: model outputs next 2-frame state, last frame is the prediction
+3. **Denormalize** each prediction: `output = output * std + mean`, then `expm1` for precipitation
+4. **Save** each step as NetCDF with geographic coordinates
+
+The recurrence state stays on GPU throughout the rollout (no CPU round-trip per step).
+
+### Note on Specific Humidity (q)
+
+The specific humidity `q` is stored and predicted in **g/kg** (grams of water vapor per kg of dry air), not the raw ERA5 kg/kg. During data preparation, ERA5 q values were multiplied by 1000 before normalization. This scaling is baked into the model's normalization statistics (`mean.nc` / `std.nc`), so denormalized output will be in g/kg. Typical near-surface values range from 0–25 g/kg.
+
+## Citation
 
 ```bibtex
 @article{chen2023fuxi,
-  title   = {FuXi: a cascade machine learning forecasting system for 15-day global weather forecast},
-  author  = {Chen, Lei and Zhong, Xiaohui and Zhang, Feng and Cheng, Yuan and Xu, Yimin and Qi, Yan and Li, Hao},
-  journal = {npj Climate and Atmospheric Science},
-  year    = {2023},
-  volume  = {6},
-  number  = {1},
-  pages   = {190}
+  title={FuXi: A cascade machine learning forecasting system for 15-day global weather forecast},
+  author={Chen, Lei and Zhong, Xiaohui and Zhang, Feng and Cheng, Yuan and Xu, Yinghui and Qi, Yuan and Li, Hao},
+  journal={npj Climate and Atmospheric Science},
+  year={2023}
 }
 ```
 
----
-
-**Code:** FuXi-1.0 — https://github.com/tpys/FuXi
+## License
 
-*© 2026 Fudan University & [SAIS](https://www.sais.com.cn/) · FuXi Weather.*
+Please refer to the model license for usage terms.