| """Create an ERA5-style HDF5 dataset consumable by OneScience ERA5Dataset.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import h5py |
| import numpy as np |
|
|
| from common import load_config, resolve_path |
| from variables import c85_from_config |
|
|
|
|
| def create_year( |
| path: Path, |
| year: int, |
| steps: int, |
| height: int, |
| width: int, |
| channels: list[str], |
| time_step_hours: int, |
| hdf5_config: dict, |
| ) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with h5py.File(path, "w") as handle: |
| fields = handle.create_dataset( |
| "fields", |
| shape=(steps, len(channels), height, width), |
| dtype="float32", |
| chunks=tuple(min(size, limit) for size, limit in zip((steps, len(channels), height, width), hdf5_config["chunks"])), |
| fillvalue=0.0, |
| compression=hdf5_config["compression"], |
| compression_opts=hdf5_config["compression_level"], |
| ) |
| fields.attrs["variables"] = np.asarray(channels, dtype=h5py.string_dtype("utf-8")) |
| fields.attrs["time_step"] = time_step_hours |
| fields.attrs["year"] = year |
| lat = np.linspace(90.0, -90.0, height, dtype=np.float32)[:, None] |
| lon = np.arange(width, dtype=np.float32)[None, :] * (360.0 / width) |
| t2m_index = channels.index("t2m") |
| for step in range(steps): |
| fields[step, t2m_index] = ( |
| np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon + step * 15.0)) |
| ) |
| handle.create_dataset("global_means", data=np.zeros((1, len(channels), 1, 1), dtype=np.float32)) |
| handle.create_dataset("global_stds", data=np.ones((1, len(channels), 1, 1), dtype=np.float32)) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default="conf/config.yaml") |
| args = parser.parse_args() |
| cfg = load_config(args.config) |
| data_cfg = cfg["data"] |
| channels, _ = c85_from_config(cfg) |
| root = resolve_path(cfg["paths"]["data_root"], cfg) |
| generated_years = set() |
| for split, split_cfg in data_cfg["splits"].items(): |
| for year in split_cfg["years"]: |
| if year in generated_years: |
| raise ValueError(f"Year {year} is assigned to more than one data split") |
| generated_years.add(year) |
| create_year( |
| root / "data" / f"{year}.h5", |
| year, |
| split_cfg["time_steps"], |
| *data_cfg["grid_size"], |
| channels, |
| data_cfg["time_step_hours"], |
| data_cfg["hdf5"], |
| ) |
| print(f"Created ERA5-compatible yearly datasets at {root / 'data'}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
