Contribution standards: ruff clean, CHANGELOG, FGNUNet docstring

- Run ruff format + fix across all fgn/ Python files
- Remove unused imports (Sequence, Callable, ShardTensor, math, torch)
- Replace assert with if/raise (S101), fix import order (I001),
  simplify loops to list-comprehension/extend (PERF401/102)
- Add noqa: E402 on intentional post-path-insert imports in stage4
- Upgrade FGNUNet docstring to MOD-003 (r-string, NumPy sections,
  Parameters/Forward/Outputs with LaTeX shapes, Examples)
- Add CHANGELOG.md entry under [2.1.0a0]

Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>

Author: kashif

CHANGELOG.md

@@ -140,6 +140,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added support for Batched radius search, which enables Domino
   and GeoTransolver with local features and batch size > 1.
 - Added the underfill recipe.
+- Adds Functional Generative Networks (FGN) weather training example
+  (`examples/weather/fgn`). Implements the latent-conditioned U-Net
+  stochastic generator from
+  `arXiv:2506.10772 <https://arxiv.org/abs/2506.10772>`_ (WeatherNext 2)
+  as a PhysicsNeMo ``Module``, trained with fair-CRPS loss on ERA5 via the
+  earth2studio ARCO data source. Supports autoregressive rollout training
+  with per-channel normalization, FSDP + ShardTensor domain parallelism,
+  deep-ensemble inference (paper §2.2.1), and validation diagnostics
+  (CRPS, RMSE, spread-skill, rank histograms, power spectra).
 
 ### Changed
 

examples/weather/fgn/datasets/arco.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta
-from typing import Any, Sequence
+from typing import Any
 
 import numpy as np
 import torch
@@ -24,7 +24,19 @@
 # Paper Table A.1 atmospheric schema
 PAPER_ATMOS_VARS: tuple[str, ...] = ("z", "q", "t", "u", "v", "w")
 PAPER_LEVELS: tuple[int, ...] = (
-    50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000,
+    50,
+    100,
+    150,
+    200,
+    250,
+    300,
+    400,
+    500,
+    600,
+    700,
+    850,
+    925,
+    1000,
 )
 PAPER_SURFACE_IN_OUT: tuple[str, ...] = ("t2m", "u10m", "v10m", "msl", "sst")
 
@@ -96,8 +108,14 @@ class ArcoFGNDataset(FGNDataset):
 
     def __init__(self, params: Any, train: bool) -> None:
         def _get(name: str, default: Any) -> Any:
