File size: 5,339 Bytes
1616901 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #!/usr/bin/env python3
"""Generate a compact deterministic ERA5-style dataset for workflow checks."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import h5py
import numpy as np
import yaml
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
parser.add_argument("--output-dir", help="Override data.data_dir from the config")
parser.add_argument("--height", type=int)
parser.add_argument("--width", type=int)
parser.add_argument("--timesteps", type=int)
parser.add_argument("--seed", type=int)
return parser.parse_args()
def _resolve(path: str, root: Path) -> Path:
candidate = Path(path).expanduser()
return candidate if candidate.is_absolute() else (root / candidate).resolve()
def generate_dataset(
output_dir: str | Path,
years: list[int],
channels: list[str],
height: int,
width: int,
timesteps: int,
time_step_hours: int = 6,
seed: int = 42,
) -> Path:
"""Write yearly HDF5 files and static fields without external data access."""
if min(height, width, timesteps) <= 0:
raise ValueError("height, width and timesteps must be positive")
output_dir = Path(output_dir)
data_dir = output_dir / "data"
data_dir.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(seed)
channels = list(channels)
channel_count = len(channels)
channel_mean = np.linspace(0.0, 0.9, channel_count, dtype=np.float32)
channel_std = np.linspace(0.8, 1.2, channel_count, dtype=np.float32)
latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[None, :, None]
longitude = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, None, :]
base_pattern = 0.15 * latitude + 0.08 * np.sin(longitude)
for year_offset, year in enumerate(sorted(set(int(item) for item in years))):
path = data_dir / f"{year}.h5"
with h5py.File(path, "w") as output:
fields = output.create_dataset(
"fields",
shape=(timesteps, channel_count, height, width),
dtype="float32",
chunks=(1, channel_count, height, width),
compression="lzf",
)
fields.attrs["variables"] = np.asarray(channels, dtype="S")
fields.attrs["time_step"] = int(time_step_hours)
for timestep in range(timesteps):
values = np.empty((channel_count, height, width), dtype=np.float32)
for channel in range(channel_count):
phase = 0.05 * (timestep + year_offset) + 0.2 * channel
smooth = base_pattern + 0.03 * np.sin(phase)
noise = rng.normal(0.0, 0.01, size=(height, width)).astype(np.float32)
values[channel] = channel_mean[channel] + channel_std[channel] * (smooth + noise)
fields[timestep] = values
output.create_dataset("global_means", data=channel_mean.reshape(1, channel_count, 1, 1))
output.create_dataset("global_stds", data=channel_std.reshape(1, channel_count, 1, 1))
print(f"Generated {path} shape=({timesteps},{channel_count},{height},{width})")
static_dir = output_dir / "static"
static_dir.mkdir(parents=True, exist_ok=True)
latitude_2d = latitude.squeeze(0)
land_mask = np.broadcast_to((latitude_2d > 0).astype(np.float32), (height, width)).copy()
orography = (0.2 * np.cos(longitude) + 0.05 * latitude).astype(np.float32).squeeze(0)
sea_ice_mask = np.broadcast_to((latitude_2d < -0.55).astype(np.float32), (height, width)).copy()
np.save(static_dir / "land_mask.npy", land_mask)
np.save(static_dir / "orography.npy", orography)
np.save(static_dir / "sea_ice_mask.npy", sea_ice_mask)
metadata = {
"format": "ERA5-HDF5-window-v1",
"years": sorted(set(int(item) for item in years)),
"channels": channels,
"shape": [timesteps, channel_count, height, width],
"time_step_hours": int(time_step_hours),
"seed": int(seed),
}
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
print(f"Generated static fields and metadata under {output_dir}")
return output_dir
def main() -> None:
args = parse_args()
with Path(args.config).open(encoding="utf-8") as source:
config = yaml.safe_load(source)
data_cfg = config["data"]
fake_cfg = config.get("fake_data", {})
data_dir = _resolve(args.output_dir or data_cfg["data_dir"], PROJECT_ROOT)
years = list(data_cfg["train_years"]) + list(data_cfg["val_years"]) + list(data_cfg["test_years"])
generate_dataset(
data_dir,
years=years,
channels=list(data_cfg["channels"]),
height=int(args.height or fake_cfg.get("height", 32)),
width=int(args.width or fake_cfg.get("width", 64)),
timesteps=int(args.timesteps or fake_cfg.get("timesteps", 12)),
time_step_hours=int(data_cfg.get("time_step_hours", 6)),
seed=int(args.seed if args.seed is not None else fake_cfg.get("seed", 42)),
)
if __name__ == "__main__":
main()
|