Publish FuXi-Ocean engineering reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,153 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- ocean-forecasting
+- eddy-resolving
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">FuXi-Ocean</span></strong>
+</p>
+
+# Model Introduction
+
+FuXi-Ocean performs autoregressive ocean-field prediction at six-hour intervals from historical ocean states and initialization-time atmospheric variables, providing an executable engineering workflow for research on high-resolution ocean forecasting methods.
+
+Paper: A deep learning global ocean forecasting model with sub-daily and eddy-resolving resolution  
+https://doi.org/10.1038/s41612-026-01444-2
+
+# Model Description
+
+The method was proposed by research teams from Tianjin University, the Shanghai Academy of Artificial Intelligence for Science, Fudan University, the Shanghai Innovation Institute, FuXi Intelligent Computing Technology Co. Ltd, Hainan Tropical Ocean University, and collaborating institutions. The paper trains on HYCOM reanalysis and analysis ocean fields together with ERA5 near-surface atmospheric variables. The model performs autoregressive prediction of temperature, salinity, zonal current, meridional current, and sea-surface height at six-hour intervals.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Autoregressive ocean-field prediction | Validate the primary task of predicting subsequent ocean states from four historical states and initialization-time atmospheric variables on synthetic sampled tiles. |
+| Core-method validation | Validate shared convolutional encoding, a spatiotemporal prior, historical-feature fusion, low-resolution attention, shared decoding, and latitude-weighted Charbonnier loss. |
+| Local engineering validation | Validate indexing, overlap crops, data contracts, inference, evaluation, and visualization interfaces for the complete `2160×4320` scientific-grid protocol; the current workflow does not produce a complete global prediction. |
+| ModelScope/OneCode execution | Validate structured data, training, inference, per-lead ocean-forecast metrics, and visualization in ModelScope or OneCode environments. |
+| Multi-GPU training | Validate distributed training and the checkpoint workflow through `torchrun`. |
+
+# Usage Instructions
+
+## 1.OneCode
+
+Experience intelligent, one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Download and Installation
+
+```bash
+hf download OneScience-Group/FuXi-Ocean --local-dir ./FuXi-Ocean
+cd FuXi-Ocean
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for connectivity validation with the default small-sample configuration.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+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
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The paper uses HYCOM ocean fields and ERA5 near-surface atmospheric variables to predict the next six-hour ocean state on a `2160×4320` global grid from four historical states. This repository uses six structured synthetic tiles while preserving 105 ocean channels, five atmospheric channels, 26 depth levels, and the complete global indexing protocol. The synthetic data validate the engineering workflow only and do not represent the real HYCOM or ERA5 distributions, training scale, or paper performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-device training, use:
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```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 uses latitude-weighted Charbonnier loss and multi-step autoregressive rollout; the default configuration reduces tiles, samples, model scale, epochs, and rollout length without reducing the real variables or global-coordinate protocol. Formal experiments require real HYCOM and ERA5 data and full-scale computing resources, with training artifacts saved to:
+
+```text
+result/checkpoints/fuxi_ocean.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not bundle official weights under `weight/`, and no confirmed official model weights were identified. The complete resource link in the paper's code-availability statement is https://doi.org/10.5281/zenodo.17412508; it must not be treated as a confirmed weight-download link. The local checkpoint is an engineering artifact trained on synthetic `sampled_tiles` and makes no compatibility claim with paper or official weights.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the locally trained checkpoint and runs three autoregressive steps on test tiles by default while retaining a record of the paper's 40-step protocol. Outputs are explicitly marked as `sampled_tiles` and `is_complete_global=false` and include their coverage fraction, so they must not be treated as a completed global prediction. Numerical outputs and metadata are saved to:
+
+```text
+result/output/predictions.npz
+result/output/metadata.json
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation computes latitude-weighted RMSE and MBE per lead for the model and persistence baseline over sampled tiles, reports S, T, U, V, and SSH unit groups separately, and retains per-channel results without averaging all 105 channels across incompatible units. Synthetic-data evaluation validates the engineering workflow only and does not represent formal paper performance. Results are saved to:
+
+```text
+result/evaluation/metrics.json
+result/evaluation/lead_metrics.png
+```
+
+# 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 engineering reproduction of the public FuXi-Ocean specifications.
+
+Use of this repository's code is subject to the Apache-2.0 license terms.
+
+The original paper is published under CC BY 4.0, and its citation and use are subject to that license.
+
+Use of HYCOM, ERA5, official Zenodo artifacts, and other third-party data or weights remains subject to the licenses and terms of their respective projects.

README_zh.md

