scripts/common.py

"""Shared config, channel and OneScience ERA5Dataset helpers."""

from __future__ import annotations

import json
import sys
from functools import lru_cache
from pathlib import Path
from typing import Any

import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[1]
SYNTHETIC_GENERATOR_VERSION = "neuralgcm-hydrostatic-v2"


def load_config(path: str | Path | None = None) -> dict[str, Any]:
    path = Path(path or PROJECT_ROOT / "conf/config.yaml")
    with path.open(encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def resolve_path(value: str | Path, config_path: str | Path | None = None) -> Path:
    path = Path(value).expanduser()
    if path.is_absolute():
        return path
    base = Path(config_path or PROJECT_ROOT / "conf/config.yaml").resolve().parent.parent
    return base / path


def channel_order(config: dict[str, Any]) -> list[str]:
    return list(config["data"]["channel_order"])


def pressure_levels(config: dict[str, Any]) -> list[int]:
    return list(config["model"]["pressure_levels_hpa"])


def as_time_major_frames(value: Any, *, name: str = "frames"):
    """Normalize OneScience ERA5Dataset output to ``(T, C, H, W)``.

    ERA5Dataset squeezes the leading time dimension when ``output_steps=1``;
    callers must restore it before indexing forecast frames. Input frames are
    allowed to remain ``(C, H, W)`` and should not use this helper.
    """
    import numpy as np

    if hasattr(value, "detach"):
        value = value.detach().cpu().numpy()
    value = np.asarray(value)
    if value.ndim == 3:
        value = value[None, ...]
    if value.ndim != 4:
        raise ValueError(
            f"{name} must have shape (T,C,H,W) or (C,H,W), got {value.shape}"
        )
    return value


def load_era5_dataset(config: dict[str, Any], years: list[int], *, input_steps: int | None = None, output_steps: int | None = None):
    """Construct the required OneScience ERA5Dataset, without replacing it."""
    try:
        from onescience.datapipes.climate import ERA5Dataset
    except Exception as exc:
        # Source-tree fallback mirrors the earth examples and keeps this
        # project usable before OneScience is installed as a wheel.
        local_src = Path("/public/home/yangzt01/onescience/src")
        if local_src.exists() and str(local_src) not in sys.path:
            sys.path.insert(0, str(local_src))
        try:
            from onescience.datapipes.climate import ERA5Dataset
        except Exception as fallback_exc:
            raise RuntimeError(
            "OneScience ERA5Dataset import failed; load OneScience and its "
            f"runtime modules first: {type(fallback_exc).__name__}: {fallback_exc}"
            ) from fallback_exc
    data_dir = resolve_path(config["data"]["data_dir"])
    return ERA5Dataset(
        dataset_dir=str(data_dir),
        used_years=years,
        used_variables=channel_order(config),
        input_steps=input_steps or int(config["data"]["input_steps"]),
        output_steps=output_steps or int(config["data"]["output_steps"]),
        normalize=bool(config["data"].get("normalize", False)),
    )


def era5_data_is_synthetic(config: dict[str, Any], years: list[int]) -> bool:
    """Return true only when every requested HDF5 file declares synthetic data."""
    import h5py

    data_dir = resolve_path(config["data"]["data_dir"]) / "data"
    paths = [data_dir / f"{year}.h5" for year in years]
    if not paths or any(not path.exists() for path in paths):
        return False
    try:
        for path in paths:
            with h5py.File(path, "r") as handle:
                fields = handle[config["data"].get("field_key", "fields")]
                if not bool(fields.attrs.get("synthetic", False)):
                    return False
    except (KeyError, OSError):
        return False
    return True


def validate_synthetic_era5_version(
    config: dict[str, Any], years: list[int]
) -> None:
    """Reject obsolete virtual fields that are known to destabilize the model."""
    import h5py

    data_dir = resolve_path(config["data"]["data_dir"]) / "data"
    for year in years:
        path = data_dir / f"{year}.h5"
        with h5py.File(path, "r") as handle:
            fields = handle[config["data"].get("field_key", "fields")]
            if not bool(fields.attrs.get("synthetic", False)):
                continue
            version = fields.attrs.get("generator_version")
            if isinstance(version, bytes):
                version = version.decode()
            if version != SYNTHETIC_GENERATOR_VERSION:
                raise RuntimeError(
                    f"Synthetic ERA5 file {path} uses obsolete generator_version="
                    f"{version!r}; expected {SYNTHETIC_GENERATOR_VERSION!r}. "
                    "Regenerate it with scripts/fake_data.py before running a "
                    "NeuralGCM rollout."
                )


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8")


def era5_sample_to_xarray(sample: Any, config: dict[str, Any], *, timestamp: Any):
    """Convert one ERA5Dataset frame to the official NeuralGCM xarray contract.

    The HDF5 loader returns flattened channels in ``[C, latitude, longitude]``;
    official NeuralGCM expects named variables with pressure ``level`` and
    explicit latitude/longitude coordinates. Spatial interpolation to the
    configured native grid is performed before the model API sees the data.
    """
    import numpy as np
    import xarray as xr

    invar = sample[0]
    if hasattr(invar, "detach"):
        invar = invar.detach().cpu().numpy()
    channels = channel_order(config)
    levels = pressure_levels(config)
    height, width = invar.shape[-2:]
    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
    # ERA5Dataset stores (latitude, longitude), while NeuralGCM's xarray API
    # expects (longitude, latitude) for horizontal fields.
    dataset = xr.Dataset(coords={"latitude": lat, "longitude": lon, "time": [np.datetime64(timestamp)]})
    grouped: dict[str, list[tuple[int, Any]]] = {}
    for index, name in enumerate(channels):
        if name in {"sea_ice_cover", "sea_surface_temperature"}:
            values = xr.DataArray(invar[index].T, dims=("longitude", "latitude"), coords={"latitude": lat, "longitude": lon})
        else:
            base, _, suffix = name.rpartition("_")
            if not suffix.isdigit() or base not in config["model"]["input_variables"] + config["model"].get("optional_input_variables", []):
                continue
            values = xr.DataArray(invar[index].T, dims=("longitude", "latitude"), coords={"latitude": lat, "longitude": lon}).expand_dims(level=[int(suffix)])
            values = values.expand_dims(time=[np.datetime64(timestamp)])
            grouped.setdefault(base, []).append((int(suffix), values))
            continue
        values = values.expand_dims(time=[np.datetime64(timestamp)])
        dataset[name] = values
    for base, entries in grouped.items():
        entries.sort(key=lambda item: levels.index(item[0]) if item[0] in levels else item[0])
        merged = xr.concat([value for _, value in entries], dim="level")
        dataset[base] = merged.transpose("time", "level", "longitude", "latitude") if "time" in merged.dims else merged.transpose("level", "longitude", "latitude")
    return dataset


def era5_frames_to_xarray(
    frames: Any,
    config: dict[str, Any],
    *,
    start_time: Any,
):
    """Vectorized ERA5 ``(T,C,H,W)`` to NeuralGCM xarray conversion.

    This is equivalent to concatenating ``era5_sample_to_xarray`` outputs, but
    constructs every multi-level variable in one operation. It avoids hundreds
    of small DataArray allocations per training window.
    """
    import numpy as np
    import xarray as xr

    frames = as_time_major_frames(frames, name="ERA5 trajectory")
    channels = channel_order(config)
    if frames.shape[1] != len(channels):
        raise ValueError(
            f"ERA5 trajectory has {frames.shape[1]} channels, expected "
            f"{len(channels)}"
        )
    n_time, _, height, width = frames.shape
    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)
    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
    step_hours = int(config["data"].get("time_step_hours", 6))
    times = np.datetime64(start_time) + np.arange(n_time) * np.timedelta64(step_hours, "h")
    coords = {"time": times, "latitude": lat, "longitude": lon}
    dataset = xr.Dataset(coords=coords)

    level_indices: dict[str, list[tuple[int, int]]] = {}
    allowed = set(config["model"]["input_variables"])
    allowed.update(config["model"].get("optional_input_variables", []))
    for channel_index, name in enumerate(channels):
        if name in {"sea_ice_cover", "sea_surface_temperature"}:
            dataset[name] = (
                ("time", "longitude", "latitude"),
                np.asarray(frames[:, channel_index]).transpose(0, 2, 1),
            )
            continue
        base, _, suffix = name.rpartition("_")
        if suffix.isdigit() and base in allowed:
            level_indices.setdefault(base, []).append((int(suffix), channel_index))

    configured_levels = pressure_levels(config)
    for base, entries in level_indices.items():
        entries.sort(
            key=lambda item: configured_levels.index(item[0])
            if item[0] in configured_levels
            else item[0]
        )
        indices = [index for _, index in entries]
        levels = [level for level, _ in entries]
        values = np.asarray(frames[:, indices]).transpose(0, 1, 3, 2)
        dataset[base] = (
            ("time", "level", "longitude", "latitude"),
            values,
        )
        dataset = dataset.assign_coords(level=np.asarray(levels))
    return dataset


