model/aifs.py

#!/usr/bin/env python
# coding: utf-8
"""
AIFS v1.1 — Onescience-style Model  (真正从零构建)
====================================================

``class AIFS(nn.Module)`` — 兼容 onescience 的 ECMWF AIFS
编码器-处理器-解码器图神经网络。

两种构建模式
------------
1. **从零构建** — 无需 ``.ckpt`` 文件，与 FengWu/Fuxi 一致::

       model = AIFS.from_scratch(device="cuda")

   从项目内的静态文件加载架构配置（``aifs_config.json``, 47 KB）和
   N320 网格坐标（``grid-n320.npz``, 4.2 MB），通过 ``anemoi-graphs``
   构建 HeteroData 图，由 ``AnemoiModelEncProcDec`` 随机初始化所有权重。

2. **加载预训练** — 向后兼容::

       model = AIFS("aifs.ckpt", device="cuda", pretrained=True)

架构 (anemoi-models 0.5.0)
--------------------------
- **Encoder**  : ``GraphTransformerForwardMapper`` — data nodes → hidden nodes
- **Processor**: ``TransformerProcessor`` × 16 layers — 在 o96 隐藏图上做
  滑动窗口注意力 (window_size=1120, flash_attention, 16 heads)
- **Decoder**  : ``GraphTransformerBackwardMapper`` — hidden → data nodes
- **Bounding** : ``ReLUBounding`` / ``HardtanhBounding`` / ``FractionBounding``
  — 强制输出变量满足物理约束

静态文件说明
------------
- ``aifs_config.json``: 从官方 checkpoint 的 ``ai-models.json`` 提取
  (model_config + data_indices + dataset), 定义完整模型架构
- ``grid-n320.npz``: N320 高斯网格经纬度坐标 (542,080 节点), 从官方
  checkpoint 的 ``latitudes.numpy`` / ``longitudes.numpy`` 提取

Reference
---------
- Lang et al., *AIFS — ECMWF's data-driven forecasting system*, arXiv:2406.01465
- anemoi-models: https://github.com/ecmwf/anemoi-models
- anemoi-graphs: https://github.com/ecmwf/anemoi-graphs
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass
from typing import Dict, List, Optional

import numpy as np
import torch
from torch import nn

from onescience.models.meta import ModelMetaData

LOG = logging.getLogger(__name__)

# ============================================================================
# 静态文件路径（相对于本文件所在目录）
# ============================================================================
_HERE = os.path.dirname(os.path.abspath(__file__))
_CONFIG_PATH = os.path.join(_HERE, "aifs_config.json")
_GRID_PATH = os.path.join(_HERE, "grid-n320.npz")


# ============================================================================
# Metadata  (onescience convention)
# ============================================================================

@dataclass
class MetaData(ModelMetaData):
    name: str = "AIFS"
    jit: bool = False
    cuda_graphs: bool = False
    amp: bool = True
    amp_gpu: bool = True
    bf16: bool = False
    onnx: bool = False              # flash_attn 不兼容 ONNX
    func_torch: bool = False
    auto_grad: bool = False
    var_dim: int = -1               # 每图节点的变量数


# ============================================================================
# Model
# ============================================================================

class AIFS(nn.Module):
    """AIFS v1.1 编码器-处理器-解码器 GNN。

    Parameters
    ----------
    checkpoint_path : str, optional
        ``pretrained=True`` 时必传——预训练 ``.ckpt`` 文件路径。
    device : str
        PyTorch 设备。默认 ``"cuda"``。
    pretrained : bool
        ``False`` (默认): 从零构建——读取本地 ``aifs_config.json`` +
        ``grid-n320.npz``, 不依赖 ``.ckpt`` 文件。
        ``True``: 从 ``.ckpt`` 加载预训练权重（向后兼容）。
    """

    metadata = MetaData()

    # ==================================================================
    # Factory: 真正从零（无需 .ckpt）
    # ==================================================================

    @classmethod
    def from_scratch(cls, device: str = "cuda") -> "AIFS":
        """从零创建 AIFS 模型——无需 ``.ckpt`` 文件。

        等价于 FengWu/Fuxi 的 ``model = Fengwu()`` 模式。
        所有权重随机初始化，架构配置和网格坐标从项目内静态文件加载。
        """
        return cls(device=device, pretrained=False)

    # ==================================================================
    # Construction
    # ==================================================================

    def __init__(
        self,
        checkpoint_path: Optional[str] = None,
        device: str = "cuda",
        pretrained: bool = False,
    ):
        super().__init__()

        if pretrained:
            if not checkpoint_path:
                raise ValueError(
                    "checkpoint_path is required when pretrained=True"
                )
            if not os.path.exists(checkpoint_path):
                raise FileNotFoundError(
                    f"Checkpoint not found: {checkpoint_path}"
                )
            self._init_from_pretrained(checkpoint_path, device)
        else:
            self._init_from_scratch(device)

    # ==================================================================
    # Mode 1: 从零构建（无需 .ckpt）
    # ==================================================================

    def _init_from_scratch(self, device: str):
        """用静态配置文件和标准网格坐标构建模型。"""
        LOG.info("Building AIFS from scratch (no checkpoint) …")

        # ---- 1. 加载静态架构配置 ----------------------------------------
        if not os.path.exists(_CONFIG_PATH):
            raise FileNotFoundError(
                f"Static config not found: {_CONFIG_PATH}. "
                "Run extract_static_config.py first."
            )
        if not os.path.exists(_GRID_PATH):
            raise FileNotFoundError(
                f"Grid file not found: {_GRID_PATH}. "
                "Run extract_static_config.py first."
            )

        with open(_CONFIG_PATH) as f:
            config_data = json.load(f)
        LOG.info(
            "Loaded static config (%d KB)",
            os.path.getsize(_CONFIG_PATH) // 1024,
        )

        grid = np.load(_GRID_PATH)
        lat = np.asarray(grid["latitudes"], dtype=np.float32)
        lon = np.asarray(grid["longitudes"], dtype=np.float32)
        LOG.info(
            "Loaded N320 grid: %d nodes", len(lat),
        )

        # ---- 2. 构建 model_config + data_indices (官方类) ---------
        # IndexCollection 需要 OmegaConf, AnemoiModelEncProcDec 需要 DotDict.
        # 先从原始 JSON dict 创建 OmegaConf, 再转为 DotDict.
        raw_mc = config_data["model_config"]

        from anemoi.utils.config import DotDict

        # ---- 兼容性: 0.5.0 config → 0.9.0 API ---------------------------
        # 1. activation 参数已从 mapper/processor 移除（内部默认 GELU）
        for section in ("encoder", "decoder", "processor"):
            raw_mc["model"][section].pop("activation", None)
        raw_mc["model"].pop("activation", None)

        # 2. layer_kernels 从顶层移入各子模块 config
        lk = raw_mc["model"].get("layer_kernels", {})
        for section in ("encoder", "decoder", "processor"):
            if section in lk:
                raw_mc["model"][section]["layer_kernels"] = lk[section]

        model_config = DotDict(raw_mc)

        all_vars = config_data["dataset"]["variables"]
        name_to_index = {name: i for i, name in enumerate(all_vars)}

        from anemoi.models.data_indices.collection import IndexCollection
        from omegaconf import OmegaConf
        data_indices = IndexCollection(
            config=OmegaConf.create(raw_mc),
            name_to_index=name_to_index,
        )
        graph_data = self._build_graph(lat, lon)

        # 提取 area_weight 作为 node_weights
        area_wt = (
            graph_data["data"]["area_weight"].cpu().numpy().squeeze()
        )

        # ---- 4. 实例化 AnemoiModelEncProcDec（随机权重）---------------
        # 训练时直接用裸模型——AnemoiModelInterface 自带 normalizer
        # 预处理，会导致训练时双重归一化。保存 checkpoint 时再包装。
        LOG.info("Instantiating AnemoiModelEncProcDec (random weights) …")
        from anemoi.models.models.encoder_processor_decoder import \
            AnemoiModelEncProcDec

        self._model = AnemoiModelEncProcDec(
            model_config=model_config,
            data_indices=data_indices,
            statistics={},
            graph_data=graph_data,
            truncation_data={},
        ).to(device)

        # 保存接口构建所需素材（_save 中构建 AnemoiModelInterface 用）
        self._interface_config = model_config
        self._interface_graph_data = graph_data
        self._interface_data_indices = data_indices

        # ---- 5. 提取变量排序和坐标 -----------------------------------
        self._meta = config_data
        self._arrays = {
            "latitudes": lat,
            "longitudes": lon,
            "area_weight": area_wt,
        }
        self._extract_metadata()

        LOG.info(
            "AIFS (from scratch) ready: %d → %d vars, %d grid pts, "
            "%.1f M params",
            len(self._input_vars),
            len(self._output_vars),
            len(lat),
            sum(p.numel() for p in self.parameters()) / 1e6,
        )

    def _build_graph(
        self,
        lat: np.ndarray,
        lon: np.ndarray,
    ) -> "HeteroData":
        """构建 AIFS 的 HeteroData 图（N320 ↔ o96）。

        完全使用 ``anemoi-graphs`` 官方 API, 不依赖任何外部数据文件。
        节点坐标来自标准 ECMWF 网格定义, 边由纯几何算法计算。
        """
        from torch_geometric.data import HeteroData

        from anemoi.graphs.edges.builders.cutoff import CutOffEdges
        from anemoi.graphs.edges.builders.knn import KNNEdges
        from anemoi.graphs.nodes.builders.from_reduced_gaussian import \
            ReducedGaussianGridNodes
        from anemoi.graphs.nodes.builders.from_vectors import LatLonNodes

        graph = HeteroData()

        # -- 数据节点 (N320: 542,080 nodes) --
        data_builder = LatLonNodes(
            latitudes=lat, longitudes=lon, name="data",
        )
        data_attrs = {
            "area_weight": {
                "_target_": "anemoi.graphs.nodes.attributes.SphericalAreaWeights",
                "fill_value": 0,
                "norm": "unit-max",
            },
        }
        graph = data_builder.update_graph(graph, attrs_config=data_attrs)

        # -- 隐藏节点 (o96: ~8,000 nodes) --
        hidden_builder = ReducedGaussianGridNodes(
            grid="o96", name="hidden",
        )
        graph = hidden_builder.update_graph(graph, attrs_config={})

        # -- data → hidden 边 (CutOffEdges, cutoff_factor=0.6) --
        d2h_attrs = {
            "edge_dirs": {
                "_target_": "anemoi.graphs.edges.attributes.EdgeDirection",
                "norm": "unit-std",
            },
            "edge_length": {
                "_target_": "anemoi.graphs.edges.attributes.EdgeLength",
                "norm": "unit-std",
            },
        }
        d2h_builder = CutOffEdges(
            source_name="data",
            target_name="hidden",
            cutoff_factor=0.6,
        )
        graph = d2h_builder.update_graph(graph, attrs_config=None)
        graph = d2h_builder.register_attributes(graph, d2h_attrs)

        # -- hidden → data 边 (KNNEdges, K=3) --
        h2d_attrs = {
            "edge_dirs": {
                "_target_": "anemoi.graphs.edges.attributes.EdgeDirection",
                "norm": "unit-std",
            },
            "edge_length": {
                "_target_": "anemoi.graphs.edges.attributes.EdgeLength",
                "norm": "unit-std",
            },
        }
        h2d_builder = KNNEdges(
            source_name="hidden",
            target_name="data",
            num_nearest_neighbours=3,
        )
        graph = h2d_builder.update_graph(graph, attrs_config=None)
        graph = h2d_builder.register_attributes(graph, h2d_attrs)

        LOG.info(
            "Graph built: data(%d nodes) ↔ hidden(%d nodes), "
            "d→h edges=%d, h→d edges=%d",
            graph["data"].num_nodes,
            graph["hidden"].num_nodes,
            graph["data", "to", "hidden"].edge_index.shape[1],
            graph["hidden", "to", "data"].edge_index.shape[1],
        )
        return graph

    # ==================================================================
    # Mode 2: 加载预训练 (pretrained=True, 向后兼容)
    # ==================================================================

    def _init_from_pretrained(self, checkpoint_path: str, device: str):
        """从 checkpoint 加载完整序列化模型（图 + 权重 + 索引）。"""
        LOG.info("Loading AIFS checkpoint from %s …", checkpoint_path)
        _cp = torch.load(
            checkpoint_path, map_location="cpu", weights_only=False,
        )
        self._model = _cp.to(device)

        # 提取元数据
        self._meta, self._arrays = self._read_checkpoint_metadata(
            checkpoint_path,
        )
        self._extract_metadata()

        LOG.info(
            "AIFS (pretrained) ready: %d → %d vars, %d grid pts, "
            "%.1f M params",
            len(self._input_vars),
            len(self._output_vars),
            len(self.latitudes),
            sum(p.numel() for p in self.parameters()) / 1e6,
        )

    # ==================================================================
    # Metadata helpers
    # ==================================================================

    @staticmethod
    def _read_checkpoint_metadata(checkpoint_path: str):
        """从 checkpoint ZIP 中读取 JSON 元数据和支持数组。"""
        import zipfile
        from anemoi.utils.checkpoints import load_supporting_arrays

        with zipfile.ZipFile(checkpoint_path, "r") as zf:
            # 支持新旧两种元数据文件名
            metadata = None
            for name in ["anemoi.json", "ai-models.json"]:
                for fname in zf.namelist():
                    if fname.endswith(name):
                        metadata = json.load(zf.open(fname, "r"))
                        break
                if metadata:
                    break

            if metadata is None:
                raise FileNotFoundError(
                    f"No metadata JSON found in {checkpoint_path}"
                )

            arrays = load_supporting_arrays(
                zf, metadata.get("supporting_arrays_paths", {}),
            )
            return metadata, arrays

    def _extract_metadata(self):
        """从已加载的元数据中提取变量排序和网格坐标。"""
        all_vars: List[str] = self._meta["dataset"]["variables"]
        di = self._meta["data_indices"]["data"]
        self._input_vars = [all_vars[i] for i in di["input"]["full"]]
        self._output_vars = [all_vars[i] for i in di["output"]["full"]]
        self._aifs_name_to_ds_idx: Dict[str, int] = {
            n: i for i, n in enumerate(all_vars)
        }

        lat = np.asarray(self._arrays["latitudes"], dtype=np.float32)
        lon = np.asarray(self._arrays["longitudes"], dtype=np.float32)
        self.register_buffer("latitudes", torch.from_numpy(lat))
        self.register_buffer("longitudes", torch.from_numpy(lon))

    # ------------------------------------------------------------------
    # Public accessors
    # ------------------------------------------------------------------

    @property
    def num_input_vars(self) -> int:
        """输入张量中的变量数 (103)。"""
        return len(self._input_vars)

    @property
    def num_output_vars(self) -> int:
        """输出张量中的变量数 (102)。"""
        return len(self._output_vars)

    @property
    def num_grid_points(self) -> int:
        """N320 网格节点数 (542,080)。"""
        return len(self.latitudes)

    @property
    def node_weights(self) -> Optional[np.ndarray]:
        """损失函数中每节点面积权重，无则为 None。"""
        for k in ["node_weights", "area_weight"]:
            if k in self._arrays:
                return np.asarray(
                    self._arrays[k], dtype=np.float32,
                ).squeeze()
        return None

    @property
    def input_variables(self) -> List[str]:
        """输入张量中 AIFS 变量名的有序列表。"""
        return self._input_vars

    @property
    def output_variables(self) -> List[str]:
        """输出张量中 AIFS 变量名的有序列表。"""
        return self._output_vars

    def aifs_name_to_dataset_index(self, name: str) -> int:
        """返回 *name* 在 115 变量数据集中的位置。"""
        return self._aifs_name_to_ds_idx.get(name, -1)

    def input_name_to_channel(self, name: str) -> int:
        """返回 *name* 在输入张量中的通道索引。"""
        try:
            return self._input_vars.index(name)
        except ValueError:
            return -1

    def output_name_to_channel(self, name: str) -> int:
        """返回 *name* 在输出张量中的通道索引。"""
        try:
            return self._output_vars.index(name)
        except ValueError:
            return -1

    # ------------------------------------------------------------------
    # Forward  (匹配官方 AnemoiModelEncProcDec.forward)
    # ------------------------------------------------------------------

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。

        等价于 ``AnemoiModelEncProcDec.forward(x, model_comm_group=None)``
        在单 GPU 训练/推理时的行为。

        Parameters
        ----------
        x : torch.Tensor
            归一化输入。
            - 4-D: ``(B, T, G, V_in)`` — 自动提升为 5-D。
            - 5-D: ``(B, T, E, G, V_in)`` — ensemble 维度保留。

        Returns
        -------
        torch.Tensor
            - 4-D 输入 → ``(B, G, V_out)``
            - 5-D 输入 → ``(B, E, G, V_out)``
        """
        if x.ndim == 4:
            x = x.unsqueeze(2)          # (B, T, G, V) → (B, T, 1, G, V)

        out = self._model(x)            # (B, E, G, V_out)

        if out.shape[1] == 1:
            out = out.squeeze(1)        # (B, 1, G, V) → (B, G, V)
        return out

    # ------------------------------------------------------------------
    # Inference helper
    # ------------------------------------------------------------------

    @torch.no_grad()
    def predict(self, x: torch.Tensor) -> torch.Tensor:
        """推理前向传播（无梯度，eval 模式，autocast fp16）。

        FlashAttention 要求 fp16/bf16 输入，用 autocast 自动转换。
        """
        was_training = self.training
        self.eval()
        with torch.amp.autocast("cuda", dtype=torch.float16):
            out = self.forward(x)
        if was_training:
            self.train()
        return out.float()

    # ------------------------------------------------------------------
    # Train / eval propagation
    # ------------------------------------------------------------------

    def train(self, mode: bool = True):
        """将 train/eval 模式传播到被封装的官方模型。"""
        super().train(mode)
        self._model.train(mode)
        return self

    def eval(self):
        return self.train(False)