scripts/data_loader.py

"""OneScience ERA5Dataset adapter for ClimODE's 32x64 global grid."""

from __future__ import annotations

from pathlib import Path
from typing import Iterable, Sequence

import h5py
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

try:
    from onescience.datapipes.climate.era5 import ERA5Dataset
except ImportError as exc:  # pragma: no cover - exercised only without OneScience
    ERA5Dataset = None
    _ERA5_IMPORT_ERROR = exc
else:
    _ERA5_IMPORT_ERROR = None


OFFICIAL_VARIABLES = ("z", "t", "t2m", "u10", "v10")


def _require_era5dataset() -> None:
    if ERA5Dataset is None:
        raise ImportError(
            "ClimODE data loading requires OneScience ERA5Dataset; "
            "activate an environment containing onescience before running."
        ) from _ERA5_IMPORT_ERROR


def _as_channel_vector(values: np.ndarray | torch.Tensor) -> torch.Tensor:
    tensor = torch.as_tensor(values, dtype=torch.float32)
    return tensor.reshape(-1)


def _regrid_periodic(frame: torch.Tensor, target_size: tuple[int, int]) -> torch.Tensor:
    """Bilinearly sample [C,H,W] on WeatherBench cell centers."""

    if frame.ndim != 3:
        raise ValueError(f"Expected [C,H,W], got {tuple(frame.shape)}")
    target_height, target_width = target_size
    if frame.shape[-2:] == target_size:
        return frame
    periodic = torch.cat([frame, frame[..., :1]], dim=-1).unsqueeze(0)
    latitude = torch.linspace(
        90.0 - 90.0 / target_height,
        -90.0 + 90.0 / target_height,
        target_height,
        device=frame.device,
        dtype=frame.dtype,
    )
    longitude = (
        torch.arange(target_width, device=frame.device, dtype=frame.dtype)
        * (360.0 / target_width)
    )
    lat2d, lon2d = torch.meshgrid(latitude, longitude, indexing="ij")
    grid = torch.stack([lon2d / 180.0 - 1.0, -lat2d / 90.0], dim=-1)
    return F.grid_sample(
        periodic,
        grid.unsqueeze(0),
        mode="bilinear",
        padding_mode="border",
        align_corners=True,
    )[0]


def _load_stats(stats_dir: str | Path, channels: int) -> tuple[torch.Tensor, torch.Tensor]:
    stats_path = Path(stats_dir)
    minimum = _as_channel_vector(np.load(stats_path / "min_values.npy"))
    maximum = _as_channel_vector(np.load(stats_path / "max_values.npy"))
    if minimum.numel() != channels or maximum.numel() != channels:
        raise ValueError(
            f"Expected {channels} channel statistics, got {minimum.numel()} and {maximum.numel()}"
        )
    if torch.any(maximum <= minimum):
        raise ValueError("All max_values must be greater than min_values")
    return minimum, maximum


def _normalize(frame: torch.Tensor, minimum: torch.Tensor, maximum: torch.Tensor) -> torch.Tensor:
    scale = (maximum - minimum).clamp_min(torch.finfo(frame.dtype).eps)
    return (frame - minimum[:, None, None]) / scale[:, None, None]


class ClimODEDataset(Dataset):
    """Return three history frames and the following target frame.

    The underlying annual files are always read by OneScience ``ERA5Dataset``.
    No direct HDF5 field access is used here, which keeps the model adapter
    compatible with the OneScience ERA5 contract.
    """

    def __init__(
        self,
        data_dir: str | Path,
        years: Sequence[int],
        used_variables: Sequence[str] = OFFICIAL_VARIABLES,
        stats_dir: str | Path | None = None,
        model_size: tuple[int, int] = (32, 64),
        normalize: bool = True,
        input_steps: int = 1,
        output_steps: int = 1,
    ) -> None:
        _require_era5dataset()
        if tuple(used_variables) != OFFICIAL_VARIABLES:
            raise ValueError(
                "ClimODE requires the exact channel order ['z','t','t2m','u10','v10']"
            )
        if input_steps != 1 or output_steps != 1:
            raise ValueError("The ClimODE adapter currently uses one input and one target step")
        if len(years) == 0:
            raise ValueError("At least one year is required")
        # ``data_dir`` is the OneScience dataset root containing data/*.h5,
        # rather than the nested data/ directory itself.
        self.data_dir = Path(data_dir)
        self.years = [int(year) for year in years]
        self.variables = tuple(used_variables)
        self.model_size = tuple(model_size)
        self.normalize = normalize
        self.era5 = ERA5Dataset(
            dataset_dir=str(self.data_dir),
            used_years=self.years,
            used_variables=list(self.variables),
            input_steps=1,
            output_steps=1,
            normalize=False,
        )
        if (self.era5.H, self.era5.W) != (721, 1440):
            raise ValueError(
                "ClimODE raw-data adapter expects (721,1440), "
                f"got ({self.era5.H},{self.era5.W})"
            )
        self.samples_per_year = self.era5.samples_per_year
        if self.samples_per_year < 3:
            raise ValueError("Each year needs at least four frames for a three-frame history and target")
        self.samples_per_year_with_history = self.samples_per_year - 2
        self.minimum, self.maximum = (
            _load_stats(stats_dir or self.data_dir / "static", len(self.variables))
            if normalize
            else (torch.zeros(len(self.variables)), torch.ones(len(self.variables)))
        )

    def __len__(self) -> int:
        return len(self.years) * self.samples_per_year_with_history

    def _frame(self, sample_index: int, target: bool = False) -> torch.Tensor:
        invar, outvar, _, _, _ = self.era5[sample_index]
        frame = outvar if target else invar
        frame = _regrid_periodic(torch.as_tensor(frame, dtype=torch.float32), self.model_size)
        if self.normalize:
            frame = _normalize(frame, self.minimum, self.maximum)
        return frame

    def __getitem__(self, index: int) -> dict[str, torch.Tensor | int | str]:
        if index < 0:
            index += len(self)
        if index < 0 or index >= len(self):
            raise IndexError(index)
        year_index = index // self.samples_per_year_with_history
        local_index = index % self.samples_per_year_with_history
        base = year_index * self.samples_per_year + local_index + 2
        history = torch.stack([self._frame(base - 2), self._frame(base - 1), self._frame(base)])
        target = self._frame(base, target=True)
        return {
            "history": history,
            "input": history[-1],
            "target": target,
            "year": self.years[year_index],
            "step_index": local_index + 2,
        }


