Add Corrdiff-era5-hrrr

CHANGELOG.md

@@ -84,7 +84,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   per step. The solvers share an extended semi-linear callback API
   (`bias_fn`, `bias_int_fn`, `slope_fn`); users can select all four
   solvers by string key through `physicsnemo.diffusion.samplers.sample`.
-
+- Adds flow matching (`training.loss.type: flow_matching`) and the CorrDiff-ERA5-HRRR
+  downscaling recipe to the `examples/weather/regional_weather_diffusion` example.
 - `ClusterTree.find_dual_interaction_pairs(source_admissible=...)` marks source
   nodes that must never be approximated by their aggregate, such as nodes
   whose normals differ too much.

examples/weather/regional_weather_diffusion/.gitignore

@@ -2,4 +2,4 @@
 *.png
 *.pt
 *.tfevents*
-*wandb/
+*wandb/*.nc

examples/weather/regional_weather_diffusion/config/corrdiff_era5_hrrr.yaml

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Defaults
+
+# CorrDiff-ERA5-HRRR: kilometer-scale downscaling of ERA5 to the 3 km HRRR grid
+# over the contiguous US with x-prediction rectified flow matching:
+#   - noisy states x_t = (1 - t) x0 + t eps with t drawn uniformly from [0, 1]
+#   - the DiT predicts the clean HRRR state x0, the loss is ||x0_hat - x0||^2
+#   - trained on random 256 x 256 patches of the 1024 x 1792 HRRR domain
+
+defaults:
+
+  # Dataset
+  - dataset/corrdiff_era5_hrrr
+
+  # Model
+  - model/corrdiff_era5_hrrr
+
+  # Training
+  - training/default
+
+  # Sampler
+  - sampler/rectified_flow
+
+  # Hydra
+  - hydra/default
+
+  - _self_
+
+training:
+  experiment_name: 'corrdiff-era5-hrrr'
+  num_data_workers: 8
+
+  checkpoint_freq: 7812 # How often to save the checkpoints (8 GB each), measured in number of training steps (7812 steps = 0.5M samples)
+  validation_freq: 1562 # How often to record the validation loss, measured in number of training steps (1562 steps = 100k samples)
+  validation_steps: 8
+  validation_plot_variables: ["u10m", "v10m", "t2m", "refc"]
+  validation_plot_background_channels: ["u10m", "t2m"]
+
+  batch_size: 64
+  batch_size_per_gpu: 2 # 2 samples per GPU use about 68 GB; gradients are accumulated on fewer than 32 GPUs
+  total_train_steps: 781250 # 781250 steps with batch size 64 = 50M samples
+  clip_grad_norm: 1.0
+
+  perf:
+    fp_optimizations: amp-bf16
+
+  optimizer:
+    name: "adamw"
+    lr: 1.0E-4
+    betas: [0.9, 0.99]
+    weight_decay: 0.0
+
+  scheduler:
+    name: "CosineAnnealingLR" # cosine decay to zero at total_train_steps after the warmup
+    lr_rampup_steps: 1000
+
+  loss:
+    type: 'flow_matching'
+    prediction_type: 'x0' # the network predicts the clean state (x-prediction)
+    t_min: 0.0 # Lowest time (clean data)
+    t_max: 1.0 # Highest time (pure noise), where the sampling starts

examples/weather/regional_weather_diffusion/config/dataset/corrdiff_era5_hrrr.yaml

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Defaults
+name: data_loader_corrdiff_era5_hrrr.CorrDiffEra5HrrrDataset # Dataset class module and name, in any module in `datasets` folder
+
+# Main dataset, see datasets/data_loader_corrdiff_era5_hrrr.py for the expected layout
+data_path: 'data' # Path to the dataset
+years: [2018, 2019, 2020, 2021, 2022, 2023, 2024, 2025] # Years of data to load
+
+# Temporal selection: data starts at the HRRR version change in 2018,
+# missing samples are automatically excluded
+time_range_train: ["2018-07-01T00:00:00", "2025-10-01T00:00:00"]
+time_range_valid: ["2025-10-01T00:00:00", "2026-01-01T00:00:00"]
+
+# Domain
+crop_window: [[17, 1041], [3, 1795]] # [[y0, y1], [x0, x1]] window of the HRRR grid, 1024 x 1792 pixels
+patch_shape: [256, 256] # Shape of the random training patches; null for the full domain
+
+# Static fields appended to the background channels
+invariant_variables: ["elev_mean", "elev_std", "lsm_mean"]
+
+# Droppable inputs: during training, each listed ERA5 variable is independently
+# replaced by its climatological mean with probability drop_prob. A presence
+# flag per variable (1 = kept, 0 = dropped) is passed to the model with the
+# scalar conditions, so the model can also run without these inputs.
+drop_variables: ["tcwv", "sp"]
+drop_prob: 0.5

examples/weather/regional_weather_diffusion/config/model/corrdiff_era5_hrrr.yaml

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Defaults
+
+# Diffusion transformer (DiT) of CorrDiff-ERA5-HRRR: 678M parameters,
+# neighborhood attention with 2D rotary position embeddings, trained with
+# x-prediction rectified flow matching (training.loss.type: flow_matching).
+# These settings match the released checkpoint.
+
+# Model name
+model_name: 'diffusion'
+
+architecture: dit
+
+# Preconditioner: the times are multiplied by time_scale before the DiT
+# timestep embedding, the network output is the prediction itself
+preconditioner: time_scaled
+time_scale: 999.0 # Flow-matching times in [0, 1] are multiplied by this
+
+# Model conditions: pure downscaling, the low-resolution ERA5 fields (with the
+# static fields appended by the dataset) condition the model
+regression_conditions: []
+diffusion_conditions: ["background"]
+
+# Model hyperparameters
+hyperparameters:
+  hidden_size: 1152 # Number of channels in the transformer blocks
+  depth: 28 # Number of transformer blocks
+  num_heads: 18 # Number of attention heads
+  patch_size: 2 # Size of the patches embedded as tokens
+  attn_kernel_size: 31 # Neighborhood size of the attention, in tokens
+  attention_backend: natten2d_rope # Neighborhood attention with 2D rotary position embeddings
+  layernorm_backend: torch # Layer normalization backend of the blocks, "torch" or "apex"
+  attn_kwargs:
+    qk_norm: true # Normalize the attention queries and keys, stabilizes bf16 training
+    norm_layer: apex # Apex LayerNorm keeps the bf16 dtype that NATTEN needs for queries, keys, and values

examples/weather/regional_weather_diffusion/config/model/stormcast.yaml

@@ -23,6 +23,9 @@ regression_conditions: ["background", "invariant"] # list consisting of "backgro
                                                     # The dataset's "background" tensor includes the previous HRRR timestep.
 diffusion_conditions: ["regression", "invariant"] # list consisting of "regression", "background", "invariant" (default from StormCast paper)
 
+# Preconditioner of the diffusion network: edm preconditioner
+preconditioner: edm
+
 # Model hyperparameters
 model_type: "SongUNetPosEmbd" # Model class to use
 # Example overrides for architecture hyperparameters (uncomment to customize)

examples/weather/regional_weather_diffusion/config/model/stormscope.yaml

@@ -25,6 +25,9 @@ regression_conditions: []
 diffusion_conditions: ["background", "invariant"] # list consisting of "regression", "background", "invariant".
                                                    # The dataset's "background" tensor includes any past-state history.
 
+# Preconditioner of the diffusion network: edm preconditioner
+preconditioner: edm
+
 # Model hyperparameters (defaults follow Stormscope)
 hyperparameters:
   depth: 16

examples/weather/regional_weather_diffusion/config/sampler/rectified_flow.yaml

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Defaults
+
+# Sampler args for physicsnemo.diffusion.samplers.sample with a rectified-flow
+# model (training.loss.type: flow_matching): integrates the rectified-flow ODE
+# dx/dt = (x_t - x0_hat) / t from pure noise (t = 1) to t = 0.
+
+name: 'Rectified flow'
+args:
+  num_steps: 18 # Number of solver steps
+  solver: heun # "heun" (2nd order, 2 model evaluations per step) or "euler" (1st order)
+  time_shift: 1.0 # Time shift a, t -> a t / (1 + (a - 1) t); use a > 1 to sample a larger domain than the training patches
+  S_churn: 0. # Set > 0 for stochastic sampling (EDM-style churn, see the EDM paper); the churn of the solvers assumes alpha = 1 and is an approximation for rectified flow
+  S_min: 0.
+  S_max: .inf
+  S_noise: 1

examples/weather/regional_weather_diffusion/config/test_flow_matching.yaml

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# Defaults
+
+# Tiny DiT downscaling model trained with flow matching on mock data, used by
+# test_training.py
+
+# Defaults
+defaults:
+
+  # Dataset
+  - dataset/mock
+
+  # Model
+  - model/corrdiff_era5_hrrr
+
+  # Training
+  - training/default
+
+  # Sampler
+  - sampler/rectified_flow
+
+  # Hydra
+  - hydra/default
+
+  - _self_
+
+training:
+  num_data_workers: 1
+  validation_plot_variables: ["state_0"]
+  log_to_wandb: True
+  wandb_mode: "offline"  # not using online logging while testing
+  log_to_tensorboard: True
+
+  print_progress_freq: 5 # How often to print progress, measured in number of training steps
+  checkpoint_freq: 5 # How often to save the checkpoints, measured in number of training steps
+  validation_freq: 5 # how often to record the validation loss, measured in number of training steps
+  batch_size: 1
+  total_train_steps: 10
+
+  optimizer:
+    name: "adamw"
+    betas: [0.9, 0.99]
+
+  scheduler:
+    name: null
+
+  loss:
+    type: 'flow_matching'
+    prediction_type: 'x0'
+
+sampler:
+  args:
+    num_steps: 2
+
+dataset:
+  num_samples: 20
+  model_type: downscaling
+
+model:
+  hyperparameters:
+    depth: 2
+    hidden_size: 32
+    num_heads: 1
+    attn_kernel_size: 7

examples/weather/regional_weather_diffusion/config/training/default.yaml

@@ -19,7 +19,7 @@
 outdir: 'rundir' # Root path under which to save training outputs
 experiment_name: 'stormcast-training' # Name for the training experiment
 run_id: '0' # Unique ID to use for this training run
-rundir: ./${training.outdir}/${training.experiment_name}/${training.run_id} # Path where experiement outputs will be saved
+rundir: ${training.outdir}/${training.experiment_name}/${training.run_id} # Path where experiment outputs will be saved
 num_data_workers: 4 # Number of dataloader worker threads per proc
 log_to_wandb: False # Whether or not to log to Weights & Biases (requires wandb account)
 wandb_mode: "online" # logging mode, "online" or "offline"
@@ -95,7 +95,10 @@ channel_loss_weights: null
 
 # Loss options
 loss:
-  type: 'regression' # Loss type; use 'regression' or 'edm' for the regression and diffusion, respectively
+  type: 'regression' # Loss type: 'regression', 'edm' (denoising score matching), or 'flow_matching' (rectified flow)
+  prediction_type: 'x0' # Network output of a diffusion model: 'x0', 'epsilon', 'score', or 'flow'; 'x0' with the EDM preconditioner
+  # t_min: 0.0 # Flow matching: lowest time (clean data)
+  # t_max: 1.0 # Flow matching: highest time (1 = pure noise)
   # sigma_distribution: "lognormal" # EDM noise sampling distribution, "lognormal" or "loguniform"
   # sigma_data: 0.5 # Standard deviation of the data (set to match your normalized data's std)
   # # Lognormal distribution parameters (used when sigma_distribution: "lognormal")