| """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() |
|
|
