File size: 2,846 Bytes
f4a39ee | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | #!/usr/bin/env python3
"""Extract official NeuralGCM static fields from released checkpoints."""
from __future__ import annotations
import argparse
import pickle
try:
from common import PROJECT_ROOT, load_config, resolve_path
except ModuleNotFoundError: # supports ``python -m scripts.prepare_static_data``
from scripts.common import PROJECT_ROOT, load_config, resolve_path
STATIC_VARIABLES = ("geopotential_at_surface", "land_sea_mask")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
parser.add_argument(
"--mode",
action="append",
dest="modes",
help="profile to extract; repeat for multiple profiles (default: all)",
)
parser.add_argument(
"--overwrite",
action="store_true",
help="replace an existing profile static NetCDF",
)
args = parser.parse_args()
config = load_config(args.config)
profiles = config["model"]["profiles"]
static_files = config["data"].get("static_files", {})
modes = args.modes or list(profiles)
import xarray as xr
for mode in modes:
if mode not in profiles:
raise ValueError(f"Unknown model profile {mode!r}")
if mode not in static_files:
raise ValueError(f"data.static_files has no path for {mode!r}")
checkpoint_path = resolve_path(
profiles[mode]["official_reference"], args.config
)
output_path = resolve_path(static_files[mode], args.config)
if output_path.exists() and not args.overwrite:
print(f"Static fields already exist: {output_path}")
continue
with checkpoint_path.open("rb") as handle:
payload = pickle.load(handle)
if not isinstance(payload, dict) or "aux_ds_dict" not in payload:
raise ValueError(
f"{checkpoint_path} is missing the official aux_ds_dict"
)
source = xr.Dataset.from_dict(payload["aux_ds_dict"])
missing = [name for name in STATIC_VARIABLES if name not in source]
if missing:
raise ValueError(
f"{checkpoint_path} is missing static variables {missing}"
)
static = source[list(STATIC_VARIABLES)]
static.attrs.update(
{
"source": str(checkpoint_path),
"neuralgcm_profile": mode,
"synthetic": "false",
}
)
output_path.parent.mkdir(parents=True, exist_ok=True)
static.to_netcdf(output_path)
print(
f"Extracted {mode} static fields to {output_path}: "
f"longitude={static.sizes['longitude']}, "
f"latitude={static.sizes['latitude']}"
)
if __name__ == "__main__":
main()
|