| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any |
|
|
| import h5py |
| import torch |
| from onescience.datapipes.climate.era5 import ERA5Dataset |
| from torch.utils.data import Dataset |
|
|
| from grid import lambert_grid |
|
|
|
|
| class StormCastDataset(Dataset): |
| """Pair OneScience ERA5 backgrounds with synchronized local state targets.""" |
|
|
| def __init__( |
| self, |
| data_root: str | Path, |
| years: list[int], |
| era5_variables: list[str], |
| state_variables: list[str], |
| invariant_variables: list[str], |
| image_size: list[int] | tuple[int, int], |
| input_steps: int = 1, |
| output_steps: int = 1, |
| normalize: bool = True, |
| ) -> None: |
| if input_steps != 1 or output_steps != 1: |
| raise ValueError("StormCast pairing currently requires one input and one target step") |
|
|
| self.data_root = Path(data_root) |
| self.years = years |
| self.era5_variables = era5_variables |
| self.state_variables = state_variables |
| self.invariant_variables = invariant_variables |
| self.image_size = tuple(image_size) |
| self.normalize = normalize |
| self.era5 = ERA5Dataset( |
| dataset_dir=str(self.data_root / "era5"), |
| used_years=years, |
| used_variables=era5_variables, |
| input_steps=input_steps, |
| output_steps=output_steps, |
| normalize=normalize, |
| ) |
| self.samples_per_year = self.era5.samples_per_year |
| self._validate_era5_grid() |
| self._validate_local_files() |
| self.invariants = self._load_invariants() |
| self._initialize_background_regrid() |
|
|
| def _validate_era5_grid(self) -> None: |
| if self.era5.H < 2 or self.era5.W < 2: |
| raise ValueError("ERA5 grid must have at least two points per dimension") |
| expected = (721, 1440) |
| if (self.era5.H, self.era5.W) != expected: |
| raise ValueError( |
| f"StormCast expects ERA5 on the global {expected} grid, " |
| f"got {(self.era5.H, self.era5.W)}" |
| ) |
|
|
| def _validate_local_files(self) -> None: |
| for year in self.years: |
| path = self.data_root / "hrrr" / "data" / f"{year}.h5" |
| if not path.is_file(): |
| raise FileNotFoundError(f"Missing local state file: {path}") |
| with h5py.File(path, "r") as handle: |
| fields = handle["fields"] |
| variables = [ |
| value.decode() if isinstance(value, bytes) else str(value) |
| for value in fields.attrs["variables"] |
| ] |
| if variables != self.state_variables: |
| raise ValueError( |
| "Local state channel order differs from data.state_variables" |
| ) |
| expected_steps = self.samples_per_year + 1 |
| if fields.shape[0] != expected_steps: |
| raise ValueError( |
| f"{path} has {fields.shape[0]} steps, expected {expected_steps}" |
| ) |
| if tuple(fields.shape[-2:]) != self.image_size: |
| raise ValueError( |
| f"Local state grid is {tuple(fields.shape[-2:])}, " |
| f"expected regional grid {self.image_size}" |
| ) |
|
|
| def _load_invariants(self) -> torch.Tensor: |
| path = self.data_root / "hrrr" / "invariants.h5" |
| with h5py.File(path, "r") as handle: |
| fields = handle["fields"] |
| variables = [ |
| value.decode() if isinstance(value, bytes) else str(value) |
| for value in fields.attrs["variables"] |
| ] |
| if variables != self.invariant_variables: |
| raise ValueError( |
| "Invariant channel order differs from data.invariant_variables" |
| ) |
| invariants = torch.as_tensor(fields[:], dtype=torch.float32) |
| if tuple(invariants.shape[-2:]) != self.image_size: |
| raise ValueError( |
| f"Invariant grid is {tuple(invariants.shape[-2:])}, " |
| f"expected {self.image_size}" |
| ) |
| return invariants |
|
|
| def _initialize_background_regrid(self) -> None: |
| with h5py.File(self.data_root / "hrrr" / "invariants.h5", "r") as handle: |
| if "lat" in handle and "lon" in handle: |
| target_lat = torch.as_tensor(handle["lat"][:], dtype=torch.float32) |
| target_lon = torch.as_tensor(handle["lon"][:], dtype=torch.float32) |
| else: |
| target_lat_np, target_lon_np = lambert_grid(self.image_size) |
| target_lat = torch.from_numpy(target_lat_np) |
| target_lon = torch.from_numpy(target_lon_np) |
| if target_lat.shape != self.image_size or target_lon.shape != self.image_size: |
| raise ValueError("StormCast target latitude/longitude grid has wrong shape") |
|
|
| lat_position = (90.0 - target_lat) / (180.0 / (self.era5.H - 1)) |
| lon_position = torch.remainder(target_lon, 360.0) / (360.0 / self.era5.W) |
| self.lat0 = lat_position.floor().long().clamp(0, self.era5.H - 2) |
| self.lat1 = self.lat0 + 1 |
| self.lon0 = lon_position.floor().long().remainder(self.era5.W) |
| self.lon1 = (self.lon0 + 1).remainder(self.era5.W) |
| self.lat_weight = lat_position - self.lat0 |
| self.lon_weight = lon_position - lon_position.floor() |
|
|
| def _regrid_background(self, background: torch.Tensor) -> torch.Tensor: |
| f00 = background[..., self.lat0, self.lon0] |
| f01 = background[..., self.lat0, self.lon1] |
| f10 = background[..., self.lat1, self.lon0] |
| f11 = background[..., self.lat1, self.lon1] |
| lon_weight = self.lon_weight.to(background.dtype) |
| lat_weight = self.lat_weight.to(background.dtype) |
| top = torch.lerp(f00, f01, lon_weight) |
| bottom = torch.lerp(f10, f11, lon_weight) |
| return torch.lerp(top, bottom, lat_weight) |
|
|
| def __len__(self) -> int: |
| return len(self.era5) |
|
|
| def __getitem__(self, index: int) -> dict[str, Any]: |
| background, _, _, step_index, time_index = self.era5[index] |
| background = self._regrid_background(background) |
| year_index = index // self.samples_per_year |
| year = self.years[year_index] |
| path = self.data_root / "hrrr" / "data" / f"{year}.h5" |
|
|
| with h5py.File(path, "r") as handle: |
| state = torch.as_tensor( |
| handle["fields"][step_index : step_index + 2], dtype=torch.float32 |
| ) |
| if self.normalize: |
| means = torch.as_tensor(handle["global_means"][:], dtype=torch.float32) |
| stds = torch.as_tensor(handle["global_stds"][:], dtype=torch.float32) |
| state = (state - means) / stds |
|
|
| return { |
| "background": background, |
| "state": (state[0], state[1]), |
| "invariant": self.invariants, |
| "step_index": step_index, |
| "time_index": time_index, |
| } |
|
|
