File size: 8,124 Bytes
807a08b | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | """Generate small, deterministic ERA5-shaped HDF5 files for workflow checks."""
from __future__ import annotations
import argparse
from pathlib import Path
from typing import Iterable
import h5py
import numpy as np
import yaml
VARIABLES = ("z", "t", "t2m", "u10", "v10")
ERA5_HEIGHT = 721
ERA5_WIDTH = 1440
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def _parse_years(value: str | Iterable[int]) -> list[int]:
if isinstance(value, str):
return [int(item.strip()) for item in value.split(",") if item.strip()]
return [int(item) for item in value]
def _synthetic_frame(
step: int,
year: int,
lat: np.ndarray,
lon: np.ndarray,
) -> np.ndarray:
"""Create smooth fields with distinct scales for the five official channels."""
lat_rad = np.deg2rad(lat)[:, None]
lon_rad = np.deg2rad(lon)[None, :]
phase = 2.0 * np.pi * (step + (year % 100)) / 1460.0
spatial = np.sin(lat_rad) + 0.35 * np.cos(lon_rad) + 0.15 * np.sin(
2.0 * lon_rad + phase
)
seasonal = np.cos(lat_rad) * np.sin(phase)
channels = np.stack(
[
5000.0 + 300.0 * spatial + 20.0 * seasonal,
260.0 + 12.0 * spatial + 2.0 * seasonal,
280.0 + 18.0 * spatial + 3.0 * seasonal,
4.0 * np.cos(lon_rad + phase) + 0.5 * spatial,
3.0 * np.sin(lon_rad - phase) - 0.5 * spatial,
],
axis=0,
)
return channels.astype(np.float32, copy=False)
def _write_static(static_dir: Path, height: int, width: int) -> None:
static_dir.mkdir(parents=True, exist_ok=True)
lat = np.linspace(
90.0 - 90.0 / height,
-90.0 + 90.0 / height,
height,
dtype=np.float32,
)
lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
orography = (1200.0 * np.maximum(np.cos(np.deg2rad(lat2d)), 0.0)).astype(
np.float32
)
lsm = (np.cos(np.deg2rad(lat2d)) > 0.25).astype(np.float32)
with h5py.File(static_dir / "constants.h5", "w") as handle:
handle.create_dataset("orography", data=orography)
handle.create_dataset("lsm", data=lsm)
handle.create_dataset("lat2d", data=lat2d)
handle.create_dataset("lon2d", data=lon2d)
handle.attrs["variables"] = np.asarray(
["orography", "lsm"], dtype=h5py.string_dtype()
)
def generate_data(
output_dir: str | Path,
years: Iterable[int],
timesteps: int,
height: int = ERA5_HEIGHT,
width: int = ERA5_WIDTH,
seed: int = 42,
overwrite: bool = False,
write_static: bool = True,
) -> dict[str, list[float]]:
"""Generate annual files and return per-channel global statistics."""
if timesteps < 4:
raise ValueError("timesteps must be at least 4 for a three-frame history")
if (height, width) != (ERA5_HEIGHT, ERA5_WIDTH):
raise ValueError(
"ClimODE virtual ERA5 data must use the raw shape "
f"({ERA5_HEIGHT}, {ERA5_WIDTH})"
)
root = Path(output_dir)
data_dir = root / "data"
static_dir = root / "static"
data_dir.mkdir(parents=True, exist_ok=True)
years = _parse_years(years)
lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
rng = np.random.default_rng(seed)
minimum = np.full(len(VARIABLES), np.inf, dtype=np.float64)
maximum = np.full(len(VARIABLES), -np.inf, dtype=np.float64)
total = np.zeros(len(VARIABLES), dtype=np.float64)
total_sq = np.zeros(len(VARIABLES), dtype=np.float64)
total_count = 0
for year in years:
path = data_dir / f"{year}.h5"
if path.exists() and not overwrite:
raise FileExistsError(f"Refusing to overwrite existing file: {path}")
with h5py.File(path, "w") as handle:
fields = handle.create_dataset(
"fields",
shape=(timesteps, len(VARIABLES), height, width),
dtype=np.float32,
chunks=(1, len(VARIABLES), height, width),
)
fields.attrs["variables"] = np.asarray(
VARIABLES, dtype=h5py.string_dtype()
)
fields.attrs["time_step"] = 6
for step in range(timesteps):
frame = _synthetic_frame(step, year, lat, lon)
# A tiny deterministic per-frame perturbation keeps years distinct
# without materializing another 20 MB random tensor per frame.
frame += np.float32(rng.normal(0.0, 1.0e-3))
fields[step] = frame
flat = frame.reshape(len(VARIABLES), -1).astype(np.float64)
minimum = np.minimum(minimum, flat.min(axis=1))
maximum = np.maximum(maximum, flat.max(axis=1))
total += flat.sum(axis=1)
total_sq += np.square(flat).sum(axis=1)
total_count += flat.shape[1]
# Placeholders are replaced with statistics over every requested year.
handle.create_dataset("global_means", shape=(1, len(VARIABLES), 1, 1), dtype=np.float32)
handle.create_dataset("global_stds", shape=(1, len(VARIABLES), 1, 1), dtype=np.float32)
means = (total / total_count).astype(np.float32)
variances = np.maximum(
total_sq / total_count - means.astype(np.float64) ** 2, 1.0e-12
)
stds = np.sqrt(variances).astype(np.float32)
for year in years:
with h5py.File(data_dir / f"{year}.h5", "r+") as handle:
handle["global_means"][:] = means.reshape(1, -1, 1, 1)
handle["global_stds"][:] = stds.reshape(1, -1, 1, 1)
static_height, static_width = 32, 64
if write_static:
_write_static(static_dir, static_height, static_width)
np.save(static_dir / "min_values.npy", minimum.astype(np.float32))
np.save(static_dir / "max_values.npy", maximum.astype(np.float32))
return {"min": minimum.tolist(), "max": maximum.tolist()}
def _load_config(path: Path) -> dict:
with path.open("r", encoding="utf-8") as handle:
return yaml.safe_load(handle)
def _resolve(path: str | Path) -> Path:
value = Path(path)
return value if value.is_absolute() else PROJECT_ROOT / value
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml"
)
parser.add_argument("--output-dir", type=Path, default=None)
parser.add_argument("--years", type=str, default=None, help="Comma-separated years")
parser.add_argument("--timesteps", type=int, default=None)
parser.add_argument("--height", type=int, default=None)
parser.add_argument("--width", type=int, default=None)
parser.add_argument("--seed", type=int, default=None)
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args()
config_path = _resolve(args.config)
config = _load_config(config_path) if config_path.exists() else {}
fake = config.get("fake_data", {})
output_dir = _resolve(
args.output_dir or config.get("data", {}).get("data_dir", "./data")
)
years = _parse_years(args.years) if args.years else fake.get("years", [2006, 2016, 2017])
stats = generate_data(
output_dir=output_dir,
years=years,
timesteps=(
args.timesteps
if args.timesteps is not None
else fake.get("timesteps", 8)
),
height=(
args.height
if args.height is not None
else fake.get("height", ERA5_HEIGHT)
),
width=(
args.width
if args.width is not None
else fake.get("width", ERA5_WIDTH)
),
seed=args.seed if args.seed is not None else fake.get("seed", 42),
overwrite=args.overwrite,
)
print(f"Generated years={years} under {Path(output_dir).resolve()}")
print(f"min={stats['min']}")
print(f"max={stats['max']}")
if __name__ == "__main__":
main()
|