File size: 2,935 Bytes
3fef26f | 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 | """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)
# ERA5Dataset always initializes these files, even when normalize=False.
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()
|