@@ -0,0 +1,154 @@
+---
+license: apache-2.0
+language:
+- zh
+- en
+tags:
+- OneScience
+- 地球科学
+- 海洋预报
+- 涡分辨率模型
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">FuXi-Ocean</span></strong>
+</p>
+
+# 模型介绍
+
+FuXi-Ocean 用于根据历史海洋状态和初始化大气变量进行 6 小时间隔的海洋场自回归预测，为高时空分辨率海洋预报方法研究提供可运行的工程验证流程。
+
+论文：A deep learning global ocean forecasting model with sub-daily and eddy-resolving resolution  
+https://doi.org/10.1038/s41612-026-01444-2
+
+# 模型描述
+
+该方法由天津大学、上海科学智能研究院、复旦大学、上海创新研究院、上海伏羲智算科技有限公司和海南热带海洋学院等机构的研究团队提出。论文使用 HYCOM 再分析与分析海洋场以及 ERA5 近地面大气变量开展训练。模型执行 6 小时间隔的温度、盐度、纬向流、经向流和海表高度自回归预测任务。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 海洋场自回归预测 | 在合成 sampled tiles 上验证由 4 帧历史海洋状态和初始化大气变量预测后续海洋状态的主任务。 |
+| 核心方法验证 | 验证共享卷积编码、时空先验、历史特征融合、低分辨率 attention、共享解码和纬度加权 Charbonnier 损失。 |
+| 本地工程验证 | 验证完整 `2160×4320` 科学网格的索引、重叠裁剪、数据契约、推理、评估和可视化接口；当前不生成完整全球预测。 |
+| ModelScope/OneCode 运行 | 在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、逐时效海洋预报指标和可视化流程。 |
+| 多卡训练 | 通过 `torchrun` 验证分布式训练和 checkpoint 流程。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/FuXi-Ocean --local_dir ./FuXi-Ocean
+cd FuXi-Ocean
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 GPU 或 DCU 运行。
+- CPU 可用于默认小样本配置的连通性验证。
+- DCU 用户需预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
+
+**DCU环境**
+
+```bash
+# 请首先激活 DTK 及 Conda
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU环境**
+
+```bash
+# 请首先激活 Conda
+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
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### 训练数据介绍
+
+论文使用 HYCOM 海洋场和 ERA5 近地面大气变量，模型根据 4 个历史时刻预测 `2160×4320` 全球网格上的下一 6 小时海洋状态。本仓库使用 6 个结构化合成 tile，保留 105 个海洋通道、5 个大气通道、26 个深度层和完整全球索引协议。虚拟数据仅用于验证工程流程，不代表 HYCOM 或 ERA5 的真实数据分布、训练规模或论文性能。
+
+```bash
+python scripts/fake_data.py
+```
+
+### 训练
+
+单卡训练可使用：
+
+```bash
+python scripts/train.py
+```
+
+多卡训练可使用：
+
+```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
+```
+
+训练使用纬度加权 Charbonnier 损失和多步自回归流程；默认配置仅缩小 tile、样本数、模型规模、训练轮数和 rollout 步数，不缩小真实变量与全局坐标协议。正式实验需要真实 HYCOM、ERA5 数据和完整计算资源，训练产物保存到：
+
+```text
+result/checkpoints/fuxi_ocean.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+本仓库不在 `weight/` 中内置官方权重，也未发现可确认的官方模型权重。论文代码可用性部分提供的完整资源链接为 https://doi.org/10.5281/zenodo.17412508 ，该链接不能视为已确认的权重下载链接。当前本地 checkpoint 是合成 `sampled_tiles` 的工程训练产物，不声明与论文或官方权重兼容。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载本地训练 checkpoint，默认在测试 tile 上执行 3 步自回归预测，同时记录论文的 40 步协议。结果明确标记为 `sampled_tiles`、`is_complete_global=false` 并保存覆盖率，不能视为已完成的全球预测；数值结果和元数据保存到：
+
+```text
+result/output/predictions.npz
+result/output/metadata.json
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估在已采样 tile ���逐时效计算模型和 persistence 基线的纬度加权 RMSE 与 MBE，按 S、T、U、V 和 SSH 单位组报告并保留逐通道结果，不对单位不同的 105 个通道求统一均值。合成数据评估仅用于验证工程流程，不代表论文正式性能；结果保存到：
+
+```text
+result/evaluation/metrics.json
+result/evaluation/lead_metrics.png
+```
+
+# OneScience官方信息
+
+| 平台 | OneScience 主仓库 | Skills 仓库 |
+| --- | --- | --- |
+| 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 |
+
+# 引用与许可证
+
+本仓库为 FuXi-Ocean 公开规格的独立工程复现版本。
+
+本仓库代码的使用应遵循 Apache-2.0 许可证条款。
+
+原论文按照 CC BY 4.0 许可证发布，引用和使用时应遵循该许可证条款。
+
+HYCOM、ERA5、官方 Zenodo 产物和其他第三方数据或权重的使用应遵循各自项目的许可证及使用条款。

conf/config.yaml

