File size: 2,119 Bytes
9191802 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | """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()
|