File size: 4,047 Bytes
006ea64 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | from __future__ import annotations
import argparse
from pathlib import Path
import h5py
import numpy as np
import yaml
from grid import lambert_grid
def load_config(path: str | Path) -> dict:
with Path(path).open("r", encoding="utf-8") as handle:
return yaml.safe_load(handle)
def write_temporal_fields(
path: Path,
variables: list[str],
num_timesteps: int,
image_size: tuple[int, int],
time_step_hours: int,
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
channels = len(variables)
height, width = image_size
means = np.zeros((1, channels, 1, 1), dtype=np.float32)
stds = np.ones((1, channels, 1, 1), dtype=np.float32)
with h5py.File(path, "w") as handle:
fields = handle.create_dataset(
"fields",
shape=(num_timesteps, channels, height, width),
dtype=np.float32,
chunks=(1, channels, height, width),
fillvalue=0.0,
)
fields.attrs["variables"] = variables
fields.attrs["time_step"] = time_step_hours
handle.create_dataset("global_means", data=means)
handle.create_dataset("global_stds", data=stds)
def write_invariants(
path: Path,
variables: list[str],
image_size: tuple[int, int],
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
height, width = image_size
target_lat, target_lon = lambert_grid(image_size)
invariants = np.zeros((len(variables), height, width), dtype=np.float32)
with h5py.File(path, "w") as handle:
fields = handle.create_dataset(
"fields",
shape=(len(variables), height, width),
dtype=np.float32,
chunks=(1, height, width),
data=invariants,
)
fields.attrs["variables"] = variables
handle.create_dataset("lat", data=target_lat)
handle.create_dataset("lon", data=target_lon)
def generate(config: dict) -> None:
data = config["data"]
root = Path(data["root_dir"])
years = sorted(
set(data["train_years"] + data["val_years"] + data["test_years"])
)
era5_image_size = tuple(data["era5_image_size"])
image_size = tuple(data["image_size"])
if era5_image_size != (721, 1440):
raise ValueError("ERA5 grid must be 721 x 1440")
if image_size != (512, 640):
raise ValueError("Regional grid must be 512 x 640")
if len(data["era5_variables"]) != 26:
raise ValueError("The configured ERA5 input must contain 26 channels")
if len(data["state_variables"]) != 99:
raise ValueError("The configured local state must contain 99 channels")
if data["invariant_variables"] != ["lsm", "orography"]:
raise ValueError("Invariant order must be [lsm, orography]")
for year in years:
write_temporal_fields(
root / "era5" / "data" / f"{year}.h5",
data["era5_variables"],
data["num_timesteps"],
era5_image_size,
data["time_step_hours"],
)
write_temporal_fields(
root / "hrrr" / "data" / f"{year}.h5",
data["state_variables"],
data["num_timesteps"],
image_size,
data["time_step_hours"],
)
write_invariants(
root / "hrrr" / "invariants.h5",
data["invariant_variables"],
image_size,
)
print(f"Generated project validation data under {root}")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate lightweight StormCast data")
parser.add_argument("--config", default="conf/config.yaml")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
config_path = Path(args.config).resolve()
config = load_config(config_path)
project_root = config_path.parent.parent
for key in ("root_dir",):
path = Path(config["data"][key])
if not path.is_absolute():
config["data"][key] = str((project_root / path).resolve())
generate(config)
|