model/fourcastnet_v2.py

"""OneScience adapter for NVIDIA's official legacy FourCastNet v2 network."""

from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Mapping

import torch
from torch import nn


def _load_model_class():
    package_dir = Path(__file__).resolve().parent / "fcnv2"
    if not (package_dir / "fcnv2_sfnonet.py").is_file():
        raise FileNotFoundError(
            f"Bundled FCNv2 source was not found at {package_dir}"
        )

    package_path = str(package_dir)
    if package_path not in sys.path:
        sys.path.insert(0, package_path)

    try:
        from fcnv2_sfnonet import FourierNeuralOperatorNet
    except ModuleNotFoundError as error:
        if error.name == "torch_harmonics":
            raise ModuleNotFoundError(
                "FourCastNet v2 requires NVIDIA torch-harmonics. Install the "
                "version pinned by this project before constructing the model."
            ) from error
        raise
    return FourierNeuralOperatorNet


def _official_params(model_config: Mapping[str, Any]) -> SimpleNamespace:
    required = {
        "img_size",
        "in_channels",
        "out_channels",
        "spectral_transform",
        "filter_type",
        "scale_factor",
        "embed_dim",
        "num_layers",
        "num_blocks",
        "normalization_layer",
        "mlp_mode",
        "spectral_layers",
        "complex_activation",
        "hard_thresholding_fraction",
        "big_skip",
    }
    missing = sorted(required.difference(model_config))
    if missing:
        raise ValueError(f"Missing FourCastNet v2 model settings: {missing}")

    height, width = model_config["img_size"]
    hidden_height = height // model_config["scale_factor"]
    hidden_width = width // model_config["scale_factor"]
    if hidden_height < 2 or hidden_width < 2:
        raise ValueError("The internal SFNO grid must have at least 2 x 2 points")

    return SimpleNamespace(
        img_crop_shape_x=int(height),
        img_crop_shape_y=int(width),
        N_in_channels=int(model_config["in_channels"]),
        N_out_channels=int(model_config["out_channels"]),
        spectral_transform=model_config["spectral_transform"],
        filter_type=model_config["filter_type"],
        scale_factor=int(model_config["scale_factor"]),
        embed_dim=int(model_config["embed_dim"]),
        num_layers=int(model_config["num_layers"]),
        num_blocks=int(model_config["num_blocks"]),
        normalization_layer=model_config["normalization_layer"],
        mlp_mode=model_config["mlp_mode"],
        spectral_layers=int(model_config["spectral_layers"]),
        complex_activation=model_config["complex_activation"],
        hard_thresholding_fraction=float(
            model_config["hard_thresholding_fraction"]
        ),
        big_skip=bool(model_config["big_skip"]),
    )


class FourCastNetV2(nn.Module):
    """Build the exact official FCNv2 network behind a stable project API."""

    def __init__(
        self,
        model_config: Mapping[str, Any],
    ) -> None:
        super().__init__()
        self.model_config = dict(model_config)
        self.expected_shape = (
            int(model_config["in_channels"]),
            int(model_config["img_size"][0]),
            int(model_config["img_size"][1]),
        )
        model_class = _load_model_class()
        self.model = model_class(_official_params(model_config))

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        if inputs.ndim != 4:
            raise ValueError(f"Expected [B,C,H,W], got {tuple(inputs.shape)}")
        if tuple(inputs.shape[1:]) != self.expected_shape:
            raise ValueError(
                f"Expected trailing shape {self.expected_shape}, "
                f"got {tuple(inputs.shape[1:])}"
            )
        return self.model(inputs)

    def no_weight_decay(self) -> set[str]:
        return {f"model.{name}" for name in self.model.no_weight_decay()}


def _unwrap_state_dict(checkpoint: Any) -> Mapping[str, torch.Tensor]:
    if not isinstance(checkpoint, Mapping):
        raise TypeError("Checkpoint must contain a mapping")
    for key in ("model_state", "model_state_dict", "state_dict"):
        candidate = checkpoint.get(key)
        if isinstance(candidate, Mapping):
            return candidate
    if checkpoint and all(isinstance(value, torch.Tensor) for value in checkpoint.values()):
        return checkpoint
    raise KeyError("Checkpoint has no recognized model state mapping")


def _normalize_state_keys(
    state_dict: Mapping[str, torch.Tensor], model: nn.Module
) -> dict[str, torch.Tensor]:
    target_keys = set(model.state_dict())
    normalized: dict[str, torch.Tensor] = {}
    for key, value in state_dict.items():
        clean_key = key
        while clean_key.startswith("module."):
            clean_key = clean_key[len("module.") :]
        if clean_key not in target_keys and f"model.{clean_key}" in target_keys:
            clean_key = f"model.{clean_key}"
        normalized[clean_key] = value
    return normalized


def load_checkpoint(
    model: nn.Module,
    checkpoint_path: str | Path,
    *,
    expected_profile: str,
    expected_variables: list[str],
    allowed_stages: set[str],
    allowed_initializations: set[str],
    strict: bool = True,
    map_location: str | torch.device = "cpu",
) -> dict[str, Any]:
    """Load a project checkpoint without changing its parameter tensors."""

    checkpoint = torch.load(
        Path(checkpoint_path).expanduser(),
        map_location=map_location,
        weights_only=False,
    )
    validate_project_checkpoint(
        checkpoint,
        expected_profile=expected_profile,
        expected_variables=expected_variables,
        allowed_stages=allowed_stages,
        allowed_initializations=allowed_initializations,
    )
    state_dict = _normalize_state_keys(_unwrap_state_dict(checkpoint), model)
    incompatible = model.load_state_dict(state_dict, strict=strict)
    return {
        "checkpoint": checkpoint,
        "missing_keys": list(incompatible.missing_keys),
        "unexpected_keys": list(incompatible.unexpected_keys),
    }


def validate_project_checkpoint(
    checkpoint: Any,
    *,
    expected_profile: str,
    expected_variables: list[str],
    allowed_stages: set[str],
    allowed_initializations: set[str],
) -> None:
    if not isinstance(checkpoint, Mapping):
        raise TypeError("Project checkpoint must contain metadata")
    expected = {
        "checkpoint_format": "fourcastnet_v2_project",
        "scratch_lineage": True,
        "model_profile": expected_profile,
        "variables": expected_variables,
    }
    for key, value in expected.items():
        if checkpoint.get(key) != value:
            raise ValueError(
                f"Checkpoint metadata {key!r} does not match the project config"
            )
    if checkpoint.get("stage") not in allowed_stages:
        raise ValueError(
            f"Checkpoint stage must be one of {sorted(allowed_stages)}"
        )
    if checkpoint.get("initialization") not in allowed_initializations:
        raise ValueError(
            "Checkpoint does not have an approved random-initialization lineage"
        )
    expected_initialization = {
        "one_step": "random",
        "finetune": "one_step_checkpoint",
    }.get(checkpoint.get("stage"))
    if checkpoint.get("initialization") != expected_initialization:
        raise ValueError(
            "Checkpoint stage and initialization metadata are inconsistent"
        )