Upload folder using huggingface_hub

.ms_upload_cache

@@ -0,0 +1 @@
+{"version": 3, "repo_id": "OneScience/PrithviWxC", "files": {"README.md|1787118190.0|6111": {"hash": "4644dd18dcfd1c681633df078cd60b87a647c9cdf9505c83e0b9f6e03af55faf", "size": 6111, "status": "c"}, "conf/config.yaml|1787189055.0|2484": {"hash": "dc0ae83fd9cb50c61ce7a15fea77481f85c1f587fece5e63a5127a82ba273a67", "size": 2484, "status": "c"}, "model/__pycache__/prithvi_wxc.cpython-311.pyc|1787129097.0|6399": {"hash": "001025caa3762a8c91555f64ca871e687d234518abc23382206b83e0b3881978", "size": 6399, "status": "c"}, "model/__pycache__/prithvi_wxc_official.cpython-311.pyc|1787129243.0|76780": {"hash": "76c3889a828a1ee60630b10ea5a61dce9e25e8ef045fd71e2d150349f77e48ed", "size": 76780, "status": "c"}, "model/prithvi_wxc_official.py|1787129237.0|61202": {"hash": "58a3712e4739889be50dcb43106e7bfa43ed3c10e1ff44422a8192d6f343f992", "size": 61202, "status": "c"}, "configuration.json|1787037409.0|38": {"hash": "97d4072fc4a7a3b71e3184609fc491e77a324a982d9df860d9379c009dff84de", "size": 38, "status": "c"}, "LICENSE|1787037359.0|1106": {"hash": "c681a18eea18d4a07e4f429ca50d3f874c9d7982bb4836fff613218dfe47818e", "size": 1106, "status": "c"}, "model/prithvi_wxc.py|1787037757.0|5712": {"hash": "9f4186ae8f902dfa1b606c4ef53731c4923fda7c212a99941167f65ce37360d3", "size": 5712, "status": "c"}, "scripts/fake_data.py|1787118238.0|4985": {"hash": "b16a395251bffa0ee2948643d66c5d40a89fb63b397c23ec8703cd90f97f2859", "size": 4985, "status": "c"}, "scripts/inference.py|1787118190.0|4127": {"hash": "07920093609f16618001ef9a138ac1b843dd8985c1abfdc8accedf18518bec04", "size": 4127, "status": "c"}, "scripts/result.py|1787129368.0|8391": {"hash": "9a5b66ed2b41bbb766b3ce26a286be904d6d2483d8fde2ada6ea6c6044e6aed0", "size": 8391, "status": "c"}, "scripts/train.py|1787118190.0|10421": {"hash": "af5547fc972472f51bfcc719da9f6d1746492428e03a73345e2f23ac41f635b8", "size": 10421, "status": "c"}, "weight/.gitkeep|1787037589.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}}}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Inter Agency Implementation and Advanced Concepts
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,180 @@
+---
+license: mit
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth science
+- Weather forecasting
+- Medium- to long-range weather forecasting
+- Foundation models
+- Vision Transformer
+- ERA5
+frameworks: PyTorch
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Prithvi WxC</span>
+  </strong>
+</p>
+
+
+# Model Introduction
+
+Prithvi WxC (Weather and Climate) was proposed by NASA-IMPACT, IBM, and other teams. It is a weather and climate foundation model based on a Vision Transformer (alternating local/global attention with Hiera and MaxViT), supporting forecasting (6-hour-step rollout) and climate simulation (internal error growth).
+
+Paper:Prithvi WxC: Foundation Model for Weather and Climate
+
+https://arxiv.org/abs/2409.13598
+
+# Model Description
+
+Prithvi WxC is a deterministic global weather foundation model: it takes atmospheric states at two consecutive 6-hour time steps, optionally with static fields, and outputs the target state. Longer lead times are obtained through autoregressive rollout.This repository is organized from the official `NASA-IMPACT/Prithvi-WxC` implementation and integrated with the OneScience data loading and training workflow.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global weather and climate foundation model research | Train or fine-tune a Vision Transformer forecasting model on ERA5 data. |
+| Long-horizon autoregressive rollout | Generate medium- to long-range forecasts autoregressively at 6-hour intervals. |
+| Local quick validation | Use synthetic data to check data loading, training, inference, and result scripts. |
+| ModelScope/OneCode execution | Download the model package, install dependencies, and run the scripts directly. |
+| Multi-card training | Launch multi-process training with `torchrun`. |
+
+
+# Usage
+
+## 1. OneCode Usage
+
+Use the OneCode online environment for intelligent one-click AI4S programming:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU is recommended.
+- CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower.
+- DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+- The paper-level configuration (`embed_dim=2560`, 25 encoder blocks, 5 decoder blocks, and approximately 2.3 billion parameters) requires substantial GPU memory.
+
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/PrithviWxC --local-dir ./PrithviWxC
+cd PrithviWxC
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# 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
+# 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
+
+The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in `conf/config.yaml` is correct:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+For a quick workflow validation, run the synthetic data script first:
+
+```bash
+python scripts/fake_data.py
+```
+
+> Note: `scripts/fake_data.py` generates the `[T, C, H, W]` HDF5 data required by the two input time steps and generates `data/static/static.npy` (currently `[4, 32, 64]`) for training and inference.
+
+### Training
+
+Single card：
+
+```bash
+python scripts/train.py
+```
+
+Multiple cards：
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training outputs:
+
+```text
+data/checkpoints/model_bak.pth
+data/checkpoints/trloss.npy
+data/checkpoints/valoss.npy
+```
+
+### Training Weights
+The `weight/` folder is reserved for model weights. Official weights with approximately 2.3 billion parameters (such as PrithviWxC_160_13b_2t_0p5d_v1.pt) are published on Hugging Face, but their structure differs from this repository's small configuration. Align the channel count and grid size before loading; weights are not provided by default, and users may train the model using the paper configuration.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth`：
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are written to:
+
+```text
+result/output/
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Outputs include：
+
+- `result/rmse.npy`
+- `result/acc.npy`
+- `result/loss.png`
+- Forecast comparison plots for the specified date and variables
+
+
+# Official Source and Reproduction Notes
+
+- The model implementation comes from the official `NASA-IMPACT/Prithvi-WxC` (MIT License). The official implementation is embedded unchanged in `model/prithvi_wxc_official.py`; `model/prithvi_wxc.py` is only a YAML-driven thin wrapper (with identity normalization parameters for small-configuration connectivity validation).
+- Commit fetched for the current case directory: `79dabfcd17abe77e2d5c696707c0164a04f2ec01` (2026-02-05).
+- `conf/config.yaml` uses a small configuration (`embed_dim=32`, `n_blocks_encoder=1`, `n_blocks_decoder=1`) for connectivity validation by default; paper-level reproduction requires a 0.5°×0.625° grid, 160 channels, `embed_dim=2560`, and 13+12 encoder blocks/3+2 decoder blocks as specified in the paper.
+- The following details are not disclosed in the paper and are assumptions in this reproduction:data normalization statistics (identity normalization is currently used; real statistics will be supplied with the data), masked-training details and pretraining schedule, and some hyperparameters (such as the relative positional encoding implementation).
+
+# Official OneScience Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- This repository is an independent organization and adaptation of Prithvi WxC. The model source is based on the official `NASA-IMPACT/Prithvi-WxC` implementation by Schmude et al. (2024) and follows the MIT License.
+- Please cite:Schmude et al. Prithvi WxC: Foundation Model for Weather and Climate. arXiv:2409.13598, 2024.

conf/config.yaml

@@ -0,0 +1,73 @@
+# Prithvi WxC 训练配置示例
+# 论文配置为 0.5°×0.625°（360/361×576）网格、160 动态变量、embed_dim=2560、
+# 编码器 13 本地+12 全局块、解码器 3 本地+2 全局块（约 23 亿参数，需大规模显存）。
+# 当前为连通性验证小配置（约百万级参数）。
+model:
+  start_epoch: 0
+  max_epoch: 100
+  lr: 1E-3
+  patience: 50
+  checkpoint_dir: "./data/checkpoints"
+
+  # Prithvi WxC 结构参数
+  in_channels: 6                 # 论文为 160（20 单层 + 10 变量 ×14 层）
+  input_size_time: 2             # 两个输入时刻
+  in_channels_static: 4          # 静态通道数（论文 MEMRA-2 为 4）
+  n_lats_px: 32                  # 数据纬度方向像元数（论文约 360）
+  n_lons_px: 64                  # 数据经度方向像元数（论文 576）
+  patch_size_px: [2, 2]          # token 尺寸（论文 2×2 像素）
+  mask_unit_size_px: [8, 8]      # 掩码单元尺寸
+  mask_ratio_inputs: 0.0         # 预训练为 0.5；预报滚动微调为 0.0
+  embed_dim: 32                  # 隐藏维度（论文 2560）
+  n_blocks_encoder: 1            # 编码器本地-全局对数量（论文 13 对 -> 25 块）
+  n_blocks_decoder: 1            # 解码器本地-全局对数量（论文 5 块）
+  mlp_multiplier: 4.0
+  n_heads: 4                     # 注意力头数（论文 16）
+  dropout: 0.0
+  drop_path: 0.0
+  parameter_dropout: 0.0
+  residual: "none"               # none / temporal / climate
+  masking_mode: "global"
+  positional_encoding: "absolute"
+  encoder_shifting: false
+  decoder_shifting: false
+
+# 整个数据读取流程
+datapipe:
+  name: "ERA5"
+  task: "weather_forecasting"
+
+  dataset:
+    type: "hdf5"
+    data_dir: './data/'
+    train_time: [1951, 1952]
+    val_time: [1953]
+    test_time: [1954]
+    img_size: [32, 64]
+    verbose: true
+    cache: false
+
+    # 气象变量（论文 160 通道配置的子集，用于连通性验证）
+    channels: ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature',
+               'mean_sea_level_pressure', 'geopotential_500', 'temperature_850']
+
+  # DataLoader 配置
+  dataloader:
+    mask_dtype: "float32"
+    batch_size: 1
+    num_workers: 1
+    pin_memory: true
+    drop_last: true
+    shuffle: false
+    prefetch_factor: 2
+    persistent_workers: true
+
+  # 分布式配置
+  distributed:
+    enabled: true
+    sampler: "DistributedSampler"
+    rank: 0
+    world_size: 2
+    shuffle: true
+    seed: 42
+    drop_last: true

config.json

@@ -0,0 +1,116 @@
+{
+  "model_name": "Prithvi WxC",
+  "model_type": "prithvi_wxc",
+  "architectures": [
+    "PrithviWxC"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "global-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/prithvi_wxc.py",
+    "scope": "YAML-driven wrapper around the vendored official Hiera-MaxViT encoder-decoder, with identity model scalers for the repository's small connectivity configuration"
+  },
+  "architecture": {
+    "family": "Hiera-MaxViT encoder-decoder Vision Transformer",
+    "attention_mechanism": "alternating local block attention and global grid attention",
+    "input_format": "BTCHW atmospheric history plus BCHW static fields",
+    "output_format": "BCHW next atmospheric state",
+    "tokenization": "Conv2d patch embedding followed by mask-unit grouping and pixel-shuffle reconstruction",
+    "activation": "GELU",
+    "normalization": "LayerNorm",
+    "repository_default_config": {
+      "purpose": "small connectivity validation with synthetic or reduced ERA5 data",
+      "in_channels": 6,
+      "input_size_time": 2,
+      "in_channels_static": 4,
+      "grid_size": [
+        32,
+        64
+      ],
+      "patch_size_px": [
+        2,
+        2
+      ],
+      "mask_unit_size_px": [
+        8,
+        8
+      ],
+      "mask_ratio_inputs": 0.0,
+      "embed_dim": 32,
+      "n_blocks_encoder": 1,
+      "n_blocks_decoder": 1,
+      "n_heads": 4,
+      "mlp_multiplier": 4.0,
+      "dropout": 0.0,
+      "drop_path": 0.0,
+      "parameter_dropout": 0.0,
+      "residual": "none",
+      "masking_mode": "global",
+      "positional_encoding": "absolute",
+      "encoder_shifting": false,
+      "decoder_shifting": false,
+      "normalization_scalers": "identity in the YAML wrapper"
+    },
+    "paper_reference_config": {
+      "grid_resolution_degrees": [
+        0.5,
+        0.625
+      ],
+      "grid_size": "approximately 360/361 x 576",
+      "dynamic_channels": 160,
+      "embed_dim": 2560,
+      "encoder_blocks": "13 local plus 12 global blocks",
+      "decoder_blocks": "3 local plus 2 global blocks",
+      "n_heads": 16,
+      "pretraining_mask_ratio_inputs": 0.5,
+      "forecast_finetuning_mask_ratio_inputs": 0.0,
+      "parameter_count": "approximately 2.3 billion",
+      "compatibility_note": "official 2.3B checkpoints do not match the repository default small configuration without aligning channels, grid size, and architecture"
+    }
+  },
+  "data": {
+    "dataset": "ERA5",
+    "storage": "HDF5 fields with shape TCHW",
+    "temporal_interval_hours": 6,
+    "input_steps": 2,
+    "output_steps": 1,
+    "forecast_lead_time_hours": 6,
+    "dynamic_channels": [
+      "10m_u_component_of_wind",
+      "10m_v_component_of_wind",
+      "2m_temperature",
+      "mean_sea_level_pressure",
+      "geopotential_500",
+      "temperature_850"
+    ],
+    "static_channels": 4,
+    "default_spatial_size": [
+      32,
+      64
+    ],
+    "default_year_splits": {
+      "train": [
+        1951,
+        1952
+      ],
+      "validation": [
+        1953
+      ],
+      "test": [
+        1954
+      ]
+    },
+    "synthetic_data_generator": "scripts/fake_data.py"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/prithvi_wxc.py",
+    "model/prithvi_wxc_official.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/fake_data.py",
+    "README.md",
+    "configuration.json"
+  ]
+}

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/prithvi_wxc.py

@@ -0,0 +1,144 @@
+# coding=utf-8
+#
+# SPDX-License-Identifier: MIT
+#
+# Copyright (c) 2024 Inter Agency Implementation and Advanced Concepts
+#
+# This module is a thin configuration wrapper around the official Prithvi WxC
+# model (Schmude et al., arXiv:2409.13598), taken verbatim from the
+# NASA-IMPACT/Prithvi-WxC repository (commit 79dabfcd17abe77e2d5c696707c0164a04f2ec01,
+# MIT License). The official implementation is vendored as
+# ``prithvi_wxc_official.py``; only scaler construction and tensor-level call
+# signatures are added here for YAML-driven usage.
+import torch
+import torch.nn as nn
+
+from model.prithvi_wxc_official import PrithviWxC as _PrithviWxC
+
+
+def _identity_scalers(n, device="cpu"):
+    return torch.zeros(n, device=device), torch.ones(n, device=device)
+
+
+class PrithviWxC(nn.Module):
+    """
+    Config-driven Prithvi WxC wrapper.
+
+    The official model is an encoder-decoder vision Transformer (Hiera + MaxViT,
+    alternating local/global attention) that maps an input state history to a
+    target state. This wrapper builds the official network with identity
+    scalers so that a small YAML config can be exercised; loading the 2.3B
+    checkpoints published on Hugging Face is possible by calling
+    ``load_state_dict`` with the official checkpoint state (the buffer names,
+    e.g. ``input_scalers_mu``, are identical).
+
+    Args:
+        in_channels: number of input/output parameters.
+        input_size_time: number of input timestamps (paper uses 2).
+        in_channels_static: number of static input channels (>= 1).
+        n_lats_px / n_lons_px: grid size in pixels.
+        patch_size_px: patch/token size in pixels.
+        mask_unit_size_px: mask unit size in pixels.
+        mask_ratio_inputs: input masking ratio (0..1, 0 disables).
+        embed_dim / n_blocks_encoder / n_blocks_decoder / mlp_multiplier / n_heads:
+            transformer architecture hyper-parameters (paper: 2560 / 25 / 5 / 4 / 16).
+        dropout / drop_path / parameter_dropout: regularization ratios.
+        residual: one of "none" / "temporal" / "climate".
+        masking_mode: "global" / "local" / "both".
+        positional_encoding: "absolute" or "fourier".
+        encoder_shifting / decoder_shifting: Swin-style shifting.
+    """
+
+    def __init__(
+        self,
+        in_channels=4,
+        input_size_time=2,
+        in_channels_static=4,
+        n_lats_px=32,
+        n_lons_px=64,
+        patch_size_px=(2, 2),
+        mask_unit_size_px=(8, 8),
+        mask_ratio_inputs=0.0,
+        embed_dim=16,
+        n_blocks_encoder=1,
+        n_blocks_decoder=1,
+        mlp_multiplier=4.0,
+        n_heads=4,
+        dropout=0.0,
+        drop_path=0.0,
+        parameter_dropout=0.0,
+        residual="none",
+        masking_mode="global",
+        positional_encoding="absolute",
+        encoder_shifting=False,
+        decoder_shifting=False,
+    ):
+        super().__init__()
+        self.in_channels = int(in_channels)
+        self.input_size_time = int(input_size_time)
+        self.in_channels_static = int(in_channels_static)
+        self.n_lats_px = int(n_lats_px)
+        self.n_lons_px = int(n_lons_px)
+        self.residual = residual
+        self.positional_encoding = positional_encoding
+
+        mu_in, sigma_in = _identity_scalers(self.in_channels)
+        mu_s, sigma_s = _identity_scalers(self.in_channels_static)
+
+        self.model = _PrithviWxC(
+            in_channels=self.in_channels,
+            input_size_time=self.input_size_time,
+            in_channels_static=self.in_channels_static,
+            input_scalers_mu=mu_in,
+            input_scalers_sigma=sigma_in,
+            input_scalers_epsilon=1e-6,
+            static_input_scalers_mu=mu_s,
+            static_input_scalers_sigma=sigma_s,
+            static_input_scalers_epsilon=1e-6,
+            output_scalers=sigma_in,
+            n_lats_px=self.n_lats_px,
+            n_lons_px=self.n_lons_px,
+            patch_size_px=tuple(patch_size_px),
+            mask_unit_size_px=tuple(mask_unit_size_px),
+            mask_ratio_inputs=float(mask_ratio_inputs),
+            mask_ratio_targets=0.0,
+            embed_dim=int(embed_dim),
+            n_blocks_encoder=int(n_blocks_encoder),
+            n_blocks_decoder=int(n_blocks_decoder),
+            mlp_multiplier=int(mlp_multiplier),
+            n_heads=int(n_heads),
+            dropout=float(dropout),
+            drop_path=float(drop_path),
+            parameter_dropout=float(parameter_dropout),
+            residual=residual,
+            masking_mode=masking_mode,
+            positional_encoding=positional_encoding,
+            encoder_shifting=bool(encoder_shifting),
+            decoder_shifting=bool(decoder_shifting),
+        )
+
+    def forward(self, x, static, input_time=None, lead_time=None, climate=None):
+        """
+        Args:
+            x: Tensor of shape [batch, time, parameter, lat, lon].
+            static: Tensor of shape [batch, static_channel, lat, lon].
+            input_time: Tensor of shape [batch] (optional, default zeros).
+            lead_time: Tensor of shape [batch] (optional, default zeros).
+            climate: Optional Tensor of shape [batch, parameter, lat, lon].
+        Returns:
+            Tensor of shape [batch, parameter, lat, lon].
+        """
+        if input_time is None:
+            input_time = torch.zeros(x.shape[0], device=x.device)
+        if lead_time is None:
+            lead_time = torch.zeros(x.shape[0], device=x.device)
+        batch = {
+            "x": x,
+            "y": x[:, -1],
+            "static": static,
+            "input_time": input_time,
+            "lead_time": lead_time,
+        }
+        if climate is not None:
+            batch["climate"] = climate
+        return self.model(batch)

model/prithvi_wxc_official.py

@@ -0,0 +1,1637 @@
+from importlib.metadata import version
+TORCH_VERSION = version('torch')
+
+from functools import cached_property
+from typing import Optional
+import numpy as np
+
+import torch
+from torch import Tensor
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+if TORCH_VERSION > '2.3.0':
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+
+# DropPath code is straight from timm
+# (https://huggingface.co/spaces/Roll20/pet_score/blame/main/lib/timm/models/layers/drop.py)
+# Primarily since we currently don't have timm in the environment.
+def drop_path(
+    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
+):
+    """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 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None, scale_by_keep=True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+
+class Mlp(nn.Module):
+    """
+    Multi layer perceptron.
+    """
+
+    def __init__(
+        self, features: int, hidden_features: int, dropout: float = 0.0
+    ) -> None:
+        """
+        Args:
+            features: Input/output dimension.
+            hidden_features: Hidden dimension.
+            dropout: Dropout.
+        """
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(features, hidden_features),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_features, features),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            Tensor of shape [..., channel]
+        Returns:
+            Tensor of same shape as x.
+        """
+        return self.net(x)
+
+
+class LayerNormPassThrough(nn.LayerNorm):
+    """
+    Normalising layer that allows the attention mask to be passed through
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> tuple[Tensor, Tensor | None]:
+        """
+        Forwards function
+        Args:
+            d: tuple of the data tensor and the attention mask
+        Returns:
+            output: normalised output data
+            attn_mask: the attention mask that was passed in
+        """
+        input, attn_mask = d
+        output = F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+        return output, attn_mask
+
+
+class MultiheadAttention(nn.Module):
+    """
+    Multihead attention layer for inputs of shape [..., sequence, features].
+
+    Uses `scaled_dot_product_attention` to obtain a memory efficient attention
+    computation (https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html).
+    This follows:
+    - Dao et la. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
+        (https://arxiv.org/abs/2205.14135)
+    - Rabe, Staats "Self-attention Does Not Need O(n2) Memory" (https://arxiv.org/abs/2112.05682)
+
+    Note: Even though the documentation page for `scaled_dot_product_attention`
+    states that tensors can have any number of dimensions as long as the shapes
+    are along the lines of `(B, ..., S, E)`, the fused and memory efficient
+    mechanisms we enforce here require a 4D input. Some experimentatino shows
+    that this should be of shape `(B, H, S, E)`, where `H` represents heads.
+    However, as of right now this is not confirmed int he documentation.
+    """
+
+    def __init__(self, features: int, n_heads: int, dropout: float) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            n_heads: Number of attention heads. Should be a factor of features.
+                (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+        """
+        super().__init__()
+
+        if not (features % n_heads) == 0:
+            raise ValueError(
+                f"Number of features {features} is not divisible by number of heads {n_heads}."
+            )
+
+        self.features = features
+        self.n_heads = n_heads
+        self.dropout = dropout
+
+        self.qkv_layer = torch.nn.Linear(features, features * 3, bias=False)
+        self.w_layer = torch.nn.Linear(features, features, bias=False)
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
+        """
+        Args:
+            d: tuple containing Tensor of shape [..., sequence, features] and
+            the attention mask
+        Returns:
+            Tensor of shape [..., sequence, features]
+        """
+        x, attn_mask = d
+
+        if not x.shape[-1] == self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension of size {self.features}."
+            )
+
+        passenger_dims = x.shape[:-2]
+        B = passenger_dims.numel()
+        S = x.shape[-2]
+        C = x.shape[-1]
+        x = x.reshape(B, S, C)
+
+        # x [B, S, C]
+        # q, k, v [B, H, S, C/H]
+        q, k, v = (
+            self.qkv_layer(x)
+            .view(B, S, self.n_heads, 3 * (C // self.n_heads))
+            .transpose(1, 2)
+            .chunk(chunks=3, dim=3)
+        )
+
+        # Let us enforce either flash (A100+) or memory efficient attention,
+        # falling back to math on CPU / DCU(昇腾) shim where fused kernels are unavailable.
+        if TORCH_VERSION > '2.3.0':
+            with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]):
+                # x [B, H, S, C//H]
+                x = F.scaled_dot_product_attention(
+                    q, k, v, dropout_p=self.dropout
+                )
+        else:
+            with torch.backends.cuda.sdp_kernel(
+                enable_flash=True, enable_math=False, enable_mem_efficient=True
+            ):
+                # x [B, H, S, C//H]
+                x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout)
+
+        # x [B, S, C]
+        x = x.transpose(1, 2).reshape(B, S, C)
+
+        # x [B, S, C]
+        x = self.w_layer(x)
+
+        # Back to input shape
+        x = x.view(*passenger_dims, S, self.features)
+        return x
+
+
+class Transformer(nn.Module):
+    """
+    Transformer for inputs of shape [..., S, features].
+    """
+
+    def __init__(
+        self,
+        features: int,
+        mlp_multiplier: int,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+    ) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            mlp_multiplier: Model will use features*mlp_multiplier hidden units.
+            n_heads: Number of attention heads. Should be a factor of features.
+                (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+            drop_path: DropPath.
+        """
+        super().__init__()
+
+        self.features = features
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.attention = nn.Sequential(
+            LayerNormPassThrough(features),
+            MultiheadAttention(features, n_heads, dropout),
+        )
+
+        self.ff = nn.Sequential(
+            nn.LayerNorm(features),
+            Mlp(
+                features=features,
+                hidden_features=features * mlp_multiplier,
+                dropout=dropout,
+            ),
+        )
+
+    def forward(self, d: tuple[Tensor, Tensor | None]) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [..., sequence, features]
+        Returns:
+            Tensor of shape [..., sequence, features]
+        """
+        x, attn_mask = d
+        if not x.shape[-1] == self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension of size {self.features}."
+            )
+
+        attention_x = self.attention(d)
+
+        x = x + self.drop_path(attention_x)
+        x = x + self.drop_path(self.ff(x))
+
+        return x
+
+class _Shift(nn.Module):
+    """
+    Private base class for the shifter. This allows some behaviour to be easily
+    handled when the shifter isn't used.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self._shifted = False
+
+    @torch.no_grad()
+    def reset(self) -> None:
+        """
+        Resets the bool tracking whether the data is shifted
+        """
+        self._shifted: bool = False
+
+    def forward(self, data: Tensor) -> tuple[Tensor, dict[bool, None]]:
+        return data, {True: None, False: None}
+
+
+class SWINShift(_Shift):
+    """
+    Handles the shifting of patches similar to how SWIN works. However if we
+    shift the latitudes then the poles will wrap and potentially that might be
+    problematic. The possition tokens should handle it but masking is safer.
+    """
+
+    def __init__(
+        self,
+        mu_shape: tuple[int, int],
+        global_shape: tuple[int, int],
+        local_shape: tuple[int, int],
+        patch_shape: tuple[int, int],
+        n_context_tokens: int = 2,
+    ) -> None:
+        """
+        Args:
+            mu_shape: the shape to the masking units
+            global_shape: number of global patches in lat and lon
+            local_shape: size of the local patches
+            patch_shape: patch size
+            n_context_token: number of additional context tokens at start of _each_ local sequence
+        """
+        super().__init__()
+
+        self._mu_shape = ms = mu_shape
+        self._g_shape = gs = global_shape
+        self._l_shape = ls = local_shape
+        self._p_shape = ps = patch_shape
+        self._lat_patch = (gs[0], ls[0], gs[1], ls[1])
+        self._n_context_tokens = n_context_tokens
+
+        self._g_shift_to = tuple(int(0.5 * x / p) for x, p in zip(ms, ps))
+        self._g_shift_from = tuple(-int(0.5 * x / p) for x, p in zip(ms, ps))
+
+        # Define the attention masks for the shifted MaxViT.
+        nglobal = global_shape[0] * global_shape[1]
+        nlocal = local_shape[0] * local_shape[1] + self._n_context_tokens  # "+ 1" for leadtime
+
+        lm = torch.ones((nglobal, 1, nlocal, nlocal), dtype=bool)
+        mwidth = int(0.5 * local_shape[1]) * local_shape[0]
+        lm[
+            : gs[1],
+            :,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+        ] = False
+        self.register_buffer("local_mask", lm)
+
+        gm = torch.ones((nlocal, 1, nglobal, nglobal), dtype=bool)
+        gm[: int(0.5 * ls[1]) * ls[0], :, : gs[1], : gs[1]] = False
+        self.register_buffer("global_mask", gm)
+
+    def _to_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the global/local setting back to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        y1 = x.view(nbatch, *self._g_shape, *self._l_shape, -1)
+        y2 = y1.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = y2.shape
+        return y2.view((nbatch, -1, s[2] * s[3], s[4] * s[5]))
+
+    def _to_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the local/global setting to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the lat/lon setting.
+        """
+        x = x.transpose(2, 1).contiguous()
+        return self._to_grid_global(x)
+
+    def _from_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the global/local
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        z1 = x.view(nbatch, -1, *self._lat_patch)
+        z2 = z1.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = z2.shape
+        return z2.view(nbatch, s[1] * s[2], s[3] * s[4], -1)
+
+    def _from_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the local/global
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the local/global setting
+        """
+        x = self._from_grid_global(x)
+        return x.transpose(2, 1).contiguous()
+
+    def _shift(self, x: Tensor) -> Tensor:
+        """
+        Shifts data in the gridded lat/lon setting by half the mask unit shape
+        Args:
+            x: data to be shifted
+        Returns:
+            x: either the hsifted or unshifted data
+        """
+        shift = self._g_shift_from if self._shifted else self._g_shift_to
+        x_shifted = torch.roll(x, shift, (-2, -1))
+
+        self._shifted = not self._shifted
+        return x_shifted
+
+    def _sep_lt(self, x: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Seperate off the leadtime from the local patches
+        Args:
+            x: data to have leadtime removed from
+        Returns:
+            lt: leadtime
+            x: data without the lead time in the local patch
+        """
+        lt_it = x[:, : self._n_context_tokens, :, :]
+        x_stripped = x[:, self._n_context_tokens :, :, :]
+
+        return lt_it, x_stripped
+
+    def forward(self, data: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Shift or unshift the the data depending on whether the data is already
+        shifted, as defined by self._shifted
+        Args:
+            data: data to be shifted
+        Returns:
+
+        """
+        lt, x = self._sep_lt(data)
+
+        x_grid = self._to_grid_local(x)
+        x_shifted = self._shift(x_grid)
+        x_patched = self._from_grid_local(x_shifted)
+
+        # Mask has to be repeated based on batch size
+        n_batch = x_grid.shape[0]
+        local_rep = [n_batch] + [1] * (self.local_mask.ndim - 1)
+        global_rep = [n_batch] + [1] * (self.global_mask.ndim - 1)
+
+        if self._shifted:
+            attn_mask = {
+                True: self.local_mask.repeat(local_rep),
+                False: self.global_mask.repeat(global_rep),
+            }
+        else:
+            attn_mask = {True: None, False: None}
+
+        return torch.cat((lt, x_patched), axis=1), attn_mask
+
+class SWINShiftNoBuffer(_Shift):
+    """
+    Handles the shifting of patches similar to how SWIN works. However if we
+    shift the latitudes then the poles will wrap and potentially that might be
+    problematic. The possition tokens should handle it but masking is safer.
+    """
+
+    def __init__(
+        self,
+        mu_shape: tuple[int, int],
+        global_shape: tuple[int, int],
+        local_shape: tuple[int, int],
+        patch_shape: tuple[int, int],
+        n_context_tokens: int = 2,
+    ) -> None:
+        """
+        Args:
+            mu_shape: the shape to the masking units
+            global_shape: number of global patches in lat and lon
+            local_shape: size of the local patches
+            patch_shape: patch size
+            n_context_token: number of additional context tokens at start of _each_ local sequence
+        """
+        super().__init__()
+
+        self._mu_shape = ms = mu_shape
+        self._g_shape = gs = global_shape
+        self._l_shape = ls = local_shape
+        self._p_shape = ps = patch_shape
+        self._lat_patch = (gs[0], ls[0], gs[1], ls[1])
+        self._n_context_tokens = n_context_tokens
+
+        self._g_shift_to = tuple(int(0.5 * x / p) for x, p in zip(ms, ps))
+        self._g_shift_from = tuple(-int(0.5 * x / p) for x, p in zip(ms, ps))
+
+        # Define the attention masks for the shifted MaxViT.
+        nglobal = global_shape[0] * global_shape[1]
+        nlocal = local_shape[0] * local_shape[1] + self._n_context_tokens  # "+ 1" for leadtime
+
+        lm = torch.ones((nglobal, 1, nlocal, nlocal), dtype=bool)
+        mwidth = int(0.5 * local_shape[1]) * local_shape[0]
+        lm[
+            : gs[1],
+            :,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+            self._n_context_tokens : mwidth + self._n_context_tokens,
+        ] = False
+        self.local_mask = lm
+
+        gm = torch.ones((nlocal, 1, nglobal, nglobal), dtype=bool)
+        gm[: int(0.5 * ls[1]) * ls[0], :, : gs[1], : gs[1]] = False
+        self.global_mask = gm
+
+    def _to_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the global/local setting back to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        y1 = x.view(nbatch, *self._g_shape, *self._l_shape, -1)
+        y2 = y1.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = y2.shape
+        return y2.view((nbatch, -1, s[2] * s[3], s[4] * s[5]))
+
+    def _to_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the local/global setting to the
+        lat/lon grid setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the lat/lon setting.
+        """
+        x = x.transpose(2, 1).contiguous()
+        return self._to_grid_global(x)
+
+    def _from_grid_global(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the global/local
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the global/local setting
+        """
+        nbatch, *other = x.shape
+
+        z1 = x.view(nbatch, -1, *self._lat_patch)
+        z2 = z1.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = z2.shape
+        return z2.view(nbatch, s[1] * s[2], s[3] * s[4], -1)
+
+    def _from_grid_local(self, x: Tensor) -> Tensor:
+        """
+        Shuffle and reshape the data from the lat/lon grid to the local/global
+        setting
+        Args:
+            x: the data tensor to be shuffled.
+        Returns:
+            x: data in the local/global setting
+        """
+        x = self._from_grid_global(x)
+        return x.transpose(2, 1).contiguous()
+
+    def _shift(self, x: Tensor) -> Tensor:
+        """
+        Shifts data in the gridded lat/lon setting by half the mask unit shape
+        Args:
+            x: data to be shifted
+        Returns:
+            x: either the hsifted or unshifted data
+        """
+        shift = self._g_shift_from if self._shifted else self._g_shift_to
+        x_shifted = torch.roll(x, shift, (-2, -1))
+
+        self._shifted = not self._shifted
+        return x_shifted
+
+    def _sep_lt(self, x: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Seperate off the leadtime from the local patches
+        Args:
+            x: data to have leadtime removed from
+        Returns:
+            lt: leadtime
+            x: data without the lead time in the local patch
+        """
+        lt_it = x[:, : self._n_context_tokens, :, :]
+        x_stripped = x[:, self._n_context_tokens :, :, :]
+
+        return lt_it, x_stripped
+
+    def forward(self, data: Tensor) -> tuple[Tensor, Tensor]:
+        """
+        Shift or unshift the the data depending on whether the data is already
+        shifted, as defined by self._shifted
+        Args:
+            data: data to be shifted
+        Returns:
+
+        """
+        lt, x = self._sep_lt(data)
+
+        if self.local_mask.device != x.device:
+            self.local_mask = self.local_mask.to(device=x.device)
+        if self.global_mask.device != x.device:
+            self.global_mask = self.global_mask.to(device=x.device)
+
+        x_grid = self._to_grid_local(x)
+        x_shifted = self._shift(x_grid)
+        x_patched = self._from_grid_local(x_shifted)
+
+        # Mask has to be repeated based on batch size
+        n_batch = x_grid.shape[0]
+        local_rep = [n_batch] + [1] * (self.local_mask.ndim - 1)
+        global_rep = [n_batch] + [1] * (self.global_mask.ndim - 1)
+
+        if self._shifted:
+            attn_mask = {
+                True: self.local_mask.repeat(local_rep),
+                False: self.global_mask.repeat(global_rep),
+            }
+        else:
+            attn_mask = {True: None, False: None}
+
+        return torch.cat((lt, x_patched), axis=1), attn_mask
+
+class LocalGlobalLocalBlock(nn.Module):
+    """
+    Applies alternating block and grid attention. Given a parameter n_blocks, the entire
+    module contains 2*n_blocks+1 transformer blocks. The first, third, ..., last apply
+    local (block) attention. The second, fourth, ... global (grid) attention.
+
+    This is heavily inspired by Tu et al. "MaxViT: Multi-Axis Vision Transformer"
+    (https://arxiv.org/abs/2204.01697).
+    """
+
+    def __init__(
+        self,
+        features: int,
+        mlp_multiplier: int,
+        n_heads: int,
+        dropout: float,
+        n_blocks: int,
+        drop_path: float,
+        shifter: nn.Module | None = None,
+        checkpoint: list[int]=[],
+    ) -> None:
+        """
+        Args:
+            features: Number of features for inputs to the layer.
+            mlp_multiplier: Model will use features*mlp_multiplier hidden units.
+            n_heads: Number of attention heads. Should be a factor of features.
+            (I.e. the layer uses features // n_heads.)
+            dropout: Dropout.
+            drop_path: DropPath.
+            n_blocks: Number of local-global transformer pairs.
+        """
+        super().__init__()
+
+        self.features = features
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.n_blocks = n_blocks
+        self._checkpoint = checkpoint
+
+        if len(checkpoint) > 0:
+            if min(checkpoint) < 0 or max(checkpoint) >= 2 * n_blocks + 1:
+                raise ValueError(f'Checkpoints should satisfy 0 <= i < 2*n_blocks+1. We have {checkpoint}.')
+
+        self.transformers = nn.ModuleList(
+            [
+                Transformer(
+                    features=features,
+                    mlp_multiplier=mlp_multiplier,
+                    n_heads=n_heads,
+                    dropout=dropout,
+                    drop_path=drop_path,
+                )
+                for _ in range(2 * n_blocks + 1)
+            ]
+        )
+
+        self.evaluator = [
+            self._checkpoint_wrapper if i in checkpoint else lambda m, x : m(x)
+            for i, _ in enumerate(self.transformers)
+        ]
+
+        self.shifter = shifter or _Shift()
+
+    @staticmethod
+    def _checkpoint_wrapper(model, data):
+        return checkpoint(model, data, use_reentrant=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch, global_sequence, local_sequence, features]
+        Returns:
+            Tensor of shape [batch, global_sequence, local_sequence, features]
+        """
+        if x.shape[-1] != self.features:
+            raise ValueError(
+                f"Expecting tensor with last dimension of size {self.features}."
+            )
+        if x.ndim != 4:
+            raise ValueError(
+                f"Expecting tensor with exactly four dimensions. Input has shape {x.shape}."
+            )
+
+        self.shifter.reset()
+        local: bool = True
+        attn_mask = {True: None, False: None}
+
+        transformer_iter = zip(self.evaluator, self.transformers)
+
+        # First local block
+        evaluator, transformer = next(transformer_iter)
+        x = evaluator(transformer, (x, attn_mask[local]))
+
+        for evaluator, transformer in transformer_iter:
+            local = not local
+            # We are making exactly 2*n_blocks transposes.
+            # So the output has the same shape as input.
+            x = x.transpose(1, 2)
+
+            x = evaluator(transformer, (x, attn_mask[local]))
+
+            if not local:
+                x, attn_mask = self.shifter(x)
+
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch embedding via 2D convolution.
+    """
+
+    def __init__(
+        self, patch_size: int | tuple[int, ...], channels: int, embed_dim: int
+    ):
+        super().__init__()
+
+        self.patch_size = patch_size
+        self.channels = channels
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(
+            channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch, channels, lat, lon].
+        Returns:
+            Tensor with shape [batch, embed_dim, lat//patch_size, lon//patch_size]
+        """
+
+        H, W = x.shape[-2:]
+
+        if W % self.patch_size[1] != 0:
+            raise ValueError(
+                f"Cannot do patch embedding for tensor of shape {x.size()}"
+                " with patch size {self.patch_size}. (Dimensions are BSCHW.)"
+            )
+        if H % self.patch_size[0] != 0:
+            raise ValueError(
+                f"Cannot do patch embedding for tensor of shape {x.size()}"
+                f" with patch size {self.patch_size}. (Dimensions are BSCHW.)"
+            )
+
+        x = self.proj(x)
+
+        return x
+
+
+class PrithviWxCEncoderDecoder(nn.Module):
+    """
+    Hiera-MaxViT encoder/decoder code.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        n_blocks: int,
+        mlp_multiplier: float,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+        shifter: nn.Module | None = None,
+        transformer_cp: list[int]=[],
+    ) -> None:
+        """
+        Args:
+            embed_dim: Embedding dimension
+            n_blocks: Number of local-global transformer pairs.
+            mlp_multiplier: MLP multiplier for hidden features in feed forward
+                networks.
+            n_heads: Number of attention heads.
+            dropout: Dropout.
+            drop_path: DropPath.
+        """
+        super().__init__()
+
+        self.embed_dim = embed_dim
+        self.n_blocks = n_blocks
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self._transformer_cp = transformer_cp
+
+        self.lgl_block = LocalGlobalLocalBlock(
+            features=embed_dim,
+            mlp_multiplier=mlp_multiplier,
+            n_heads=n_heads,
+            dropout=dropout,
+            drop_path=drop_path,
+            n_blocks=n_blocks,
+            shifter=shifter,
+            checkpoint=transformer_cp,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape [batch, global sequence, local sequence, embed_dim]
+        Returns:
+            Tensor of shape [batch, mask_unit_sequence, local_sequence, embed_dim].
+            Identical in shape to the input x.
+        """
+
+        x = self.lgl_block(x)
+
+        return x
+
+
+class PrithviWxC(nn.Module):
+    """
+    Encoder-decoder fusing Hiera with MaxViT. See
+    - Ryali et al. "Hiera: A Hierarchical Vision Transformer without the
+        Bells-and-Whistles" (https://arxiv.org/abs/2306.00989)
+    - Tu et al. "MaxViT: Multi-Axis Vision Transformer"
+        (https://arxiv.org/abs/2204.01697)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        input_size_time: int,
+        in_channels_static: int,
+        input_scalers_mu: Tensor,
+        input_scalers_sigma: Tensor,
+        input_scalers_epsilon: float,
+        static_input_scalers_mu: Tensor,
+        static_input_scalers_sigma: Tensor,
+        static_input_scalers_epsilon: float,
+        output_scalers: Tensor,
+        n_lats_px: int,
+        n_lons_px: int,
+        patch_size_px: tuple[int],
+        mask_unit_size_px: tuple[int],
+        mask_ratio_inputs: float,
+        mask_ratio_targets: float,
+        embed_dim: int,
+        n_blocks_encoder: int,
+        n_blocks_decoder: int,
+        mlp_multiplier: float,
+        n_heads: int,
+        dropout: float,
+        drop_path: float,
+        parameter_dropout: float,
+        residual: str,
+        masking_mode: str,
+        positional_encoding: str,
+        encoder_shifting: bool = False,
+        decoder_shifting: bool = False,
+        checkpoint_encoder: list[int]=[],
+        checkpoint_decoder: list[int]=[],
+    ) -> None:
+        """
+        Args:
+            in_channels: Number of input channels.
+            input_size_time: Number of timestamps in input.
+            in_channels_static: Number of input channels for static data.
+            input_scalers_mu: Tensor of size (in_channels,). Used to rescale
+                input.
+            input_scalers_sigma: Tensor of size (in_channels,). Used to rescale
+                input.
+            input_scalers_epsilon: Float. Used to rescale input.
+            static_input_scalers_mu: Tensor of size (in_channels_static). Used
+                to rescale static inputs.
+            static_input_scalers_sigma: Tensor of size (in_channels_static).
+                Used to rescale static inputs.
+            static_input_scalers_epsilon: Float. Used to rescale static inputs.
+            output_scalers: Tensor of shape (in_channels,). Used to rescale
+                output.
+            n_lats_px: Total latitudes in data. In pixels.
+            n_lons_px: Total longitudes in data. In pixels.
+            patch_size_px: Patch size for tokenization. In pixels lat/lon.
+            mask_unit_size_px: Size of each mask unit. In pixels lat/lon.
+            mask_ratio_inputs: Masking ratio for inputs. 0 to 1.
+            mask_ratio_targets: Masking ratio for targets. 0 to 1.
+            embed_dim: Embedding dimension
+            n_blocks_encoder: Number of local-global transformer pairs in
+                encoder.
+            n_blocks_decoder: Number of local-global transformer pairs in
+                decoder.
+            mlp_multiplier: MLP multiplier for hidden features in feed forward
+                networks.
+            n_heads: Number of attention heads.
+            dropout: Dropout.
+            drop_path: DropPath.
+            parameter_dropout: Dropout applied to parameters.
+            residual: Indicates whether and how model should work as residual
+                model. Accepted values are 'climate', 'temporal' and 'none'
+            positional_encoding: possible values are ['absolute' (default), 'fourier'].
+                'absolute'  lat lon encoded in 3 dimensions using sine and cosine
+                'fourier' lat/lon to be encoded using various frequencies
+            masking_mode: String ['local', 'global', 'both'] that controls the
+                type of masking used.
+            checkpoint_encoder: List of integers controlling if gradient checkpointing is used on encoder.
+                Format: [] for no gradient checkpointing. [3, 7] for checkpointing after 4th and 8th layer etc.
+            checkpoint_decoder: List of integers controlling if gradient checkpointing is used on decoder.
+                Format: See `checkpoint_encoder`.
+            masking_mode: The type of masking to use {'global', 'local', 'both'}
+            encoder_shifting: Whether to use swin shifting in the encoder.
+            decoder_shifting: Whether to use swin shifting in the decoder.
+        """
+        super().__init__()
+
+        if mask_ratio_targets > 0.0:
+            raise NotImplementedError("Target masking is not implemented.")
+
+        self.in_channels = in_channels
+        self.input_size_time = input_size_time
+        self.in_channels_static = in_channels_static
+        self.n_lats_px = n_lats_px
+        self.n_lons_px = n_lons_px
+        self.patch_size_px = patch_size_px
+        self.mask_unit_size_px = mask_unit_size_px
+        self.mask_ratio_inputs = mask_ratio_inputs
+        self.mask_ratio_targets = mask_ratio_targets
+        self.embed_dim = embed_dim
+        self.n_blocks_encoder = n_blocks_encoder
+        self.n_blocks_decoder = n_blocks_decoder
+        self.mlp_multiplier = mlp_multiplier
+        self.n_heads = n_heads
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.residual = residual
+        self._encoder_shift = encoder_shifting
+        self._decoder_shift = decoder_shifting
+        self.positional_encoding = positional_encoding
+        self._checkpoint_encoder = checkpoint_encoder
+        self._checkpoint_decoder = checkpoint_decoder
+
+        assert self.n_lats_px % self.mask_unit_size_px[0] == 0
+        assert self.n_lons_px % self.mask_unit_size_px[1] == 0
+        assert self.mask_unit_size_px[0] % self.patch_size_px[0] == 0
+        assert self.mask_unit_size_px[1] % self.patch_size_px[1] == 0
+
+        if self.patch_size_px[0] != self.patch_size_px[1]:
+            raise NotImplementedError(
+                "Current pixel shuffle implementation assumes same patch size along both dimensions."
+            )
+
+        self.local_shape_mu = (
+            self.mask_unit_size_px[0] // self.patch_size_px[0],
+            self.mask_unit_size_px[1] // self.patch_size_px[1],
+        )
+        self.global_shape_mu = (
+            self.n_lats_px // self.mask_unit_size_px[0],
+            self.n_lons_px // self.mask_unit_size_px[1],
+        )
+
+        assert input_scalers_mu.shape == (in_channels,)
+        assert input_scalers_sigma.shape == (in_channels,)
+        assert output_scalers.shape == (in_channels,)
+
+        if self.positional_encoding != 'fourier':
+            assert static_input_scalers_mu.shape == (in_channels_static,)
+            assert static_input_scalers_sigma.shape == (in_channels_static,)
+
+        # Input shape [batch, time, parameter, lat, lon]
+        self.input_scalers_epsilon = input_scalers_epsilon
+        self.register_buffer('input_scalers_mu', input_scalers_mu.reshape(1, 1, -1, 1, 1))
+        self.register_buffer('input_scalers_sigma', input_scalers_sigma.reshape(1, 1, -1, 1, 1))
+
+        # Static inputs shape [batch, parameter, lat, lon]
+        self.static_input_scalers_epsilon = static_input_scalers_epsilon
+        self.register_buffer('static_input_scalers_mu', static_input_scalers_mu.reshape(1, -1, 1, 1))
+        self.register_buffer('static_input_scalers_sigma', static_input_scalers_sigma.reshape(1, -1, 1, 1))
+
+        # Output shape [batch, parameter, lat, lon]
+        self.register_buffer('output_scalers', output_scalers.reshape(1, -1, 1, 1))
+
+        self.parameter_dropout = nn.Dropout2d(p=parameter_dropout)
+
+        self.patch_embedding = PatchEmbed(
+            patch_size=patch_size_px,
+            channels=in_channels * input_size_time,
+            embed_dim=embed_dim,
+        )
+
+        if self.residual == "climate":
+            self.patch_embedding_static = PatchEmbed(
+                patch_size=patch_size_px,
+                channels=in_channels + in_channels_static,
+                embed_dim=embed_dim,
+            )
+        else:
+            self.patch_embedding_static = PatchEmbed(
+                patch_size=patch_size_px,
+                channels=in_channels_static,
+                embed_dim=embed_dim,
+            )
+
+        self.input_time_embedding = nn.Linear(1, embed_dim//4, bias=True)
+        self.lead_time_embedding = nn.Linear(1, embed_dim//4, bias=True)
+
+        self.mask_token = nn.Parameter(torch.randn(1, 1, 1, self.embed_dim))
+        self._nglobal_mu = np.prod(self.global_shape_mu)
+        self._global_idx = torch.arange(self._nglobal_mu)
+
+        self._nlocal_mu = np.prod(self.local_shape_mu)
+        self._local_idx = torch.arange(self._nlocal_mu)
+
+        if self._encoder_shift:
+            self.encoder_shifter = e_shifter = SWINShiftNoBuffer(
+                self.mask_unit_size_px,
+                self.global_shape_mu,
+                self.local_shape_mu,
+                self.patch_size_px,
+                n_context_tokens=0,
+            )
+        else:
+            self.encoder_shifter = e_shifter = None
+        self.encoder = PrithviWxCEncoderDecoder(
+            embed_dim=embed_dim,
+            n_blocks=n_blocks_encoder,
+            mlp_multiplier=mlp_multiplier,
+            n_heads=n_heads,
+            dropout=dropout,
+            drop_path=drop_path,
+            shifter=e_shifter,
+            transformer_cp=checkpoint_encoder,
+        )
+
+        if n_blocks_decoder != 0:
+            if self._decoder_shift:
+                self.decoder_shifter = d_shifter = SWINShift(
+                    self.mask_unit_size_px,
+                    self.global_shape_mu,
+                    self.local_shape_mu,
+                    self.patch_size_px,
+                    n_context_tokens=0,
+                )
+            else:
+                self.decoder_shifter = d_shifter = None
+
+            self.decoder = PrithviWxCEncoderDecoder(
+                embed_dim=embed_dim,
+                n_blocks=n_blocks_decoder,
+                mlp_multiplier=mlp_multiplier,
+                n_heads=n_heads,
+                dropout=dropout,
+                drop_path=0.,
+                shifter=d_shifter,
+                transformer_cp=checkpoint_decoder,
+            )
+
+            self.unembed = nn.Linear(
+                self.embed_dim,
+                self.in_channels * self.patch_size_px[0] * self.patch_size_px[1],
+                bias=True,
+            )
+
+        self.masking_mode = masking_mode.lower()
+        match self.masking_mode:
+            case "local":
+                self.generate_mask = self._gen_mask_local
+            case "global":
+                self.generate_mask = self._gen_mask_global
+            case "both":
+                self._mask_both_local: bool = True
+                self.generate_mask = self._gen_mask_both
+            case _:
+                raise ValueError(f"Masking mode '{masking_mode}' not supported")
+
+    def swap_masking(self) -> None:
+        if hasattr(self, '_mask_both_local'):
+            self._mask_both_local = not self._mask_both_local
+
+    @cached_property
+    def n_masked_global(self):
+        return int(self.mask_ratio_inputs * np.prod(self.global_shape_mu))
+
+    @cached_property
+    def n_masked_local(self):
+        return int(self.mask_ratio_inputs * np.prod(self.local_shape_mu))
+
+    @staticmethod
+    def _shuffle_along_axis(a, axis):
+        # https://stackoverflow.com/questions/5040797/shuffling-numpy-array-along-a-given-axis
+        idx = torch.argsort(input=torch.rand(*a.shape), dim=axis)
+        return torch.gather(a, dim=axis, index=idx)
+
+    def _gen_mask_local(self, sizes: tuple[int]) -> tuple[Tensor]:
+        """
+        Args:
+            batch_size: Number of elements in batch
+        Returns:
+            Tuple of torch tensors. [indices masked, indices unmasked].
+            Each of these is a tensor of shape (batch, global sequene)
+        """
+        # We identifies which indices (values) should be masked
+
+        maskable_indices = self._local_idx.view(1, -1).expand(*sizes[:2], -1)
+
+        maskable_indices = self._shuffle_along_axis(maskable_indices, 2)
+
+        # `...` cannot be jit'd :-(
+        indices_masked = maskable_indices[:, :, : self.n_masked_local]
+        indices_unmasked = maskable_indices[:, :, self.n_masked_local :]
+
+        return indices_masked, indices_unmasked
+
+    def _gen_mask_global(self, sizes: tuple[int]) -> tuple[Tensor]:
+        """
+        Args:
+            batch_size: Number of elements in batch
+        Returns:
+            Tuple of torch tensors. [indices masked, indices unmasked].
+            Each of these is a tensor of shape (batch, global sequene)
+        """
+        # We identifies which indices (values) should be masked
+
+        maskable_indices = self._global_idx.view(1, -1).expand(*sizes[:1], -1)
+
+        maskable_indices = self._shuffle_along_axis(maskable_indices, 1)
+
+        indices_masked = maskable_indices[:, : self.n_masked_global]
+        indices_unmasked = maskable_indices[:, self.n_masked_global :]
+
+        return indices_masked, indices_unmasked
+
+    def _gen_mask_both(self, sizes: tuple[int]) -> tuple[Tensor]:
+        if self._mask_both_local:
+            return self._gen_mask_local(sizes)
+        else:
+            return self._gen_mask_global(sizes)
+
+    @staticmethod
+    def reconstruct_batch(
+        idx_masked: Tensor,
+        idx_unmasked: Tensor,
+        data_masked: Tensor,
+        data_unmasked: Tensor,
+    ) -> Tensor:
+        """
+        Reconstructs a tensor along the mask unit dimension. Batched version.
+
+        Args:
+            idx_masked: Tensor of shape `batch, mask unit sequence`.
+            idx_unmasked: Tensor of shape `batch, mask unit sequence`.
+            data_masked: Tensor of shape `batch, mask unit sequence, ...`.
+                Should have same size along mask unit sequence dimension as
+                idx_masked. Dimensions beyond the first two, marked here as ...
+                will typically be `local_sequence, channel` or `channel, lat, lon`.
+                  These dimensions should agree with data_unmasked.
+            data_unmasked: Tensor of shape `batch, mask unit sequence, ...`.
+                Should have same size along mask unit sequence dimension as
+                idx_unmasked. Dimensions beyond the first two, marked here as
+                ... will typically be `local_sequence, channel` or `channel,
+                lat, lon`. These dimensions should agree with data_masked.
+        Returns:
+            Tensor of same shape as inputs data_masked and data_unmasked. I.e.
+            `batch, mask unit sequence, ...`. Index for the total data composed
+            of the masked and the unmasked part
+        """
+        dim: int = idx_masked.ndim
+
+        idx_total = torch.argsort(torch.cat([idx_masked, idx_unmasked], dim=-1), dim=-1)
+        idx_total = idx_total.view(*idx_total.shape, *[1] * (data_unmasked.ndim - dim))
+        idx_total = idx_total.expand(*idx_total.shape[:dim], *data_unmasked.shape[dim:])
+
+        data = torch.cat([data_masked, data_unmasked], dim=dim - 1)
+        data = torch.gather(data, dim=dim - 1, index=idx_total)
+
+        return data, idx_total
+
+    def fourier_pos_encoding(self, x_static):
+        """
+        Args
+            x_static: B x C x H x W. first two channels are lat, and lon respectively
+        Returns
+            Tensor of shape B x E x H x W where E is the embedding dimension.
+        """
+
+        # B x C x H x W -> B x 1 x H/P x W/P
+        latitudes_patch = F.avg_pool2d(x_static[:, [0]], kernel_size=self.patch_size_px, stride=self.patch_size_px)
+        longitudes_patch = F.avg_pool2d(x_static[:, [1]], kernel_size=self.patch_size_px, stride=self.patch_size_px)
+
+        modes = torch.arange(self.embed_dim//4, device=x_static.device).view(1, -1, 1, 1) + 1.
+        pos_encoding = torch.cat(
+            (
+                torch.sin(latitudes_patch*modes),
+                torch.sin(longitudes_patch*modes),
+                torch.cos(latitudes_patch*modes),
+                torch.cos(longitudes_patch*modes),
+            ),
+            axis=1
+        )
+
+        return pos_encoding # B x E x H/P x W/P
+
+    def time_encoding(self, input_time, lead_time):
+        '''
+        Args:
+            input_time: Tensor of shape [batch].
+            lead_time: Tensor of shape [batch].
+        Returns:
+            Tensor of shape [batch, embed_dim, 1, 1]
+        '''
+        input_time = self.input_time_embedding(input_time.view(-1, 1, 1, 1))
+        lead_time = self.lead_time_embedding(lead_time.view(-1, 1, 1, 1))
+
+        time_encoding = torch.cat(
+            (
+                torch.cos(input_time),
+                torch.cos(lead_time),
+                torch.sin(input_time),
+                torch.sin(lead_time),
+            ),
+            axis=3
+        )
+        return time_encoding
+
+    def to_patching(self, x: Tensor) -> Tensor:
+        """Transform data from lat/lon space to two axis patching
+
+        Args: ->
+            x: Tesnor in lat/lon space (N, C, Nlat//P_0, Nlon//P_1)
+
+        Returns:
+            Tensor in patch space (N, G, L, C)
+        """
+        n_batch = x.shape[0]
+
+        x = x.view(
+            n_batch,
+            self.embed_dim,
+            self.global_shape_mu[0],
+            self.local_shape_mu[0],
+            self.global_shape_mu[1],
+            self.local_shape_mu[1],
+        )
+        x = x.permute(0, 2, 4, 3, 5, 1).contiguous()
+
+        s = x.shape
+        return x.view(n_batch, s[1] * s[2], s[3] * s[4], -1)
+
+    def from_patching(self, x: Tensor) -> Tensor:
+        """Transform data from two axis patching to lat/lon space
+
+        Args:
+            x: Tensor in patch space with shape (N, G, L, C*P_0*P_1)
+
+        Returns:
+            Tensor in lat/lon space (N, C*P_0*P_1, Nlat//P_0, Nlon // P_1)
+        """
+        n_batch = x.shape[0]
+
+        x = x.view(
+            n_batch,
+            self.global_shape_mu[0],
+            self.global_shape_mu[1],
+            self.local_shape_mu[0],
+            self.local_shape_mu[1],
+            -1,
+        )
+        x = x.permute(0, 5, 1, 3, 2, 4).contiguous()
+
+        s = x.shape
+        return x.view(n_batch, -1, s[2]*s[3], s[4]*s[5])
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+        """
+        Args:
+            batch: Dictionary containing the keys 'x', 'y', 'input_time',
+                'lead_time' and 'static'. The associated torch tensors have the
+                following shapes:
+                x: Tensor of shape [batch, time, parameter, lat, lon]
+                y: Tensor of shape [batch, parameter, lat, lon]
+                static: Tensor of shape [batch, channel_static, lat, lon]
+                climate: Optional tensor of shape [batch, parameter, lat, lon]
+                input_time: Tensor of shape [batch]. Or none.
+                lead_time: Tensor of shape [batch]. Or none.
+        Returns:
+            Tensor of shape [batch, parameter, lat, lon].
+        """
+        assert batch["x"].shape[2] == self.in_channels
+        assert batch["x"].shape[3] == self.n_lats_px
+        assert batch["x"].shape[4] == self.n_lons_px
+        assert batch["y"].shape[1] == self.in_channels
+        assert batch["y"].shape[2] == self.n_lats_px
+        assert batch["y"].shape[3] == self.n_lons_px
+        if self.positional_encoding == 'fourier':
+            # the first two features (lat, lon) are encoded separately
+            assert batch['static'].shape[1] - 2 == self.in_channels_static, "When setting self.positional_encoding to fourier, the number of static params change in the dataset. So, in the config, reduce num_static_channels (e.g., 4 instead of 7)."
+        else:
+            assert batch['static'].shape[1] == self.in_channels_static
+        assert batch["static"].shape[2] == self.n_lats_px
+        assert batch["static"].shape[3] == self.n_lons_px
+
+        x_rescaled = (batch["x"] - self.input_scalers_mu) / (
+            self.input_scalers_sigma + self.input_scalers_epsilon
+        )
+        batch_size = x_rescaled.shape[0]
+
+        if self.positional_encoding == 'fourier':
+            x_static_pos = self.fourier_pos_encoding(batch['static']) # B, embed_dim, lat / patch_size, lon / patch_size
+            x_static = (batch['static'][:, 2:] - self.static_input_scalers_mu[:, 3:]) / ( # The first two channels in batch['static'] are used in positional encoding
+                self.static_input_scalers_sigma[:, 3:] + self.static_input_scalers_epsilon # This translates to the first three channels in 'static_input_scalers_mu'
+            )
+        else:
+            x_static = (batch["static"] - self.static_input_scalers_mu) / (
+                self.static_input_scalers_sigma + self.static_input_scalers_epsilon
+            )
+
+        if self.residual == "temporal":
+            # We create a residual of same shape as y
+            index = torch.where(batch["lead_time"] > 0, batch["x"].shape[1] - 1, 0)
+            index = index.view(-1, 1, 1, 1, 1)
+            index = index.expand(batch_size, 1, *batch["x"].shape[2:])
+            x_hat = torch.gather(batch["x"], dim=1, index=index)
+            x_hat = x_hat.squeeze(1)
+            assert (
+                batch["y"].shape == x_hat.shape
+            ), f'Shapes {batch["y"].shape} and {x_hat.shape} do not agree.'
+        elif self.residual == "climate":
+            climate_scaled = (
+                batch["climate"] - self.input_scalers_mu.view(1, -1, 1, 1)
+            ) / (
+                self.input_scalers_sigma.view(1, -1, 1, 1) + self.input_scalers_epsilon
+            )
+
+        # [batch, time, parameter, lat, lon] -> [batch, time x parameter, lat, lon]
+        x_rescaled = x_rescaled.flatten(1, 2)
+        # Parameter dropout
+        x_rescaled = self.parameter_dropout(x_rescaled)
+
+        x_embedded = self.patch_embedding(x_rescaled)
+        assert x_embedded.shape[1] == self.embed_dim
+
+        if self.residual == "climate":
+            static_embedded = self.patch_embedding_static(
+                torch.cat((x_static, climate_scaled), dim=1)
+            )
+        else:
+            static_embedded = self.patch_embedding_static(x_static)
+        assert static_embedded.shape[1] == self.embed_dim
+
+        if self.positional_encoding == 'fourier':
+            static_embedded += x_static_pos
+
+        x_embedded = self.to_patching(x_embedded)
+        static_embedded = self.to_patching(static_embedded)
+
+        time_encoding = self.time_encoding(batch['input_time'], batch['lead_time'])
+
+        tokens = x_embedded + static_embedded + time_encoding
+
+        # Now we generate masks based on masking_mode
+        indices_masked, indices_unmasked = self.generate_mask(
+            (batch_size, self._nglobal_mu)
+        )
+        indices_masked = indices_masked.to(device=tokens.device)
+        indices_unmasked = indices_unmasked.to(device=tokens.device)
+        maskdim: int = indices_masked.ndim
+
+        # Unmasking
+        unmask_view = (*indices_unmasked.shape, *[1] * (tokens.ndim - maskdim))
+        unmasked = torch.gather(
+            tokens,
+            dim=maskdim - 1,
+            index=indices_unmasked.view(*unmask_view).expand(
+                *indices_unmasked.shape, *tokens.shape[maskdim:]
+            ),
+        )
+
+        # Encoder
+        x_encoded = self.encoder(unmasked)
+
+        # Generate and position encode the mask tokens
+        # (1, 1, 1, embed_dim) -> (batch, global_seq_masked, local seq, embed_dim)
+        mask_view = (*indices_masked.shape, *[1] * (tokens.ndim - maskdim))
+        masking = self.mask_token.repeat(*static_embedded.shape[:3], 1)
+        masked = masking + static_embedded
+        masked = torch.gather(
+            masked,
+            dim=maskdim - 1,
+            index=indices_masked.view(*mask_view).expand(
+                *indices_masked.shape, *tokens.shape[maskdim:]
+            ),
+        )
+
+        recon, _ = self.reconstruct_batch(
+            indices_masked, indices_unmasked, masked, x_encoded
+        )
+
+        x_decoded = self.decoder(recon)
+
+        # Output: (batch, global sequence, local sequence, in_channels * patch_size[0] * patch_size[1])
+        x_unembed = self.unembed(x_decoded)
+
+        # Reshape to (batch, global_lat, global_lon, local_lat, local_lon, in_channels * patch_size[0] * patch_size[1])
+        assert x_unembed.shape[0] == batch_size
+        assert x_unembed.shape[1] == self.global_shape_mu[0] * self.global_shape_mu[1]
+        assert x_unembed.shape[2] == self.local_shape_mu[0] * self.local_shape_mu[1]
+        assert (
+            x_unembed.shape[3]
+            == self.in_channels * self.patch_size_px[0] * self.patch_size_px[1]
+        )
+
+        x_out = self.from_patching(x_unembed)
+
+        # Pixel shuffle to (batch, in_channels, lat, lon)
+        x_out = F.pixel_shuffle(x_out, self.patch_size_px[0])
+
+        if self.residual == "temporal":
+            x_out = self.output_scalers * x_out + x_hat
+        elif self.residual == "climate":
+            x_out = self.output_scalers * x_out + batch["climate"]
+        elif self.residual == "none":
+            x_out = self.output_scalers * x_out + self.input_scalers_mu.reshape(
+                1, -1, 1, 1
+            )
+
+        return x_out
+
+
+def masked_inference(
+        model,
+        batch: dict[str, torch.Tensor],
+        mask_global: Optional[torch.Tensor] = None,
+        mask_local: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Perform masked inference with Prithvi-WxC model
+
+    Args:
+        model: The PrithviWxC model to perform the inference with.
+        batch: Dictionary containing the keys 'x', 'y', 'input_time',
+            'lead_time' and 'static'. The associated torch tensors have the
+            following shapes:
+            x: Tensor of shape [batch, time, parameter, lat, lon]
+            y: Tensor of shape [batch, parameter, lat, lon]
+            static: Tensor of shape [batch, channel_static, lat, lon]
+            climate: Optional tensor of shape [batch, parameter, lat, lon]
+            input_time: Tensor of shape [batch]. Or none.
+            lead_time: Tensor of shape [batch]. Or none.
+        mask_global: A 2D tensor of shape [n_lats_px / mask_unit_size_px[0], n_lons_px / patch_size_px[1]]
+            identifying the global patches to mask.
+        mask_local: A 2D tensor of shape [mask_unit_size_px[0] / patch_size_px[0], mask_unit_size_px[1] / patch_size_px[1]]
+            identifying the global patches to mask.
+    Returns:
+        Tensor of shape [batch, parameter, lat, lon].
+    """
+    assert batch["x"].shape[2] == model.in_channels
+    assert batch["x"].shape[3] == model.n_lats_px
+    assert batch["x"].shape[4] == model.n_lons_px
+    assert batch["y"].shape[1] == model.in_channels
+    assert batch["y"].shape[2] == model.n_lats_px
+    assert batch["y"].shape[3] == model.n_lons_px
+    if model.positional_encoding == 'fourier':
+        # the first two features (lat, lon) are encoded separately
+        assert batch['static'].shape[1] - 2 == model.in_channels_static, "When setting model.positional_encoding to fourier, the number of static params change in the dataset. So, in the config, reduce num_static_channels (e.g., 4 instead of 7)."
+    else:
+        assert batch['static'].shape[1] == model.in_channels_static
+    assert batch["static"].shape[2] == model.n_lats_px
+    assert batch["static"].shape[3] == model.n_lons_px
+
+    x_rescaled = (batch["x"] - model.input_scalers_mu) / (
+        model.input_scalers_sigma + model.input_scalers_epsilon
+    )
+    batch_size = x_rescaled.shape[0]
+
+    if model.positional_encoding == 'fourier':
+        x_static_pos = model.fourier_pos_encoding(batch['static']) # B, embed_dim, lat / patch_size, lon / patch_size
+        x_static = (batch['static'][:, 2:] - model.static_input_scalers_mu[:, 3:]) / ( # The first two channels in batch['static'] are used in positional encoding
+            model.static_input_scalers_sigma[:, 3:] + model.static_input_scalers_epsilon # This translates to the first three channels in 'static_input_scalers_mu'
+        )
+    else:
+        x_static = (batch["static"] - model.static_input_scalers_mu) / (
+            model.static_input_scalers_sigma + model.static_input_scalers_epsilon
+        )
+
+    if model.residual == "temporal":
+        # We create a residual of same shape as y
+        index = torch.where(batch["lead_time"] > 0, batch["x"].shape[1] - 1, 0)
+        index = index.view(-1, 1, 1, 1, 1)
+        index = index.expand(batch_size, 1, *batch["x"].shape[2:])
+        x_hat = torch.gather(batch["x"], dim=1, index=index)
+        x_hat = x_hat.squeeze(1)
+        assert (
+            batch["y"].shape == x_hat.shape
+        ), f'Shapes {batch["y"].shape} and {x_hat.shape} do not agree.'
+    elif model.residual == "climate":
+        climate_scaled = (
+            batch["climate"] - model.input_scalers_mu.view(1, -1, 1, 1)
+        ) / (
+            model.input_scalers_sigma.view(1, -1, 1, 1) + model.input_scalers_epsilon
+        )
+
+    # [batch, time, parameter, lat, lon] -> [batch, time x parameter, lat, lon]
+    x_rescaled = x_rescaled.flatten(1, 2)
+    # Parameter dropout
+    x_rescaled = model.parameter_dropout(x_rescaled)
+
+    x_embedded = model.patch_embedding(x_rescaled)
+    assert x_embedded.shape[1] == model.embed_dim
+
+    if model.residual == "climate":
+        static_embedded = model.patch_embedding_static(
+            torch.cat((x_static, climate_scaled), dim=1)
+        )
+    else:
+        static_embedded = model.patch_embedding_static(x_static)
+    assert static_embedded.shape[1] == model.embed_dim
+
+    if model.positional_encoding == 'fourier':
+        static_embedded += x_static_pos
+
+    x_embedded = model.to_patching(x_embedded)
+    static_embedded = model.to_patching(static_embedded)
+
+    time_encoding = model.time_encoding(batch['input_time'], batch['lead_time'])
+
+    tokens = x_embedded + static_embedded + time_encoding
+
+    if mask_global is not None:
+        indices_masked = torch.where(mask_global.flatten())[0]
+        indices_masked = torch.repeat_interleave(indices_masked[None], tokens.shape[0], dim=0)
+        indices_unmasked = torch.where(~mask_global.flatten())[0]
+        indices_unmasked = torch.repeat_interleave(indices_unmasked[None], tokens.shape[0], dim=0)
+    elif mask_local is not None:
+        indices_masked = torch.where(mask_local.flatten())[0]
+        indices_masked = indices_masked[None, None].expand(tokens.shape[0], tokens.shape[1], -1)
+        indices_unmasked = torch.where(~mask_local.flatten())[0]
+        indices_unmasked = indices_unmasked[None, None].expand(tokens.shape[0], tokens.shape[1], -1)
+        print(indices_unmasked.shape)
+    else:
+        mask_global = torch.zeros(model.global_shape_mu, dtype=torch.bool)
+        indices_masked = torch.where(mask_global.flatten())[0]
+        indices_masked = torch.repeat_interleave(indices_masked[None], tokens.shape[0], dim=0)
+        indices_unmasked = torch.where(~mask_global.flatten())[0]
+        indices_unmasked = torch.repeat_interleave(indices_unmasked[None], tokens.shape[0], dim=0)
+
+    maskdim: int = indices_masked.ndim
+
+    # Unmasking
+    unmask_view = (*indices_unmasked.shape, *[1] * (tokens.ndim - maskdim))
+    unmasked = torch.gather(
+        tokens,
+        dim=maskdim - 1,
+        index=indices_unmasked.view(*unmask_view).expand(
+            *indices_unmasked.shape, *tokens.shape[maskdim:]
+        ),
+    )
+
+    # Encoder
+    x_encoded = model.encoder(unmasked)
+
+    # Generate and position encode the mask tokens
+    # (1, 1, 1, embed_dim) -> (batch, global_seq_masked, local seq, embed_dim)
+    mask_view = (*indices_masked.shape, *[1] * (tokens.ndim - maskdim))
+    masking = model.mask_token.repeat(*static_embedded.shape[:3], 1)
+    masked = masking + static_embedded
+    masked = torch.gather(
+        masked,
+        dim=maskdim - 1,
+        index=indices_masked.view(*mask_view).expand(
+            *indices_masked.shape, *tokens.shape[maskdim:]
+        ),
+    )
+
+    recon, _ = model.reconstruct_batch(
+        indices_masked, indices_unmasked, masked, x_encoded
+    )
+
+    x_decoded = model.decoder(recon)
+
+    # Output: (batch, global sequence, local sequence, in_channels * patch_size[0] * patch_size[1])
+    x_unembed = model.unembed(x_decoded)
+
+    # Reshape to (batch, global_lat, global_lon, local_lat, local_lon, in_channels * patch_size[0] * patch_size[1])
+    assert x_unembed.shape[0] == batch_size
+    assert x_unembed.shape[1] == model.global_shape_mu[0] * model.global_shape_mu[1]
+    assert x_unembed.shape[2] == model.local_shape_mu[0] * model.local_shape_mu[1]
+    assert (
+        x_unembed.shape[3]
+        == model.in_channels * model.patch_size_px[0] * model.patch_size_px[1]
+    )
+
+    x_out = model.from_patching(x_unembed)
+
+    # Pixel shuffle to (batch, in_channels, lat, lon)
+    x_out = F.pixel_shuffle(x_out, model.patch_size_px[0])
+
+    if model.residual == "temporal":
+        x_out = model.output_scalers * x_out + x_hat
+    elif model.residual == "climate":
+        x_out = model.output_scalers * x_out + batch["climate"]
+    elif model.residual == "none":
+        x_out = model.output_scalers * x_out + model.input_scalers_mu.reshape(
+            1, -1, 1, 1
+        )
+
+    return x_out

scripts/fake_data.py

@@ -0,0 +1,118 @@
+import os
+import h5py
+import numpy as np
+import xarray as xr
+from onescience.utils.YParams import YParams
+
+
+# Prithvi WxC 输入两个时刻、预测一个时刻，并额外接收 4 通道静态场。
+def get_dims(cfg_model, cfg_data):
+    H, W = int(cfg_model.n_lats_px), int(cfg_model.n_lons_px)
+    if tuple(map(int, cfg_data.dataset.img_size)) != (H, W):
+        raise ValueError("model grid and datapipe.dataset.img_size must match")
+    input_steps, output_steps = int(cfg_model.input_size_time), 1
+    samples = max(int(cfg_data.dataloader.batch_size), 2)
+    T = input_steps + output_steps + samples - 1
+    return {
+        "T": T, "H": H, "W": W, "time_step": 6,
+        "input_steps": input_steps, "output_steps": output_steps,
+        "static_channels": int(cfg_model.in_channels_static),
+    }
+
+
+def generate_fake_h5(data_dir, var_names, years, dims):
+    """
+    为每个年份生成一个空 h5 文件。
+    利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性，
+    文件实际只含元数据，极小，但 shape 与真实数据完全一致。
+    均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。
+    """
+    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)
+    T, C = dims["T"], len(var_names)
+    H, W = dims["H"], dims["W"]
+
+    means = np.zeros((1, C, 1, 1), dtype=np.float32)
+    stds = np.ones((1, C, 1, 1), dtype=np.float32)
+
+    for year in years:
+        path = os.path.join(data_dir, "data", f"{year}.h5")
+        with h5py.File(path, "w") as f:
+            ds = f.create_dataset(
+                "fields",
+                shape=(T, C, H, W),
+                dtype="float32",
+                chunks=(1, C, H, W),
+                fillvalue=0.0,
+            )
+            ds.attrs["variables"] = var_names
+            ds.attrs["time_step"] = dims["time_step"]
+            f.create_dataset("global_means", data=means)
+            f.create_dataset("global_stds", data=stds)
+
+        size_kb = os.path.getsize(path) / 1024
+        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "
+              f"logical={T*C*H*W*4/1024**3:.1f}GB  actual={size_kb:.1f}KB")
+
+
+def get_static(data_dir, H, W, channels):
+    os.makedirs(data_dir, exist_ok=True)
+    lat = np.linspace(90, -90, H, dtype=np.float32)
+    lon = np.linspace(0, 360 - 360 / W, W, dtype=np.float32)
+    lat_grid = np.broadcast_to(lat[:, None], (H, W)) / 90.0
+    lon_grid = np.broadcast_to(lon[None, :], (H, W)) / 180.0 - 1.0
+    land_mask = (np.sin(np.deg2rad(lat_grid * 90)) > 0).astype(np.float32)
+    topography = np.cos(np.deg2rad(lat_grid * 90)).astype(np.float32)
+    base = [lat_grid, lon_grid, land_mask, topography]
+    static = np.stack((base * ((channels + 3) // 4))[:channels]).astype(np.float32)
+
+    ds = xr.Dataset(
+        data_vars={
+            "z": (("valid_time", "latitude", "longitude"), static[-1:]),
+            "lsm": (("valid_time", "latitude", "longitude"), static[min(2, channels - 1):min(2, channels - 1) + 1]),
+        },
+        coords={
+            "valid_time": ["2015-12-31"],
+            "latitude": lat.astype(np.float64),
+            "longitude": lon.astype(np.float64),
+            "number": 0,
+            "expver": "",
+        },
+        attrs={
+            "GRIB_centre": "ecmf",
+            "GRIB_centreDescription": "European Centre for Medium-Range Weather Forecasts",
+            "GRIB_subCentre": "0",
+            "Conventions": "CF-1.7",
+            "institution": "European Centre for Medium-Range Weather Forecasts",
+            "history": "Generated manually",
+        }
+    )
+
+    ds[["z"]].to_netcdf(f"{data_dir}/geopotential.nc")
+    ds[["lsm"]].to_netcdf(f"{data_dir}/land_sea_mask.nc")
+    np.save(f'{data_dir}/static.npy', static)
+    np.save(f'{data_dir}/land_mask.npy', land_mask)
+    np.save(f'{data_dir}/soil_type.npy', np.zeros((H, W), dtype=np.float32))
+    np.save(f'{data_dir}/topography.npy', topography)
+    print(f"✅ Static data: {static.shape}, dtype: {static.dtype}, save to {data_dir}")
+
+
+if __name__ == "__main__":
+    cfg_model = YParams("conf/config.yaml", "model")
+    cfg_datapipe = YParams("conf/config.yaml", "datapipe")
+
+    if cfg_datapipe.dataset.data_dir.startswith("/public/") or cfg_datapipe.dataset.data_dir.startswith("/work2/"):
+        print("请检查 config，确保各 *_dir 指向本地测试路径而非生产路径。")
+        exit()
+
+    years = cfg_datapipe.dataset.train_time + cfg_datapipe.dataset.val_time + cfg_datapipe.dataset.test_time
+    atm_vars = cfg_datapipe.dataset.channels
+    if len(atm_vars) != int(cfg_model.in_channels):
+        raise ValueError("channel count must match model.in_channels")
+
+    dims = get_dims(cfg_model, cfg_datapipe)
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, dims)
+
+    static_dir = os.path.join(cfg_datapipe.dataset.data_dir, "static")
+    get_static(static_dir, dims["H"], dims["W"], dims["static_channels"])
+
+    print("\n✅ Fake datasets generated.")

scripts/inference.py

@@ -0,0 +1,104 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from model.prithvi_wxc import PrithviWxC
+from onescience.utils.YParams import YParams
+from onescience.datapipes.climate import ERA5Datapipe
+
+
+def get_stats(data_dir, channels):
+    """从新版 h5 中读取变量列表与归一化参数（均值/标准差）"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        mu = f["global_means"][:]   # [1, C, 1, 1]
+        std = f["global_stds"][:]
+
+    channel_indices = [all_variables.index(v) for v in channels]
+    means = mu[:, channel_indices, :, :]
+    stds = std[:, channel_indices, :, :]
+    return means, stds
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)
+
+    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
+    cfg['N_out_channels'] = len(cfg_data.dataset.channels)
+
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.test_time,
+        distributed=False,
+        input_steps=2,
+        batch_size=1,
+        num_workers=4,
+    )
+    test_dataloader, _ = datapipe.get_dataloader("test")
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
+    model = PrithviWxC(
+        in_channels=cfg['N_in_channels'],
+        input_size_time=cfg.input_size_time,
+        in_channels_static=cfg.in_channels_static,
+        n_lats_px=cfg.n_lats_px,
+        n_lons_px=cfg.n_lons_px,
+        patch_size_px=cfg.patch_size_px,
+        mask_unit_size_px=cfg.mask_unit_size_px,
+        mask_ratio_inputs=0.0,
+        embed_dim=cfg.embed_dim,
+        n_blocks_encoder=cfg.n_blocks_encoder,
+        n_blocks_decoder=cfg.n_blocks_decoder,
+        mlp_multiplier=cfg.mlp_multiplier,
+        n_heads=cfg.n_heads,
+        dropout=cfg.dropout,
+        drop_path=cfg.drop_path,
+        parameter_dropout=cfg.parameter_dropout,
+        residual=cfg.residual,
+        masking_mode=cfg.masking_mode,
+        positional_encoding=cfg.positional_encoding,
+        encoder_shifting=cfg.encoder_shifting,
+        decoder_shifting=cfg.decoder_shifting,
+    ).to(device)
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    model.eval()
+    os.makedirs('result/output/', exist_ok=True)
+    print(f"📂 infer results will be generated to './result/output/'")
+    H, W = int(cfg.n_lats_px), int(cfg.n_lons_px)
+    static_path = os.path.join(cfg_data.dataset.data_dir, "static", "static.npy")
+    static_base = torch.from_numpy(np.load(static_path)).to(device=device, dtype=torch.float32).unsqueeze(0)
+    expected_static = (1, int(cfg.in_channels_static), H, W)
+    if tuple(static_base.shape) != expected_static:
+        raise ValueError(f"static data shape {tuple(static_base.shape)} != expected {expected_static}")
+    with torch.no_grad():
+        for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
+            invar = data[0].to(device, dtype=torch.float32)      # [1, 2, C, H, W]
+            filename = data[4][-1][0]
+            B = invar.shape[0]
+            static = static_base.expand(B, -1, -1, -1)
+            lead_time = torch.full((B,), 6.0, device=device)
+            pred_var = model(invar, static, lead_time=lead_time).cpu().numpy()
+            pred_var = pred_var * stds + means
+            np.save(f"result/output/{filename}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,200 @@
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import os
+import sys
+import glob
+import h5py
+from datetime import datetime
+from tqdm import tqdm
+from onescience.utils.YParams import YParams
+from matplotlib import rcParams
+
+# rcParams['font.family'] = 'serif'
+# rcParams['font.serif'] = ['DejaVu Serif']
+rcParams['mathtext.fontset'] = 'stix'
+rcParams['axes.linewidth'] = 0.9
+rcParams['xtick.major.width'] = 0.9
+rcParams['ytick.major.width'] = 0.9
+
+
+def get_metadata(data_dir, channels):
+    """从新版 h5 attrs 中读取变量列表和 time_step"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        time_step = int(ds.attrs["time_step"])
+
+    channel_indices = [all_variables.index(v) for v in channels]
+
+    total_files = [f for f in os.listdir('./result/output/') if f.endswith('.npy')]
+    total_files.sort()
+    return total_files, channel_indices, time_step
+
+
+def filename_to_index(filename, time_step):
+    """将 YYYYMMDDHH 格式的文件名转换为年度 h5 文件中的时间步索引"""
+    dt = datetime.strptime(filename, "%Y%m%d%H")
+    year_start = datetime(dt.year, 1, 1)
+    hours = (dt - year_start).total_seconds() / 3600
+    return int(hours / time_step)
+
+
+def get_result(total_files, channel_indices, time_step, data_dir, clim_mean):
+    channel_rmse = np.zeros(len(channel_indices))
+    channel_acc = np.zeros(len(channel_indices))
+    clim_mean = clim_mean[0, :, :, :]
+    if not os.path.exists('./result/rmse.npy') or not os.path.exists('result/acc.npy'):
+        numerator = np.zeros(len(channel_indices))
+        pred_sq_sum = np.zeros(len(channel_indices))
+        label_sq_sum = np.zeros(len(channel_indices))
+        for file in tqdm(total_files, unit="files"):
+            fname = file[:-4]  # 去掉 .npy
+            year = fname[:4]
+            t_idx = filename_to_index(fname, time_step)
+            with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+                label = f["fields"][t_idx]  # [C, H, W]
+                label = label[channel_indices]
+            pred = np.load(f'result/output/{file}').squeeze()
+
+            label_anom = label - clim_mean
+            pred_anom = pred - clim_mean
+            # 累加
+            numerator += np.sum(pred_anom * label_anom, axis=(1, 2))
+            pred_sq_sum += np.sum(pred_anom ** 2, axis=(1, 2))
+            label_sq_sum += np.sum(label_anom ** 2, axis=(1, 2))
+
+            channel_rmse += np.sqrt(np.mean((label - pred) ** 2, axis=(1, 2)))
+        channel_rmse /= len(total_files)
+        channel_acc = numerator / (np.sqrt(pred_sq_sum * label_sq_sum) + 1e-8)
+        np.save('./result/acc.npy', channel_acc)
+        np.save('./result/rmse.npy', channel_rmse)
+
+
+def show_result():
+    channel_rmse = np.load('./result/rmse.npy')
+    channel_acc = np.load('./result/acc.npy')
+
+    channels = [cfg_data.dataset.channels[i] for i in range(len(channel_indices))]
+    w = 24  # 最长 channel 名宽度
+
+    # 表头
+    print(f"┌{'─' * (w + 2)}┬{'─' * 14}┬{'─' * 14}┐")
+    print(f"│ {'Channel':<{w}} │ {'RMSE':>12} │ {'ACC':>12} │")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    # 数据行
+    for i, ch in enumerate(channels):
+        print(f"│ {ch:<{w}} │ {channel_rmse[i]:>12.4f} | {channel_acc[i]:>12.4f} |")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    print(f"│ {'Average':<{w}} │ {np.mean(channel_rmse):>12.4f} │ {np.mean(channel_acc):>12.4f} │")
+    print(f"└{'─' * (w + 2)}┴{'─' * 14}┴{'─' * 14}┘")
+
+
+def plot(label, pred, var, filename):
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+    xtick_labels = ['180°W', '90°W', '0°', '90°E', '180°E']
+    ytick_labels = ['90°S', '45°S', '0°', '45°N', '90°N']
+    xticks = np.linspace(0, label.shape[-1] - 1, 5)
+    yticks = np.linspace(0, label.shape[-2] - 1, 5)
+
+    vmin = min(label.min(), pred.min())
+    vmax = max(label.max(), pred.max())
+
+    diff = label - pred
+    rmse = np.sqrt(np.mean(diff ** 2))
+    diff_abs_max = np.abs(diff).max()
+
+    plot_configs = [
+        {'data': label, 'title': 'Truth', 'cmap': 'viridis', 'vmin': vmin, 'vmax': vmax},
+        {'data': pred,  'title': 'Prediction', 'cmap': 'viridis', 'vmin': vmin, 'vmax': vmax},
+        {'data': diff,  'title': f'Difference (RMSE={rmse:.2f})', 'cmap': 'RdBu_r', 'vmin': -diff_abs_max, 'vmax': diff_abs_max},
+    ]
+
+    for ax, cfg in zip(axes, plot_configs):
+        im = ax.imshow(cfg['data'], cmap=cfg['cmap'], vmin=cfg['vmin'], vmax=cfg['vmax'])
+        ax.set_title(cfg['title'], fontsize=12, pad=4)
+        ax.set_xlabel('Longitude')
+        ax.set_ylabel('Latitude')
+        ax.set_xticks(xticks)
+        ax.set_xticklabels(xtick_labels)
+        ax.set_yticks(yticks)
+        ax.set_yticklabels(ytick_labels)
+        plt.colorbar(im, ax=ax, orientation='horizontal')
+
+    fig.suptitle(var, fontsize=14, fontweight='bold', y=0.98)
+    plt.savefig(filename, dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+def plot_loss(train_loss, valid_loss):
+    mask = ~(np.isnan(train_loss) | np.isnan(valid_loss))
+    train_loss = train_loss[mask]
+    valid_loss = valid_loss[mask]
+
+    fig, ax = plt.subplots(figsize=(5, 3.5))
+    colors = {'train': '#2563EB', 'valid': '#EA580C'}
+    epochs = np.arange(1, len(train_loss) + 1)
+
+    ax.plot(epochs, train_loss, color=colors['train'], linewidth=1.5, label='Train')
+    ax.plot(epochs, valid_loss, color=colors['valid'], linewidth=1.5, label='Valid', linestyle='--')
+    min_idx = np.argmin(valid_loss)
+    ax.scatter(epochs[min_idx], valid_loss[min_idx],
+               color=colors['valid'], s=40, zorder=5, edgecolors='white')
+    ax.annotate(f'Best: {valid_loss[min_idx]:.3f}',
+                xy=(epochs[min_idx], valid_loss[min_idx]),
+                xytext=(10, 10), textcoords='offset points', fontsize=8, color=colors['valid'],
+                arrowprops=dict(arrowstyle='-', color=colors['valid'], lw=0.5))
+
+    ax.set(xlabel='Epoch', ylabel='Loss', xlim=(0, len(train_loss) + 1))
+    ax.legend(frameon=False, loc='upper right')
+    ax.grid(True, linestyle='--', alpha=0.3)
+    ax.spines[['top', 'right']].set_visible(False)
+
+    plt.tight_layout()
+    plt.savefig('./result/loss.png', dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    config_file_path = os.path.join(current_path, 'conf/config.yaml')
+    cfg = YParams(config_file_path, 'model')
+    cfg_data = YParams(config_file_path, "datapipe")
+
+    train_loss = np.load('./data/checkpoints/trloss.npy')
+    valid_loss = np.load('./data/checkpoints/valoss.npy')
+    plot_loss(train_loss, valid_loss)
+
+    data_dir = cfg_data.dataset.data_dir
+    total_files, channel_indices, time_step = get_metadata(data_dir, cfg_data.dataset.channels)
+
+    # Load data & Compute RMSE/ACC per channel
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        mu = f["global_means"][:]
+    clim_mean = mu[:, channel_indices, :, :]
+    get_result(total_files, channel_indices, time_step, data_dir, clim_mean)
+    show_result()
+
+    ##### 默认绘制 result/output 中第一个预测结果，用户可自行指定日期和变量 #####
+    eg_files = [total_files[0][:-4]]
+    channel_index = [cfg_data.dataset.channels.index(v) for v in ['2m_temperature', 'geopotential_500', 'temperature_850']]
+
+    selected_var = [cfg_data.dataset.channels[int(i)] for i in channel_index]
+    print(f"seleted date: {eg_files}")
+    print(f"selected channels: {selected_var}")
+    for file in eg_files:
+        year = file[:4]
+        t_idx = filename_to_index(file, time_step)
+        with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+            label = f["fields"][t_idx]  # [C, H, W]
+            label = label[channel_indices]
+        pred = np.load(f'result/output/{file}.npy').squeeze()
+        for i in range(len(selected_var)):
+            filename = f'./result/{file}_{selected_var[i]}.png'
+            plot(label[channel_index[i]], pred[channel_index[i]], selected_var[i], filename)
+            print(f'✅plot {filename}')

scripts/train.py

@@ -0,0 +1,243 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import numpy as np
+import torch.distributed as dist
+import logging
+import time
+
+from model.prithvi_wxc import PrithviWxC
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.utils.fcn.darcy_loss import LpLoss
+
+try:
+    from apex import optimizers
+    _FUSED_ADAM = True
+except Exception:
+    _FUSED_ADAM = False
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    logger = logging.getLogger()
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## Distributed config init
+    cfg.world_size = 1
+    if "WORLD_SIZE" in os.environ:
+        cfg.world_size = int(os.environ["WORLD_SIZE"])
+    world_rank = 0
+    local_rank = 0
+    if cfg.world_size > 1 and torch.cuda.is_available():
+        dist.init_process_group(backend="nccl", init_method="env://")
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_rank = dist.get_rank()
+    device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
+    cfg['N_out_channels'] = len(cfg_data.dataset.channels)
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.train_time,
+        distributed=dist.is_initialized(),
+        input_steps=2,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers,
+    )
+    train_dataloader, train_sampler = datapipe.get_dataloader("train")
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.val_time,
+        distributed=dist.is_initialized(),
+        input_steps=2,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers,
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    # Model init
+    model = PrithviWxC(
+        in_channels=cfg['N_in_channels'],
+        input_size_time=cfg.input_size_time,
+        in_channels_static=cfg.in_channels_static,
+        n_lats_px=cfg.n_lats_px,
+        n_lons_px=cfg.n_lons_px,
+        patch_size_px=cfg.patch_size_px,
+        mask_unit_size_px=cfg.mask_unit_size_px,
+        mask_ratio_inputs=cfg.mask_ratio_inputs,
+        embed_dim=cfg.embed_dim,
+        n_blocks_encoder=cfg.n_blocks_encoder,
+        n_blocks_decoder=cfg.n_blocks_decoder,
+        mlp_multiplier=cfg.mlp_multiplier,
+        n_heads=cfg.n_heads,
+        dropout=cfg.dropout,
+        drop_path=cfg.drop_path,
+        parameter_dropout=cfg.parameter_dropout,
+        residual=cfg.residual,
+        masking_mode=cfg.masking_mode,
+        positional_encoding=cfg.positional_encoding,
+        encoder_shifting=cfg.encoder_shifting,
+        decoder_shifting=cfg.decoder_shifting,
+    ).to(device)
+
+    if _FUSED_ADAM:
+        optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.lr)
+    else:
+        optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, mode='min')
+    loss_obj = LpLoss()
+
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/valoss.npy"
+    best_valid_loss = 1.0e6
+    best_loss_epoch = 0
+    train_losses = np.empty((0,), dtype=np.float32)
+    valid_losses = np.empty((0,), dtype=np.float32)
+
+    ## Get model params count
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    ## Load model weight if there exist well-trained model
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a model weight, load and continue training...")
+            print(f'If you want to train a new model, ensure there is no *.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+        best_valid_loss = ckpt["best_valid_loss"]
+        best_loss_epoch = ckpt["best_loss_epoch"]
+        train_losses = np.load(train_loss_file)
+        valid_losses = np.load(valid_loss_file)
+
+    ## Distributed model
+    if dist.is_initialized():
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+    world_rank == 0 and logger.info(f"start training ...")
+
+    H, W = int(cfg.n_lats_px), int(cfg.n_lons_px)
+    static_path = os.path.join(cfg_data.dataset.data_dir, "static", "static.npy")
+    static_base = torch.from_numpy(np.load(static_path)).to(device=device, dtype=torch.float32).unsqueeze(0)
+    expected_static = (1, int(cfg.in_channels_static), H, W)
+    if tuple(static_base.shape) != expected_static:
+        raise ValueError(f"static data shape {tuple(static_base.shape)} != expected {expected_static}")
+    for epoch in range(cfg.max_epoch):
+        if dist.is_initialized():
+            train_sampler.set_epoch(epoch)
+            val_sampler.set_epoch(epoch)
+        model.train()
+        train_loss = 0
+        start_time = time.time()
+        for j, data in enumerate(train_dataloader):
+            invar = data[0].to(device, dtype=torch.float32)      # [B, 2, C, H, W]
+            outvar = data[1].to(device, dtype=torch.float32)     # [B, C, H, W]
+            B = invar.shape[0]
+            static = static_base.expand(B, -1, -1, -1)
+            lead_time = torch.full((B,), 6.0, device=device)
+            outvar_pred = model(invar, static, lead_time=lead_time)
+            loss = loss_obj(outvar, outvar_pred)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+            if world_rank == 0:
+                logger.info(f'Train: Epoch {epoch}-{j+1}/{len(train_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                            f'loss:{train_loss / (j+1): .04f}')
+
+        train_loss /= len(train_dataloader)
+
+        model.eval()
+        valid_loss = 0
+        with torch.no_grad():
+            start_time = time.time()
+            for j, data in enumerate(val_dataloader):
+                invar = data[0].to(device, dtype=torch.float32)
+                outvar = data[1].to(device, dtype=torch.float32)
+                B = invar.shape[0]
+                static = static_base.expand(B, -1, -1, -1)
+                lead_time = torch.full((B,), 6.0, device=device)
+                outvar_pred = model(invar, static, lead_time=lead_time)
+                loss = loss_obj(outvar, outvar_pred)
+
+                if dist.is_initialized():
+                    loss_tensor = loss.detach().to(device)
+                    dist.all_reduce(loss_tensor)
+                    loss = loss_tensor.item() / cfg.world_size
+                    valid_loss += loss
+                else:
+                    valid_loss += loss.item()
+                if world_rank == 0:
+                    logger.info(f'Valid: Epoch {epoch}-{j+1}/{len(val_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                            f'loss:{valid_loss / (j+1): .04f}')
+
+        valid_loss /= len(val_dataloader)
+        is_save_ckp = False
+        if valid_loss < best_valid_loss:
+            best_valid_loss = valid_loss
+            best_loss_epoch = epoch
+            world_rank == 0 and save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir)
+            is_save_ckp = True
+        scheduler.step(valid_loss)
+
+        if world_rank == 0:
+            logger.info(f"Epoch [{epoch + 1}/{cfg.max_epoch}], "
+                        f"Train Loss: {train_loss:.4f}, "
+                        f"Valid Loss: {valid_loss:.4f}, "
+                        f"Best loss at Epoch: {best_loss_epoch + 1}"
+                        + (", saving checkpoint" if is_save_ckp else "")
+                        )
+            train_losses = np.append(train_losses, train_loss)
+            valid_losses = np.append(valid_losses, valid_loss)
+            np.save(train_loss_file, train_losses)
+            np.save(valid_loss_file, valid_losses)
+
+        if epoch - best_loss_epoch > cfg.patience:
+            print(f"Loss has not decrease in {cfg.patience} epochs, stopping training...")
+            exit()
+
+
+def save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, model_path):
+    model_to_save = model.module if hasattr(model, "module") else model
+    state = {"model_state_dict": model_to_save.state_dict(),
+             "optimizer_state_dict": optimizer.state_dict(),
+             "scheduler_state_dict": scheduler.state_dict(),
+             "best_valid_loss": best_valid_loss,
+             "best_loss_epoch": best_loss_epoch,
+            }
+    torch.save(state, f"{model_path}/model.pth")
+    ### the weight file saving may interrupted due to DCU queue limit, get a backup to ensure there at least has one model
+    os.system(f"mv {model_path}/model.pth {model_path}/model_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()