scripts/build_static.py

"""Build project-defined FuXi static fields and C85 validity mask from scratch."""

from __future__ import annotations

import argparse

import numpy as np

from common import load_config, resolve_path
from variables import c85_from_config


def build_static_fields(height: int, width: int) -> np.ndarray:
    latitude = np.deg2rad(np.linspace(90.0, -90.0, height, dtype=np.float32))[:, None]
    longitude = np.deg2rad(np.arange(width, dtype=np.float32) * (360.0 / width))[None, :]
    lat = np.broadcast_to(latitude, (height, width))
    lon = np.broadcast_to(longitude, (height, width))
    geopotential = np.zeros_like(lat)
    land_sea = np.ones_like(lat)
    return np.stack((geopotential, land_sea, np.cos(lat), np.sin(lat), np.cos(lon), np.sin(lon)))


def build_channel_mask(channels: list[str], height: int, width: int) -> np.ndarray:
    # The reconstruction has no external missing-channel metadata; all C85 fields are valid.
    return np.ones((len(channels), height, width), dtype=np.float32)


def ensure_static_resources(cfg: dict, force: bool = False) -> tuple:
    channels, _ = c85_from_config(cfg)
    profile = cfg["model"]["profiles"]["full"]
    height, width = profile["grid_size"]
    static_path = resolve_path(cfg["model"]["static_fields_file"], cfg)
    mask_path = resolve_path(cfg["model"]["channel_mask_file"], cfg)
    static_path.parent.mkdir(parents=True, exist_ok=True)
    mask_path.parent.mkdir(parents=True, exist_ok=True)
    if force or not static_path.is_file():
        np.save(static_path, build_static_fields(height, width))
    if force or not mask_path.is_file():
        np.save(mask_path, build_channel_mask(channels, height, width))
    return static_path, mask_path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="conf/config.yaml")
    args = parser.parse_args()
    cfg = load_config(args.config)
    static_path, mask_path = ensure_static_resources(cfg, force=True)
    print(f"Saved from-scratch static resources to {static_path} and {mask_path}")


if __name__ == "__main__":
    main()