scripts/fake_data.py

"""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()