scripts/fake_data.py

#!/usr/bin/env python3
"""Generate OneScience-compatible synthetic ERA5 HDF5 data."""

from __future__ import annotations

import argparse
from pathlib import Path

import h5py
import numpy as np

try:
    from common import PROJECT_ROOT, SYNTHETIC_GENERATOR_VERSION, channel_order, load_config, resolve_path, write_json
except ModuleNotFoundError:  # supports ``python -m scripts.fake_data``
    from scripts.common import PROJECT_ROOT, SYNTHETIC_GENERATOR_VERSION, channel_order, load_config, resolve_path, write_json


def make_field(name: str, step: int, year: int, lat: np.ndarray, lon: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    latr, lonr = np.deg2rad(lat)[:, None], np.deg2rad(lon)[None, :]
    wave = np.cos(latr) * np.sin(lonr + step * 0.05)
    noise = rng.normal(0, 1, wave.shape).astype(np.float32)
    if name == "sea_ice_cover":
        return np.clip(0.55 - 0.45 * np.cos(latr) + 0.03 * noise, 0, 1).astype(np.float32)
    if name == "sea_surface_temperature":
        return (273.15 + 26 * np.cos(latr) + 1.5 * wave + 0.1 * noise).astype(np.float32)
    base, _, suffix = name.rpartition("_")
    level = int(suffix)
    ratio = level / 1000.0
    if base == "geopotential":
        # Hydrostatic log-pressure profile calibrated against the bundled ERA5
        # statistics (about 464 km2/s2 at 1 hPa and 0.7 km2/s2 at 1000 hPa).
        vertical = 733.0 + 67000.0 * np.log(1000.0 / level)
        return (vertical + 1800.0 * (1.0 - ratio) * np.sin(latr) ** 2 + 120.0 * wave + noise).astype(np.float32)
    if base == "temperature":
        # Piecewise standard-atmosphere profile captures stratospheric warming;
        # a monotone 220->288 K profile is not physically valid above 10 hPa.
        anchors_hpa = np.asarray([1.0, 10.0, 100.0, 500.0, 1000.0])
        anchors_k = np.asarray([261.0, 229.0, 207.0, 253.0, 281.0])
        vertical = np.interp(np.log(level), np.log(anchors_hpa), anchors_k)
        return (vertical + 4.0 * wave + 0.1 * noise).astype(np.float32)
    if base == "specific_humidity":
        vertical = 4.0e-6 + 7.0e-3 * ratio**3
        return np.maximum((vertical * (0.8 + 0.2 * np.cos(latr)) + 1e-7 * noise), 1e-8).astype(np.float32)
    if base == "u_component_of_wind":
        return (12 * (1 - ratio) * wave + 0.2 * noise).astype(np.float32)
    if base == "v_component_of_wind":
        return (8 * (1 - ratio) * np.sin(2 * latr) * np.cos(lonr) + 0.2 * noise).astype(np.float32)
    if name.startswith("specific_cloud_ice_water_content_"):
        profile = 8e-6 * np.exp(-0.5 * ((level - 450.0) / 180.0) ** 2)
        return np.maximum(profile * (0.7 + 0.3 * np.cos(latr)) + 1e-7 * noise, 0).astype(np.float32)
    if name.startswith("specific_cloud_liquid_water_content_"):
        profile = 1.2e-5 * np.exp(-0.5 * ((level - 750.0) / 160.0) ** 2)
        return np.maximum(profile * (0.7 + 0.3 * np.cos(latr)) + 1e-7 * noise, 0).astype(np.float32)
    raise ValueError(f"Unsupported channel: {name}")


def generate_year(
    path: Path,
    channels: list[str],
    timesteps: int,
    height: int,
    width: int,
    year: int,
    seed: int,
    time_step_hours: int,
) -> dict:
    lat = np.linspace(90, -90, height, dtype=np.float32)
    lon = np.linspace(0, 360, width, endpoint=False, dtype=np.float32)
    path.parent.mkdir(parents=True, exist_ok=True)
    sums = np.zeros(len(channels), dtype=np.float64)
    sq = np.zeros(len(channels), dtype=np.float64)
    with h5py.File(path, "w") as handle:
        fields = handle.create_dataset("fields", shape=(timesteps, len(channels), height, width), dtype="f4", chunks=(1, 1, min(height, 32), min(width, 64)), compression="gzip", compression_opts=1)
        fields.attrs["variables"] = np.asarray(channels, dtype=h5py.string_dtype("utf-8"))
        fields.attrs["time_step"] = time_step_hours
        fields.attrs["synthetic"] = True
        fields.attrs["generator_version"] = SYNTHETIC_GENERATOR_VERSION
        for t in range(timesteps):
            rng = np.random.default_rng(seed + year * 1009 + t)
            for i, name in enumerate(channels):
                value = make_field(name, t, year, lat, lon, rng)
                fields[t, i] = value
                sums[i] += value.sum(dtype=np.float64)
                sq[i] += np.square(value, dtype=np.float64).sum(dtype=np.float64)
        count = timesteps * height * width
        means = sums / count
        stds = np.sqrt(np.maximum(sq / count - means**2, 1e-12))
        handle.create_dataset("global_means", data=means[None, :, None, None].astype("f4"))
        handle.create_dataset("global_stds", data=stds[None, :, None, None].astype("f4"))
    return {"path": str(path), "shape": [timesteps, len(channels), height, width], "mean_std_min": float(stds.min())}


def generate_static(path: Path, height: int, width: int) -> dict:
    """Generate deterministic ERA5-like static topography and land mask.

    These variables are auxiliary model inputs rather than flattened dynamic
    channels.  Values use the physical units expected by the official Gin
    profiles: geopotential at surface in m² s⁻² and land-sea mask in [0, 1].
    """
    import xarray as xr

    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
    latr, lonr = np.deg2rad(lat)[:, None], np.deg2rad(lon)[None, :]
    # Smooth continent-like mask and non-negative terrain height.  This is
    # intentionally synthetic; real ERA5 static fields should be preferred.
    mask = (0.5 + 0.5 * np.sin(2.0 * latr) * np.cos(3.0 * lonr) > 0.52).astype(np.float32)
    elevation_m = np.maximum(
        0.0,
        2200.0 * mask * (0.35 + 0.65 * np.cos(latr) ** 2)
        + 250.0 * np.sin(latr) ** 2,
    ).astype(np.float32)
    geopotential = (9.80665 * elevation_m).astype(np.float32)
    ds = xr.Dataset(
        {
            "geopotential_at_surface": (("longitude", "latitude"), geopotential.T),
            "land_sea_mask": (("longitude", "latitude"), mask.T),
        },
        coords={"latitude": lat, "longitude": lon},
        attrs={"synthetic": "true", "source": "neuralgcm_develop.fake_data"},
    )
    ds["geopotential_at_surface"].attrs["units"] = "m**2 s**-2"
    ds["land_sea_mask"].attrs["units"] = "dimensionless"
    path.parent.mkdir(parents=True, exist_ok=True)
    ds.to_netcdf(path)
    return {"path": str(path), "shape": [height, width], "variables": list(ds.data_vars)}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
    parser.add_argument("--output-dir")
    parser.add_argument("--years", nargs="*", type=int)
    parser.add_argument("--timesteps", type=int)
    parser.add_argument(
        "--forecast-steps",
        type=int,
        help="number of future 6-hour frames; defaults to data.virtual.forecast_steps",
    )
    parser.add_argument("--height", type=int)
    parser.add_argument("--width", type=int)
    args = parser.parse_args()
    config = load_config(args.config)
    out = resolve_path(args.output_dir or config["data"]["data_dir"], args.config)
    virtual = config["data"]["virtual"]
    time_step_hours = int(config["data"].get("time_step_hours", 6))
    if time_step_hours <= 0:
        raise ValueError("data.time_step_hours must be positive")
    years = args.years or sorted(set(config["data"]["train_years"] + config["data"]["val_years"] + config["data"]["test_years"]))
    channels = channel_order(config)
    # The explicit value always wins, followed by the virtual-data setting,
    # then the inference default. The project defaults to a compact two-day
    # window; 60 six-hour steps exercise the full official 15-day capability.
    forecast_steps = int(
        args.forecast_steps
        if args.forecast_steps is not None
        else virtual.get("forecast_steps", config.get("inference", {}).get("prediction_steps", 8))
    )
    if forecast_steps <= 0:
        raise ValueError("--forecast-steps must be positive")
    input_steps = int(config["data"].get("input_steps", 1))
    configured_timesteps = int(virtual["timesteps_per_year"])
    # The synthetic file contains one initial frame plus the requested future
    # frames. Explicit --timesteps remains available for tiny smoke tests.
    explicit_timesteps = args.timesteps is not None
    timesteps = int(args.timesteps) if explicit_timesteps else max(
        configured_timesteps, input_steps + forecast_steps
    )
    if timesteps <= 0:
        raise ValueError("--timesteps must be positive")
    if not explicit_timesteps and timesteps < input_steps + forecast_steps:
        raise ValueError(
            f"timesteps={timesteps} is too short for input_steps={input_steps} "
            f"and forecast_steps={forecast_steps}; need at least {input_steps + forecast_steps}"
        )
    complete_window = timesteps >= input_steps + forecast_steps
    forecast_horizon_days = forecast_steps * time_step_hours / 24
    if explicit_timesteps and not complete_window:
        print(
            f"Warning: timesteps={timesteps} provides only a smoke window; "
            f"{input_steps + forecast_steps} frames are required for the full "
            f"{forecast_horizon_days:g}-day configured forecast."
        )
    records = [
        generate_year(
            out / "data" / f"{year}.h5",
            channels,
            timesteps,
            args.height or virtual["height"],
            args.width or virtual["width"],
            year,
            int(virtual["seed"]),
            time_step_hours,
        )
        for year in years
    ]
    static_record = generate_static(out / "static.nc", args.height or virtual["height"], args.width or virtual["width"])
    write_json(out / "metadata" / "dataset_card.json", {"name": "neuralgcm-synthetic-era5", "format": "OneScience ERA5Dataset HDF5", "channels": channels, "files": records, "static_file": static_record, "native_model_grid": config["model"]["grid_shape"], "input_steps": input_steps, "forecast_steps": forecast_steps, "time_step_hours": time_step_hours, "forecast_horizon_hours": forecast_steps * time_step_hours, "forecast_horizon_days": forecast_horizon_days, "official_forecast_capability_days": [2, 15], "official_15_day_window_complete": forecast_horizon_days >= 15, "forecast_window_complete": complete_window})
    print(f"Generated {len(records)} years under {out}")


if __name__ == "__main__":
    main()