Upload folder using huggingface_hub

README.md

@@ -0,0 +1,164 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Short-to-Medium-Range Weather Forecast
+- ERA5
+- FourCastNet
+- SFNO
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FourCastNet_v2</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+FourCastNet v2 is a global weather forecast model based on the Spherical Fourier Neural Operator (SFNO), proposed by NVIDIA and its collaborators.
+
+Paper: Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere
+
+https://arxiv.org/abs/2306.03838
+
+# Model Description
+
+The key architectural change from v1 is replacing the Adaptive Fourier Neural Operator (AFNO) with the Spherical Fourier Neural Operator (SFNO).
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global weather forecast training | Train an SFNO-style FourCastNet v2 model with 73-channel ERA5 HDF5 data. |
+| Local quick validation | Use synthetic ERA5 files to check the training, inference, and result-visualization pipeline. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch PyTorch DDP 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/FourCastNet_v2 --local-dir ./FourCastNet_v2
+cd FourCastNet_v2
+```
+
+### 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
+```
+
+### Training Data Introduction
+
+The OneScience community provides an ERA5 data slice that can be downloaded as follows:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+Real HDF5 annual files must contain `fields`, variable attributes, `time_step`, `global_means`, and `global_stds`. When real data is unavailable, first generate synthetic files for pipeline validation:
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+The default checkpoint is saved to `data/checkpoint/one_step/model_bak.pt`.
+
+### Fine-tuning
+
+Single GPU:
+
+```bash
+python scripts/train.py --stage finetune
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py --stage finetune
+```
+
+The checkpoint is saved to `data/checkpoint/<stage>/model_bak.pt` by default.
+
+### Training Weights
+
+This repository provides weights trained on ERA5 reanalysis data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are written to `result/output/` by default.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The default output includes latitude-weighted RMSE/ACC metrics and `result/figures/t2m_forecast.png`.
+
+# 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
+
+- The SFNO numerical implementation of FourCastNet v2 follows the design of NVIDIA Earth2MIP and related official implementations. The upstream code and model licenses and copyright notices must be retained.

conf/config.yaml

@@ -0,0 +1,188 @@
+project:
+  name: FourCastNet_v2
+  seed: 42
+  output_dir: ./result
+  checkpoint_dir: ./data/checkpoint
+
+data:
+  dataset_dir: ./data/era5_fake
+  train_years: [2014,2015]
+  val_years: [2016]
+  test_years: [2018]
+  official_splits:
+    train_years: [2014, 2015]
+    val_years: [2016, 2017]
+    test_years: [2018]
+  input_steps: 1
+  output_steps: 1
+  time_step_hours: 6
+  grid_shape: [721, 1440]
+  normalize: true
+  variables:
+    - u10m
+    - v10m
+    - u100m
+    - v100m
+    - t2m
+    - sp
+    - msl
+    - tcwv
+    - u50
+    - u100
+    - u150
+    - u200
+    - u250
+    - u300
+    - u400
+    - u500
+    - u600
+    - u700
+    - u850
+    - u925
+    - u1000
+    - v50
+    - v100
+    - v150
+    - v200
+    - v250
+    - v300
+    - v400
+    - v500
+    - v600
+    - v700
+    - v850
+    - v925
+    - v1000
+    - z50
+    - z100
+    - z150
+    - z200
+    - z250
+    - z300
+    - z400
+    - z500
+    - z600
+    - z700
+    - z850
+    - z925
+    - z1000
+    - t50
+    - t100
+    - t150
+    - t200
+    - t250
+    - t300
+    - t400
+    - t500
+    - t600
+    - t700
+    - t850
+    - t925
+    - t1000
+    - r50
+    - r100
+    - r150
+    - r200
+    - r250
+    - r300
+    - r400
+    - r500
+    - r600
+    - r700
+    - r850
+    - r925
+    - r1000
+
+fake_data:
+  time_steps_per_year: 20
+  chunk_time_steps: 1
+  fill_value: 0.0
+  materialize_pattern: true
+
+model:
+  profile: smoke
+  profiles:
+    full_resolution:
+      img_size: [721, 1440]
+      in_channels: 73
+      out_channels: 73
+      spectral_transform: sht
+      filter_type: non-linear
+      scale_factor: 6
+      embed_dim: 256
+      num_layers: 12
+      num_blocks: 8
+      normalization_layer: instance_norm
+      mlp_mode: distributed
+      spectral_layers: 3
+      complex_activation: real
+      hard_thresholding_fraction: 1.0
+      big_skip: true
+    smoke:
+      img_size: [16, 32]
+      in_channels: 73
+      out_channels: 73
+      spectral_transform: sht
+      filter_type: linear
+      scale_factor: 4
+      embed_dim: 8
+      num_layers: 2
+      num_blocks: 1
+      normalization_layer: instance_norm
+      mlp_mode: serial
+      spectral_layers: 1
+      complex_activation: real
+      hard_thresholding_fraction: 0.5
+      big_skip: true
+      # The smoke profile keeps the same equations but reduces the grid/model.
+
+checkpoint:
+  initialize_from: scratch
+  prefix: model_bak
+  finetune_from: ./data/checkpoint/one_step/model_bak.pt
+  strict: true
+
+training:
+  stage: one_step
+  epochs: 3
+  batch_size: 1
+  num_workers: 0
+  learning_rate: 0.0006
+  weight_decay: 0.0
+  optimizer_betas: [0.9, 0.95]
+  max_grad_norm: 32.0
+  scheduler: cosine
+  amp: false
+  max_train_batches: null
+  max_val_batches: null
+  finetune:
+    autoregressive_steps: 2
+    epochs: 3
+    learning_rate: 0.0001
+
+distributed:
+  backend: nccl
+  master_addr: 127.0.0.1
+  master_port: 29500
+
+inference:
+  checkpoint_path: ./data/checkpoint/finetune/model_bak.pt
+  rollout_steps: 1
+  max_samples: 1
+  save_normalized: false
+  output_dir: ./result/output
+
+visualization:
+  variable: t2m
+  sample_index: 0
+  output_dir: ./result/figures
+  cmap: coolwarm
+
+slurm:
+  job_name: fcnv2_train
+  nodes: 1
+  gpus_per_node: 8
+  cpus_per_task: 8
+  time: "24:00:00"
+  partition: null
+  conda_env: fourcastnetv2_develop

config.json

@@ -0,0 +1,47 @@
+{
+  "model_name": "FourCastNet v2",
+  "model_type": "fourcastnet_v2",
+  "architectures": [
+    "FourCastNetV2"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "global-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/fourcastnet_v2.py",
+    "scope": "adapter for the bundled NVIDIA FourCastNet v2 SFNO network"
+  },
+  "architecture": {
+    "family": "spherical Fourier neural operator",
+    "grid_shape": [
+      721,
+      1440
+    ],
+    "input_channels": 73,
+    "output_channels": 73,
+    "spectral_transform": "sht",
+    "filter_type": "non-linear",
+    "scale_factor": 6,
+    "embedding_size": 256,
+    "layers": 12,
+    "blocks_per_layer": 8,
+    "normalization": "instance_norm",
+    "mlp_mode": "distributed",
+    "spectral_layers": 3,
+    "complex_activation": "real",
+    "hard_thresholding_fraction": 1.0
+  },
+  "data": {
+    "dataset": "ERA5",
+    "spatial_resolution_degrees": 0.25,
+    "time_step_hours": 6,
+    "input_steps": 1,
+    "output_steps": 1,
+    "protocol": "synthetic_era5"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/fourcastnet_v2.py",
+    "model/fcnv2"
+  ]
+}

configuration.json

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

model/fcnv2/fcnv2_activations.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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.
+
+import torch
+from torch import nn
+
+
+class ComplexReLU(nn.Module):
+    def __init__(self, negative_slope=0.0, mode="cartesian", bias_shape=None):
+        super(ComplexReLU, self).__init__()
+
+        # store parameters
+        self.mode = mode
+        if self.mode in ["modulus", "halfplane"]:
+            if bias_shape is not None:
+                self.bias = nn.Parameter(torch.zeros(bias_shape, dtype=torch.float32))
+            else:
+                self.bias = nn.Parameter(torch.zeros((1), dtype=torch.float32))
+        else:
+            bias = torch.zeros((1), dtype=torch.float32)
+            self.register_buffer("bias", bias)
+
+        self.negative_slope = negative_slope
+        self.act = nn.LeakyReLU(negative_slope=negative_slope)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        if self.mode == "cartesian":
+            zr = torch.view_as_real(z)
+            za = self.act(zr)
+            out = torch.view_as_complex(za)
+        elif self.mode == "modulus":
+            zabs = torch.sqrt(torch.square(z.real) + torch.square(z.imag))
+            out = self.act(zabs + self.bias) * torch.exp(1.0j * z.angle())
+        elif self.mode == "halfplane":
+            # bias is an angle parameter in this case
+            modified_angle = torch.angle(z) - self.bias
+            condition = torch.logical_and(
+                (0.0 <= modified_angle), (modified_angle < torch.pi / 2.0)
+            )
+            out = torch.where(condition, z, self.negative_slope * z)
+        elif self.mode == "real":
+            zr = torch.view_as_real(z)
+            outr = torch.stack((self.act(zr[..., 0]), zr[..., 1]), dim=-1)
+            out = torch.view_as_complex(outr)
+        else:
+            # identity
+            out = z
+
+        return out
+
+
+class ComplexActivation(nn.Module):
+    def __init__(self, activation, mode="cartesian", bias_shape=None):
+        super(ComplexActivation, self).__init__()
+
+        # store parameters
+        self.mode = mode
+        if self.mode == "modulus":
+            if bias_shape is not None:
+                self.bias = nn.Parameter(torch.zeros(bias_shape, dtype=torch.float32))
+            else:
+                self.bias = nn.Parameter(torch.zeros((1), dtype=torch.float32))
+        else:
+            bias = torch.zeros((1), dtype=torch.float32)
+            self.register_buffer("bias", bias)
+
+        # real valued activation
+        self.act = activation
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        if self.mode == "cartesian":
+            zr = torch.view_as_real(z)
+            za = self.act(zr)
+            out = torch.view_as_complex(za)
+        elif self.mode == "modulus":
+            zabs = torch.sqrt(torch.square(z.real) + torch.square(z.imag))
+            out = self.act(zabs + self.bias) * torch.exp(1.0j * z.angle())
+        else:
+            # identity
+            out = z
+
+        return out

model/fcnv2/fcnv2_contractions.py

