| """Generate native-grid virtual ERA5 data for OneScience ERA5Dataset.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from config import load_config, project_path |
|
|
| def generate(dataset_dir, year, time_steps, height, width, variables, time_step_hours, seed, dtype): |
| import h5py |
|
|
| dataset_dir = Path(dataset_dir) |
| data_dir = dataset_dir / "data" |
| stats_dir = dataset_dir / "stats" |
| data_dir.mkdir(parents=True, exist_ok=True) |
| stats_dir.mkdir(parents=True, exist_ok=True) |
| path = data_dir / f"{year}.h5" |
| rng = np.random.default_rng(seed + year) |
| lat = np.linspace(90.0, -90.0, height, dtype=np.float32)[:, None] |
| lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)[None, :] |
| with h5py.File(path, "w") as handle: |
| fields = handle.create_dataset( |
| "fields", shape=(time_steps, len(variables), height, width), dtype=dtype, |
| chunks=(1, 1, min(height, 64), min(width, 128)), |
| ) |
| fields.attrs["variables"] = np.asarray(variables, dtype=h5py.string_dtype()) |
| fields.attrs["time_step"] = time_step_hours |
| for step in range(time_steps): |
| phase = 2.0 * np.pi * (step / max(time_steps, 1)) |
| base = np.sin(np.deg2rad(lat)) + 0.25 * np.cos(np.deg2rad(lon) + phase) |
| noise = rng.normal(0.0, 0.01, size=(len(variables), height, width)).astype(np.float32) |
| fields[step] = np.asarray([(index + 1) * base + noise[index] for index in range(len(variables))], dtype=dtype) |
| |
| np.save(stats_dir / "global_means.npy", np.zeros((1, len(variables), 1, 1), dtype=np.float32)) |
| np.save(stats_dir / "global_stds.npy", np.ones((1, len(variables), 1, 1), dtype=np.float32)) |
| return path |
|
|
|
|
| def main(): |
| config = load_config() |
| data_config = config["data"] |
| runtime_config = config["runtime"] |
| inference_config = config["inference"] |
| parser = argparse.ArgumentParser(description="Generate native-grid virtual ERA5 HDF5 data.") |
| parser.add_argument("--dataset-dir", default=str(project_path(data_config["virtual_dir"]))) |
| parser.add_argument("--year", type=int, default=inference_config["year"]) |
| parser.add_argument("--time-steps", type=int, default=data_config["virtual_time_steps"]) |
| parser.add_argument("--height", type=int, default=data_config["raw_height"]) |
| parser.add_argument("--width", type=int, default=data_config["raw_width"]) |
| parser.add_argument("--seed", type=int, default=runtime_config["seed"]) |
| parser.add_argument("--dtype", choices=["float16", "float32"], default=data_config["virtual_dtype"]) |
| args = parser.parse_args() |
| print(generate(args.dataset_dir, args.year, args.time_steps, args.height, args.width, |
| data_config["variables"], data_config["time_step_hours"], args.seed, args.dtype)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