class ClimODESeriesDataset(Dataset):
    """Official-style sequence batches with years as the inner batch axis.

    Each item contains a contiguous sequence for every requested year.  The
    outer DataLoader should use ``batch_size=1``; the sequence length plays the
    role of the official training batch of time points.
    """

    def __init__(
        self,
        data_dir: str | Path,
        years: Sequence[int],
        used_variables: Sequence[str] = OFFICIAL_VARIABLES,
        stats_dir: str | Path | None = None,
        model_size: tuple[int, int] = (32, 64),
        sequence_length: int = 8,
        normalize: bool = True,
    ) -> None:
        _require_era5dataset()
        if tuple(used_variables) != OFFICIAL_VARIABLES:
            raise ValueError("ClimODE requires the exact channel order ['z','t','t2m','u10','v10']")
        if sequence_length < 1:
            raise ValueError("sequence_length must be positive")
        self.data_dir = Path(data_dir)
        self.years = [int(year) for year in years]
        self.variables = tuple(used_variables)
        self.model_size = tuple(model_size)
        self.sequence_length = int(sequence_length)
        self.normalize = normalize
        self.era5 = ERA5Dataset(
            dataset_dir=str(self.data_dir),
            used_years=self.years,
            used_variables=list(self.variables),
            input_steps=1,
            output_steps=1,
            normalize=False,
        )
        if (self.era5.H, self.era5.W) != (721, 1440):
            raise ValueError(
                "ClimODE raw-data adapter expects (721,1440), "
                f"got ({self.era5.H},{self.era5.W})"
            )
        self.samples_per_year = self.era5.samples_per_year
        self.frames_per_year = self.era5.T
        first_start = 2
        # The official DataLoader keeps its final, possibly shorter batch.
        self.starts = list(range(first_start, self.frames_per_year, self.sequence_length))
        if not self.starts:
            raise ValueError(
                f"Not enough frames ({self.era5.T}) for sequence_length={sequence_length} "
                "and a three-frame history"
            )
        self.minimum, self.maximum = (
            _load_stats(stats_dir or self.data_dir / "static", len(self.variables))
            if normalize
            else (torch.zeros(len(self.variables)), torch.ones(len(self.variables)))
        )

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

    def _frame(self, year_index: int, frame_index: int) -> torch.Tensor:
        if frame_index < 0 or frame_index >= self.frames_per_year:
            raise IndexError(frame_index)
        sample_index = year_index * self.samples_per_year + min(
            frame_index, self.samples_per_year - 1
        )
        invar, outvar, _, _, _ = self.era5[sample_index]
        # ERA5Dataset's final input index is T-2; its paired target is frame T-1.
        frame = outvar if frame_index == self.frames_per_year - 1 else invar
        frame = _regrid_periodic(torch.as_tensor(frame, dtype=torch.float32), self.model_size)
        if self.normalize:
            frame = _normalize(frame, self.minimum, self.maximum)
        return frame

    def __getitem__(self, index: int) -> dict[str, torch.Tensor | int]:
        start = self.starts[index]
        history_per_year = []
        sequence_per_year = []
        for year_index in range(len(self.years)):
            history_per_year.append(
                torch.stack(
                    [
                        self._frame(year_index, start - 2),
                        self._frame(year_index, start - 1),
                        self._frame(year_index, start),
                    ]
                )
            )
            stop = min(start + self.sequence_length, self.frames_per_year)
            sequence_per_year.append(
                torch.stack(
                    [self._frame(year_index, step) for step in range(start, stop)]
                )
            )
        stop = min(start + self.sequence_length, self.frames_per_year)
        return {
            "history": torch.stack(history_per_year, dim=0),
            "observations": torch.stack(sequence_per_year, dim=1),
            "time_steps": torch.arange(start, stop, dtype=torch.float32),
            "sequence_index": index,
        }


def load_constants(
    static_file: str | Path,
    expected_size: tuple[int, int] = (32, 64),
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Load [orography, lsm], latitude and longitude from constants.h5."""

    with h5py.File(static_file, "r") as handle:
        constants = torch.stack(
            [
                torch.as_tensor(handle["orography"][:], dtype=torch.float32),
                torch.as_tensor(handle["lsm"][:], dtype=torch.float32),
            ]
        ).unsqueeze(0)
        lat2d = torch.as_tensor(handle["lat2d"][:], dtype=torch.float32)
        lon2d = torch.as_tensor(handle["lon2d"][:], dtype=torch.float32)
    if tuple(constants.shape[-2:]) != expected_size:
        raise ValueError(f"Static constants have shape {tuple(constants.shape[-2:])}")
    return constants, lat2d, lon2d


def make_dataloader(
    dataset: Dataset,
    batch_size: int,
    shuffle: bool,
    num_workers: int = 0,
    pin_memory: bool = False,
) -> DataLoader:
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=False,
    )