@@ -0,0 +1,74 @@
+seed: 2026
+data:
+  path: data/fuxi_ocean_tiles.npz
+  format_version: fuxi_ocean_sampled_tiles_v2
+  input_shape: [4, 105, 2160, 4320]
+  atmosphere_shape: [5, 2160, 4320]
+  output_shape: [105, 2160, 4320]
+  global_height: 2160
+  global_width: 4320
+  history_steps: 4
+  ocean_channels: 105
+  atmosphere_channels: 5
+  time_step_hours: 6
+  tile_height: 16
+  tile_width: 16
+  tile_overlap: 4
+  train_samples: 4
+  test_samples: 2
+  depth_levels_m: [0, 2, 4, 6, 10, 15, 20, 30, 40, 50, 60, 70, 80, 90, 100, 125, 150, 200, 300, 400, 500, 600, 800, 1000, 1250, 1500]
+model:
+  embed_dim: 24
+  patch_size: 2
+  encoder_levels: 2
+  attention_layers: 2
+  attention_heads: 4
+  epsilon: 0.000001
+training:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.0005
+  weight_decay: 0.00001
+  multistep_rollout: 2
+runtime:
+  device: auto
+  ddp_backend: gloo
+  num_threads: 2
+inference:
+  rollout_steps: 3
+paper_model:
+  paper_determined:
+    input_shape: [4, 105, 2160, 4320]
+    atmosphere_shape: [5, 2160, 4320]
+    output_shape: [105, 2160, 4320]
+    ocean_channel_groups: {S: 26, T: 26, U: 26, V: 26, SSH: 1}
+    atmosphere_variables: [U10M, V10M, MSL, T2M, D2M]
+    shared_convolutional_patch_encoder: true
+    spatiotemporal_prior_network: true
+    history_fusion: four encodings concatenated then 1x1 convolution
+    low_resolution_attention: stacked attention with AdaLN and FFN
+    shared_transposed_convolution_decoder: true
+    skip_connections: true
+    depth_dependent_land_sea_mask: true
+    loss: latitude-weighted Charbonnier
+    training_stages: single-step pretraining then multi-step fine-tuning
+    metrics: [latitude-weighted RMSE, MBE, persistence]
+  missing_architecture:
+    patch_size: MISSING
+    hidden_width: MISSING
+    attention_layer_count: MISSING
+    attention_head_count: MISSING
+    exact_prior_network: MISSING
+    normalization_statistics: MISSING
+    optimizer_and_schedule: MISSING
+    batch_size_and_epochs: MISSING
+    multistep_finetuning_horizon: MISSING
+  rollout_steps: 40
+  time_step_hours: 6
+paths:
+  checkpoint: result/checkpoints/fuxi_ocean.pt
+  training_metrics: result/training/metrics.json
+  inference: result/output/predictions.npz
+  inference_metadata: result/output/metadata.json
+  evaluation_metrics: result/evaluation/metrics.json
+  evaluation_plot: result/evaluation/lead_metrics.png

config.json