-            return getattr(params, name, default) if hasattr(params, name) else (
-                params[name] if isinstance(params, dict) and name in params else default
+            return (
+                getattr(params, name, default)
+                if hasattr(params, name)
+                else (
+                    params[name]
+                    if isinstance(params, dict) and name in params
+                    else default
+                )
             )
 
         state = _get("state_variables", None)
@@ -232,9 +250,7 @@ def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
         # slot from _fetch_tp_accumulation below.
         if self._tp_channel_idx is not None:
             ci = self._tp_channel_idx
-            arco_vars = [
-                v for i, v in enumerate(self._state_variables) if i != ci
-            ]
+            arco_vars = [v for i, v in enumerate(self._state_variables) if i != ci]
         else:
             ci = None
             arco_vars = list(self._state_variables)
@@ -368,7 +384,10 @@ def _fetch_tp_accumulation(self, frame_times: list[datetime]) -> np.ndarray:
         fetch all distinct hourly stamps required by any frame in a single
         earth2studio call to minimise GCS round-trips.
         """
-        assert self.tp_accumulation_hours is not None
+        if self.tp_accumulation_hours is None:
+            raise RuntimeError(
+                "_fetch_tp_accumulation called without tp_accumulation_hours set"
+            )
         N = self.tp_accumulation_hours
 
         # Union of hours we need across all frames, sorted.
@@ -387,9 +406,7 @@ def _fetch_tp_accumulation(self, frame_times: list[datetime]) -> np.ndarray:
         if self.stride > 1:
             hourly = hourly[:, :: self.stride, :: self.stride]
 
-        acc = np.zeros(
-            (len(frame_times), self.height, self.width), dtype=np.float32
-        )
+        acc = np.zeros((len(frame_times), self.height, self.width), dtype=np.float32)
         for k, hours_k in enumerate(per_frame_hours):
             acc[k] = sum(hourly[hour_to_idx[h]] for h in hours_k)
         return acc
@@ -418,9 +435,7 @@ def _load_stats(self, stats_path: str) -> None:
         self._mean = mean
         self._std = std
 
-    def _broadcast_stats_for(
-        self, x: np.ndarray | torch.Tensor
-    ) -> tuple[Any, Any]:
+    def _broadcast_stats_for(self, x: np.ndarray | torch.Tensor) -> tuple[Any, Any]:
         """Reshape `(V,)` stats to broadcast along the channel axis of ``x``.
 
         Supports ``x`` of shape ``(V, H, W)``, ``(T, V, H, W)``, or
@@ -435,8 +450,12 @@ def _broadcast_stats_for(
         else:
             raise ValueError(f"unsupported state tensor ndim {x.ndim}")
         if isinstance(x, torch.Tensor):
-            mean = torch.as_tensor(self._mean, dtype=x.dtype, device=x.device).reshape(shape)
-            std = torch.as_tensor(self._std, dtype=x.dtype, device=x.device).reshape(shape)
+            mean = torch.as_tensor(self._mean, dtype=x.dtype, device=x.device).reshape(
+                shape
+            )
+            std = torch.as_tensor(self._std, dtype=x.dtype, device=x.device).reshape(
+                shape
+            )
         else:
             mean = self._mean.reshape(shape)
             std = self._std.reshape(shape)

examples/weather/fgn/inference.py

@@ -14,15 +14,14 @@
 from pathlib import Path
 
 import hydra
-from omegaconf import DictConfig
 import torch
+from datasets import dataset_classes
+from omegaconf import DictConfig
+from utils.trainer import find_latest_model_checkpoint
 
 from physicsnemo.core import Module
 from physicsnemo.distributed import DistributedManager
 
-from datasets import dataset_classes
-from utils.trainer import find_latest_model_checkpoint
-
 
 def _resolve_checkpoints(cfg: DictConfig) -> list[str]:
     """Resolve the inference config to an ordered list of checkpoint paths.
@@ -31,9 +30,11 @@ def _resolve_checkpoints(cfg: DictConfig) -> list[str]:
     back to ``inference.checkpoint`` (single path or ``"latest"``). Single-
     model inference is just the length-1 deep-ensemble case.
     """
-    checkpoints = cfg.inference.get("checkpoints", None) if hasattr(
-        cfg.inference, "get"
-    ) else getattr(cfg.inference, "checkpoints", None)
+    checkpoints = (
+        cfg.inference.get("checkpoints", None)
+        if hasattr(cfg.inference, "get")
+        else getattr(cfg.inference, "checkpoints", None)
+    )
     if checkpoints:
         return [str(c) for c in checkpoints]
 
@@ -139,19 +140,25 @@ def run_inference(cfg: DictConfig) -> dict[str, float | str | int | list[int]]:
     # only uses the first step, so collapse to (B, C, H, W).
     if target.ndim == 5:
         target = target[:, 0]
-    background = sample["background"].unsqueeze(0).to(device=device, dtype=torch.float32)
+    background = (
+        sample["background"].unsqueeze(0).to(device=device, dtype=torch.float32)
+    )
 
     invariants = dataset.get_invariants()
     if invariants is not None:
-        invariants = torch.from_numpy(invariants).unsqueeze(0).to(
-            device=device, dtype=torch.float32
+        invariants = (
+            torch.from_numpy(invariants)
+            .unsqueeze(0)
+            .to(device=device, dtype=torch.float32)
         )
 
     all_trajectories: list[torch.Tensor] = []
     num_steps = int(cfg.inference.num_steps)
     output_only = dataset.output_only_channels()
     with torch.no_grad():
-        for ckpt_path, n_members in zip(checkpoint_paths, members_per_model, strict=True):
+        for ckpt_path, n_members in zip(
+            checkpoint_paths, members_per_model, strict=True
+        ):
             if n_members <= 0:
                 continue
             model = Module.from_checkpoint(ckpt_path).to(device).eval()

examples/weather/fgn/scripts/compute_arco_stats.py

@@ -31,7 +31,19 @@
 # Paper Table A.1 defaults -- keep in sync with datasets.arco.DEFAULT_STATE.
 DEFAULT_ATMOS_VARS = ("z", "q", "t", "u", "v", "w")
 DEFAULT_LEVELS = (
-    50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000,
+    50,
+    100,
+    150,
+    200,
+    250,
+    300,
+    400,
+    500,
+    600,
+    700,
+    850,
+    925,
+    1000,
 )
 DEFAULT_SURFACE = ("t2m", "u10m", "v10m", "msl", "sst")
 DEFAULT_STATE = tuple(
@@ -50,43 +62,61 @@ def parse_args() -> argparse.Namespace:
         "arrays to be laid out. Default: paper Table A.1 (83 channels).",
     )
     p.add_argument(
-        "--start", type=str, default="1979-01-01",
+        "--start",
+        type=str,
+        default="1979-01-01",
         help="Earliest sample time (ISO date).",
     )
     p.add_argument(
-        "--end", type=str, default="2018-01-01",
+        "--end",
+        type=str,
+        default="2018-01-01",
         help="Latest sample time (ISO date, exclusive).",
     )
     p.add_argument(
-        "--step-hours", type=int, default=6,
+        "--step-hours",
+        type=int,
+        default=6,
         help="Sampling cadence; restricts timestamps to a 6h grid by default.",
     )
     p.add_argument(
-        "--samples", type=int, default=256,
+        "--samples",
+        type=int,
+        default=256,
         help="Number of random timestamps to average over.",
     )
     p.add_argument(
-        "--stride", type=int, default=1,
+        "--stride",
+        type=int,
+        default=1,
         help="Spatial stride applied to the 721x1440 grid to cut fetch cost.",
     )
     p.add_argument(
-        "--tp-accumulation-hours", type=int, default=None,
+        "--tp-accumulation-hours",
+        type=int,
+        default=None,
         help="If set to N, any variable named tp{N:02d} (e.g. tp06 for N=6) "
-             "in --variables is treated as a paper §3 N-hour accumulation of "
-             "ARCO hourly ``tp``, not a native ARCOLexicon key. Matches "
-             "ArcoFGNDataset.tp_accumulation_hours so stats and training see "
-             "the same representation.",
+        "in --variables is treated as a paper §3 N-hour accumulation of "
+        "ARCO hourly ``tp``, not a native ARCOLexicon key. Matches "
+        "ArcoFGNDataset.tp_accumulation_hours so stats and training see "
+        "the same representation.",
     )
     p.add_argument("--seed", type=int, default=0)
     p.add_argument(
-        "--output", type=Path, required=True,
+        "--output",
+        type=Path,
+        required=True,
         help="Destination .npz path.",
     )
     return p.parse_args()
 
 
 def iter_sample_times(
-    start: datetime, end: datetime, step_hours: int, samples: int, rng: np.random.Generator
+    start: datetime,
+    end: datetime,
+    step_hours: int,
+    samples: int,
+    rng: np.random.Generator,
 ):
     total_hours = int((end - start).total_seconds() // 3600)
     max_offset = total_hours // step_hours

examples/weather/fgn/scripts/prefetch_arco.py

@@ -36,9 +36,9 @@
 from __future__ import annotations
 
 import argparse
+import sys
 from datetime import datetime, timedelta
 from pathlib import Path
-import sys
 
 _EXAMPLE_DIR = Path(__file__).resolve().parents[1]
 if str(_EXAMPLE_DIR) not in sys.path:
@@ -47,12 +47,18 @@
 
 def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(description=__doc__)
-    p.add_argument("--start", default="2024-01-01", help="Window start (ISO, inclusive).")
+    p.add_argument(
+        "--start", default="2024-01-01", help="Window start (ISO, inclusive)."
+    )
     p.add_argument("--end", default="2025-01-01", help="Window end (ISO, exclusive).")
     p.add_argument("--step-hours", type=int, default=6)
     p.add_argument("--history-frames", type=int, default=2)
-    p.add_argument("--tp-accumulation-hours", type=int, default=6,
-                   help="N for tp{N:02d} accumulation. 0 = skip tp fetch.")
+    p.add_argument(
+        "--tp-accumulation-hours",
+        type=int,
+        default=6,
+        help="N for tp{N:02d} accumulation. 0 = skip tp fetch.",
+    )
     p.add_argument(
         "--variables",
         nargs="+",
@@ -65,20 +71,28 @@ def parse_args() -> argparse.Namespace:
         help="Date for one-off invariants fetch (z, lsm).",
     )
     p.add_argument(
-        "--batch-days", type=int, default=31,
+        "--batch-days",
+        type=int,
+        default=31,
         help="Days of data per time-batch. Default: 31.",
     )
     p.add_argument(
-        "--var-group-size", type=int, default=10,
+        "--var-group-size",
+        type=int,
+        default=10,
         help="Variables per sub-batch. Limits concurrent GCS requests to "
-             "batch_days_timestamps × var_group_size. Default: 10.",
+        "batch_days_timestamps × var_group_size. Default: 10.",
     )
     p.add_argument(
-        "--no-tp", dest="include_tp", action="store_false",
+        "--no-tp",
+        dest="include_tp",
+        action="store_false",
         help="Skip hourly tp prefetch.",
     )
     p.add_argument(
-        "--no-invariants", dest="include_invariants", action="store_false",
+        "--no-invariants",
+        dest="include_invariants",
+        action="store_false",
         help="Skip invariants prefetch.",
     )
     p.set_defaults(include_tp=True, include_invariants=True)
@@ -94,7 +108,9 @@ def _window_times(start: datetime, end: datetime, step_hours: int) -> list[datet
     return out
 
 
-def _batch(times: list[datetime], batch_days: int, step_hours: int = 1) -> list[list[datetime]]:
+def _batch(
+    times: list[datetime], batch_days: int, step_hours: int = 1
+) -> list[list[datetime]]:
     """Split a time list into chunks covering at most batch_days of real time."""
     n = max(1, batch_days * 24 // step_hours)
     return [times[i : i + n] for i in range(0, len(times), n)]
@@ -104,7 +120,9 @@ def main() -> int:
     args = parse_args()
     # Silence earth2studio's per-fetch DEBUG lines — they flood log files at scale.
     import sys
+
     from loguru import logger
+
     logger.remove()
     logger.add(sys.stderr, level="INFO")
 
@@ -169,7 +187,7 @@ def main() -> int:
                 f"[state {ti}/{n_tb} vars {vi}/{n_vg}]  "
                 f"{batch_times[0].date()} → {batch_times[-1].date()}"
                 f"  ({len(batch_times)} steps × {len(vgroup)} vars"
-                f" = {len(batch_times)*len(vgroup)} requests)"
+                f" = {len(batch_times) * len(vgroup)} requests)"
             )
             arco(time=batch_times, variable=vgroup)
 

examples/weather/fgn/scripts/stage4_ar_schedule.py

@@ -53,8 +53,8 @@
 if str(_EXAMPLE_DIR) not in sys.path:
     sys.path.insert(0, str(_EXAMPLE_DIR))
 
-from hydra import compose, initialize
-from omegaconf import DictConfig, OmegaConf
+from hydra import compose, initialize  # noqa: E402
+from omegaconf import DictConfig, OmegaConf  # noqa: E402
 
 # Stage 4 of paper Table A.2.
 PAPER_STAGES: list[dict] = [
@@ -143,7 +143,6 @@ def _seed_from_prev_stage(
     stage_checkpoint_dir.mkdir(parents=True, exist_ok=True)
     shutil.copy2(last_mdlus, stage_checkpoint_dir / last_mdlus.name)
     # Copy optimizer/scheduler state if present so the resume is exact.
-    opt_pt = last_mdlus.with_suffix("").with_suffix(".pt")
     # physicsnemo names these ``checkpoint.{mp_rank}.{epoch}.pt``; find the
     # matching epoch by filename suffix (``.<epoch>.mdlus``).
     epoch_suffix = last_mdlus.stem.split(".")[-1]