@@ -0,0 +1,181 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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.
+
+import torch
+
+# Helper routines for FNOs
+
+
+@torch.jit.script
+def compl_contract2d_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    tmp = torch.einsum("bixys,kixyr->srbkxy", a, b)
+    res = torch.stack(
+        [tmp[0, 0, ...] - tmp[1, 1, ...], tmp[1, 0, ...] + tmp[0, 1, ...]], dim=-1
+    )
+    return res
+
+
+@torch.jit.script
+def compl_contract2d_fwd_c(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    ac = torch.view_as_complex(a)
+    bc = torch.view_as_complex(b)
+    res = torch.einsum("bixy,kixy->bkxy", ac, bc)
+    return torch.view_as_real(res)
+
+
+@torch.jit.script
+def compl_contract_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    tmp = torch.einsum("bins,kinr->srbkn", a, b)
+    res = torch.stack(
+        [tmp[0, 0, ...] - tmp[1, 1, ...], tmp[1, 0, ...] + tmp[0, 1, ...]], dim=-1
+    )
+    return res
+
+
+@torch.jit.script
+def compl_contract_fwd_c(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    ac = torch.view_as_complex(a)
+    bc = torch.view_as_complex(b)
+    res = torch.einsum("bin,kin->bkn", ac, bc)
+    return torch.view_as_real(res)
+
+
+@torch.jit.script
+def compl_ttc1_c_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    ac = torch.view_as_complex(a)
+    bc = torch.view_as_complex(b)
+    res = torch.einsum("jt,bct->jbct", ac, bc)
+    return torch.view_as_real(res)
+
+
+@torch.jit.script
+def compl_ttc2_c_fwd(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+    ac = torch.view_as_complex(a)
+    bc = torch.view_as_complex(b)
+    cc = torch.view_as_complex(c)
+    res = torch.einsum("oi,icj,jbct->bot", ac, bc, cc)
+    return torch.view_as_real(res)
+
+
+def contract_tt(x, w):
+    y = compl_ttc1_c_fwd(w[2], x)
+    return compl_ttc2_c_fwd(w[0], w[1], y)
+
+
+# Helper routines for spherical MLPs
+@torch.jit.script
+def compl_mul1d_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    tmp = torch.einsum("bixs,ior->srbox", a, b)
+    res = torch.stack(
+        [tmp[0, 0, ...] - tmp[1, 1, ...], tmp[1, 0, ...] + tmp[0, 1, ...]], dim=-1
+    )
+    return res
+
+
+@torch.jit.script
+def compl_mul1d_fwd_c(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    ac = torch.view_as_complex(a)
+    bc = torch.view_as_complex(b)
+    resc = torch.einsum("bix,io->box", ac, bc)
+    res = torch.view_as_real(resc)
+    return res
+
+
+@torch.jit.script
+def compl_muladd1d_fwd(
+    a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+) -> torch.Tensor:
+    res = compl_mul1d_fwd(a, b) + c
+    return res
+
+
+@torch.jit.script
+def compl_muladd1d_fwd_c(
+    a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+) -> torch.Tensor:
+    tmpcc = torch.view_as_complex(compl_mul1d_fwd_c(a, b))
+    cc = torch.view_as_complex(c)
+    return torch.view_as_real(tmpcc + cc)
+
+
+# for the real-valued case:
+
+
+@torch.jit.script
+def compl_mul1d_fwd_r(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    res = torch.einsum("bix,io->box", a, b)
+    return res
+
+
+@torch.jit.script
+def compl_muladd1d_fwd_r(
+    a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+) -> torch.Tensor:
+    tmp = compl_mul1d_fwd_r(a, b)
+    return tmp + c
+
+
+# Helper routines for FFT MLPs
+
+
+@torch.jit.script
+def compl_mul2d_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    tmp = torch.einsum("bixys,ior->srboxy", a, b)
+    res = torch.stack(
+        [tmp[0, 0, ...] - tmp[1, 1, ...], tmp[1, 0, ...] + tmp[0, 1, ...]], dim=-1
+    )
+    return res
+
+
+@torch.jit.script
+def compl_mul2d_fwd_c(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    ac = torch.view_as_complex(a)
+    bc = torch.view_as_complex(b)
+    resc = torch.einsum("bixy,io->boxy", ac, bc)
+    res = torch.view_as_real(resc)
+    return res
+
+
+@torch.jit.script
+def compl_muladd2d_fwd(
+    a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+) -> torch.Tensor:
+    res = compl_mul2d_fwd(a, b) + c
+    return res
+
+
+@torch.jit.script
+def compl_muladd2d_fwd_c(
+    a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+) -> torch.Tensor:
+    tmpcc = torch.view_as_complex(compl_mul2d_fwd_c(a, b))
+    cc = torch.view_as_complex(c)
+    return torch.view_as_real(tmpcc + cc)
+
+
+# for the real-valued case:
+@torch.jit.script
+def compl_mul2d_fwd_r(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    res = torch.einsum("bixy,io->boxy", a, b)
+    return res
+
+
+@torch.jit.script
+def compl_muladd2d_fwd_r(
+    a: torch.Tensor, b: torch.Tensor, c: torch.Tensor
+) -> torch.Tensor:
+    tmp = compl_mul2d_fwd_c(a, b)
+    return torch.view_as_real(tmp + c)

model/fcnv2/fcnv2_layers.py

@@ -0,0 +1,662 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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.
+
+import math
+import warnings
+
+import torch
+import torch.fft
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda import amp
+from torch.utils.checkpoint import checkpoint
+from torch_harmonics import *  # noqa
+
+from fcnv2_activations import ComplexReLU  # noqa
+from fcnv2_contractions import (
+    compl_contract2d_fwd,
+    compl_contract2d_fwd_c,
+    compl_contract_fwd,
+    compl_contract_fwd_c,
+    compl_mul2d_fwd,
+    compl_mul2d_fwd_c,
+    compl_muladd2d_fwd,
+    compl_muladd2d_fwd_c,
+    contract_tt,
+)
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)  # noqa
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+    tensor: an n-dimensional `torch.Tensor`
+    mean: the mean of the normal distribution
+    std: the standard deviation of the normal distribution
+    a: the minimum cutoff value
+    b: the maximum cutoff value
+    Examples:
+    >>> w = torch.empty(3, 5)
+    >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+@torch.jit.script
+def drop_path(
+    x: torch.Tensor, drop_prob: float = 0.0, training: bool = False
+) -> torch.Tensor:
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1.0 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2d ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(
+        self, img_size=(224, 224), patch_size=(16, 16), in_chans=3, embed_dim=768
+    ):
+        super(PatchEmbed, self).__init__()
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x):
+        # gather input
+        B, C, H, W = x.shape
+        assert (  # noqa
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        # new: B, C, H*W
+        x = self.proj(x).flatten(2)
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        output_bias=True,
+        drop_rate=0.0,
+        checkpointing=False,
+    ):
+        super(MLP, self).__init__()
+        self.checkpointing = checkpointing
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        fc1 = nn.Conv2d(in_features, hidden_features, 1, bias=True)
+        act = act_layer()
+        fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=output_bias)
+        if drop_rate > 0.0:
+            drop = nn.Dropout(drop_rate)
+            self.fwd = nn.Sequential(fc1, act, drop, fc2, drop)
+        else:
+            self.fwd = nn.Sequential(fc1, act, fc2)
+
+    @torch.jit.ignore
+    def checkpoint_forward(self, x):
+        return checkpoint(self.fwd, x)
+
+    def forward(self, x):
+        if self.checkpointing:
+            return self.checkpoint_forward(x)
+        else:
+            return self.fwd(x)
+
+
+class RealFFT2(nn.Module):
+    """
+    Helper routine to wrap FFT similarly to the SHT
+    """
+
+    def __init__(self, nlat, nlon, lmax=None, mmax=None):
+        super(RealFFT2, self).__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.lmax = lmax or self.nlat
+        self.mmax = mmax or self.nlon // 2 + 1
+        self.num_batches = 1
+
+        assert self.lmax % 2 == 0  # noqa
+
+    def forward(self, x):
+        # do batched FFT
+        xs = torch.split(x, x.shape[1] // self.num_batches, dim=1)
+
+        ys = []
+        for xt in xs:
+            yt = torch.fft.rfft2(xt, dim=(-2, -1), norm="ortho")
+            ys.append(
+                torch.cat(
+                    (
+                        yt[..., : math.ceil(self.lmax / 2), : self.mmax],
+                        yt[..., -math.floor(self.lmax / 2) :, : self.mmax],
+                    ),
+                    dim=-2,
+                )
+            )
+
+        # connect
+        y = torch.cat(ys, dim=1).contiguous()
+
+        # y = torch.fft.rfft2(x, dim=(-2, -1), norm="ortho")
+        # y = torch.cat((y[..., :math.ceil(self.lmax/2), :self.mmax], y[..., -math.floor(self.lmax/2):, :self.mmax]), dim=-2)
+        return y
+
+
+class InverseRealFFT2(nn.Module):
+    """
+    Helper routine to wrap FFT similarly to the SHT
+    """
+
+    def __init__(self, nlat, nlon, lmax=None, mmax=None):
+        super(InverseRealFFT2, self).__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.lmax = lmax or self.nlat
+        self.mmax = mmax or self.nlon // 2 + 1
+        self.num_batches = 1
+
+    def forward(self, x):
+        # do batched FFT
+        xs = torch.split(x, x.shape[1] // self.num_batches, dim=1)
+
+        ys = []
+        for xt in xs:
+            ys.append(
+                torch.fft.irfft2(
+                    xt, dim=(-2, -1), s=(self.nlat, self.nlon), norm="ortho"
+                )
+            )
+        out = torch.cat(ys, dim=1).contiguous()
+
+        # out = torch.fft.irfft2(x, dim=(-2, -1), s=(self.nlat, self.nlon), norm="ortho")
+        return out
+
+
+class SpectralConv2d(nn.Module):
+    """
+    Spectral Convolution as utilized in
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        hidden_size,
+        sparsity_threshold=0.0,
+        hard_thresholding_fraction=1,
+        use_complex_kernels=False,
+        compression=None,
+        rank=0,
+        bias=False,
+    ):
+        super(SpectralConv2d, self).__init__()
+
+        self.hidden_size = hidden_size
+        self.sparsity_threshold = sparsity_threshold
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.scale = 1 / hidden_size**2
+        self.contract_handle = (
+            compl_contract2d_fwd_c if use_complex_kernels else compl_contract2d_fwd
+        )
+
+        self.forward_transform = forward_transform
+        self.inverse_transform = inverse_transform
+
+        self.output_dims = (self.inverse_transform.nlat, self.inverse_transform.nlon)
+        modes_lat = self.inverse_transform.lmax
+        modes_lon = self.inverse_transform.mmax
+        self.modes_lat = int(modes_lat * self.hard_thresholding_fraction)
+        self.modes_lon = int(modes_lon * self.hard_thresholding_fraction)
+
+        # new simple linear layer
+        self.w = nn.Parameter(
+            self.scale
+            * torch.randn(
+                self.hidden_size, self.hidden_size, self.modes_lat, self.modes_lon, 2
+            )
+        )
+        # optional bias
+        if bias:
+            self.b = nn.Parameter(
+                self.scale * torch.randn(1, self.hidden_size, *self.output_dims)
+            )
+
+    def forward(self, x):
+        dtype = x.dtype
+        # x = x.float()
+        B, C, H, W = x.shape
+
+        with amp.autocast(enabled=False):
+            x = x.to(torch.float32)
+            x = self.forward_transform(x)
+            x = torch.view_as_real(x)
+            x = x.to(dtype)
+
+        # do spectral conv
+        modes = torch.zeros(x.shape, device=x.device)
+
+        # modes[:, :, :self.modes_lat,  :self.modes_lon, :] = self.contract_handle(x[:, :, :self.modes_lat,  :self.modes_lon, :], self.wh)
+        # modes[:, :, -self.modes_lat:, :self.modes_lon, :] = self.contract_handle(x[:, :, -self.modes_lat:, :self.modes_lon, :], self.wl)
+        modes = self.contract_handle(x, self.w)
+
+        # finalize
+        x = F.softshrink(modes, lambd=self.sparsity_threshold)
+        x = torch.view_as_complex(x)
+
+        with amp.autocast(enabled=False):
+            x = x.to(torch.float32)
+            x = torch.view_as_complex(x)
+            x = self.inverse_transform(x)
+            x = x.to(dtype)
+
+        if hasattr(self, "b"):
+            x = x + self.b
+
+        return x
+
+
+class SpectralConvS2(nn.Module):
+    """
+    Spectral Convolution as utilized in
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        hidden_size,
+        sparsity_threshold=0.0,
+        use_complex_kernels=False,
+        compression=None,
+        rank=128,
+        bias=False,
+    ):
+        super(SpectralConvS2, self).__init__()
+
+        self.hidden_size = hidden_size
+        self.sparsity_threshold = sparsity_threshold
+        self.scale = 0.02
+
+        self.forward_transform = forward_transform
+        self.inverse_transform = inverse_transform
+
+        self.modes_lat = self.forward_transform.lmax
+        self.modes_lon = self.forward_transform.mmax
+
+        assert self.inverse_transform.lmax == self.modes_lat  # noqa
+        assert self.inverse_transform.mmax == self.modes_lon  # noqa
+
+        # remember the lower triangular indices
+        ii, jj = torch.tril_indices(self.modes_lat, self.modes_lon)
+        self.register_buffer("ii", ii)
+        self.register_buffer("jj", jj)
+
+        if compression == "tt":
+            self.rank = rank
+            # tensortrain coefficients
+            g1 = nn.Parameter(self.scale * torch.randn(self.hidden_size, self.rank, 2))
+            g2 = nn.Parameter(
+                self.scale * torch.randn(self.rank, self.hidden_size, self.rank, 2)
+            )
+            g3 = nn.Parameter(self.scale * torch.randn(self.rank, len(ii), 2))
+            self.w = nn.ParameterList([g1, g2, g3])
+
+            self.contract_handle = (
+                contract_tt  # if use_complex_kernels else raise(NotImplementedError)
+            )
+        else:
+            self.w = nn.Parameter(
+                self.scale * torch.randn(self.hidden_size, self.hidden_size, len(ii), 2)
+            )
+            self.contract_handle = (
+                compl_contract_fwd_c if use_complex_kernels else compl_contract_fwd
+            )
+
+        if bias:
+            self.b = nn.Parameter(
+                self.scale * torch.randn(1, self.hidden_size, *self.output_dims)
+            )
+
+    def forward(self, x):
+        dtype = x.dtype
+        # x = x.float()
+        B, C, H, W = x.shape
+
+        with amp.autocast(enabled=False):
+            x = x.to(torch.float32)
+            x = self.forward_transform(x)
+            x = torch.view_as_real(x)
+            x = x.to(dtype)
+
+        # Populate the sparse spectral grid without an in-place write. The
+        # latter breaks autograd under multi-process DDP on some HIP builds.
+        spectral_height, spectral_width = x.shape[2:4]
+        contracted = self.contract_handle(
+            x[:, :, self.ii, self.jj, :], self.w
+        )
+        spectral_indices = self.ii * spectral_width + self.jj
+        modes = torch.zeros_like(x).reshape(
+            B, C, spectral_height * spectral_width, 2
+        ).index_copy(
+            2, spectral_indices, contracted
+        )
+        modes = modes.view_as(x)
+
+        # finalize
+        x = F.softshrink(modes, lambd=self.sparsity_threshold)
+
+        with amp.autocast(enabled=False):
+            x = x.to(torch.float32)
+            x = torch.view_as_complex(x)
+            x = self.inverse_transform(x)
+            x = x.to(dtype)
+
+        if hasattr(self, "b"):
+            x = x + self.b
+
+        return x
+
+
+class SpectralAttention2d(nn.Module):
+    """
+    2d Spectral Attention layer
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        embed_dim,
+        sparsity_threshold=0.0,
+        hidden_size_factor=2,
+        use_complex_network=True,
+        use_complex_kernels=False,
+        complex_activation="real",
+        bias=False,
+        spectral_layers=1,
+        drop_rate=0.0,
+    ):
+        super(SpectralAttention2d, self).__init__()
+
+        self.embed_dim = embed_dim
+        self.sparsity_threshold = sparsity_threshold
+        self.hidden_size = int(hidden_size_factor * self.embed_dim)
+        self.scale = 0.02
+        self.spectral_layers = spectral_layers
+        self.mul_add_handle = (
+            compl_muladd2d_fwd_c if use_complex_kernels else compl_muladd2d_fwd
+        )
+        self.mul_handle = compl_mul2d_fwd_c if use_complex_kernels else compl_mul2d_fwd
+
+        self.modes_lat = forward_transform.lmax
+        self.modes_lon = forward_transform.mmax
+
+        # only storing the forward handle to be able to call it
+        self.forward_transform = forward_transform.forward
+        self.inverse_transform = inverse_transform.forward
+
+        assert inverse_transform.lmax == self.modes_lat  # noqa
+        assert inverse_transform.mmax == self.modes_lon  # noqa
+
+        # weights
+        w = [self.scale * torch.randn(self.embed_dim, self.hidden_size, 2)]
+        # w = [self.scale * torch.randn(self.embed_dim + 2*self.embed_freqs, self.hidden_size, 2)]
+        # w = [self.scale * torch.randn(self.embed_dim + 4*self.embed_freqs, self.hidden_size, 2)]
+        for l in range(1, self.spectral_layers):
+            w.append(self.scale * torch.randn(self.hidden_size, self.hidden_size, 2))
+        self.w = nn.ParameterList(w)
+
+        if bias:
+            self.b = nn.ParameterList(
+                [
+                    self.scale * torch.randn(self.hidden_size, 1, 2)
+                    for _ in range(self.spectral_layers)
+                ]
+            )
+
+        self.wout = nn.Parameter(
+            self.scale * torch.randn(self.hidden_size, self.embed_dim, 2)
+        )
+
+        self.drop = nn.Dropout(drop_rate) if drop_rate > 0.0 else nn.Identity()
+
+        self.activation = ComplexReLU(
+            mode=complex_activation, bias_shape=(self.hidden_size, 1, 1)
+        )
+
+    def forward_mlp(self, xr):
+        for l in range(self.spectral_layers):
+            if hasattr(self, "b"):
+                xr = self.mul_add_handle(
+                    xr, self.w[l].to(xr.dtype), self.b[l].to(xr.dtype)
+                )
+            else:
+                xr = self.mul_handle(xr, self.w[l].to(xr.dtype))
+            xr = torch.view_as_complex(xr)
+            xr = self.activation(xr)
+            xr = self.drop(xr)
+            xr = torch.view_as_real(xr)
+
+        xr = self.mul_handle(xr, self.wout)
+
+        return xr
+
+    def forward(self, x):
+        dtype = x.dtype
+        # x = x.to(torch.float32)
+
+        # FWD transform
+        with amp.autocast(enabled=False):
+            x = x.to(torch.float32)
+            x = self.forward_transform(x)
+            x = torch.view_as_real(x)
+
+        # MLP
+        x = self.forward_mlp(x)
+
+        # BWD transform
+        with amp.autocast(enabled=False):
+            x = torch.view_as_complex(x)
+            x = self.inverse_transform(x)
+            x = x.to(dtype)
+
+        return x
+
+
+class SpectralAttentionS2(nn.Module):
+    """
+    geometrical Spectral Attention layer
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        embed_dim,
+        sparsity_threshold=0.0,
+        hidden_size_factor=2,
+        use_complex_network=True,
+        use_complex_kernels=False,
+        complex_activation="real",
+        bias=False,
+        spectral_layers=1,
+        drop_rate=0.0,
+    ):
+        super(SpectralAttentionS2, self).__init__()
+
+        self.embed_dim = embed_dim
+        self.sparsity_threshold = sparsity_threshold
+        self.hidden_size = int(hidden_size_factor * self.embed_dim)
+        self.scale = 0.02
+        # self.mul_add_handle = compl_muladd1d_fwd_c if use_complex_kernels else compl_muladd1d_fwd
+        self.mul_add_handle = (
+            compl_muladd2d_fwd_c if use_complex_kernels else compl_muladd2d_fwd
+        )
+        # self.mul_handle = compl_mul1d_fwd_c if use_complex_kernels else compl_mul1d_fwd
+        self.mul_handle = compl_mul2d_fwd_c if use_complex_kernels else compl_mul2d_fwd
+        self.spectral_layers = spectral_layers
+
+        self.modes_lat = forward_transform.lmax
+        self.modes_lon = forward_transform.mmax
+
+        # only storing the forward handle to be able to call it
+        self.forward_transform = forward_transform.forward
+        self.inverse_transform = inverse_transform.forward
+
+        assert inverse_transform.lmax == self.modes_lat  # noqa
+        assert inverse_transform.mmax == self.modes_lon  # noqa
+
+        # weights
+        w = [self.scale * torch.randn(self.embed_dim, self.hidden_size, 2)]
+        # w = [self.scale * torch.randn(self.embed_dim + 4*self.embed_freqs, self.hidden_size, 2)]
+        for l in range(1, self.spectral_layers):
+            w.append(self.scale * torch.randn(self.hidden_size, self.hidden_size, 2))
+        self.w = nn.ParameterList(w)
+
+        if bias:
+            self.b = nn.ParameterList(
+                [
+                    self.scale * torch.randn(2 * self.hidden_size, 1, 1, 2)
+                    for _ in range(self.spectral_layers)
+                ]
+            )
+
+        self.wout = nn.Parameter(
+            self.scale * torch.randn(self.hidden_size, self.embed_dim, 2)
+        )
+
+        self.drop = nn.Dropout(drop_rate) if drop_rate > 0.0 else nn.Identity()
+
+        self.activation = ComplexReLU(
+            mode=complex_activation, bias_shape=(self.hidden_size, 1, 1)
+        )
+
+    def forward_mlp(self, xr):
+        for l in range(self.spectral_layers):
+            if hasattr(self, "b"):
+                xr = self.mul_add_handle(
+                    xr, self.w[l].to(xr.dtype), self.b[l].to(xr.dtype)
+                )
+            else:
+                xr = self.mul_handle(xr, self.w[l].to(xr.dtype))
+            xr = torch.view_as_complex(xr)
+            xr = self.activation(xr)
+            xr = self.drop(xr)
+            xr = torch.view_as_real(xr)
+
+        # final MLP
+        xr = self.mul_handle(xr, self.wout)
+
+        return xr
+
+    def forward(self, x):
+        dtype = x.dtype
+        # x = x.to(torch.float32)
+
+        # FWD transform
+        with amp.autocast(enabled=False):
+            x = x.to(torch.float32)
+            x = self.forward_transform(x)
+            x = torch.view_as_real(x)
+
+        # MLP
+        x = self.forward_mlp(x)
+
+        # BWD transform
+        with amp.autocast(enabled=False):
+            x = torch.view_as_complex(x)
+            x = self.inverse_transform(x)
+            x = x.to(dtype)
+
+        return x

model/fcnv2/fcnv2_sfnonet.py

@@ -0,0 +1,615 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 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.
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch_harmonics as harmonics
+from apex.normalization import FusedLayerNorm
+
+# helpers
+# to fake the sht module with ffts
+from fcnv2_layers import (
+    MLP,
+    DropPath,
+    InverseRealFFT2,
+    RealFFT2,
+    SpectralAttention2d,
+    SpectralAttentionS2,
+    SpectralConv2d,
+    SpectralConvS2,
+    trunc_normal_,
+)
+
+
+class SafeRealSHT(harmonics.RealSHT):
+    """RealSHT variant that avoids in-place writes during autograd."""
+
+    def forward(self, x):
+        if x.dim() < 2:
+            raise ValueError("Expected tensor with at least 2 dimensions")
+        if x.shape[-2:] != (self.nlat, self.nlon):
+            raise ValueError(
+                f"Expected spatial shape {(self.nlat, self.nlon)}, got {tuple(x.shape[-2:])}"
+            )
+
+        transformed = torch.view_as_real(
+            2.0 * torch.pi * torch.fft.rfft(x, dim=-1, norm="forward")
+        )
+        weights = self.weights.to(dtype=transformed.dtype)
+        real = torch.einsum(
+            "...km,mlk->...lm", transformed[..., : self.mmax, 0], weights
+        )
+        imag = torch.einsum(
+            "...km,mlk->...lm", transformed[..., : self.mmax, 1], weights
+        )
+        return torch.view_as_complex(torch.stack((real, imag), dim=-1).contiguous())
+
+
+class SafeInverseRealSHT(harmonics.InverseRealSHT):
+    """InverseRealSHT variant that avoids in-place writes during autograd."""
+
+    def forward(self, x):
+        if x.dim() < 2:
+            raise ValueError("Expected tensor with at least 2 dimensions")
+        if x.shape[-2:] != (self.lmax, self.mmax):
+            raise ValueError(
+                f"Expected spectral shape {(self.lmax, self.mmax)}, got {tuple(x.shape[-2:])}"
+            )
+
+        spectral = torch.view_as_real(x)
+        spatial = torch.einsum(
+            "...lmr,mlk->...kmr", spectral, self.pct.to(dtype=spectral.dtype)
+        )
+        spatial = torch.view_as_complex(spatial.contiguous())
+
+        # torch_harmonics clears these components in-place. Build the same
+        # result functionally so autograd can retain the spectral tensor.
+        longitude = torch.arange(
+            self.mmax, device=spatial.device, dtype=spatial.real.dtype
+        )
+        constrained = (longitude != 0).to(spatial.real.dtype)
+        if self.nlon % 2 == 0 and self.nlon // 2 < self.mmax:
+            constrained = constrained * (
+                longitude != self.nlon // 2
+            ).to(spatial.real.dtype)
+        spatial = torch.complex(spatial.real, spatial.imag * constrained)
+        return torch.fft.irfft(spatial, n=self.nlon, dim=-1, norm="forward")
+
+
+class SpectralFilterLayer(nn.Module):
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        embed_dim,
+        filter_type="linear",
+        sparsity_threshold=0.0,
+        use_complex_kernels=True,
+        hidden_size_factor=2,
+        compression=None,
+        rank=128,
+        complex_network=True,
+        complex_activation="real",
+        spectral_layers=1,
+        drop_rate=0.0,
+    ):
+        super(SpectralFilterLayer, self).__init__()
+
+        if filter_type == "non-linear" and isinstance(
+            forward_transform, harmonics.RealSHT
+        ):
+            self.filter = SpectralAttentionS2(
+                forward_transform,
+                inverse_transform,
+                embed_dim,
+                sparsity_threshold,
+                use_complex_network=complex_network,
+                use_complex_kernels=use_complex_kernels,
+                hidden_size_factor=hidden_size_factor,
+                complex_activation=complex_activation,
+                spectral_layers=spectral_layers,
+                drop_rate=drop_rate,
+                bias=False,
+            )
+        elif filter_type == "non-linear" and isinstance(
+            forward_transform, harmonics.RealFFT2
+        ):
+            self.filter = SpectralAttention2d(
+                forward_transform,
+                inverse_transform,
+                embed_dim,
+                sparsity_threshold,
+                use_complex_kernels=use_complex_kernels,
+                hidden_size_factor=hidden_size_factor,
+                complex_activation=complex_activation,
+                spectral_layers=spectral_layers,
+                drop_rate=drop_rate,
+                bias=False,
+            )
+
+        elif filter_type == "linear" and isinstance(
+            forward_transform, harmonics.RealSHT
+        ):
+            self.filter = SpectralConvS2(
+                forward_transform,
+                inverse_transform,
+                embed_dim,
+                sparsity_threshold,
+                use_complex_kernels=use_complex_kernels,
+                compression=compression,
+                rank=rank,
+                bias=False,
+            )
+
+        elif filter_type == "linear" and isinstance(forward_transform, RealFFT2):
+            self.filter = SpectralConv2d(
+                forward_transform,
+                inverse_transform,
+                embed_dim,
+                sparsity_threshold,
+                use_complex_kernels=use_complex_kernels,
+                compression=compression,
+                rank=rank,
+                bias=False,
+            )
+
+        else:
+            raise (NotImplementedError)
+
+    def forward(self, x):
+        return self.filter(x)
+
+
+class FourierNeuralOperatorBlock(nn.Module):
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        embed_dim,
+        filter_type="linear",
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=(nn.LayerNorm, nn.LayerNorm),
+        # num_blocks = 8,
+        sparsity_threshold=0.0,
+        use_complex_kernels=True,
+        compression=None,
+        rank=128,
+        inner_skip="linear",
+        outer_skip=None,  # None, nn.linear or nn.Identity
+        concat_skip=False,
+        mlp_mode="none",
+        complex_network=True,
+        complex_activation="real",
+        spectral_layers=1,
+        checkpointing=False,
+    ):
+        super(FourierNeuralOperatorBlock, self).__init__()
+
+        # norm layer
+        self.norm0 = norm_layer[0]()  # ((h,w))
+
+        # convolution layer
+        self.filter_layer = SpectralFilterLayer(
+            forward_transform,
+            inverse_transform,
+            embed_dim,
+            filter_type,
+            sparsity_threshold,
+            use_complex_kernels=use_complex_kernels,
+            hidden_size_factor=mlp_ratio,
+            compression=compression,
+            rank=rank,
+            complex_network=complex_network,
+            complex_activation=complex_activation,
+            spectral_layers=spectral_layers,
+            drop_rate=drop_rate,
+        )
+
+        if inner_skip == "linear":
+            self.inner_skip = nn.Conv2d(embed_dim, embed_dim, 1, 1)
+        elif inner_skip == "identity":
+            self.inner_skip = nn.Identity()
+
+        self.concat_skip = concat_skip
+
+        if concat_skip and inner_skip is not None:
+            self.inner_skip_conv = nn.Conv2d(2 * embed_dim, embed_dim, 1, bias=False)
+
+        if filter_type == "linear":
+            self.act_layer = act_layer()
+
+        # dropout
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        # norm layer
+        self.norm1 = norm_layer[1]()  # ((h,w))
+
+        if mlp_mode != "none":
+            mlp_hidden_dim = int(embed_dim * mlp_ratio)
+            self.mlp = MLP(
+                in_features=embed_dim,
+                hidden_features=mlp_hidden_dim,
+                act_layer=act_layer,
+                drop_rate=drop_rate,
+                checkpointing=checkpointing,
+            )
+
+        if outer_skip == "linear":
+            self.outer_skip = nn.Conv2d(embed_dim, embed_dim, 1, 1)
+        elif outer_skip == "identity":
+            self.outer_skip = nn.Identity()
+
+        if concat_skip and outer_skip is not None:
+            self.outer_skip_conv = nn.Conv2d(2 * embed_dim, embed_dim, 1, bias=False)
+
+    def forward(self, x):
+        residual = x
+
+        x = self.norm0(x)
+        x = self.filter_layer(x).contiguous()
+
+        if hasattr(self, "inner_skip"):
+            if self.concat_skip:
+                x = torch.cat((x, self.inner_skip(residual)), dim=1)
+                x = self.inner_skip_conv(x)
+            else:
+                x = x + self.inner_skip(residual)
+
+        if hasattr(self, "act_layer"):
+            x = self.act_layer(x)
+
+        x = self.norm1(x)
+
+        if hasattr(self, "mlp"):
+            x = self.mlp(x)
+
+        x = self.drop_path(x)
+
+        if hasattr(self, "outer_skip"):
+            if self.concat_skip:
+                x = torch.cat((x, self.outer_skip(residual)), dim=1)
+                x = self.outer_skip_conv(x)
+            else:
+                x = x + self.outer_skip(residual)
+
+        return x
+
+    # @torch.jit.ignore
+    # def checkpoint_forward(self, x):
+    #     return checkpoint(self._forward, x)
+
+    # def forward(self, x):
+    #     if self.checkpointing:
+    #         return self.checkpoint_forward(x)
+    #     else:
+    #         return self._forward(x)
+
+
+class FourierNeuralOperatorNet(nn.Module):
+    def __init__(
+        self,
+        params,
+        spectral_transform="sht",
+        filter_type="non-linear",
+        img_size=(721, 1440),
+        scale_factor=16,
+        in_chans=2,
+        out_chans=2,
+        embed_dim=256,
+        num_layers=12,
+        mlp_mode="none",
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path_rate=0.0,
+        num_blocks=16,
+        sparsity_threshold=0.0,
+        normalization_layer="instance_norm",
+        hard_thresholding_fraction=1.0,
+        use_complex_kernels=True,
+        big_skip=True,
+        compression=None,
+        rank=128,
+        complex_network=True,
+        complex_activation="real",
+        spectral_layers=3,
+        laplace_weighting=False,
+        checkpointing=False,
+    ):
+        super(FourierNeuralOperatorNet, self).__init__()
+
+        self.params = params
+        self.spectral_transform = (
+            params.spectral_transform
+            if hasattr(params, "spectral_transform")
+            else spectral_transform
+        )
+        self.filter_type = (
+            params.filter_type if hasattr(params, "filter_type") else filter_type
+        )
+        self.img_size = (params.img_crop_shape_x, params.img_crop_shape_y)
+        self.scale_factor = (
+            params.scale_factor if hasattr(params, "scale_factor") else scale_factor
+        )
+        self.in_chans = (
+            params.N_in_channels if hasattr(params, "N_in_channels") else in_chans
+        )
+        self.out_chans = (
+            params.N_out_channels if hasattr(params, "N_out_channels") else out_chans
+        )
+        self.embed_dim = self.num_features = (
+            params.embed_dim if hasattr(params, "embed_dim") else embed_dim
+        )
+        self.num_layers = (
+            params.num_layers if hasattr(params, "num_layers") else num_layers
+        )
+        self.num_blocks = (
+            params.num_blocks if hasattr(params, "num_blocks") else num_blocks
+        )
+        self.hard_thresholding_fraction = (
+            params.hard_thresholding_fraction
+            if hasattr(params, "hard_thresholding_fraction")
+            else hard_thresholding_fraction
+        )
+        self.normalization_layer = (
+            params.normalization_layer
+            if hasattr(params, "normalization_layer")
+            else normalization_layer
+        )
+        self.mlp_mode = params.mlp_mode if hasattr(params, "mlp_mode") else mlp_mode
+        self.big_skip = params.big_skip if hasattr(params, "big_skip") else big_skip
+        self.compression = (
+            params.compression if hasattr(params, "compression") else compression
+        )
+        self.rank = params.rank if hasattr(params, "rank") else rank
+        self.complex_network = (
+            params.complex_network
+            if hasattr(params, "complex_network")
+            else complex_network
+        )
+        self.complex_activation = (
+            params.complex_activation
+            if hasattr(params, "complex_activation")
+            else complex_activation
+        )
+        self.spectral_layers = (
+            params.spectral_layers
+            if hasattr(params, "spectral_layers")
+            else spectral_layers
+        )
+        self.laplace_weighting = (
+            params.laplace_weighting
+            if hasattr(params, "laplace_weighting")
+            else laplace_weighting
+        )
+        self.checkpointing = (
+            params.checkpointing if hasattr(params, "checkpointing") else checkpointing
+        )
+
+        # compute downsampled image size
+        self.h = self.img_size[0] // self.scale_factor
+        self.w = self.img_size[1] // self.scale_factor
+
+        # dropout
+        self.pos_drop = nn.Dropout(p=drop_rate) if drop_rate > 0.0 else nn.Identity()
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.num_layers)]
+
+        # pick norm layer
+        if self.normalization_layer == "layer_norm":
+            norm_layer0 = partial(
+                nn.LayerNorm,
+                normalized_shape=(self.img_size[0], self.img_size[1]),
+                eps=1e-6,
+            )
+            norm_layer1 = partial(
+                nn.LayerNorm, normalized_shape=(self.h, self.w), eps=1e-6
+            )
+        elif self.normalization_layer == "instance_norm":
+            norm_layer0 = partial(
+                nn.InstanceNorm2d,
+                num_features=self.embed_dim,
+                eps=1e-6,
+                affine=True,
+                track_running_stats=False,
+            )
+            norm_layer1 = norm_layer0
+        # elif self.normalization_layer == "batch_norm":
+        #     norm_layer = partial(nn.InstanceNorm2d, num_features=self.embed_dim, eps=1e-6, affine=True, track_running_stats=False)
+        else:
+            raise NotImplementedError(
+                f"Error, normalization {self.normalization_layer} not implemented."
+            )
+
+        # ENCODER is just an MLP?
+        encoder_hidden_dim = self.embed_dim
+        encoder_act = nn.GELU
+
+        # encoder0 = nn.Conv2d(self.in_chans, encoder_hidden_dim, 1, bias=True)
+        # encoder1 = nn.Conv2d(encoder_hidden_dim, self.embed_dim, 1, bias=False)
+        # encoder_act = nn.GELU()
+        # self.encoder = nn.Sequential(encoder0, encoder_act, encoder1, norm_layer0())
+
+        self.encoder = MLP(
+            in_features=self.in_chans,
+            hidden_features=encoder_hidden_dim,
+            out_features=self.embed_dim,
+            output_bias=False,
+            act_layer=encoder_act,
+            drop_rate=0.0,
+            checkpointing=checkpointing,
+        )
+
+        # self.input_encoding = nn.Conv2d(self.in_chans, self.embed_dim, 1)
+        # self.pos_embed = nn.Parameter(torch.zeros(1, self.pos_embed_dim, self.img_size[0], self.img_size[1]))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, self.embed_dim, self.img_size[0], self.img_size[1])
+        )
+
+        # prepare the SHT
+        modes_lat = int(self.h * self.hard_thresholding_fraction)
+        modes_lon = int((self.w // 2 + 1) * self.hard_thresholding_fraction)
+
+        if self.spectral_transform == "sht":
+            self.trans_down = SafeRealSHT(
+                *self.img_size, lmax=modes_lat, mmax=modes_lon, grid="equiangular"
+            ).float()
+            self.itrans_up = SafeInverseRealSHT(
+                *self.img_size, lmax=modes_lat, mmax=modes_lon, grid="equiangular"
+            ).float()
+            self.trans = SafeRealSHT(
+                self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid="legendre-gauss"
+            ).float()
+            self.itrans = SafeInverseRealSHT(
+                self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid="legendre-gauss"
+            ).float()
+
+            # we introduce some ad-hoc rescaling of the weights to aid gradient computation:
+            sht_rescaling_factor = 1e5
+            self.trans_down.weights = self.trans_down.weights * sht_rescaling_factor
+            self.itrans_up.pct = self.itrans_up.pct / sht_rescaling_factor
+            self.trans.weights = self.trans.weights * sht_rescaling_factor
+            self.itrans.pct = self.itrans.pct / sht_rescaling_factor
+
+        elif self.spectral_transform == "fft":
+            self.trans_down = RealFFT2(
+                *self.img_size, lmax=modes_lat, mmax=modes_lon
+            ).float()
+            self.itrans_up = InverseRealFFT2(
+                *self.img_size, lmax=modes_lat, mmax=modes_lon
+            ).float()
+            self.trans = RealFFT2(
+                self.h, self.w, lmax=modes_lat, mmax=modes_lon
+            ).float()
+            self.itrans = InverseRealFFT2(
+                self.h, self.w, lmax=modes_lat, mmax=modes_lon
+            ).float()
+        else:
+            raise (ValueError("Unknown spectral transform"))
+
+        self.blocks = nn.ModuleList([])
+        for i in range(self.num_layers):
+            first_layer = i == 0
+            last_layer = i == self.num_layers - 1
+
+            forward_transform = self.trans_down if first_layer else self.trans
+            inverse_transform = self.itrans_up if last_layer else self.itrans
+
+            inner_skip = "linear" if 0 < i < self.num_layers - 1 else None
+            outer_skip = "identity" if 0 < i < self.num_layers - 1 else None
+            mlp_mode = self.mlp_mode if not last_layer else "none"
+
+            if first_layer:
+                norm_layer = (norm_layer0, norm_layer1)
+            elif last_layer:
+                norm_layer = (norm_layer1, norm_layer0)
+            else:
+                norm_layer = (norm_layer1, norm_layer1)
+
+            block = FourierNeuralOperatorBlock(
+                forward_transform,
+                inverse_transform,
+                self.embed_dim,
+                filter_type=self.filter_type,
+                mlp_ratio=mlp_ratio,
+                drop_rate=drop_rate,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                sparsity_threshold=sparsity_threshold,
+                use_complex_kernels=use_complex_kernels,
+                inner_skip=inner_skip,
+                outer_skip=outer_skip,
+                mlp_mode=mlp_mode,
+                compression=self.compression,
+                rank=self.rank,
+                complex_network=self.complex_network,
+                complex_activation=self.complex_activation,
+                spectral_layers=self.spectral_layers,
+                checkpointing=self.checkpointing,
+            )
+
+            self.blocks.append(block)
+
+        # DECODER is also an MLP
+        decoder_hidden_dim = self.embed_dim
+        decoder_act = nn.GELU
+
+        # decoder0 = nn.Conv2d(self.embed_dim + self.big_skip*self.in_chans, decoder_hidden_dim, 1, bias=True)
+        # decoder1 = nn.Conv2d(decoder_hidden_dim, self.out_chans, 1, bias=False)
+        # decoder_act = nn.GELU()
+        # self.decoder = nn.Sequential(decoder0, decoder_act, decoder1)
+
+        self.decoder = MLP(
+            in_features=self.embed_dim + self.big_skip * self.in_chans,
+            hidden_features=decoder_hidden_dim,
+            out_features=self.out_chans,
+            output_bias=False,
+            act_layer=decoder_act,
+            drop_rate=0.0,
+            checkpointing=checkpointing,
+        )
+
+        trunc_normal_(self.pos_embed, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
+            trunc_normal_(m.weight, std=0.02)
+            # nn.init.normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm) or isinstance(m, FusedLayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def forward_features(self, x):
+        # x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        return x
+
+    def forward(self, x):
+        # save big skip
+        if self.big_skip:
+            residual = x
+
+        # encoder
+        x = self.encoder(x)
+
+        # do positional embedding
+        x = x + self.pos_embed
+
+        # forward features
+        x = self.forward_features(x)
+
+        # concatenate the big skip
+        if self.big_skip:
+            x = torch.cat((x, residual), dim=1)
+
+        # decoder
+        x = self.decoder(x)
+
+        return x

model/fourcastnet_v2.py

@@ -0,0 +1,215 @@
+"""OneScience adapter for NVIDIA's official legacy FourCastNet v2 network."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Mapping
+
+import torch
+from torch import nn
+
+
+def _load_model_class():
+    package_dir = Path(__file__).resolve().parent / "fcnv2"
+    if not (package_dir / "fcnv2_sfnonet.py").is_file():
+        raise FileNotFoundError(
+            f"Bundled FCNv2 source was not found at {package_dir}"
+        )
+
+    package_path = str(package_dir)
+    if package_path not in sys.path:
+        sys.path.insert(0, package_path)
+
+    try:
+        from fcnv2_sfnonet import FourierNeuralOperatorNet
+    except ModuleNotFoundError as error:
+        if error.name == "torch_harmonics":
+            raise ModuleNotFoundError(
+                "FourCastNet v2 requires NVIDIA torch-harmonics. Install the "
+                "version pinned by this project before constructing the model."
+            ) from error
+        raise
+    return FourierNeuralOperatorNet
+
+
+def _official_params(model_config: Mapping[str, Any]) -> SimpleNamespace:
+    required = {
+        "img_size",
+        "in_channels",
+        "out_channels",
+        "spectral_transform",
+        "filter_type",
+        "scale_factor",
+        "embed_dim",
+        "num_layers",
+        "num_blocks",
+        "normalization_layer",
+        "mlp_mode",
+        "spectral_layers",
+        "complex_activation",
+        "hard_thresholding_fraction",
+        "big_skip",
+    }
+    missing = sorted(required.difference(model_config))
+    if missing:
+        raise ValueError(f"Missing FourCastNet v2 model settings: {missing}")
+
+    height, width = model_config["img_size"]
+    hidden_height = height // model_config["scale_factor"]
+    hidden_width = width // model_config["scale_factor"]
+    if hidden_height < 2 or hidden_width < 2:
+        raise ValueError("The internal SFNO grid must have at least 2 x 2 points")
+
+    return SimpleNamespace(
+        img_crop_shape_x=int(height),
+        img_crop_shape_y=int(width),
+        N_in_channels=int(model_config["in_channels"]),
+        N_out_channels=int(model_config["out_channels"]),
+        spectral_transform=model_config["spectral_transform"],
+        filter_type=model_config["filter_type"],
+        scale_factor=int(model_config["scale_factor"]),
+        embed_dim=int(model_config["embed_dim"]),
+        num_layers=int(model_config["num_layers"]),
+        num_blocks=int(model_config["num_blocks"]),
+        normalization_layer=model_config["normalization_layer"],
+        mlp_mode=model_config["mlp_mode"],
+        spectral_layers=int(model_config["spectral_layers"]),
+        complex_activation=model_config["complex_activation"],
+        hard_thresholding_fraction=float(
+            model_config["hard_thresholding_fraction"]
+        ),
+        big_skip=bool(model_config["big_skip"]),
+    )
+
+
+class FourCastNetV2(nn.Module):
+    """Build the exact official FCNv2 network behind a stable project API."""
+
+    def __init__(
+        self,
+        model_config: Mapping[str, Any],
+    ) -> None:
+        super().__init__()
+        self.model_config = dict(model_config)
+        self.expected_shape = (
+            int(model_config["in_channels"]),
+            int(model_config["img_size"][0]),
+            int(model_config["img_size"][1]),
+        )
+        model_class = _load_model_class()
+        self.model = model_class(_official_params(model_config))
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        if inputs.ndim != 4:
+            raise ValueError(f"Expected [B,C,H,W], got {tuple(inputs.shape)}")
+        if tuple(inputs.shape[1:]) != self.expected_shape:
+            raise ValueError(
+                f"Expected trailing shape {self.expected_shape}, "
+                f"got {tuple(inputs.shape[1:])}"
+            )
+        return self.model(inputs)
+
+    def no_weight_decay(self) -> set[str]:
+        return {f"model.{name}" for name in self.model.no_weight_decay()}
+
+
+def _unwrap_state_dict(checkpoint: Any) -> Mapping[str, torch.Tensor]:
+    if not isinstance(checkpoint, Mapping):
+        raise TypeError("Checkpoint must contain a mapping")
+    for key in ("model_state", "model_state_dict", "state_dict"):
+        candidate = checkpoint.get(key)
+        if isinstance(candidate, Mapping):
+            return candidate
+    if checkpoint and all(isinstance(value, torch.Tensor) for value in checkpoint.values()):
+        return checkpoint
+    raise KeyError("Checkpoint has no recognized model state mapping")
+
+
+def _normalize_state_keys(
+    state_dict: Mapping[str, torch.Tensor], model: nn.Module
+) -> dict[str, torch.Tensor]:
+    target_keys = set(model.state_dict())
+    normalized: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        clean_key = key
+        while clean_key.startswith("module."):
+            clean_key = clean_key[len("module.") :]
+        if clean_key not in target_keys and f"model.{clean_key}" in target_keys:
+            clean_key = f"model.{clean_key}"
+        normalized[clean_key] = value
+    return normalized
+
+
+def load_checkpoint(
+    model: nn.Module,
+    checkpoint_path: str | Path,
+    *,
+    expected_profile: str,
+    expected_variables: list[str],
+    allowed_stages: set[str],
+    allowed_initializations: set[str],
+    strict: bool = True,
+    map_location: str | torch.device = "cpu",
+) -> dict[str, Any]:
+    """Load a project checkpoint without changing its parameter tensors."""
+
+    checkpoint = torch.load(
+        Path(checkpoint_path).expanduser(),
+        map_location=map_location,
+        weights_only=False,
+    )
+    validate_project_checkpoint(
+        checkpoint,
+        expected_profile=expected_profile,
+        expected_variables=expected_variables,
+        allowed_stages=allowed_stages,
+        allowed_initializations=allowed_initializations,
+    )
+    state_dict = _normalize_state_keys(_unwrap_state_dict(checkpoint), model)
+    incompatible = model.load_state_dict(state_dict, strict=strict)
+    return {
+        "checkpoint": checkpoint,
+        "missing_keys": list(incompatible.missing_keys),
+        "unexpected_keys": list(incompatible.unexpected_keys),
+    }
+
+
+def validate_project_checkpoint(
+    checkpoint: Any,
+    *,
+    expected_profile: str,
+    expected_variables: list[str],
+    allowed_stages: set[str],
+    allowed_initializations: set[str],
+) -> None:
+    if not isinstance(checkpoint, Mapping):
+        raise TypeError("Project checkpoint must contain metadata")
+    expected = {
+        "checkpoint_format": "fourcastnet_v2_project",
+        "scratch_lineage": True,
+        "model_profile": expected_profile,
+        "variables": expected_variables,
+    }
+    for key, value in expected.items():
+        if checkpoint.get(key) != value:
+            raise ValueError(
+                f"Checkpoint metadata {key!r} does not match the project config"
+            )
+    if checkpoint.get("stage") not in allowed_stages:
+        raise ValueError(
+            f"Checkpoint stage must be one of {sorted(allowed_stages)}"
+        )
+    if checkpoint.get("initialization") not in allowed_initializations:
+        raise ValueError(
+            "Checkpoint does not have an approved random-initialization lineage"
+        )
+    expected_initialization = {
+        "one_step": "random",
+        "finetune": "one_step_checkpoint",
+    }.get(checkpoint.get("stage"))
+    if checkpoint.get("initialization") != expected_initialization:
+        raise ValueError(
+            "Checkpoint stage and initialization metadata are inconsistent"
+        )

scripts/common.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
+
+
+def load_config(path: str | Path = DEFAULT_CONFIG) -> dict[str, Any]:
+    config_path = Path(path).expanduser().resolve()
+    with config_path.open("r", encoding="utf-8") as stream:
+        config = yaml.safe_load(stream)
+    config["_config_path"] = str(config_path)
+    config["_project_root"] = str(PROJECT_ROOT)
+    validate_config(config)
+    return config
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    variables = config["data"]["variables"]
+    if len(variables) != 73 or len(set(variables)) != 73:
+        raise ValueError("FourCastNet v2 requires 73 unique variables")
+
+    profile_name = config["model"]["profile"]
+    profiles = config["model"]["profiles"]
+    if profile_name not in profiles:
+        raise ValueError(f"Unknown model profile: {profile_name}")
+
+    profile = profiles[profile_name]
+    if profile["in_channels"] != len(variables):
+        raise ValueError("Model input channels do not match the variable ledger")
+    if profile["out_channels"] != len(variables):
+        raise ValueError("Model output channels do not match the variable ledger")
+
+    if config["data"]["input_steps"] != 1:
+        raise ValueError("FourCastNet v2 expects exactly one input time step")
+    if config["data"]["output_steps"] != 1:
+        raise ValueError("One-step pretraining expects data.output_steps=1")
+    if config["training"]["finetune"]["autoregressive_steps"] < 2:
+        raise ValueError("Fine-tuning requires at least two autoregressive steps")
+    if config["inference"]["rollout_steps"] < 1:
+        raise ValueError("inference.rollout_steps must be positive")
+
+    if config["training"]["stage"] not in {"one_step", "finetune"}:
+        raise ValueError("training.stage must be 'one_step' or 'finetune'")
+    if config["checkpoint"]["initialize_from"] != "scratch":
+        raise ValueError("checkpoint.initialize_from must be 'scratch'")
+    if not config["checkpoint"].get("finetune_from"):
+        raise ValueError("checkpoint.finetune_from must name a one-step checkpoint")
+    prefix = config["checkpoint"].get("prefix", "model_bak")
+    if not prefix or Path(prefix).name != prefix:
+        raise ValueError("checkpoint.prefix must be a non-empty file name")
+
+
+def resolve_path(config: dict[str, Any], value: str | Path) -> Path:
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path
+    return Path(config["_project_root"]) / path
+
+
+def active_model_config(config: dict[str, Any]) -> dict[str, Any]:
+    return dict(config["model"]["profiles"][config["model"]["profile"]])
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)

scripts/data.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as functional
+from torch.utils.data import DataLoader, Dataset
+
+
+def _import_era5_dataset(onescience_source_dir: str | None = None):
+    if onescience_source_dir:
+        source_dir = str(Path(onescience_source_dir).expanduser().resolve())
+        if source_dir not in sys.path:
+            sys.path.insert(0, source_dir)
+    from onescience.datapipes.climate import ERA5Dataset
+
+    return ERA5Dataset
+
+
+class SpatialAdapter(Dataset):
+    """Resize OneScience ERA5 samples only for reduced smoke profiles."""
+
+    def __init__(self, dataset: Dataset, output_size: tuple[int, int]) -> None:
+        self.dataset = dataset
+        self.output_size = output_size
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def _resize(self, tensor: torch.Tensor) -> torch.Tensor:
+        if tuple(tensor.shape[-2:]) == self.output_size:
+            return tensor
+        leading_shape = tensor.shape[:-2]
+        resized = functional.interpolate(
+            tensor.reshape(-1, 1, *tensor.shape[-2:]),
+            size=self.output_size,
+            mode="bilinear",
+            align_corners=False,
+        )
+        return resized.reshape(*leading_shape, *self.output_size)
+
+    def __getitem__(self, index: int):
+        inputs, targets, cos_zenith, step_idx, time_index = self.dataset[index]
+        inputs = self._resize(inputs)
+        targets = self._resize(targets)
+        cos_zenith = self._resize(cos_zenith)
+        return inputs, targets, cos_zenith, step_idx, time_index
+
+
+def build_dataset(
+    config: dict[str, Any],
+    years: list[int],
+    *,
+    output_steps: int = 1,
+) -> Dataset:
+    from common import active_model_config, resolve_path
+
+    era5_dataset = _import_era5_dataset(config["project"].get("onescience_source_dir"))
+    data_config = config["data"]
+    dataset = era5_dataset(
+        dataset_dir=str(resolve_path(config, data_config["dataset_dir"])),
+        used_years=years,
+        used_variables=data_config["variables"],
+        input_steps=data_config["input_steps"],
+        output_steps=output_steps,
+        normalize=data_config["normalize"],
+    )
+    model_size = tuple(active_model_config(config)["img_size"])
+    data_size = tuple(data_config["grid_shape"])
+    if model_size != data_size:
+        dataset = SpatialAdapter(dataset, model_size)
+    return dataset
+
+
+def build_loader(
+    config: dict[str, Any],
+    years: list[int],
+    *,
+    train: bool,
+    distributed: bool,
+    output_steps: int = 1,
+) -> tuple[DataLoader, torch.utils.data.Sampler | None]:
+    dataset = build_dataset(config, years, output_steps=output_steps)
+    sampler = None
+    if distributed:
+        sampler = torch.utils.data.distributed.DistributedSampler(
+            dataset, shuffle=train
+        )
+    loader = DataLoader(
+        dataset,
+        batch_size=config["training"]["batch_size"],
+        shuffle=train and sampler is None,
+        sampler=sampler,
+        num_workers=config["training"]["num_workers"],
+        pin_memory=True,
+        drop_last=False,
+    )
+    return loader, sampler
+
+
+def load_statistics(config: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
+    import h5py
+    import numpy as np
+
+    from common import resolve_path
+
+    data_config = config["data"]
+    year = data_config["test_years"][0]
+    path = resolve_path(config, data_config["dataset_dir"]) / "data" / f"{year}.h5"
+    with h5py.File(path, "r") as handle:
+        fields = handle["fields"]
+        all_variables = [
+            item.decode() if isinstance(item, bytes) else str(item)
+            for item in fields.attrs["variables"]
+        ]
+        indices = [all_variables.index(name) for name in data_config["variables"]]
+        if "global_means" in handle:
+            means = handle["global_means"][:]
+            stds = handle["global_stds"][:]
+        else:
+            stats_dir = path.parents[1] / "stats"
+            means = np.load(stats_dir / "global_means.npy")
+            stds = np.load(stats_dir / "global_stds.npy")
+    return torch.from_numpy(means[:, indices]), torch.from_numpy(stds[:, indices])

scripts/fake_data.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+import h5py
+import numpy as np
+
+from common import DEFAULT_CONFIG, load_config, resolve_path
+
+
+def generate_year(
+    path: Path,
+    variables: list[str],
+    *,
+    time_steps: int,
+    height: int,
+    width: int,
+    time_step_hours: int,
+    chunk_time_steps: int,
+    fill_value: float,
+    materialize_pattern: bool,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    channels = len(variables)
+    with h5py.File(path, "w") as handle:
+        fields = handle.create_dataset(
+            "fields",
+            shape=(time_steps, channels, height, width),
+            dtype=np.float32,
+            chunks=(chunk_time_steps, 1, height, width),
+            fillvalue=np.float32(fill_value),
+            compression="lzf",
+        )
+        fields.attrs["variables"] = np.asarray(variables, dtype=h5py.string_dtype())
+        fields.attrs["time_step"] = time_step_hours
+        handle.create_dataset(
+            "global_means", data=np.zeros((1, channels, 1, 1), dtype=np.float32)
+        )
+        handle.create_dataset(
+            "global_stds", data=np.ones((1, channels, 1, 1), dtype=np.float32)
+        )
+
+        if materialize_pattern:
+            latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None]
+            longitude = np.linspace(
+                0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32
+            )
+            base = latitude + np.sin(longitude)[None, :]
+            # Two frames are enough to exercise non-zero input and target reads.
+            for time_index in range(min(time_steps, 2)):
+                for channel_index in range(channels):
+                    fields[time_index, channel_index] = (
+                        base + channel_index / channels + time_index * 0.01
+                    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Generate ERA5-compatible FCNv2 data")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
+    parser.add_argument("--no-pattern", action="store_true")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    data = config["data"]
+    fake = config["fake_data"]
+    output_dir = resolve_path(config, data["dataset_dir"])
+    years = sorted(set(data["train_years"] + data["val_years"] + data["test_years"]))
+    height, width = data["grid_shape"]
+    for year in years:
+        path = output_dir / "data" / f"{year}.h5"
+        generate_year(
+            path,
+            data["variables"],
+            time_steps=fake["time_steps_per_year"],
+            height=height,
+            width=width,
+            time_step_hours=data["time_step_hours"],
+            chunk_time_steps=fake["chunk_time_steps"],
+            fill_value=fake["fill_value"],
+            materialize_pattern=fake["materialize_pattern"] and not args.no_pattern,
+        )
+        logical_gib = (
+            fake["time_steps_per_year"] * len(data["variables"]) * height * width * 4
+        ) / 1024**3
+        print(f"{path}: logical={logical_gib:.2f} GiB, actual={path.stat().st_size / 1024**2:.2f} MiB")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+MODEL_DIR = SCRIPT_DIR.parent / "model"
+for module_dir in (SCRIPT_DIR, MODEL_DIR):
+    if str(module_dir) not in sys.path:
+        sys.path.insert(0, str(module_dir))
+
+import numpy as np
+import torch
+
+from common import DEFAULT_CONFIG, active_model_config, load_config, resolve_path
+from data import build_loader, load_statistics
+from fourcastnet_v2 import (
+    FourCastNetV2,
+    load_checkpoint,
+)
+
+
+def choose_device() -> torch.device:
+    return torch.device("cuda", 0) if torch.cuda.is_available() else torch.device("cpu")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run FourCastNet v2 inference")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
+    parser.add_argument("--checkpoint")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    inference = config["inference"]
+
+    device = choose_device()
+    model = FourCastNetV2(active_model_config(config)).to(device)
+    if args.checkpoint:
+        checkpoint_path = resolve_path(config, args.checkpoint)
+    else:
+        checkpoint_path = resolve_path(config, inference["checkpoint_path"])
+    result = load_checkpoint(
+        model,
+        checkpoint_path,
+        expected_profile=config["model"]["profile"],
+        expected_variables=config["data"]["variables"],
+        allowed_stages={"one_step", "finetune"},
+        allowed_initializations={"random", "one_step_checkpoint"},
+        strict=config["checkpoint"]["strict"],
+        map_location=device,
+    )
+    if result["missing_keys"] or result["unexpected_keys"]:
+        print(
+            f"missing_keys={result['missing_keys']} "
+            f"unexpected_keys={result['unexpected_keys']}"
+        )
+
+    loader, _ = build_loader(
+        config,
+        config["data"]["test_years"],
+        train=False,
+        distributed=False,
+        output_steps=inference["rollout_steps"],
+    )
+    means, stds = load_statistics(config)
+    means = means.numpy()
+    stds = stds.numpy()
+
+    output_dir = resolve_path(config, inference["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    model.eval()
+    with torch.no_grad():
+        for sample_index, batch in enumerate(loader):
+            if sample_index >= inference["max_samples"]:
+                break
+            inputs = batch[0].to(device)
+            targets = batch[1]
+            state = inputs
+            predictions = []
+            for _ in range(inference["rollout_steps"]):
+                state = model(state)
+                predictions.append(state.cpu())
+            prediction = torch.stack(predictions, dim=1).numpy()
+            if targets.ndim == 4:
+                targets = targets.unsqueeze(1)
+            target = targets.numpy()
+            input_array = inputs.cpu().numpy()
+            if not inference["save_normalized"]:
+                prediction = prediction * stds[:, None] + means[:, None]
+                target = target * stds[:, None] + means[:, None]
+                input_array = input_array * stds + means
+            path = output_dir / f"sample_{sample_index:04d}.npz"
+            np.savez_compressed(
+                path,
+                input=input_array,
+                prediction=prediction,
+                target=target,
+                variables=np.asarray(config["data"]["variables"]),
+                time_index=np.asarray(batch[4], dtype=str).T,
+            )
+            print(path)
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+from common import DEFAULT_CONFIG, load_config, resolve_path
+
+
+def latitude_weights(height: int) -> np.ndarray:
+    latitude = np.linspace(np.pi / 2, -np.pi / 2, height)
+    weights = np.cos(latitude).clip(min=0)
+    return weights / weights.mean()
+
+
+def compute_metrics(prediction: np.ndarray, target: np.ndarray) -> dict[str, list[float]]:
+    weights = latitude_weights(target.shape[-2])[None, None, None, :, None]
+    error = prediction - target
+    rmse = np.sqrt(np.mean(error**2 * weights, axis=(0, 1, 3, 4)))
+    spatial_weight = weights / (
+        weights.sum(axis=(-2, -1), keepdims=True) * target.shape[-1]
+    )
+    pred_mean = np.sum(prediction * spatial_weight, axis=(-2, -1), keepdims=True)
+    target_mean = np.sum(target * spatial_weight, axis=(-2, -1), keepdims=True)
+    pred_anomaly = prediction - pred_mean
+    target_anomaly = target - target_mean
+    numerator = np.sum(pred_anomaly * target_anomaly * weights, axis=(0, 1, 3, 4))
+    denominator = np.sqrt(
+        np.sum(pred_anomaly**2 * weights, axis=(0, 1, 3, 4))
+        * np.sum(target_anomaly**2 * weights, axis=(0, 1, 3, 4))
+    )
+    acc = numerator / np.maximum(denominator, 1e-12)
+    return {"rmse": rmse.tolist(), "acc": acc.tolist()}
+
+
+def plot_sample(
+    prediction: np.ndarray,
+    target: np.ndarray,
+    variable: str,
+    channel_index: int,
+    cmap: str,
+    output_path: Path,
+) -> None:
+    predicted = prediction[0, 0, channel_index]
+    expected = target[0, 0, channel_index]
+    error = predicted - expected
+    value_min = min(predicted.min(), expected.min())
+    value_max = max(predicted.max(), expected.max())
+    error_limit = max(abs(error.min()), abs(error.max()), 1e-12)
+    extent = (0, 360, -90, 90)
+
+    figure, axes = plt.subplots(3, 1, figsize=(12, 10), constrained_layout=True)
+    image = axes[0].imshow(
+        expected, origin="upper", extent=extent, aspect="auto", cmap=cmap,
+        vmin=value_min, vmax=value_max,
+    )
+    axes[0].set_title(f"Target {variable}")
+    figure.colorbar(image, ax=axes[0], orientation="vertical")
+    image = axes[1].imshow(
+        predicted, origin="upper", extent=extent, aspect="auto", cmap=cmap,
+        vmin=value_min, vmax=value_max,
+    )
+    axes[1].set_title(f"Prediction {variable}")
+    figure.colorbar(image, ax=axes[1], orientation="vertical")
+    image = axes[2].imshow(
+        error, origin="upper", extent=extent, aspect="auto", cmap="RdBu_r",
+        vmin=-error_limit, vmax=error_limit,
+    )
+    axes[2].set_title(f"Error {variable}")
+    figure.colorbar(image, ax=axes[2], orientation="vertical")
+    for axis in axes:
+        axis.set_xlabel("Longitude")
+        axis.set_ylabel("Latitude")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    figure.savefig(output_path, dpi=160)
+    plt.close(figure)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate and plot FCNv2 output")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
+    parser.add_argument("--input")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    inference_dir = resolve_path(config, config["inference"]["output_dir"])
+    input_path = Path(args.input).expanduser().resolve() if args.input else None
+    files = [input_path] if input_path else sorted(inference_dir.glob("sample_*.npz"))
+    if not files:
+        raise FileNotFoundError(f"No inference outputs found in {inference_dir}")
+
+    predictions = []
+    targets = []
+    for path in files:
+        with np.load(path) as data:
+            predictions.append(data["prediction"])
+            targets.append(data["target"])
+    prediction = np.concatenate(predictions)
+    target = np.concatenate(targets)
+    metrics = compute_metrics(prediction, target)
+
+    output_dir = resolve_path(config, config["visualization"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    (output_dir / "metrics.json").write_text(
+        json.dumps(metrics, indent=2), encoding="utf-8"
+    )
+    variable = config["visualization"]["variable"]
+    channel_index = config["data"]["variables"].index(variable)
+    sample_index = config["visualization"]["sample_index"]
+    if not 0 <= sample_index < prediction.shape[0]:
+        raise IndexError(
+            f"visualization.sample_index={sample_index} is outside "
+            f"the available range [0, {prediction.shape[0] - 1}]"
+        )
+    plot_sample(
+        prediction[sample_index : sample_index + 1],
+        target[sample_index : sample_index + 1],
+        variable,
+        channel_index,
+        config["visualization"]["cmap"],
+        output_dir / f"{variable}_forecast.png",
+    )
+    print(output_dir / "metrics.json")
+    print(output_dir / f"{variable}_forecast.png")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import tempfile
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any, Iterable
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+MODEL_DIR = SCRIPT_DIR.parent / "model"
+for module_dir in (SCRIPT_DIR, MODEL_DIR):
+    if str(module_dir) not in sys.path:
+        sys.path.insert(0, str(module_dir))
+
+from common import (
+    DEFAULT_CONFIG,
+    active_model_config,
+    load_config,
+    resolve_path,
+    seed_everything,
+)
+from data import build_loader
+from fourcastnet_v2 import (
+    FourCastNetV2,
+    load_checkpoint,
+)
+
+
+def initialize_distributed(backend: str) -> tuple[torch.device, int, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if torch.cuda.is_available():
+        device_count = torch.cuda.device_count()
+        if not 0 <= local_rank < device_count:
+            raise RuntimeError(
+                f"LOCAL_RANK={local_rank} is not available; "
+                f"this process can see {device_count} CUDA devices "
+                f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '<unset>')})"
+            )
+        # Select the rank-local device before NCCL initialization.  Otherwise
+        # every process starts with the default device (usually cuda:0).
+        torch.cuda.set_device(local_rank)
+        device = torch.device("cuda", local_rank)
+    else:
+        device = torch.device("cpu")
+
+    if world_size > 1 and not dist.is_initialized():
+        dist.init_process_group(backend=backend, init_method="env://")
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    return device, rank, local_rank, world_size
+
+
+def spherical_relative_l2(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    height = target.shape[-2]
+    latitude = torch.linspace(
+        torch.pi / 2,
+        -torch.pi / 2,
+        height,
+        device=target.device,
+        dtype=target.dtype,
+    )
+    weights = torch.cos(latitude).clamp_min(0)
+    weights = weights / weights.mean()
+    weights = weights.view(1, 1, height, 1)
+    error = ((prediction - target).square() * weights).sum(dim=(-2, -1))
+    reference = (target.square() * weights).sum(dim=(-2, -1)).clamp_min(1e-12)
+    return torch.sqrt(error / reference).mean()
+
+
+def autoregressive_loss(
+    model: torch.nn.Module,
+    inputs: torch.Tensor,
+    targets: torch.Tensor,
+    steps: int,
+) -> torch.Tensor:
+    if steps == 1:
+        targets = targets if targets.ndim == 4 else targets[:, 0]
+    elif targets.ndim != 5 or targets.shape[1] != steps:
+        raise ValueError(f"Expected targets [B,{steps},C,H,W], got {targets.shape}")
+
+    state = inputs
+    losses = []
+    for step in range(steps):
+        state = model(state)
+        target = targets if steps == 1 else targets[:, step]
+        losses.append(spherical_relative_l2(state, target))
+    return torch.stack(losses).mean()
+
+
+def limited_batches(loader: Iterable, maximum: int | None):
+    for index, batch in enumerate(loader):
+        if maximum is not None and index >= maximum:
+            break
+        yield batch
+
+
+def reduce_average(total: float, count: int, device: torch.device) -> float:
+    values = torch.tensor([total, count], dtype=torch.float64, device=device)
+    if dist.is_initialized():
+        dist.all_reduce(values, op=dist.ReduceOp.SUM)
+    return (values[0] / values[1].clamp_min(1)).item()
+
+
+def run_epoch(
+    model: torch.nn.Module,
+    loader,
+    device: torch.device,
+    *,
+    steps: int,
+    optimizer: torch.optim.Optimizer | None,
+    amp: bool,
+    max_batches: int | None,
+    max_grad_norm: float,
+) -> float:
+    training = optimizer is not None
+    model.train(training)
+    total = 0.0
+    count = 0
+    context = nullcontext if training else torch.no_grad
+    with context():
+        for batch in limited_batches(loader, max_batches):
+            inputs = batch[0].to(device, non_blocking=True)
+            targets = batch[1].to(device, non_blocking=True)
+            if training:
+                optimizer.zero_grad(set_to_none=True)
+            with torch.autocast(
+                device_type=device.type,
+                dtype=torch.float16,
+                enabled=amp and device.type == "cuda",
+            ):
+                loss = autoregressive_loss(model, inputs, targets, steps)
+            if training:
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
+                optimizer.step()
+            total += loss.detach().item()
+            count += 1
+    return reduce_average(total, count, device)
+
+
+def save_checkpoint_atomic(path: Path, state: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as stream:
+        temporary_path = Path(stream.name)
+    try:
+        torch.save(state, temporary_path)
+        os.replace(temporary_path, path)
+    finally:
+        temporary_path.unlink(missing_ok=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train FourCastNet v2")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
+    parser.add_argument("--stage", choices=("one_step", "finetune"))
+    parser.add_argument("--resume")
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    seed_everything(config["project"]["seed"])
+    training = config["training"]
+    stage = args.stage or training["stage"]
+    if stage == "one_step" and args.resume:
+        raise ValueError("One-step training always starts from random initialization")
+    steps = 1 if stage == "one_step" else training["finetune"]["autoregressive_steps"]
+    epochs = training["epochs"] if stage == "one_step" else training["finetune"]["epochs"]
+    learning_rate = (
+        training["learning_rate"]
+        if stage == "one_step"
+        else training["finetune"]["learning_rate"]
+    )
+
+    device, rank, local_rank, world_size = initialize_distributed(
+        config["distributed"]["backend"]
+    )
+    print(
+        f"rank={rank}/{world_size} local_rank={local_rank} "
+        f"device={device} visible_devices={torch.cuda.device_count()}",
+        flush=True,
+    )
+    train_loader, train_sampler = build_loader(
+        config,
+        config["data"]["train_years"],
+        train=True,
+        distributed=world_size > 1,
+        output_steps=steps,
+    )
+    val_loader, val_sampler = build_loader(
+        config,
+        config["data"]["val_years"],
+        train=False,
+        distributed=world_size > 1,
+        output_steps=steps,
+    )
+
+    model = FourCastNetV2(active_model_config(config)).to(device)
+    if stage == "finetune":
+        resume_path = args.resume or config["checkpoint"]["finetune_from"]
+        result = load_checkpoint(
+            model,
+            resolve_path(config, resume_path),
+            expected_profile=config["model"]["profile"],
+            expected_variables=config["data"]["variables"],
+            allowed_stages={"one_step"},
+            allowed_initializations={"random"},
+            strict=config["checkpoint"]["strict"],
+            map_location=device,
+        )
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=learning_rate,
+        betas=tuple(training["optimizer_betas"]),
+        weight_decay=training["weight_decay"],
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+    if world_size > 1:
+        ddp_options = (
+            {
+                "device_ids": [device.index],
+                "output_device": device.index,
+                # SFNO's SHT buffers are immutable coefficients.  DDP's
+                # per-forward buffer broadcast mutates them in-place and
+                # invalidates the graph in multi-step autoregressive training.
+                "broadcast_buffers": False,
+            }
+            if device.type == "cuda"
+            else {"broadcast_buffers": False}
+        )
+        model = DistributedDataParallel(model, **ddp_options)
+
+    checkpoint_dir = (
+        resolve_path(config, config["project"]["checkpoint_dir"]) / stage
+    )
+    checkpoint_prefix = config["checkpoint"].get("prefix", "model_bak")
+    best_loss = float("inf")
+    history = []
+    for epoch in range(epochs):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        if val_sampler is not None:
+            val_sampler.set_epoch(epoch)
+        train_loss = run_epoch(
+            model,
+            train_loader,
+            device,
+            steps=steps,
+            optimizer=optimizer,
+            amp=training["amp"],
+            max_batches=training["max_train_batches"],
+            max_grad_norm=training["max_grad_norm"],
+        )
+        val_loss = run_epoch(
+            model,
+            val_loader,
+            device,
+            steps=steps,
+            optimizer=None,
+            amp=training["amp"],
+            max_batches=training["max_val_batches"],
+            max_grad_norm=training["max_grad_norm"],
+        )
+        scheduler.step()
+        history.append({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss})
+        if rank == 0:
+            print(
+                f"epoch={epoch + 1}/{epochs} train_loss={train_loss:.6f} "
+                f"val_loss={val_loss:.6f}"
+            )
+            raw_model = model.module if hasattr(model, "module") else model
+            state = {
+                "checkpoint_format": "fourcastnet_v2_project",
+                "scratch_lineage": True,
+                "model_state_dict": raw_model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "scheduler_state_dict": scheduler.state_dict(),
+                "epoch": epoch,
+                "stage": stage,
+                "initialization": (
+                    "random" if stage == "one_step" else "one_step_checkpoint"
+                ),
+                "model_profile": config["model"]["profile"],
+                "variables": config["data"]["variables"],
+            }
+            save_checkpoint_atomic(checkpoint_dir / f"{checkpoint_prefix}_last.pt", state)
+            if val_loss < best_loss:
+                best_loss = val_loss
+                save_checkpoint_atomic(checkpoint_dir / f"{checkpoint_prefix}.pt", state)
+            checkpoint_dir.mkdir(parents=True, exist_ok=True)
+            (checkpoint_dir / "history.json").write_text(
+                json.dumps(history, indent=2), encoding="utf-8"
+            )
+
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()