| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| SCRIPT_DIR = Path(__file__).resolve().parent |
| if str(SCRIPT_DIR) not in sys.path: |
| sys.path.insert(0, str(SCRIPT_DIR)) |
|
|
| import h5py |
| import numpy as np |
|
|
| from common import DEFAULT_CONFIG, load_config, resolve_path |
|
|
|
|
| def generate_year( |
| path: Path, |
| variables: list[str], |
| *, |
| time_steps: int, |
| height: int, |
| width: int, |
| time_step_hours: int, |
| chunk_time_steps: int, |
| fill_value: float, |
| materialize_pattern: bool, |
| ) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| channels = len(variables) |
| with h5py.File(path, "w") as handle: |
| fields = handle.create_dataset( |
| "fields", |
| shape=(time_steps, channels, height, width), |
| dtype=np.float32, |
| chunks=(chunk_time_steps, 1, height, width), |
| fillvalue=np.float32(fill_value), |
| compression="lzf", |
| ) |
| fields.attrs["variables"] = np.asarray(variables, dtype=h5py.string_dtype()) |
| fields.attrs["time_step"] = time_step_hours |
| handle.create_dataset( |
| "global_means", data=np.zeros((1, channels, 1, 1), dtype=np.float32) |
| ) |
| handle.create_dataset( |
| "global_stds", data=np.ones((1, channels, 1, 1), dtype=np.float32) |
| ) |
|
|
| if materialize_pattern: |
| latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None] |
| longitude = np.linspace( |
| 0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32 |
| ) |
| base = latitude + np.sin(longitude)[None, :] |
| |
| for time_index in range(min(time_steps, 2)): |
| for channel_index in range(channels): |
| fields[time_index, channel_index] = ( |
| base + channel_index / channels + time_index * 0.01 |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Generate ERA5-compatible FCNv2 data") |
| parser.add_argument("--config", default=str(DEFAULT_CONFIG)) |
| parser.add_argument("--no-pattern", action="store_true") |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
| data = config["data"] |
| fake = config["fake_data"] |
| output_dir = resolve_path(config, data["dataset_dir"]) |
| years = sorted(set(data["train_years"] + data["val_years"] + data["test_years"])) |
| height, width = data["grid_shape"] |
| for year in years: |
| path = output_dir / "data" / f"{year}.h5" |
| generate_year( |
| path, |
| data["variables"], |
| time_steps=fake["time_steps_per_year"], |
| height=height, |
| width=width, |
| time_step_hours=data["time_step_hours"], |
| chunk_time_steps=fake["chunk_time_steps"], |
| fill_value=fake["fill_value"], |
| materialize_pattern=fake["materialize_pattern"] and not args.no_pattern, |
| ) |
| logical_gib = ( |
| fake["time_steps_per_year"] * len(data["variables"]) * height * width * 4 |
| ) / 1024**3 |
| print(f"{path}: logical={logical_gib:.2f} GiB, actual={path.stat().st_size / 1024**2:.2f} MiB") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
