OneScience-Group/OneForecast

1"""Generate metadata-compatible ERA5 HDF5 fixtures at the native 0.25 degree grid."""2 3from __future__ import annotations4 5import argparse6from pathlib import Path7 8import h5py9import numpy as np10 11 12OFFICIAL_VARIABLES = (13    [f"Z{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]14    + [f"Q{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]15    + [f"T{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]16    + [f"U{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]17    + [f"V{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]18    + ["U10M", "V10M", "T2M", "MSLP"]19)20 21VARIABLE_ALIASES = {22    **{f"Z{x}": f"geopotential_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},23    **{f"Q{x}": f"specific_humidity_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},24    **{f"T{x}": f"temperature_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},25    **{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)},26    **{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)},27    "U10M": "10m_u_component_of_wind",28    "V10M": "10m_v_component_of_wind",29    "T2M": "2m_temperature",30    "MSLP": "mean_sea_level_pressure",31}32REAL_VARIABLES = tuple(VARIABLE_ALIASES[name] for name in OFFICIAL_VARIABLES)33 34 35def _synthetic_field(time_index: int, channel: int, height: int, width: int) -> np.ndarray:36    lat = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None]37    lon = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, :]38    phase = np.float32(channel * 0.17)39    field = np.sin(lon + phase) + 0.4 * np.cos(np.float32(time_index / 3.0) + phase) + 0.2 * lat40    return np.asarray(field, dtype=np.float32)41 42 43def generate_fake_h5(output_dir: Path, years: list[int], stats_years: set[int], time_steps: int,44                     height: int, width: int, seed: int) -> None:45    data_dir = output_dir / "data"46    data_dir.mkdir(parents=True, exist_ok=True)47    sums = np.zeros(len(OFFICIAL_VARIABLES), dtype=np.float64)48    squared_sums = np.zeros_like(sums)49    value_count = 050 51    for offset, year in enumerate(years):52        path = data_dir / f"{year}.h5"53        with h5py.File(path, "w") as handle:54            dataset = handle.create_dataset(55                "fields", shape=(time_steps, len(OFFICIAL_VARIABLES), height, width),56                dtype="float32", chunks=(1, 1, height, width),57            )58            dataset.attrs["variables"] = list(REAL_VARIABLES)59            dataset.attrs["time_step"] = 660            for time_index in range(time_steps):61                for channel in range(len(OFFICIAL_VARIABLES)):62                    field = _synthetic_field(time_index + offset, channel, height, width)63                    dataset[time_index, channel] = field64                    if year in stats_years:65                        sums[channel] += field.sum(dtype=np.float64)66                        squared_sums[channel] += np.square(field, dtype=np.float64).sum()67            if year in stats_years:68                value_count += time_steps * height * width69        print(f"{path}: fields={(time_steps, len(OFFICIAL_VARIABLES), height, width)}, variables={len(OFFICIAL_VARIABLES)}")70 71    means = (sums / value_count).reshape(1, -1, 1, 1)72    variances = squared_sums / value_count - np.square(means.reshape(-1))73    stds = np.sqrt(np.maximum(variances, 1e-12)).reshape(1, -1, 1, 1)74    stats_dir = output_dir / "stats"75    stats_dir.mkdir(parents=True, exist_ok=True)76    np.save(stats_dir / "global_means.npy", means)77    np.save(stats_dir / "global_stds.npy", stds)78 79 80def main() -> None:81    parser = argparse.ArgumentParser()82    parser.add_argument("--output-dir", type=Path, default=Path("./data"))83    parser.add_argument("--years", nargs="+", type=int, default=[2000, 2001, 2002])84    parser.add_argument("--stats-years", nargs="+", type=int, default=None)85    parser.add_argument("--time-steps", type=int, default=3)86    parser.add_argument("--height", type=int, default=721)87    parser.add_argument("--width", type=int, default=1440)88    parser.add_argument("--seed", type=int, default=42)89    args = parser.parse_args()90    if len(args.years) < 1 or min(args.time_steps, args.height, args.width) < 1:91        parser.error("years, time-steps, height, and width must be positive")92    if (args.height, args.width) != (721, 1440):93        parser.error("ERA5 fixtures must use the native 0.25 degree grid 721x1440")94    stats_years = set(args.stats_years or args.years[:1])95    if not stats_years.issubset(args.years):96        parser.error("stats-years must be included in years")97    generate_fake_h5(args.output_dir, args.years, stats_years, args.time_steps,98                     args.height, args.width, args.seed)99 100 101if __name__ == "__main__":102    main()103