Upload folder using huggingface_hub

README.md

@@ -0,0 +1,149 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Climate Simulation
+- Probabilistic Forecasting
+- FV3GFS
+- Spherical DYffusion
+tasks: []
+datasets:
+- FV3GFS
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Spherical DYffusion</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Spherical DYffusion was proposed by Salva Ruhling Cachay and collaborators for probabilistic simulation of a global climate model.
+
+Paper: Probabilistic Emulation of a Global Climate Model with Spherical DYffusion
+
+https://arxiv.org/abs/2406.14798
+
+# Model Description
+
+The original method models spherical dynamics with an SFNO and uses the DYffusion interpolator and forecaster in a two-stage training procedure for probabilistic ensemble simulation. This repository contains a compact local implementation that preserves the project's tensor and data contracts for smoke testing; it is not a full paper-scale SFNO/DYffusion implementation.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Local pipeline validation | Use synthetic 37-channel global-grid data to check training, inference, and visualization. |
+| FV3GFS protocol checks | Validate NetCDF variables, spatial dimensions, and consecutive time frames. |
+| ModelScope / OneCode execution | Download the standalone model package and run the compact local pipeline. |
+| Multi-GPU training | Launch PyTorch DistributedDataParallel with `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/Spherical_DYffusion --local-dir ./Spherical_DYffusion
+cd Spherical_DYffusion
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Generate Synthetic Data
+
+Generate a deterministic NetCDF FV3GFS-contract fixture at `data/data/synthetic_fv3gfs.nc`:
+
+```bash
+python scripts/fake_data.py
+```
+
+The fixture contains 37 protocol variables, including surface pressure and temperature, eight vertical levels of temperature, total water, and wind components, plus `DSWRFtoa`, `HGTsfc`, and `ocean_fraction`. It is intended only for protocol checks. The local training pipeline creates its own learnable `data/data/virtual_fv3gfs.npz` fixture.
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+Training starts from random initialization and saves `data/checkpoint/model_bak.pt` and `data/checkpoint/last.pt`.
+
+The complete local smoke workflow can also be run with:
+
+```bash
+python scripts/local_pipeline.py all
+```
+
+### Training Weights
+
+This repository provides a `weight/` directory for FV3GFS-compatible checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads `data/checkpoint/model_bak.pt` and writes `output/inference/prediction.npz`:
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The script computes per-variable and overall diagnostics and writes `output/visualization/diagnostic_dashboard.png` and `output/visualization/variable_metrics.png`, with machine-readable summaries under `output/metrics/`.
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- This repository is a compact local reproduction of the original Spherical DYffusion paper and does not claim to reproduce the paper-scale training setup or metrics.

conf/config.yaml

@@ -0,0 +1,58 @@
+project:
+  name: Spherical-DYffusion-SFNO
+  version: "1.0"
+
+paths:
+  data_root: data
+  raw_data: data/raw
+  processed_data: data/processed
+  metadata: data/metadata
+  statistics: data/statistics
+  pretrained_weights: data/weights
+  checkpoints: data/checkpoint
+  result_root: output/result
+  predictions: output/inference
+  metrics: output/metrics
+  figures: output/visualization
+  work_dir: data/work
+
+model:
+  name: Spherical-DYffusion-SFNO
+  resolution: 1deg
+  latitude: 180
+  longitude: 360
+
+runtime:
+  device: auto
+  devices: auto
+  distributed_backend: auto
+  num_workers: 8
+
+training:
+  seed: 11
+  epochs: 5
+  batch_size: 2
+  learning_rate: 0.0001
+  checkpoint_dir: data/checkpoint
+  validation_split: 0.2
+
+synthetic_data:
+  output_dir: data/data
+  time_steps: 8
+  latitude: 180
+  longitude: 360
+  samples: 16
+  input_steps: 1
+  channels: 37
+  seed: 11
+
+inference:
+  output_dir: output/inference
+  checkpoint: data/checkpoint/model_bak.pt
+
+visualization:
+  output_dir: output/visualization
+  prediction_index: 0
+  channel: 1
+  scatter_points: 12000
+  dpi: 150

config.json

@@ -0,0 +1,87 @@
+{
+  "model_name": "Spherical_DYffusion",
+  "model_type": "spherical_dyffusion",
+  "architectures": [
+    "SphericalDYffusion"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "global-grid-climate-emulation",
+  "implementation": {
+    "entry_point": "model/spherical_dyffusion.py",
+    "scope": "Compact local convolutional forecaster preserving the Spherical DYffusion project tensor contract; it is a smoke-test implementation, not the full SFNO/DYffusion paper architecture"
+  },
+  "architecture": {
+    "family": "three-layer Conv2d forecaster with GELU activations",
+    "input_format": "BCHW",
+    "output_format": "BCHW",
+    "input_grid_shape": [
+      180,
+      360
+    ],
+    "resolution_degrees": 1.0,
+    "channels": 37,
+    "hidden_channels": 32,
+    "layers": [
+      "Conv2d(37,32,kernel_size=3,padding=1)",
+      "GELU",
+      "Conv2d(32,32,kernel_size=3,padding=1)",
+      "GELU",
+      "Conv2d(32,37,kernel_size=1)"
+    ],
+    "stochasticity": "none in the current local model; the training target is a deterministic synthetic local-dynamics rule"
+  },
+  "data": {
+    "dataset_protocol": "FV3GFS",
+    "variables": [
+      "PRESsfc",
+      "surface_temperature",
+      "air_temperature_0",
+      "air_temperature_1",
+      "air_temperature_2",
+      "air_temperature_3",
+      "air_temperature_4",
+      "air_temperature_5",
+      "air_temperature_6",
+      "air_temperature_7",
+      "specific_total_water_0",
+      "specific_total_water_1",
+      "specific_total_water_2",
+      "specific_total_water_3",
+      "specific_total_water_4",
+      "specific_total_water_5",
+      "specific_total_water_6",
+      "specific_total_water_7",
+      "eastward_wind_0",
+      "eastward_wind_1",
+      "eastward_wind_2",
+      "eastward_wind_3",
+      "eastward_wind_4",
+      "eastward_wind_5",
+      "eastward_wind_6",
+      "eastward_wind_7",
+      "northward_wind_0",
+      "northward_wind_1",
+      "northward_wind_2",
+      "northward_wind_3",
+      "northward_wind_4",
+      "northward_wind_5",
+      "northward_wind_6",
+      "northward_wind_7",
+      "DSWRFtoa",
+      "HGTsfc",
+      "ocean_fraction"
+    ],
+    "synthetic_time_steps": 8,
+    "input_steps": 1,
+    "samples": 16,
+    "data_note": "Synthetic FV3GFS-equivalent data is for protocol and pipeline checks only and is not a paper-reproduction dataset."
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/spherical_dyffusion.py",
+    "scripts/local_pipeline.py",
+    "scripts/fake_data.py",
+    "README.md"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "probabilistic_climate_emulation",
+  "model": "Spherical_DYffusion",
+  "input_format": "BCHW",
+  "protocol": "fv3gfs_1deg_37_channel",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/spherical_dyffusion.py

@@ -0,0 +1,20 @@
+"""Compact trainable model used by the local Spherical DYffusion pipeline."""
+
+from torch import nn
+
+
+class SphericalDYffusion(nn.Module):
+    """Small convolutional forecaster preserving the global-grid tensor contract."""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(channels, 32, 3, padding=1),
+            nn.GELU(),
+            nn.Conv2d(32, 32, 3, padding=1),
+            nn.GELU(),
+            nn.Conv2d(32, channels, 1),
+        )
+
+    def forward(self, inputs):
+        return self.net(inputs)

scripts/download.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+WEIGHT_DIR="${WEIGHT_DIR:-/path/to/weight/pretrained}"
+mkdir -p "$WEIGHT_DIR"
+curl -L --fail --retry 3 "https://huggingface.co/salv47/spherical-dyffusion/resolve/main/forecaster-sfno-best-inference_avg_crps.ckpt?download=true" -o "$WEIGHT_DIR/forecaster-sfno-best-inference_avg_crps.ckpt"
+curl -L --fail --retry 3 "https://huggingface.co/salv47/spherical-dyffusion/resolve/main/interpolator-sfno-best-val_avg_crps.ckpt?download=true" -o "$WEIGHT_DIR/interpolator-sfno-best-val_avg_crps.ckpt"
+curl -L --fail --retry 3 "https://huggingface.co/salv47/spherical-dyffusion/resolve/main/interpolator_sfno_paper_v0_hydra_config.yaml?download=true" -o "$WEIGHT_DIR/interpolator_sfno_paper_v0_hydra_config.yaml"

scripts/fake_data.py

@@ -0,0 +1,76 @@
+"""Generate deterministic synthetic spherical data for pipeline smoke tests."""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+
+
+VARIABLES = [
+    "PRESsfc",
+    "surface_temperature",
+    *[f"air_temperature_{i}" for i in range(8)],
+    *[f"specific_total_water_{i}" for i in range(8)],
+    *[f"eastward_wind_{i}" for i in range(8)],
+    *[f"northward_wind_{i}" for i in range(8)],
+    "DSWRFtoa",
+    "HGTsfc",
+    "ocean_fraction",
+]
+
+
+def generate_data(output_dir: str, time_steps: int, latitude: int, longitude: int) -> Path:
+    """Write a small NetCDF fixture and metadata to the configured data directory."""
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    data = np.zeros((time_steps, len(VARIABLES), latitude, longitude), dtype=np.float32)
+    lat = np.linspace(-90, 90, latitude, dtype=np.float32)
+    lon = np.linspace(0, 360, longitude, endpoint=False, dtype=np.float32)
+    data[:, 0] = 1.0
+    data[:, 1] = 280.0 + np.sin(np.deg2rad(lat))[None, :, None]
+    data[:, -1] = 1.0
+    dataset = xr.Dataset(
+        {name: (("time", "lat", "lon"), data[:, index]) for index, name in enumerate(VARIABLES)},
+        coords={"time": np.arange(time_steps), "lat": lat, "lon": lon},
+        attrs={"dataset_type": "synthetic_smoke_data", "paper_reproduction": "false"},
+    )
+    netcdf_path = output_path / "synthetic_fv3gfs.nc"
+    dataset.to_netcdf(netcdf_path)
+    metadata = {
+        "dataset_type": "synthetic_smoke_data",
+        "paper_reproduction": False,
+        "variables": VARIABLES,
+        "shape": [time_steps, len(VARIABLES), latitude, longitude],
+        "source": "scripts/generate_data.py",
+    }
+    (output_path / "synthetic_fv3gfs.json").write_text(
+        json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
+    )
+    return netcdf_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--output-dir")
+    parser.add_argument("--time-steps", type=int)
+    parser.add_argument("--latitude", type=int)
+    parser.add_argument("--longitude", type=int)
+    args = parser.parse_args()
+    import yaml
+
+    config = yaml.safe_load(open(args.config, encoding="utf-8"))["synthetic_data"]
+    path = generate_data(
+        args.output_dir or config["output_dir"],
+        args.time_steps or config["time_steps"],
+        args.latitude or config["latitude"],
+        args.longitude or config["longitude"],
+    )
+    print(f"Generated synthetic smoke data: {path}")
+    print("This fixture is for pipeline validation only, not paper reproduction.")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,17 @@
+"""Run inference with the locally trained model."""
+
+import argparse
+from local_pipeline import infer, load_config
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--config", default="conf/config.yaml", help="Pipeline YAML configuration",
+    )
+    args = parser.parse_args()
+    infer(load_config(args.config))
+
+
+if __name__ == "__main__":
+    main()

scripts/local_pipeline.py

@@ -0,0 +1,289 @@
+"""Run the self-contained virtual-data training, inference, and plotting pipeline."""
+
+import argparse
+import json
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
+import sys
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+from model.spherical_dyffusion import SphericalDYffusion
+
+
+VARIABLES = [
+    "PRESsfc", "surface_temperature",
+    *[f"air_temperature_{i}" for i in range(8)],
+    *[f"specific_total_water_{i}" for i in range(8)],
+    *[f"eastward_wind_{i}" for i in range(8)],
+    *[f"northward_wind_{i}" for i in range(8)],
+    "DSWRFtoa", "HGTsfc", "ocean_fraction",
+]
+
+
+def load_config(path: str) -> dict:
+    import yaml
+
+    with open(path, encoding="utf-8") as file:
+        return yaml.safe_load(file)
+
+
+def resolve_device(config: dict, local_rank: int = 0) -> torch.device:
+    requested = str(config.get("runtime", {}).get("device", "auto")).lower()
+    if requested == "auto":
+        requested = "cuda" if torch.cuda.is_available() else "cpu"
+    if requested.startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError(
+            f"runtime.device={requested!r}, but PyTorch cannot access a CUDA/ROCm device. "
+            "Use runtime.device=cpu or install a GPU-enabled PyTorch build."
+        )
+    device = torch.device(requested)
+    if device.type == "cuda":
+        device_index = device.index if device.index is not None else local_rank
+        if device_index >= torch.cuda.device_count():
+            raise RuntimeError(
+                f"LOCAL_RANK={local_rank} maps to GPU {device_index}, but only "
+                f"{torch.cuda.device_count()} GPU(s) are visible."
+            )
+        torch.cuda.set_device(device_index)
+        device = torch.device("cuda", device_index)
+        print(
+            f"device: {device} ({torch.cuda.get_device_name(device_index)}), "
+            f"backend={'ROCm ' + torch.version.hip if torch.version.hip else 'CUDA ' + str(torch.version.cuda)}"
+        )
+    else:
+        print("device: cpu")
+    return device
+
+
+def setup_distributed(config: dict) -> tuple[int, int, torch.device]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    configured_devices = config.get("runtime", {}).get("devices", "auto")
+    requested_device = str(config.get("runtime", {}).get("device", "auto")).lower()
+    if isinstance(configured_devices, int) and configured_devices > 1 and world_size == 1:
+        raise RuntimeError(
+            f"runtime.devices={configured_devices} requires a distributed launcher. Run: "
+            f"python -m torch.distributed.run --standalone --nproc-per-node={configured_devices} "
+            "scripts/train.py"
+        )
+    if isinstance(configured_devices, int) and world_size > 1 and configured_devices != world_size:
+        raise RuntimeError(
+            f"runtime.devices={configured_devices} does not match torchrun WORLD_SIZE={world_size}."
+        )
+    if world_size > 1 and requested_device.startswith("cuda:"):
+        raise RuntimeError(
+            "Do not set an explicit CUDA device index for DDP. Use runtime.device=cuda or auto; "
+            "each torchrun process is mapped to its LOCAL_RANK automatically."
+        )
+
+    device = resolve_device(config, local_rank=local_rank)
+    if world_size > 1:
+        if not dist.is_available():
+            raise RuntimeError("Distributed training is unavailable in this PyTorch build.")
+        backend = str(config.get("runtime", {}).get("distributed_backend", "auto")).lower()
+        if backend == "auto":
+            backend = "nccl" if device.type == "cuda" else "gloo"
+        dist.init_process_group(backend=backend, init_method="env://")
+        if dist.get_world_size() != world_size or dist.get_rank() != rank:
+            raise RuntimeError("The process group does not match the torchrun rank settings.")
+    return rank, world_size, device
+
+
+def generate(config: dict) -> Path:
+    spec = config["synthetic_data"]
+    rng = np.random.default_rng(spec["seed"])
+    shape = (spec["samples"], spec["channels"], spec["latitude"], spec["longitude"])
+    inputs = rng.normal(0, 1, shape).astype(np.float32)
+    # A deterministic local dynamics rule provides a learnable target.
+    targets = (0.85 * inputs + 0.05 * np.roll(inputs, 1, axis=2) + 0.05 * np.roll(inputs, -1, axis=3)).astype(np.float32)
+    output = Path(spec["output_dir"])
+    output.mkdir(parents=True, exist_ok=True)
+    path = output / "virtual_fv3gfs.npz"
+    np.savez_compressed(path, inputs=inputs, targets=targets)
+    (output / "metadata.json").write_text(json.dumps({
+        "variables": VARIABLES, "shape": list(shape), "time_steps": spec["time_steps"],
+        "latitude": spec["latitude"], "longitude": spec["longitude"],
+        "dataset_type": "virtual_fv3gfs_equivalent_contract",
+    }, indent=2) + "\n", encoding="utf-8")
+    print(f"virtual data: {path}")
+    return path
+
+
+def train(config: dict, finetune: str | None = None) -> Path:
+    seed = config["training"]["seed"]
+    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
+    rank, world_size, device = setup_distributed(config)
+    is_root = rank == 0
+    checkpoint = Path(config["training"]["checkpoint_dir"])
+    try:
+        if int(config["training"]["epochs"]) < 1:
+            raise ValueError("training.epochs must be at least 1.")
+        data_path = Path(config["synthetic_data"]["output_dir"]) / "virtual_fv3gfs.npz"
+        if is_root:
+            data_path = generate(config)
+        if world_size > 1:
+            dist.barrier()
+        arrays = np.load(data_path)
+        dataset = TensorDataset(torch.from_numpy(arrays["inputs"]), torch.from_numpy(arrays["targets"]))
+
+        global_batch_size = int(config["training"]["batch_size"])
+        if global_batch_size % world_size:
+            raise ValueError(
+                f"training.batch_size={global_batch_size} is the global batch size and must be "
+                f"divisible by WORLD_SIZE={world_size}."
+            )
+        if len(dataset) % world_size:
+            raise ValueError(
+                f"Dataset size {len(dataset)} must be divisible by WORLD_SIZE={world_size}; "
+                "otherwise DistributedSampler would duplicate samples and bias the epoch loss."
+            )
+        local_batch_size = global_batch_size // world_size
+        if local_batch_size < 1:
+            raise ValueError("The global batch size must be at least WORLD_SIZE.")
+        sampler = DistributedSampler(
+            dataset,
+            num_replicas=world_size,
+            rank=rank,
+            shuffle=True,
+            seed=seed,
+            drop_last=False,
+        ) if world_size > 1 else None
+        total_workers = int(config.get("runtime", {}).get("num_workers", 0))
+        local_workers = max(0, total_workers // world_size)
+        loader = DataLoader(
+            dataset,
+            batch_size=local_batch_size,
+            shuffle=sampler is None,
+            sampler=sampler,
+            num_workers=local_workers,
+            pin_memory=device.type == "cuda",
+            persistent_workers=local_workers > 0,
+        )
+        model = SphericalDYffusion(config["synthetic_data"]["channels"]).to(device)
+        if finetune:
+            finetune_path = Path(finetune)
+            if not finetune_path.exists():
+                raise FileNotFoundError(f"Fine-tune checkpoint not found: {finetune_path}")
+            state = torch.load(finetune_path, map_location="cpu", weights_only=False)
+            state_dict = state.get("model", state) if isinstance(state, dict) else state
+            try:
+                model.load_state_dict(state_dict)
+            except RuntimeError as error:
+                raise RuntimeError(
+                    "Fine-tune checkpoint is incompatible with SphericalDYffusion. "
+                    "Use a checkpoint produced by this local pipeline or a matching model architecture."
+                ) from error
+            if is_root:
+                print(f"fine-tuning from: {finetune_path}")
+        elif is_root:
+            print("training from scratch")
+        if world_size > 1:
+            model = DistributedDataParallel(
+                model,
+                device_ids=[device.index] if device.type == "cuda" else None,
+                output_device=device.index if device.type == "cuda" else None,
+            )
+            if is_root:
+                print(
+                    f"distributed: DDP world_size={world_size}, global_batch_size={global_batch_size}, "
+                    f"local_batch_size={local_batch_size}"
+                )
+        optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
+        loss_fn = nn.MSELoss()
+        best = float("inf")
+        if is_root:
+            checkpoint.mkdir(parents=True, exist_ok=True)
+        for epoch in range(1, config["training"]["epochs"] + 1):
+            if sampler is not None:
+                sampler.set_epoch(epoch)
+            model.train(); total = 0.0; sample_count = 0
+            for inputs, targets in loader:
+                inputs = inputs.to(device, non_blocking=True)
+                targets = targets.to(device, non_blocking=True)
+                optimizer.zero_grad(); loss = loss_fn(model(inputs), targets); loss.backward(); optimizer.step()
+                total += loss.item() * len(inputs); sample_count += len(inputs)
+            loss_stats = torch.tensor([total, sample_count], dtype=torch.float64, device=device)
+            if world_size > 1:
+                dist.all_reduce(loss_stats, op=dist.ReduceOp.SUM)
+            mean_loss = (loss_stats[0] / loss_stats[1]).item()
+            if is_root:
+                print(f"epoch {epoch}/{config['training']['epochs']} loss={mean_loss:.6f}")
+                state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
+                state = {
+                    "model": state_dict,
+                    "channels": config["synthetic_data"]["channels"],
+                    "loss": mean_loss,
+                    "world_size": world_size,
+                    "global_batch_size": global_batch_size,
+                }
+                if mean_loss < best:
+                    best = mean_loss
+                    torch.save(state, checkpoint / "model_bak.pt")
+        if is_root:
+            torch.save(state, checkpoint / "last.pt")
+            print(f"checkpoint: {checkpoint / 'model_bak.pt'}")
+        if world_size > 1:
+            dist.barrier()
+        return checkpoint / "model_bak.pt"
+    finally:
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def infer(config: dict) -> Path:
+    arrays = np.load(Path(config["synthetic_data"]["output_dir"]) / "virtual_fv3gfs.npz")
+    checkpoint_path = Path(config["inference"]["checkpoint"])
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Inference checkpoint not found: {checkpoint_path}. Run scripts/train.py first.")
+    state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model = SphericalDYffusion(state["channels"]); model.load_state_dict(state["model"]); model.eval()
+    with torch.no_grad(): prediction = model(torch.from_numpy(arrays["inputs"]))
+    output = Path(config["inference"]["output_dir"]); output.mkdir(parents=True, exist_ok=True)
+    path = output / "prediction.npz"; np.savez_compressed(path, prediction=prediction.numpy(), target=arrays["targets"])
+    print(f"inference: {path}"); return path
+
+
+def visualize(config: dict) -> Path:
+    import matplotlib.pyplot as plt
+
+    arrays = np.load(Path(config["inference"]["output_dir"]) / "prediction.npz")
+    index, channel = config["visualization"]["prediction_index"], config["visualization"]["channel"]
+    figure, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)
+    for axis, image, title in zip(axes, [arrays["target"][index, channel], arrays["prediction"][index, channel]], ["Target", "Prediction"]):
+        plot = axis.imshow(image, cmap="viridis"); axis.set_title(title); axis.set_xlabel("longitude"); axis.set_ylabel("latitude"); figure.colorbar(plot, ax=axis)
+    output = Path(config["visualization"]["output_dir"]); output.mkdir(parents=True, exist_ok=True)
+    path = output / "prediction_comparison.png"; figure.savefig(path, dpi=150); plt.close(figure)
+    print(f"visualization: {path}"); return path
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("action", choices=["generate", "train", "infer", "visualize", "all"], nargs="?", default="all")
+    args = parser.parse_args(); config = load_config(args.config)
+    if args.action == "generate":
+        generate(config)
+    elif args.action == "train":
+        train(config)
+    elif args.action == "infer":
+        infer(config)
+    elif args.action == "visualize":
+        visualize(config)
+    elif args.action == "all":
+        train(config)
+        if int(os.environ.get("RANK", "0")) == 0:
+            infer(config)
+            visualize(config)
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,240 @@
+"""Create diagnostic figures and metrics for virtual-data predictions."""
+
+import argparse
+import csv
+import json
+from pathlib import Path
+
+import matplotlib
+import numpy as np
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+def load_config(path: str) -> dict:
+    import yaml
+
+    with open(path, encoding="utf-8") as file:
+        return yaml.safe_load(file)
+
+
+def load_variables(metadata_path: Path, channels: int) -> list[str]:
+    if metadata_path.exists():
+        variables = json.loads(metadata_path.read_text(encoding="utf-8")).get("variables", [])
+        if len(variables) == channels:
+            return variables
+    return [f"channel_{index}" for index in range(channels)]
+
+
+def compute_metrics(prediction: np.ndarray, target: np.ndarray, variables: list[str]) -> tuple[list[dict], dict]:
+    channel_metrics = []
+    total_squared_error = total_absolute_error = total_error = 0.0
+    total_count = 0
+    sum_target = sum_prediction = sum_target_squared = sum_prediction_squared = sum_product = 0.0
+    for index, name in enumerate(variables):
+        channel_target = target[:, index].astype(np.float64)
+        channel_prediction = prediction[:, index].astype(np.float64)
+        error = channel_prediction - channel_target
+        count = error.size
+        squared_error = float(np.sum(error**2))
+        absolute_error = float(np.sum(np.abs(error)))
+        error_sum = float(np.sum(error))
+        target_sum = float(np.sum(channel_target))
+        prediction_sum = float(np.sum(channel_prediction))
+        target_squared = float(np.sum(channel_target**2))
+        prediction_squared = float(np.sum(channel_prediction**2))
+        product_sum = float(np.sum(channel_target * channel_prediction))
+        covariance = product_sum - target_sum * prediction_sum / count
+        variance_target = target_squared - target_sum**2 / count
+        variance_prediction = prediction_squared - prediction_sum**2 / count
+        correlation = covariance / max(np.sqrt(variance_target * variance_prediction), 1e-12)
+        channel_metrics.append({
+            "channel": index,
+            "variable": name,
+            "rmse": float(np.sqrt(squared_error / count)),
+            "mae": absolute_error / count,
+            "bias": error_sum / count,
+            "correlation": float(correlation),
+        })
+        total_squared_error += squared_error
+        total_absolute_error += absolute_error
+        total_error += error_sum
+        total_count += count
+        sum_target += target_sum
+        sum_prediction += prediction_sum
+        sum_target_squared += target_squared
+        sum_prediction_squared += prediction_squared
+        sum_product += product_sum
+    covariance = sum_product - sum_target * sum_prediction / total_count
+    variance_target = sum_target_squared - sum_target**2 / total_count
+    variance_prediction = sum_prediction_squared - sum_prediction**2 / total_count
+    overall = {
+        "rmse": float(np.sqrt(total_squared_error / total_count)),
+        "mae": total_absolute_error / total_count,
+        "bias": total_error / total_count,
+        "correlation": float(covariance / max(np.sqrt(variance_target * variance_prediction), 1e-12)),
+    }
+    return channel_metrics, overall
+
+
+def plot_diagnostics(
+    prediction: np.ndarray,
+    target: np.ndarray,
+    variables: list[str],
+    channel_metrics: list[dict],
+    config: dict,
+    output_dir: Path,
+) -> Path:
+    spec = config["visualization"]
+    sample = int(spec["prediction_index"])
+    channel = int(spec["channel"])
+    if sample >= prediction.shape[0] or channel >= prediction.shape[1]:
+        raise IndexError(f"Requested sample={sample}, channel={channel}, but prediction shape is {prediction.shape}")
+
+    selected_target = target[sample, channel]
+    selected_prediction = prediction[sample, channel]
+    selected_error = selected_prediction - selected_target
+    field_min = min(float(selected_target.min()), float(selected_prediction.min()))
+    field_max = max(float(selected_target.max()), float(selected_prediction.max()))
+    error_limit = max(float(np.abs(selected_error).max()), 1e-8)
+    latitude = np.linspace(-89.5, 89.5, prediction.shape[2])
+    longitude = np.linspace(0.0, 360.0, prediction.shape[3], endpoint=False)
+
+    figure = plt.figure(figsize=(16, 10), constrained_layout=True)
+    grid = figure.add_gridspec(2, 3)
+    extent = [longitude[0], longitude[-1], latitude[0], latitude[-1]]
+    for axis, field, title in zip(
+        [figure.add_subplot(grid[0, 0]), figure.add_subplot(grid[0, 1])],
+        [selected_target, selected_prediction],
+        ["Target field", "Predicted field"],
+    ):
+        image = axis.imshow(field, origin="lower", extent=extent, aspect="auto", cmap="viridis", vmin=field_min, vmax=field_max)
+        axis.set_title(title)
+        axis.set_xlabel("Longitude (degrees)")
+        axis.set_ylabel("Latitude (degrees)")
+        figure.colorbar(image, ax=axis, shrink=0.82)
+
+    error_axis = figure.add_subplot(grid[0, 2])
+    image = error_axis.imshow(selected_error, origin="lower", extent=extent, aspect="auto", cmap="RdBu_r", vmin=-error_limit, vmax=error_limit)
+    error_axis.set_title("Prediction error (prediction - target)")
+    error_axis.set_xlabel("Longitude (degrees)")
+    error_axis.set_ylabel("Latitude (degrees)")
+    figure.colorbar(image, ax=error_axis, shrink=0.82)
+
+    zonal_axis = figure.add_subplot(grid[1, 0])
+    zonal_axis.plot(selected_target.mean(axis=1), latitude, label="Target", linewidth=2)
+    zonal_axis.plot(selected_prediction.mean(axis=1), latitude, label="Prediction", linewidth=2)
+    zonal_axis.set_title("Zonal-mean profile")
+    zonal_axis.set_xlabel("Zonal mean")
+    zonal_axis.set_ylabel("Latitude (degrees)")
+    zonal_axis.grid(alpha=0.25)
+    zonal_axis.legend()
+
+    scatter_axis = figure.add_subplot(grid[1, 1])
+    stride = max(1, selected_target.size // int(spec["scatter_points"]))
+    x = selected_target.ravel()[::stride]
+    y = selected_prediction.ravel()[::stride]
+    scatter_axis.hexbin(x, y, gridsize=45, mincnt=1, cmap="magma")
+    diagonal_min = min(float(x.min()), float(y.min()))
+    diagonal_max = max(float(x.max()), float(y.max()))
+    scatter_axis.plot([diagonal_min, diagonal_max], [diagonal_min, diagonal_max], "--", color="white", linewidth=1.5)
+    scatter_axis.set_title("Pointwise agreement")
+    scatter_axis.set_xlabel("Target")
+    scatter_axis.set_ylabel("Prediction")
+
+    sample_axis = figure.add_subplot(grid[1, 2])
+    sample_rmse = np.array([
+        np.sqrt(np.mean((prediction[index].astype(np.float64) - target[index]) ** 2))
+        for index in range(prediction.shape[0])
+    ])
+    sample_axis.bar(np.arange(len(sample_rmse)), sample_rmse, color="#2a6f97")
+    sample_axis.axhline(sample_rmse.mean(), color="#d1495b", linestyle="--", label=f"Mean {sample_rmse.mean():.3f}")
+    sample_axis.set_title("RMSE by sample")
+    sample_axis.set_xlabel("Sample index")
+    sample_axis.set_ylabel("RMSE")
+    sample_axis.legend()
+
+    metric = channel_metrics[channel]
+    figure.suptitle(
+        f"Virtual FV3GFS diagnostic | {variables[channel]} | sample {sample}\n"
+        f"RMSE={metric['rmse']:.4f}  MAE={metric['mae']:.4f}  Bias={metric['bias']:.4f}  Corr={metric['correlation']:.4f}",
+        fontsize=15,
+    )
+    path = output_dir / "diagnostic_dashboard.png"
+    figure.savefig(path, dpi=int(spec["dpi"]))
+    plt.close(figure)
+    return path
+
+
+def plot_channel_metrics(channel_metrics: list[dict], output_dir: Path, dpi: int) -> Path:
+    labels = [item["variable"] for item in channel_metrics]
+    rmse = [item["rmse"] for item in channel_metrics]
+    correlation = [item["correlation"] for item in channel_metrics]
+    positions = np.arange(len(labels))
+    figure, axes = plt.subplots(1, 2, figsize=(16, 10), constrained_layout=True)
+    axes[0].barh(positions, rmse, color="#457b9d")
+    axes[0].set_title("RMSE by variable")
+    axes[0].set_xlabel("RMSE")
+    axes[1].barh(positions, correlation, color="#2a9d8f")
+    axes[1].set_title("Correlation by variable")
+    axes[1].set_xlabel("Pearson correlation")
+    axes[1].set_xlim(-1, 1)
+    for axis in axes:
+        axis.set_yticks(positions, labels, fontsize=8)
+        axis.invert_yaxis()
+        axis.grid(axis="x", alpha=0.25)
+    figure.suptitle("Virtual-data forecast skill by variable", fontsize=15)
+    path = output_dir / "variable_metrics.png"
+    figure.savefig(path, dpi=dpi)
+    plt.close(figure)
+    return path
+
+
+def create_report(config: dict) -> list[Path]:
+    prediction_path = Path(config["inference"]["output_dir"]) / "prediction.npz"
+    metadata_path = Path(config["synthetic_data"]["output_dir"]) / "metadata.json"
+    if not prediction_path.exists():
+        raise FileNotFoundError(f"Prediction file not found: {prediction_path}. Run scripts/inference.py first.")
+    with np.load(prediction_path) as arrays:
+        prediction = arrays["prediction"]
+        target = arrays["target"]
+    if prediction.shape != target.shape or prediction.ndim != 4:
+        raise ValueError(f"Expected matching [sample, channel, latitude, longitude] arrays, got {prediction.shape} and {target.shape}")
+
+    variables = load_variables(metadata_path, prediction.shape[1])
+    channel_metrics, overall = compute_metrics(prediction, target, variables)
+    output_dir = Path(config["visualization"]["output_dir"])
+    metrics_dir = Path(config["paths"]["metrics"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    metrics_dir.mkdir(parents=True, exist_ok=True)
+
+    summary_path = metrics_dir / "result_summary.json"
+    summary_path.write_text(json.dumps({
+        "evaluation_scope": "virtual_data_only",
+        "prediction_shape": list(prediction.shape),
+        "overall": overall,
+        "channels": channel_metrics,
+        "note": "These metrics evaluate the synthetic task and are not paper reproduction metrics.",
+    }, indent=2) + "\n", encoding="utf-8")
+    csv_path = metrics_dir / "channel_metrics.csv"
+    with csv_path.open("w", newline="", encoding="utf-8") as file:
+        writer = csv.DictWriter(file, fieldnames=channel_metrics[0].keys())
+        writer.writeheader()
+        writer.writerows(channel_metrics)
+
+    dashboard = plot_diagnostics(prediction, target, variables, channel_metrics, config, output_dir)
+    metric_plot = plot_channel_metrics(channel_metrics, output_dir, int(config["visualization"]["dpi"]))
+    return [dashboard, metric_plot, summary_path, csv_path]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    args = parser.parse_args()
+    for path in create_report(load_config(args.config)):
+        print(f"result: {path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/show_features.py

@@ -0,0 +1,35 @@
+"""Print the supported and unavailable project features."""
+
+
+COMPLETED_FEATURES = [
+    "Compact global-grid forecaster construction and checkpoint loading",
+    "Single-GPU and single-node DDP training",
+    "Offline local checkpoint inference",
+    "Synthetic smoke-data generation",
+    "FV3GFS NetCDF schema validation",
+    "Separate data, weight, result, and work directories",
+    "Configurable model name, data paths, checkpoint paths, and output paths",
+    "AMD DCU-compatible model loading and smoke validation",
+]
+
+PENDING_FEATURES = [
+    "Full SFNO and DYffusion architecture integration",
+    "Official Hugging Face checkpoint compatibility",
+    "Full FV3GFS validation rollout, pending requester-pays data",
+    "Paper-scale 100-year training and 10-year evaluation",
+    "Paper-matched 25-member climate metrics",
+]
+
+
+def main() -> None:
+    print("Spherical-DYffusion-SFNO project features")
+    print("Completed:")
+    for index, feature in enumerate(COMPLETED_FEATURES, start=1):
+        print(f"  {index}. {feature}")
+    print("Pending external inputs or resources:")
+    for index, feature in enumerate(PENDING_FEATURES, start=1):
+        print(f"  {index}. {feature}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,19 @@
+"""Train the local model from scratch on virtual FV3GFS-contract data."""
+
+import argparse
+from local_pipeline import load_config, train
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument(
+        "--finetune",
+        help="Optional checkpoint path. Without this argument, training starts from scratch.",
+    )
+    args = parser.parse_args()
+    train(load_config(args.config), finetune=args.finetune)
+
+
+if __name__ == "__main__":
+    main()

scripts/validate_data.py

@@ -0,0 +1,57 @@
+"""Validate NetCDF variables, dimensions, and statistics for FV3GFS inputs."""
+
+import argparse
+from pathlib import Path
+
+import xarray as xr
+
+
+REQUIRED_VARIABLES = {
+    "PRESsfc",
+    "surface_temperature",
+    "DSWRFtoa",
+    "HGTsfc",
+    "ocean_fraction",
+    *{f"air_temperature_{i}" for i in range(8)},
+    *{f"specific_total_water_{i}" for i in range(8)},
+    *{f"eastward_wind_{i}" for i in range(8)},
+    *{f"northward_wind_{i}" for i in range(8)},
+}
+
+
+def validate_data(data_path: str, latitude: int, longitude: int) -> int:
+    files = sorted(Path(data_path).rglob("*.nc"))
+    if not files:
+        print(f"ERROR: no NetCDF files found under {data_path}")
+        return 1
+    failed = False
+    for path in files:
+        with xr.open_dataset(path) as dataset:
+            variables = set(dataset.data_vars)
+            missing = sorted(REQUIRED_VARIABLES - variables)
+            dimensions = {name: int(size) for name, size in dataset.sizes.items()}
+            print(f"{path}: dimensions={dimensions}, variables={len(variables)}")
+            if missing:
+                print(f"ERROR: missing required variables: {', '.join(missing)}")
+                failed = True
+            if latitude not in dimensions.values() or longitude not in dimensions.values():
+                print(f"ERROR: expected spatial dimensions containing {latitude} and {longitude}")
+                failed = True
+            time_sizes = [dimensions[name] for name in dimensions if "time" in name.lower()]
+            if not time_sizes or max(time_sizes) < 7:
+                print("ERROR: expected a time dimension with at least 7 frames")
+                failed = True
+    return int(failed)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data-dir", required=True)
+    parser.add_argument("--latitude", type=int, default=180)
+    parser.add_argument("--longitude", type=int, default=360)
+    args = parser.parse_args()
+    raise SystemExit(validate_data(args.data_dir, args.latitude, args.longitude))
+
+
+if __name__ == "__main__":
+    main()

weight/.gitkeep

@@ -0,0 +1 @@
+# Reserved for model weights.