| """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: |
| |
| 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() |
|
|