@@ -0,0 +1,31 @@
+{
+  "format_version": "1.0",
+  "model_name": "FuXi-Ocean",
+  "model_type": "tile-streamed-global-ocean-transformer",
+  "architectures": ["FuXiOcean"],
+  "framework": "PyTorch",
+  "domain": "earth-science-ocean",
+  "license": "Apache-2.0",
+  "paper": {
+    "title": "A deep learning global ocean forecasting model with sub-daily and eddy-resolving resolution",
+    "doi": "10.1038/s41612-026-01444-2",
+    "paper_license": "CC-BY-4.0"
+  },
+  "implementation": {
+    "entry_point": "model/fuxi_ocean.py",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "scientific_grid": {
+    "history": 4,
+    "ocean_channels": 105,
+    "atmosphere_channels": 5,
+    "height": 2160,
+    "width": 4320,
+    "hours_per_step": 6,
+    "forecast_steps": 40
+  },
+  "execution_protocol": "Tiles retain all 105 channels, physical coordinates, and indices into the unchanged 2160x4320 global grid. Tile dimensions are not the scientific grid dimensions."
+}

configuration.json

@@ -0,0 +1,16 @@
+{
+  "format_version": "1.0",
+  "framework": "PyTorch",
+  "task": "global-ocean-autoregressive-forecasting",
+  "model": "FuXi-Ocean",
+  "input_format": "sampled_tiles with deterministic full-grid enumeration and overlap crops",
+  "global_ocean_shape": [4, 105, 2160, 4320],
+  "atmosphere_shape": [5, 2160, 4320],
+  "output_shape": [105, 2160, 4320],
+  "time_step_hours": 6,
+  "paper_rollout_steps": 40,
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py"
+}

model/fuxi_ocean.py

@@ -0,0 +1,205 @@
+"""Executable tile-streamed FuXi-Ocean reconstruction."""
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+import torch.distributed as dist
+
+
+FORMAT_VERSION = "fuxi_ocean_sampled_tiles_v2"
+GLOBAL_SHAPE = (2160, 4320)
+HISTORY_STEPS = 4
+OCEAN_CHANNELS = 105
+ATMOSPHERE_CHANNELS = 5
+DEPTHS = (0, 2, 4, 6, 10, 15, 20, 30, 40, 50, 60, 70, 80, 90, 100, 125,
+           150, 200, 300, 400, 500, 600, 800, 1000, 1250, 1500)
+
+
+def enumerate_global_tiles(global_shape, tile_shape, overlap=0):
+    """Enumerate all windows and non-overlapping ownership crops in row-major order."""
+    height, width = map(int, global_shape)
+    tile_h, tile_w = map(int, tile_shape)
+    overlap = int(overlap)
+    if min(height, width, tile_h, tile_w) <= 0 or overlap < 0 or overlap >= min(tile_h, tile_w):
+        raise ValueError("invalid global/tile shape or overlap")
+
+    def starts(size, tile):
+        values = list(range(0, max(size - tile, 0) + 1, tile - overlap))
+        last = max(size - tile, 0)
+        if not values or values[-1] != last:
+            values.append(last)
+        return values
+
+    ys, xs = starts(height, tile_h), starts(width, tile_w)
+    records = []
+    for row, y0 in enumerate(ys):
+        for column, x0 in enumerate(xs):
+            y1, x1 = min(y0 + tile_h, height), min(x0 + tile_w, width)
+            top = 0 if row == 0 else (ys[row - 1] + tile_h + y0) // 2 - y0
+            bottom = y1 - y0 if row == len(ys) - 1 else (y1 + ys[row + 1]) // 2 - y0
+            left = 0 if column == 0 else (xs[column - 1] + tile_w + x0) // 2 - x0
+            right = x1 - x0 if column == len(xs) - 1 else (x1 + xs[column + 1]) // 2 - x0
+            records.append((y0, y1, x0, x1, top, bottom, left, right))
+    return np.asarray(records, dtype=np.int32)
+
+
+def assemble_tiles(tiles, tile_records, global_shape, require_complete=True):
+    """Crop overlap by ownership records and assemble a dense field for verification/use."""
+    tiles = np.asarray(tiles)
+    records = np.asarray(tile_records)
+    height, width = map(int, global_shape)
+    output = np.zeros((*tiles.shape[1:-2], height, width), dtype=tiles.dtype)
+    coverage = np.zeros((height, width), dtype=np.uint8)
+    for tile, record in zip(tiles, records):
+        y0, _y1, x0, _x1, top, bottom, left, right = map(int, record)
+        output[..., y0 + top:y0 + bottom, x0 + left:x0 + right] = tile[..., top:bottom, left:right]
+        coverage[y0 + top:y0 + bottom, x0 + left:x0 + right] = 1
+    is_complete = bool(coverage.all())
+    if require_complete and not is_complete:
+        raise ValueError("tile set does not completely cover the global grid")
+    return output, coverage
+
+
+def ocean_channel_depth_indices():
+    """Channel order: S26, T26, U26, V26, SSH."""
+    return np.asarray(list(range(26)) * 4 + [0], dtype=np.int64)
+
+
+def channel_mask(depth_mask):
+    indices = torch.as_tensor(ocean_channel_depth_indices(), device=depth_mask.device)
+    return depth_mask.index_select(1, indices)
+
+
+class ChannelLayerNorm(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.norm = nn.LayerNorm(channels)
+
+    def forward(self, x):
+        return self.norm(x.movedim(1, -1)).movedim(-1, 1)
+
+
+class PriorNetwork(nn.Module):
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.spatial = nn.Conv2d(3, embed_dim, 1)
+        self.temporal = nn.Sequential(nn.Linear(6, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim))
+
+    def forward(self, coordinates, bathymetry, time_features):
+        spatial = self.spatial(torch.cat((coordinates, bathymetry), dim=1))
+        hour, day, step = time_features.unbind(dim=1)
+        temporal = torch.stack((torch.sin(2 * math.pi * hour / 24), torch.cos(2 * math.pi * hour / 24),
+                                torch.sin(2 * math.pi * day / 366), torch.cos(2 * math.pi * day / 366),
+                                torch.sin(step / 40), torch.cos(step / 40)), dim=1)
+        return spatial, self.temporal(temporal)
+
+
+class SharedPatchEncoder(nn.Module):
+    def __init__(self, in_channels, embed_dim, patch_size, levels):
+        super().__init__()
+        layers = []
+        channels = in_channels
+        self.skip_channels = []
+        for level in range(levels):
+            out = embed_dim * (2 ** level)
+            layers += [nn.Conv2d(channels, out, patch_size, patch_size), ChannelLayerNorm(out), nn.GELU()]
+            self.skip_channels.append(out)
+            channels = out
+        self.layers = nn.ModuleList(layers)
+        self.levels = levels
+
+    def forward(self, x, prior):
+        skips = []
+        for level in range(self.levels):
+            conv, norm, activation = self.layers[level * 3:level * 3 + 3]
+            x = conv(x)
+            local_prior = torch.nn.functional.interpolate(prior, size=x.shape[-2:], mode="bilinear", align_corners=False)
+            repeats = math.ceil(x.shape[1] / local_prior.shape[1])
+            modulation = local_prior.repeat(1, repeats, 1, 1)[:, :x.shape[1]]
+            x = activation(norm(x * (1.0 + torch.tanh(modulation))))
+            skips.append(x)
+        return x, skips
+
+
+class AdaLNBlock(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
+        self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 4))
+        self.attention = nn.MultiheadAttention(dim, heads, batch_first=True)
+        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim))
+
+    def forward(self, x, context):
+        scale1, shift1, scale2, shift2 = self.modulation(context).chunk(4, dim=-1)
+        q = self.norm1(x) * (1 + scale1[:, None]) + shift1[:, None]
+        x = x + self.attention(q, q, q, need_weights=False)[0]
+        q = self.norm2(x) * (1 + scale2[:, None]) + shift2[:, None]
+        return x + self.ffn(q)
+
+
+class FuXiOcean(nn.Module):
+    def __init__(self, embed_dim=24, patch_size=2, encoder_levels=2, attention_layers=2, attention_heads=4):
+        super().__init__()
+        self.prior = PriorNetwork(embed_dim)
+        self.encoder = SharedPatchEncoder(OCEAN_CHANNELS + ATMOSPHERE_CHANNELS, embed_dim, patch_size, encoder_levels)
+        latent = embed_dim * 2 ** (encoder_levels - 1)
+        self.fusion = nn.Sequential(nn.Conv2d(HISTORY_STEPS * latent, latent, 1), ChannelLayerNorm(latent), nn.GELU())
+        self.blocks = nn.ModuleList(AdaLNBlock(latent, attention_heads) for _ in range(attention_layers))
+        self.context_projection = nn.Linear(embed_dim, latent)
+        decoder = []
+        for level in reversed(range(encoder_levels)):
+            skip = embed_dim * 2 ** level
+            out = embed_dim * 2 ** max(level - 1, 0)
+            decoder.append(nn.Sequential(nn.ConvTranspose2d(latent + skip, out, patch_size, patch_size),
+                                         ChannelLayerNorm(out), nn.GELU()))
+            latent = out
+        self.decoder = nn.ModuleList(decoder)
+        self.output = nn.Conv2d(latent, OCEAN_CHANNELS, 1)
+
+    def forward(self, ocean_history, atmosphere, coordinates, bathymetry, depth_mask, time_features):
+        if ocean_history.ndim != 5 or ocean_history.shape[1:3] != (HISTORY_STEPS, OCEAN_CHANNELS):
+            raise ValueError("ocean_history must be [B,4,105,tile_h,tile_w]")
+        spatial_prior, temporal_prior = self.prior(coordinates, bathymetry, time_features)
+        encoded, all_skips = [], []
+        for index in range(HISTORY_STEPS):
+            state = torch.cat((ocean_history[:, index], atmosphere), dim=1)
+            latent, skips = self.encoder(state, spatial_prior)
+            encoded.append(latent)
+            all_skips.append(skips)
+        x = self.fusion(torch.cat(encoded, dim=1))
+        shape = x.shape
+        tokens = x.flatten(2).transpose(1, 2)
+        context = self.context_projection(temporal_prior)
+        for block in self.blocks:
+            tokens = block(tokens, context)
+        x = tokens.transpose(1, 2).reshape(shape)
+        for decoder_index, decoder in enumerate(self.decoder):
+            level = len(self.decoder) - 1 - decoder_index
+            skip = torch.stack([values[level] for values in all_skips]).mean(dim=0)
+            x = decoder(torch.cat((x, skip), dim=1))
+        prediction = self.output(x)
+        return prediction * channel_mask(depth_mask).to(prediction.dtype)
+
+
+def latitude_weighted_charbonnier(prediction, target, latitude_deg, valid_mask, epsilon=1e-6):
+    weights = torch.cos(torch.deg2rad(latitude_deg)).clamp_min(0)[:, None, :, None]
+    mask = valid_mask.to(prediction.dtype)
+    numerator = (torch.sqrt((prediction - target) ** 2 + epsilon ** 2) * weights * mask).sum()
+    denominator = (weights * mask).sum()
+    if dist.is_available() and dist.is_initialized():
+        global_numerator = numerator.detach().clone()
+        global_denominator = denominator.detach().clone()
+        dist.all_reduce(global_numerator, op=dist.ReduceOp.SUM)
+        dist.all_reduce(global_denominator, op=dist.ReduceOp.SUM)
+        # DDP averages gradients, so this yields the gradient of the global ratio.
+        return numerator * dist.get_world_size() / global_denominator.clamp_min(1.0)
+    return numerator / denominator.clamp_min(1.0)
+
+
+def weighted_error_sums(prediction, target, latitude_deg, valid_mask):
+    weights = np.cos(np.deg2rad(latitude_deg))[:, None, :, None] * valid_mask
+    error = prediction - target
+    return ((weights * error ** 2).sum(axis=(0, 2, 3)),
+            (weights * error).sum(axis=(0, 2, 3)), weights.sum(axis=(0, 2, 3)))