def _target_grid(mode: str):
    from dinosaur import spherical_harmonic

    targets = {
        "weather_forecast": spherical_harmonic.Grid.TL255,
        "climate_scale": spherical_harmonic.Grid.TL127,
        "forecast_2_8_deg": spherical_harmonic.Grid.TL63,
        "stochastic_1_4_deg": spherical_harmonic.Grid.TL127,
    }
    try:
        return targets[mode]()
    except KeyError as exc:
        raise ValueError(f"Unknown model mode {mode!r}") from exc


@lru_cache(maxsize=16)
def _profile_regridder(
    height: int,
    width: int,
    mode: str,
    latitude_spacing: str,
    longitude_offset: float,
):
    """Construct and cache the profile's conservative regridder."""
    from dinosaur import horizontal_interpolation, spherical_harmonic

    source_grid = spherical_harmonic.Grid(
        latitude_nodes=height,
        longitude_nodes=width,
        latitude_spacing=latitude_spacing,
        longitude_offset=longitude_offset,
    )
    return horizontal_interpolation.ConservativeRegridder(
        source_grid, _target_grid(mode), skipna=True
    )


def regrid_for_neuralgcm(dataset: Any, official_model: Any):
    """Conservatively regrid ERA5 fields to the checkpoint's Gaussian grid."""
    from dinosaur import horizontal_interpolation
    from dinosaur import spherical_harmonic
    from dinosaur import xarray_utils

    source_grid = spherical_harmonic.Grid(
        latitude_nodes=dataset.sizes["latitude"],
        longitude_nodes=dataset.sizes["longitude"],
        latitude_spacing=xarray_utils.infer_latitude_spacing(dataset.latitude),
        longitude_offset=xarray_utils.infer_longitude_offset(dataset.longitude),
    )
    regridder = horizontal_interpolation.ConservativeRegridder(
        source_grid, official_model.data_coords.horizontal, skipna=True
    )
    regridded = xarray_utils.regrid(dataset, regridder)
    return xarray_utils.fill_nan_with_nearest(regridded)


