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