FuXi_v21 / scripts /variables.py
yzt15806542928's picture
Upload folder using huggingface_hub
9191802 verified
Raw
History Blame Contribute Delete
1.8 kB
from __future__ import annotations
PRESSURE_LEVELS = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
PRESSURE_VARIABLES = ["z", "t", "u", "v", "q"]
SURFACE_VARIABLES = [
"msl", "t2m", "d2m", "sst", "ws10m", "ws100m", "u10m", "v10m",
"u100m", "v100m", "lcc", "mcc", "hcc", "tcc", "ssr", "ssrd",
"fdir", "ttr", "tcw", "tp",
]
C85_CHANNEL_NAMES = [f"{name}{level}" for name in PRESSURE_VARIABLES for level in PRESSURE_LEVELS]
C85_CHANNEL_NAMES += SURFACE_VARIABLES
DIAGNOSTIC_CHANNELS = ["ssr", "ssrd", "fdir", "ttr", "tp"]
assert len(C85_CHANNEL_NAMES) == 85
def c85_from_config(config: dict) -> tuple[list[str], list[str]]:
"""Build and validate the fixed C85 channel contract from configuration."""
variables = config["variables"]
mapping = variables["mapping"]
if mapping != {
"pressure_order": "variable_major",
"pressure_name": "{variable}{level}",
"surface_name": "{variable}",
"source_dataset": "fields",
"units": "source_native",
"transform": "identity",
}:
raise ValueError("variables.mapping must preserve the C85 source and ordering contract")
channels = [
mapping["pressure_name"].format(variable=name, level=level)
for name in variables["pressure"]
for level in variables["pressure_levels"]
]
channels.extend(mapping["surface_name"].format(variable=name) for name in variables["surface"])
diagnostics = list(variables["diagnostic"])
if channels != C85_CHANNEL_NAMES:
raise ValueError("Configured variables do not match the required C85 channel order")
if diagnostics != DIAGNOSTIC_CHANNELS:
raise ValueError("Configured diagnostic variables do not match the required C85 contract")
return channels, diagnostics