model/era5_adapter.py

"""OneScience ERA5 adapter for the official OneForecast 69-channel contract."""

from __future__ import annotations

from pathlib import Path
import tempfile
from typing import Any, Iterable

import numpy as np

SOURCE_GRID = (721, 1440)
ONEFORECAST_FILE_GRID = (121, 240)
SPATIAL_STRIDE = 6

OFFICIAL_VARIABLES = tuple(
    [f"Z{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
    + [f"Q{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
    + [f"T{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
    + [f"U{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
    + [f"V{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
    + ["U10M", "V10M", "T2M", "MSLP"]
)

VARIABLE_ALIASES = {
    **{f"Z{x}": f"geopotential_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
    **{f"Q{x}": f"specific_humidity_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
    **{f"T{x}": f"temperature_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
    **{f"U{x}": f"u_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
    **{f"V{x}": f"v_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
    "U10M": "10m_u_component_of_wind",
    "V10M": "10m_v_component_of_wind",
    "T2M": "2m_temperature",
    "MSLP": "mean_sea_level_pressure",
}


def _decode_variables(values: Iterable[Any]) -> list[str]:
    return [value.decode() if isinstance(value, bytes) else str(value) for value in values]


class OneForecastERA5Adapter:
    """Validate files and construct OneScience's ERA5 DataLoader."""

    def __init__(self, dataset_dir: str | Path, years: Iterable[int], batch_size: int = 1,
                 input_steps: int = 1, output_steps: int = 1, normalize: bool = True,
                 num_workers: int = 0, distributed: bool = False) -> None:
        self.dataset_dir = Path(dataset_dir).expanduser().resolve()
        self.years = [int(year) for year in years]
        self.batch_size = batch_size
        self.input_steps = input_steps
        self.output_steps = output_steps
        self.normalize = normalize
        self.num_workers = num_workers
        self.distributed = distributed
        self.source_variables: list[str] = []
        self.channel_indices: list[int] = []
        self.global_means: np.ndarray | None = None
        self.global_stds: np.ndarray | None = None
        self.time_step_hours: int | None = None
        self.source_grid: tuple[int, int] | None = None
        self._external_stats: tuple[Path, Path] | None = None
        self._layout_dir: tempfile.TemporaryDirectory[str] | None = None
        self._validate_files()

    def _year_path(self, year: int) -> Path:
        for path in (self.dataset_dir / "data" / f"{year}.h5", self.dataset_dir / f"{year}.h5"):
            if path.is_file():
                return path
        raise FileNotFoundError(f"ERA5 file for year {year} was not found below {self.dataset_dir}")

    def _validate_files(self) -> None:
        try:
            import h5py
        except ImportError as exc:
            raise RuntimeError("h5py is required to validate ERA5 HDF5 files") from exc
        if not self.years:
            raise ValueError("At least one ERA5 year is required")
        reference_variables: list[str] | None = None
        reference_indices: list[int] | None = None
        for year in self.years:
            path = self._year_path(year)
            with h5py.File(path, "r") as handle:
                if "fields" not in handle:
                    raise ValueError(f"{path} does not contain a fields dataset")
                fields = handle["fields"]
                if len(fields.shape) != 4:
                    raise ValueError(f"{path}: fields must have shape [T, C, H, W], got {fields.shape}")
                variables = _decode_variables(fields.attrs.get("variables", []))
                source_variables = [
                    name if name in variables else VARIABLE_ALIASES[name]
                    for name in OFFICIAL_VARIABLES
                    if name in variables or VARIABLE_ALIASES[name] in variables
                ]
                missing = [
                    name for name in OFFICIAL_VARIABLES
                    if name not in variables and VARIABLE_ALIASES[name] not in variables
                ]
                if missing:
                    raise ValueError(f"{path}: missing official variables: {missing}")
                indices = [variables.index(name) for name in source_variables]
                if reference_variables is not None and variables != reference_variables:
                    raise ValueError(f"{path}: complete variable metadata differs between yearly files")
                if reference_indices is not None and indices != reference_indices:
                    raise ValueError(f"{path}: official channel indices differ between yearly files")
                reference_variables, reference_indices = variables, indices
                self.source_variables = source_variables
                self.channel_indices = indices
                if fields.shape[1] != len(variables):
                    raise ValueError(f"{path}: variables metadata does not match channel dimension")
                if fields.shape[1] != 69 or tuple(fields.shape[2:]) not in (SOURCE_GRID, ONEFORECAST_FILE_GRID):
                    raise ValueError(
                        f"{path}: expected fields [T, 69, 721, 1440] or [T, 69, 121, 240], got {fields.shape}"
                    )
                grid = tuple(fields.shape[2:])
                if self.source_grid is not None and grid != self.source_grid:
                    raise ValueError(f"{path}: spatial grid differs between yearly files")
                self.source_grid = grid
                if fields.shape[0] < self.input_steps + self.output_steps:
                    raise ValueError(f"{path}: not enough time steps for configured window")
                if "time_step" not in fields.attrs:
                    raise ValueError(f"{path}: fields.attrs['time_step'] is required by ERA5Datapipe")
                time_step = int(fields.attrs["time_step"])
                if time_step != 6 or (self.time_step_hours is not None and time_step != self.time_step_hours):
                    raise ValueError(f"{path}: expected a consistent 6-hour time_step, got {time_step}")
                self.time_step_hours = time_step
                if "global_means" in handle and "global_stds" in handle:
                    means = np.asarray(handle["global_means"])
                    stds = np.asarray(handle["global_stds"])
                else:
                    candidates = (
                        (self.dataset_dir / "stats" / "global_means.npy",
                         self.dataset_dir / "stats" / "global_stds.npy"),
                        (self.dataset_dir / "mean.npy", self.dataset_dir / "std.npy"),
                        (self.dataset_dir.parent / "mean.npy", self.dataset_dir.parent / "std.npy"),
                    )
                    stats_paths = next(((mean, std) for mean, std in candidates
                                        if mean.is_file() and std.is_file()), None)
                    if stats_paths is None:
                        raise ValueError(f"{path}: embedded or external ERA5 statistics are required")
                    self._external_stats = stats_paths
                    means, stds = (np.load(item) for item in stats_paths)
                expected_shape = (1, len(variables), 1, 1)
                if means.shape != expected_shape or stds.shape != expected_shape:
                    raise ValueError(f"{path}: statistics must have shape {expected_shape}")
                if not np.isfinite(means).all() or not np.isfinite(stds).all() or not (stds > 0).all():
                    raise ValueError(f"{path}: statistics must be finite and standard deviations positive")
                if self.global_means is not None and not np.array_equal(means, self.global_means):
                    raise ValueError(f"{path}: global_means differ between yearly files")
                if self.global_stds is not None and not np.array_equal(stds, self.global_stds):
                    raise ValueError(f"{path}: global_stds differ between yearly files")
                self.global_means, self.global_stds = means, stds

    def _onescience_dataset_dir(self) -> Path:
        if self._layout_dir is not None:
            return Path(self._layout_dir.name)
        self._layout_dir = tempfile.TemporaryDirectory(prefix="oneforecast_era5_")
        root = Path(self._layout_dir.name)
        data_dir = root / "data"
        data_dir.mkdir()
        for year in self.years:
            source_path = self._year_path(year)
            target_path = data_dir / f"{year}.h5"
            if self.source_grid == SOURCE_GRID:
                import h5py

                with h5py.File(source_path, "r") as source_handle:
                    source_fields = source_handle["fields"]
                    layout = h5py.VirtualLayout(
                        shape=(source_fields.shape[0], source_fields.shape[1], *ONEFORECAST_FILE_GRID),
                        dtype=source_fields.dtype,
                    )
                    virtual_source = h5py.VirtualSource(str(source_path), "fields", shape=source_fields.shape)
                    layout[:] = virtual_source[:, :, ::SPATIAL_STRIDE, ::SPATIAL_STRIDE]
                    with h5py.File(target_path, "w", libver="latest") as target_handle:
                        fields = target_handle.create_virtual_dataset("fields", layout)
                        for name, value in source_fields.attrs.items():
                            fields.attrs[name] = value
            else:
                target_path.symlink_to(source_path)
        if self._external_stats is not None:
            stats_dir = root / "stats"
            stats_dir.mkdir()
            (stats_dir / "global_means.npy").symlink_to(self._external_stats[0])
            (stats_dir / "global_stds.npy").symlink_to(self._external_stats[1])

        return root

    def get_dataloader(self, mode: str):
        """Delegate loading to OneScience, then align native ERA5 to OneForecast's grid."""
        try:
            from onescience.datapipes.climate.era5 import ERA5Datapipe
        except ImportError as exc:
            raise RuntimeError("OneScience ERA5Datapipe is required for data loading") from exc
        datapipe = ERA5Datapipe(
            dataset_dir=str(self._onescience_dataset_dir()), used_years=self.years,
            used_variables=self.source_variables, distributed=self.distributed,
            input_steps=self.input_steps, output_steps=self.output_steps,
            normalize=self.normalize, batch_size=self.batch_size, num_workers=self.num_workers,
        )
        loader, sampler = datapipe.get_dataloader(mode=mode)
        return _SpatiallyAdaptedLoader(loader, self.source_grid), sampler

    def inspect(self) -> dict[str, Any]:
        try:
            import h5py
        except ImportError as exc:
            raise RuntimeError("h5py is required to inspect ERA5 HDF5 files") from exc
        path = self._year_path(self.years[0])
        with h5py.File(path, "r") as handle:
            fields = handle["fields"]
            variables = _decode_variables(fields.attrs["variables"])
            indices = [variables.index(name) for name in self.source_variables]
            return {"path": str(path), "fields_shape": list(fields.shape),
                     "source_grid": list(fields.shape[2:]),
                     "oneforecast_file_grid": list(ONEFORECAST_FILE_GRID),
                     "oneforecast_model_grid": [120, 240],
                     "spatial_transform": "identity" if tuple(fields.shape[2:]) == ONEFORECAST_FILE_GRID else "stride_6",
                     "time_step_hours": int(fields.attrs["time_step"]),
                     "variable_count": len(variables), "official_channel_indices": indices,
                     "source_variables": self.source_variables,
                     "statistics_shape": list(self.global_means.shape),
                     "statistics_shared_across_years": True,
                     "official_variables_match": len(indices) == len(OFFICIAL_VARIABLES)}

    def selected_statistics(self) -> tuple[np.ndarray, np.ndarray]:
        """Return normalization statistics in the model's 69-channel order."""
        if self.global_means is None or self.global_stds is None:
            raise RuntimeError("ERA5 statistics have not been validated")
        return self.global_means[:, self.channel_indices], self.global_stds[:, self.channel_indices]


def _adapt_spatial(value: Any, source_grid: tuple[int, int] | None) -> Any:
    if not hasattr(value, "shape") or len(value.shape) < 2:
        return value
    if tuple(value.shape[-2:]) == ONEFORECAST_FILE_GRID:
        return value
    if tuple(value.shape[-2:]) != SOURCE_GRID or source_grid != SOURCE_GRID:
        return value
    return value[..., ::SPATIAL_STRIDE, ::SPATIAL_STRIDE]


class _SpatiallyAdaptedLoader:
    """Preserve the DataLoader interface while adapting fields after ERA5Datapipe."""

    def __init__(self, loader: Any, source_grid: tuple[int, int] | None) -> None:
        self.loader = loader
        self.source_grid = source_grid

    def __len__(self) -> int:
        return len(self.loader)

    def __iter__(self):
        for batch in self.loader:
            yield tuple(_adapt_spatial(value, self.source_grid) for value in batch)