File size: 5,053 Bytes
04fe5fd | 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 | """Generate metadata-compatible ERA5 HDF5 fixtures at the native 0.25 degree grid."""
from __future__ import annotations
import argparse
from pathlib import Path
import h5py
import numpy as np
OFFICIAL_VARIABLES = (
[f"Z{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+ [f"Q{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+ [f"T{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+ [f"U{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+ [f"V{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+ ["U10M", "V10M", "T2M", "MSLP"]
)
VARIABLE_ALIASES = {
**{f"Z{x}": f"geopotential_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
**{f"Q{x}": f"specific_humidity_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
**{f"T{x}": f"temperature_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
**{f"U{x}": f"u_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
**{f"V{x}": f"v_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
"U10M": "10m_u_component_of_wind",
"V10M": "10m_v_component_of_wind",
"T2M": "2m_temperature",
"MSLP": "mean_sea_level_pressure",
}
REAL_VARIABLES = tuple(VARIABLE_ALIASES[name] for name in OFFICIAL_VARIABLES)
def _synthetic_field(time_index: int, channel: int, height: int, width: int) -> np.ndarray:
lat = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None]
lon = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, :]
phase = np.float32(channel * 0.17)
field = np.sin(lon + phase) + 0.4 * np.cos(np.float32(time_index / 3.0) + phase) + 0.2 * lat
return np.asarray(field, dtype=np.float32)
def generate_fake_h5(output_dir: Path, years: list[int], stats_years: set[int], time_steps: int,
height: int, width: int, seed: int) -> None:
data_dir = output_dir / "data"
data_dir.mkdir(parents=True, exist_ok=True)
sums = np.zeros(len(OFFICIAL_VARIABLES), dtype=np.float64)
squared_sums = np.zeros_like(sums)
value_count = 0
for offset, year in enumerate(years):
path = data_dir / f"{year}.h5"
with h5py.File(path, "w") as handle:
dataset = handle.create_dataset(
"fields", shape=(time_steps, len(OFFICIAL_VARIABLES), height, width),
dtype="float32", chunks=(1, 1, height, width),
)
dataset.attrs["variables"] = list(REAL_VARIABLES)
dataset.attrs["time_step"] = 6
for time_index in range(time_steps):
for channel in range(len(OFFICIAL_VARIABLES)):
field = _synthetic_field(time_index + offset, channel, height, width)
dataset[time_index, channel] = field
if year in stats_years:
sums[channel] += field.sum(dtype=np.float64)
squared_sums[channel] += np.square(field, dtype=np.float64).sum()
if year in stats_years:
value_count += time_steps * height * width
print(f"{path}: fields={(time_steps, len(OFFICIAL_VARIABLES), height, width)}, variables={len(OFFICIAL_VARIABLES)}")
means = (sums / value_count).reshape(1, -1, 1, 1)
variances = squared_sums / value_count - np.square(means.reshape(-1))
stds = np.sqrt(np.maximum(variances, 1e-12)).reshape(1, -1, 1, 1)
stats_dir = output_dir / "stats"
stats_dir.mkdir(parents=True, exist_ok=True)
np.save(stats_dir / "global_means.npy", means)
np.save(stats_dir / "global_stds.npy", stds)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--output-dir", type=Path, default=Path("./data"))
parser.add_argument("--years", nargs="+", type=int, default=[2000, 2001, 2002])
parser.add_argument("--stats-years", nargs="+", type=int, default=None)
parser.add_argument("--time-steps", type=int, default=3)
parser.add_argument("--height", type=int, default=721)
parser.add_argument("--width", type=int, default=1440)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
if len(args.years) < 1 or min(args.time_steps, args.height, args.width) < 1:
parser.error("years, time-steps, height, and width must be positive")
if (args.height, args.width) != (721, 1440):
parser.error("ERA5 fixtures must use the native 0.25 degree grid 721x1440")
stats_years = set(args.stats_years or args.years[:1])
if not stats_years.issubset(args.years):
parser.error("stats-years must be included in years")
generate_fake_h5(args.output_dir, args.years, stats_years, args.time_steps,
args.height, args.width, args.seed)
if __name__ == "__main__":
main()
|