scripts/fake_data.py

@@ -0,0 +1,76 @@
+"""Create small physical-coordinate tiles without allocating a global dense sample."""
+
+import json
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+import sys
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.fuxi_ocean import enumerate_global_tiles
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = config["data"]
+    rng = np.random.default_rng(config["seed"])
+    count = data["train_samples"] + data["test_samples"]
+    th, tw = data["tile_height"], data["tile_width"]
+    history, channels, atmosphere_channels = data["history_steps"], data["ocean_channels"], data["atmosphere_channels"]
+    global_records = enumerate_global_tiles((data["global_height"], data["global_width"]), (th, tw), data["tile_overlap"])
+    selected_indices = np.linspace(0, len(global_records) - 1, count, dtype=np.int64)
+    selected_records = global_records[selected_indices]
+    origins = selected_records[:, (0, 2)]
+    ocean = np.empty((count, history, channels, th, tw), dtype=np.float32)
+    targets = np.empty((count, channels, th, tw), dtype=np.float32)
+    atmosphere = np.empty((count, atmosphere_channels, th, tw), dtype=np.float32)
+    latitude = np.empty((count, th), dtype=np.float32)
+    longitude = np.empty((count, tw), dtype=np.float32)
+    bathymetry = np.empty((count, 1, th, tw), dtype=np.float32)
+    depth_mask = np.empty((count, 26, th, tw), dtype=np.float32)
+    time_features = np.empty((count, 3), dtype=np.float32)
+    yy, xx = np.meshgrid(np.arange(th), np.arange(tw), indexing="ij")
+    channel_scale = np.linspace(0.2, 1.0, channels, dtype=np.float32)[:, None, None]
+    for sample, (y0, x0) in enumerate(origins):
+        lat = 90.0 - (y0 + np.arange(th) + 0.5) * 180.0 / data["global_height"]
+        lon = (x0 + np.arange(tw) + 0.5) * 360.0 / data["global_width"]
+        latitude[sample], longitude[sample] = lat, lon
+        wave = np.sin(np.deg2rad(lat))[:, None] + 0.5 * np.cos(np.deg2rad(lon))[None, :]
+        for history_index in range(history):
+            ocean[sample, history_index] = channel_scale * (wave + 0.025 * history_index) + rng.normal(0, 0.005, (channels, th, tw))
+        atmosphere[sample] = np.stack([wave + 0.03 * k for k in range(atmosphere_channels)])
+        bathymetry[sample, 0] = 200 + 1400 * (0.5 + 0.5 * np.sin((xx + x0) / 12))
+        for depth, depth_m in enumerate(data["depth_levels_m"]):
+            depth_mask[sample, depth] = bathymetry[sample, 0] >= depth_m
+        targets[sample] = ocean[sample, -1] + channel_scale * (0.015 * atmosphere[sample, 0] - 0.006 * atmosphere[sample, 1])
+        time_features[sample] = ((sample * 6) % 24, 100 + sample // 4, 0)
+    channel_depth = np.asarray(list(range(26)) * 4 + [0])
+    targets *= depth_mask[:, channel_depth]
+    path = ROOT / data["path"]
+    path.parent.mkdir(parents=True, exist_ok=True)
+    input_shape = np.asarray([history, channels, data["global_height"], data["global_width"]])
+    atmosphere_shape = np.asarray([atmosphere_channels, data["global_height"], data["global_width"]])
+    output_shape = np.asarray([channels, data["global_height"], data["global_width"]])
+    owned_pixels = sum((r[5] - r[4]) * (r[7] - r[6]) for r in selected_records)
+    coverage_fraction = owned_pixels / (data["global_height"] * data["global_width"])
+    np.savez_compressed(path, format_version=data["format_version"], input_shape=input_shape,
+                        atmosphere_shape=atmosphere_shape, output_shape=output_shape,
+                        global_tile_records=global_records, selected_tile_indices=selected_indices,
+                        selected_tile_records=selected_records, tile_origins=origins,
+                        tile_record_fields=np.asarray(["y0", "y1", "x0", "x1", "crop_top", "crop_bottom", "crop_left", "crop_right"]),
+                        tile_order="row-major", overlap_crop_semantics="midpoint ownership; crops partition global grid exactly",
+                        coverage_fraction=coverage_fraction, is_complete_global=False, synthetic=True,
+                        ocean=ocean, atmosphere=atmosphere, targets=targets, latitude_deg=latitude,
+                        longitude_deg=longitude, bathymetry_m=bathymetry, depth_mask=depth_mask,
+                        time_features=time_features, train_count=data["train_samples"])
+    print(json.dumps({"path": str(path.relative_to(ROOT)), "input_shape": input_shape.tolist(),
+                      "tile_shape": list(ocean.shape), "global_tile_count": len(global_records),
+                      "coverage_fraction": coverage_fraction, "bytes": path.stat().st_size}))
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,81 @@
+"""Autoregressive tile inference with the unchanged scientific global-grid contract."""
+
+import json
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "scripts"))
+from model.fuxi_ocean import FORMAT_VERSION, FuXiOcean, assemble_tiles
+from train import TileDataset, validate_data_contract
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.set_num_threads(config["runtime"]["num_threads"])
+    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
+    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
+    required = {"model", "model_config", "format_version"}
+    if not required.issubset(checkpoint):
+        raise ValueError(f"checkpoint missing keys: {sorted(required - checkpoint.keys())}")
+    model_config = checkpoint["model_config"]
+    model_args = model_config["architecture"]
+    model = FuXiOcean(**model_args).to(device)
+    model.load_state_dict(checkpoint["model"]); model.eval()
+    data = np.load(ROOT / config["data"]["path"])
+    validate_data_contract(data, config)
+    shape_keys = ("input_shape", "atmosphere_shape", "output_shape")
+    if checkpoint["format_version"] != FORMAT_VERSION or model_config["data_format_version"] != str(data["format_version"]) or any(model_config[key] != data[key].tolist() for key in shape_keys):
+        raise ValueError("checkpoint, model configuration, and data contract are incompatible")
+    predictions, truths, initial, latitudes, masks, origins = [], [], [], [], [], []
+    with torch.no_grad():
+        for index in range(int(data["train_count"]), len(data["ocean"])):
+            batch = TileDataset(data, [index])[0]
+            ocean, atmosphere, coordinates, bathymetry, mask, time_info, target, latitude = [value[None].to(device) for value in batch]
+            history = ocean; sample_predictions, sample_truths = [], []
+            for lead in range(config["inference"]["rollout_steps"]):
+                time_info[:, 2] = lead
+                prediction = model(history, atmosphere, coordinates, bathymetry, mask, time_info)
+                truth = target + lead * 0.005
+                sample_predictions.append(prediction.cpu().numpy()[0]); sample_truths.append(truth.cpu().numpy()[0])
+                history = torch.cat((history[:, 1:], prediction[:, None]), dim=1)
+            predictions.append(sample_predictions); truths.append(sample_truths); initial.append(ocean[0, -1].cpu().numpy())
+            latitudes.append(latitude[0].cpu().numpy()); masks.append(mask[0].cpu().numpy()); origins.append(data["tile_origins"][index])
+    output = ROOT / config["paths"]["inference"]; output.parent.mkdir(parents=True, exist_ok=True)
+    records = data["selected_tile_records"][int(data["train_count"]):]
+    # Exercise overlap-crop assembly on a bounded canvas without allocating the global field.
+    test_tile = np.ones((1, 1, config["data"]["tile_height"], config["data"]["tile_width"]), dtype=np.float32)
+    _, stitch_coverage = assemble_tiles(test_tile, [[0, config["data"]["tile_height"], 0, config["data"]["tile_width"], 0,
+                                                     config["data"]["tile_height"], 0, config["data"]["tile_width"]]],
+                                                 (config["data"]["tile_height"], config["data"]["tile_width"]))
+    coverage_fraction = float(sum((r[5] - r[4]) * (r[7] - r[6]) for r in records) /
+                              (config["data"]["global_height"] * config["data"]["global_width"]))
+    checkpoint_source = str(config["paths"]["checkpoint"])
+    np.savez_compressed(output, output_kind="sampled_tiles", format_version=FORMAT_VERSION,
+                         checkpoint_source=checkpoint_source, sample_count=len(predictions), synthetic=True,
+                         coverage_fraction=coverage_fraction, is_complete_global=False,
+                         variable_groups=json.dumps({"S": {"channels": [0, 26], "unit": "psu"}, "T": {"channels": [26, 52], "unit": "degC"}, "U": {"channels": [52, 78], "unit": "m s-1"}, "V": {"channels": [78, 104], "unit": "m s-1"}, "SSH": {"channels": [104, 105], "unit": "m"}}),
+                         prediction=np.asarray(predictions), truth=np.asarray(truths), initial=np.asarray(initial),
+                         latitude_deg=np.asarray(latitudes), depth_mask=np.asarray(masks), tile_origins=np.asarray(origins),
+                         tile_records=records, output_shape=data["output_shape"],
+                         lead_hours=config["data"]["time_step_hours"] * np.arange(1, config["inference"]["rollout_steps"] + 1))
+    metadata = {"output_kind": "sampled_tiles", "format_version": FORMAT_VERSION,
+                 "checkpoint_source": checkpoint_source, "sample_count": len(predictions), "synthetic": True,
+                 "output_shape": data["output_shape"].tolist(), "paper_rollout_steps": config["paper_model"]["rollout_steps"],
+                 "executed_rollout_steps": config["inference"]["rollout_steps"], "tile_streaming": True,
+                 "tile_origins": np.asarray(origins).tolist(), "coverage_fraction": coverage_fraction,
+                 "is_complete_global": False, "stitch_interface_verified": bool(stitch_coverage.all()),
+                 "coverage_semantics": "fraction of global cells owned after deterministic overlap crops",
+                 "variable_groups": json.loads(str(np.load(output)["variable_groups"]))}
+    metadata_path = ROOT / config["paths"]["inference_metadata"]
+    metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")
+    print(f"predictions={output.relative_to(ROOT)} shape={np.asarray(predictions).shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,61 @@
+"""Compute per-lead latitude-weighted model and persistence RMSE/MBE."""
+
+import json
+from pathlib import Path
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.fuxi_ocean import ocean_channel_depth_indices, weighted_error_sums
+
+
+def aggregate(prediction, truth, latitude, depth_mask):
+    channel_depth = ocean_channel_depth_indices()
+    mask = np.take(depth_mask, channel_depth, axis=1)
+    squared, bias, weight = weighted_error_sums(prediction, truth, latitude, mask)
+    return np.sqrt(squared / np.maximum(weight, 1)), bias / np.maximum(weight, 1)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / config["paths"]["inference"])
+    model_rmse, model_mbe, persistence_rmse, persistence_mbe = [], [], [], []
+    for lead in range(data["prediction"].shape[1]):
+        rmse, mbe = aggregate(data["prediction"][:, lead], data["truth"][:, lead], data["latitude_deg"], data["depth_mask"])
+        prmse, pmbe = aggregate(data["initial"], data["truth"][:, lead], data["latitude_deg"], data["depth_mask"])
+        model_rmse.append(rmse); model_mbe.append(mbe); persistence_rmse.append(prmse); persistence_mbe.append(pmbe)
+    arrays = [np.asarray(value) for value in (model_rmse, model_mbe, persistence_rmse, persistence_mbe)]
+    lead_hours = data["lead_hours"].tolist()
+    groups = json.loads(str(data["variable_groups"]))
+    def grouped(values, i):
+        return {name: {"unit": spec["unit"], "rmse": float(np.mean(values[0][i, slice(*spec["channels"])])),
+                       "mbe": float(np.mean(values[1][i, slice(*spec["channels"])]))} for name, spec in groups.items()}
+    metrics = {"output_kind": str(data["output_kind"]), "format_version": str(data["format_version"]),
+               "checkpoint_source": str(data["checkpoint_source"]), "sample_count": int(data["sample_count"]),
+               "coverage_fraction": float(data["coverage_fraction"]), "is_complete_global": bool(data["is_complete_global"]),
+               "synthetic": bool(data["synthetic"]), "aggregation": "latitude-weighted sampled tiles; never across unit groups",
+               "output_shape": data["output_shape"].tolist(), "variable_groups": groups,
+               "per_lead": [{"lead_hours": hour, "model_by_group": grouped(arrays[:2], i),
+                              "persistence_by_group": grouped(arrays[2:], i),
+                              "model_by_channel": {"rmse": arrays[0][i].tolist(), "mbe": arrays[1][i].tolist()},
+                              "persistence_by_channel": {"rmse": arrays[2][i].tolist(), "mbe": arrays[3][i].tolist()}}
+                             for i, hour in enumerate(lead_hours)]}
+    path = ROOT / config["paths"]["evaluation_metrics"]; path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(metrics, indent=2) + "\n")
+    figure, axes = plt.subplots(1, 5, figsize=(18, 3.5))
+    for axis, (name, spec) in zip(axes, groups.items()):
+        channel_slice = slice(*spec["channels"])
+        axis.plot(lead_hours, arrays[0][:, channel_slice].mean(1), "o-", label="model")
+        axis.plot(lead_hours, arrays[2][:, channel_slice].mean(1), "s--", label="persistence")
+        axis.set(title=name, xlabel="Lead (h)", ylabel=f"RMSE ({spec['unit']})"); axis.legend()
+    figure.tight_layout()
+    plot = ROOT / config["paths"]["evaluation_plot"]; figure.savefig(plot, dpi=160); plt.close(figure)
+    print(f"metrics={path.relative_to(ROOT)} plot={plot.relative_to(ROOT)} leads={lead_hours}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,118 @@
+"""Train FuXi-Ocean on globally indexed tiles, with optional torchrun DDP."""
+
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+import numpy as np
+import torch
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.fuxi_ocean import FORMAT_VERSION, FuXiOcean, channel_mask, latitude_weighted_charbonnier
+
+
+class TileDataset(Dataset):
+    def __init__(self, data, indices):
+        self.data, self.indices = data, list(indices)
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, item):
+        index = self.indices[item]
+        lat, lon = self.data["latitude_deg"][index], self.data["longitude_deg"][index]
+        coordinates = np.stack(np.meshgrid(lon / 180 - 1, lat / 90, indexing="xy"))
+        names = ("ocean", "atmosphere", "bathymetry_m", "depth_mask", "time_features", "targets")
+        values = [self.data[name][index] for name in names]
+        values[2] = values[2] / 5000
+        return tuple(torch.as_tensor(value, dtype=torch.float32) for value in (*values[:2], coordinates, *values[2:], lat))
+
+
+def validate_data_contract(data, config):
+    expected = config["data"]
+    checks = {"format_version": str(data["format_version"]) == config["data"]["format_version"],
+              "input_shape": data["input_shape"].tolist() == expected["input_shape"],
+              "atmosphere_shape": data["atmosphere_shape"].tolist() == expected["atmosphere_shape"],
+              "output_shape": data["output_shape"].tolist() == expected["output_shape"]}
+    if not all(checks.values()):
+        raise ValueError(f"data contract mismatch: {checks}")
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.set_num_threads(config["runtime"]["num_threads"])
+    random.seed(config["seed"]); np.random.seed(config["seed"]); torch.manual_seed(config["seed"])
+    world = int(os.environ.get("WORLD_SIZE", "1")); distributed = world > 1
+    if distributed:
+        torch.distributed.init_process_group(config["runtime"]["ddp_backend"])
+    rank = torch.distributed.get_rank() if distributed else 0
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    available_devices = torch.cuda.device_count() if torch.cuda.is_available() else 0
+    requested_auto_gpu = config["runtime"]["device"] != "cpu" and available_devices >= world
+    use_cuda = requested_auto_gpu and local_rank < available_devices
+    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+    data = np.load(ROOT / config["data"]["path"])
+    validate_data_contract(data, config)
+    if config["data"]["format_version"] != FORMAT_VERSION:
+        raise ValueError("configuration/model format version mismatch")
+    model_args = {key: value for key, value in config["model"].items() if key != "epsilon"}
+    model = FuXiOcean(**model_args).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"],
+                                  weight_decay=config["training"]["weight_decay"])
+    losses = []
+    train_count = int(data["train_count"])
+    dataset = TileDataset(data, range(train_count))
+    sampler = DistributedSampler(dataset, num_replicas=world, rank=rank, shuffle=True, seed=config["seed"]) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler,
+                        shuffle=sampler is None, drop_last=False)
+    for epoch in range(config["training"]["epochs"]):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        for batch in loader:
+            ocean, atmosphere, coordinates, bathymetry, mask, time_info, target, latitude = [value.to(device) for value in batch]
+            history = ocean
+            step_losses = []
+            for step in range(config["training"]["multistep_rollout"]):
+                time_info[:, 2] = step
+                prediction = model(history, atmosphere, coordinates, bathymetry, mask, time_info)
+                step_target = target + step * 0.005
+                step_losses.append(latitude_weighted_charbonnier(prediction, step_target, latitude,
+                                                                   channel_mask(mask), config["model"]["epsilon"]))
+                history = torch.cat((history[:, 1:], prediction[:, None]), dim=1)
+            loss = torch.stack(step_losses).mean()
+            optimizer.zero_grad(); loss.backward(); optimizer.step()
+            losses.append(float(loss.detach()))
+    local = torch.tensor([sum(losses), len(losses)], dtype=torch.float64, device=device)
+    if distributed:
+        torch.distributed.all_reduce(local)
+    if rank == 0:
+        raw_model = model.module if distributed else model
+        checkpoint_path = ROOT / config["paths"]["checkpoint"]
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        model_config = {"architecture": model_args, "data_format_version": FORMAT_VERSION,
+                        "input_shape": data["input_shape"].tolist(), "atmosphere_shape": data["atmosphere_shape"].tolist(),
+                        "output_shape": data["output_shape"].tolist()}
+        torch.save({"model": raw_model.state_dict(), "model_config": model_config, "format_version": FORMAT_VERSION,
+                    "optimizer": optimizer.state_dict()}, checkpoint_path)
+        metrics_path = ROOT / config["paths"]["training_metrics"]
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.write_text(json.dumps({"mean_loss": local[0].item() / local[1].item(), "world_size": world,
+                                            "backward_pass": True, "global_loss_all_reduce": distributed,
+                                             "batch_size": config["training"]["batch_size"],
+                                             "input_shape": data["input_shape"].tolist(), "synthetic": True}, indent=2) + "\n")
+        print(f"checkpoint={checkpoint_path.relative_to(ROOT)} loss={local[0].item() / local[1].item():.6f}")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()