def regrid_for_profile(dataset: Any, mode: str):
    """Regrid to the Gaussian data grid selected by an official Gin profile."""
    from dinosaur import xarray_utils

    regridder = _profile_regridder(
        dataset.sizes["latitude"],
        dataset.sizes["longitude"],
        mode,
        xarray_utils.infer_latitude_spacing(dataset.latitude),
        float(xarray_utils.infer_longitude_offset(dataset.longitude)),
    )
    return xarray_utils.fill_nan_with_nearest(xarray_utils.regrid(dataset, regridder))


@lru_cache(maxsize=16)
def _load_static_features(
    path_text: str,
    mode: str | None,
    target_height: int,
    target_width: int,
):
    """Load and, only when necessary, regrid a reusable static dataset."""
    import xarray as xr

    with xr.open_dataset(path_text) as source:
        static = source[["geopotential_at_surface", "land_sea_mask"]].load()
    source_shape = (
        static.sizes.get("latitude"),
        static.sizes.get("longitude"),
    )
    if source_shape != (target_height, target_width):
        if mode is None:
            return None
        static = regrid_for_profile(static, mode)
    if (
        static.sizes.get("latitude"),
        static.sizes.get("longitude"),
    ) != (target_height, target_width):
        return None
    return static


def add_static_features(
    dataset: Any,
    config: dict[str, Any] | None = None,
    *,
    mode: str | None = None,
    prefer_profile: bool = True,
):
    """Attach official profile static fields, with a synthetic fallback.

    Callers attach fields after regridding the dynamic ERA5 trajectory. This
    preserves the exact Gaussian-grid topography and land/sea mask bundled in
    the official checkpoints. ``data.static_file`` remains a source-grid
    fallback for installations that do not carry the released checkpoints.
    """
    import numpy as np

    required = ("geopotential_at_surface", "land_sea_mask")
    if config is not None and not set(required).issubset(dataset):
        data_cfg = config.get("data", {})
        profile_path = (
            data_cfg.get("static_files", {}).get(mode) if mode else None
        )
        fallback_path = data_cfg.get("static_file")
        candidates = []
        if not prefer_profile and fallback_path:
            candidates.append(fallback_path)
        if mode:
            if profile_path:
                candidates.append(profile_path)
        if prefer_profile and fallback_path:
            candidates.append(fallback_path)
        for value in candidates:
            static_path = resolve_path(value)
            if not static_path.exists():
                continue
            static = _load_static_features(
                str(static_path.resolve()),
                mode,
                int(dataset.sizes["latitude"]),
                int(dataset.sizes["longitude"]),
            )
            if static is None:
                continue
            for name in required:
                if name not in dataset:
                    # Both arrays are on the same profile Gaussian grid. Assign
                    # by position rather than xarray label alignment: checkpoint
                    # coordinates are float64 while regridded ERA5 coordinates
                    # can be float32, and exact-label alignment would inject NaN.
                    values = static[name].transpose("longitude", "latitude")
                    dataset[name] = (
                        ("longitude", "latitude"),
                        np.asarray(values.values),
                    )
                    dataset[name].attrs.update(values.attrs)
            dataset.attrs["static_features_source"] = str(static_path)
            break
    if "geopotential_at_surface" not in dataset:
        dataset["geopotential_at_surface"] = (("longitude", "latitude"), np.zeros((dataset.sizes["longitude"], dataset.sizes["latitude"]), np.float32))
    if "land_sea_mask" not in dataset:
        dataset["land_sea_mask"] = (("longitude", "latitude"), np.zeros((dataset.sizes["longitude"], dataset.sizes["latitude"]), np.float32))
    # Gin FloatDataFeatures parses units from these static fields exactly as in
    # the official ERA5 pipeline.
    dataset["geopotential_at_surface"].attrs.setdefault("units", "m**2 s**-2")
    dataset["land_sea_mask"].attrs.setdefault("units", "dimensionless")
    return dataset