OneScience-Group/OneForecast

1"""OneScience ERA5 adapter for the official OneForecast 69-channel contract."""2 3from __future__ import annotations4 5from pathlib import Path6import tempfile7from typing import Any, Iterable8 9import numpy as np10 11SOURCE_GRID = (721, 1440)12ONEFORECAST_FILE_GRID = (121, 240)13SPATIAL_STRIDE = 614 15OFFICIAL_VARIABLES = tuple(16    [f"Z{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]17    + [f"Q{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]18    + [f"T{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]19    + [f"U{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]20    + [f"V{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]21    + ["U10M", "V10M", "T2M", "MSLP"]22)23 24VARIABLE_ALIASES = {25    **{f"Z{x}": f"geopotential_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},26    **{f"Q{x}": f"specific_humidity_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},27    **{f"T{x}": f"temperature_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},28    **{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)},29    **{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)},30    "U10M": "10m_u_component_of_wind",31    "V10M": "10m_v_component_of_wind",32    "T2M": "2m_temperature",33    "MSLP": "mean_sea_level_pressure",34}35 36 37def _decode_variables(values: Iterable[Any]) -> list[str]:38    return [value.decode() if isinstance(value, bytes) else str(value) for value in values]39 40 41class OneForecastERA5Adapter:42    """Validate files and construct OneScience's ERA5 DataLoader."""43 44    def __init__(self, dataset_dir: str | Path, years: Iterable[int], batch_size: int = 1,45                 input_steps: int = 1, output_steps: int = 1, normalize: bool = True,46                 num_workers: int = 0, distributed: bool = False) -> None:47        self.dataset_dir = Path(dataset_dir).expanduser().resolve()48        self.years = [int(year) for year in years]49        self.batch_size = batch_size50        self.input_steps = input_steps51        self.output_steps = output_steps52        self.normalize = normalize53        self.num_workers = num_workers54        self.distributed = distributed55        self.source_variables: list[str] = []56        self.channel_indices: list[int] = []57        self.global_means: np.ndarray | None = None58        self.global_stds: np.ndarray | None = None59        self.time_step_hours: int | None = None60        self.source_grid: tuple[int, int] | None = None61        self._external_stats: tuple[Path, Path] | None = None62        self._layout_dir: tempfile.TemporaryDirectory[str] | None = None63        self._validate_files()64 65    def _year_path(self, year: int) -> Path:66        for path in (self.dataset_dir / "data" / f"{year}.h5", self.dataset_dir / f"{year}.h5"):67            if path.is_file():68                return path69        raise FileNotFoundError(f"ERA5 file for year {year} was not found below {self.dataset_dir}")70 71    def _validate_files(self) -> None:72        try:73            import h5py74        except ImportError as exc:75            raise RuntimeError("h5py is required to validate ERA5 HDF5 files") from exc76        if not self.years:77            raise ValueError("At least one ERA5 year is required")78        reference_variables: list[str] | None = None79        reference_indices: list[int] | None = None80        for year in self.years:81            path = self._year_path(year)82            with h5py.File(path, "r") as handle:83                if "fields" not in handle:84                    raise ValueError(f"{path} does not contain a fields dataset")85                fields = handle["fields"]86                if len(fields.shape) != 4:87                    raise ValueError(f"{path}: fields must have shape [T, C, H, W], got {fields.shape}")88                variables = _decode_variables(fields.attrs.get("variables", []))89                source_variables = [90                    name if name in variables else VARIABLE_ALIASES[name]91                    for name in OFFICIAL_VARIABLES92                    if name in variables or VARIABLE_ALIASES[name] in variables93                ]94                missing = [95                    name for name in OFFICIAL_VARIABLES96                    if name not in variables and VARIABLE_ALIASES[name] not in variables97                ]98                if missing:99                    raise ValueError(f"{path}: missing official variables: {missing}")100                indices = [variables.index(name) for name in source_variables]101                if reference_variables is not None and variables != reference_variables:102                    raise ValueError(f"{path}: complete variable metadata differs between yearly files")103                if reference_indices is not None and indices != reference_indices:104                    raise ValueError(f"{path}: official channel indices differ between yearly files")105                reference_variables, reference_indices = variables, indices106                self.source_variables = source_variables107                self.channel_indices = indices108                if fields.shape[1] != len(variables):109                    raise ValueError(f"{path}: variables metadata does not match channel dimension")110                if fields.shape[1] != 69 or tuple(fields.shape[2:]) not in (SOURCE_GRID, ONEFORECAST_FILE_GRID):111                    raise ValueError(112                        f"{path}: expected fields [T, 69, 721, 1440] or [T, 69, 121, 240], got {fields.shape}"113                    )114                grid = tuple(fields.shape[2:])115                if self.source_grid is not None and grid != self.source_grid:116                    raise ValueError(f"{path}: spatial grid differs between yearly files")117                self.source_grid = grid118                if fields.shape[0] < self.input_steps + self.output_steps:119                    raise ValueError(f"{path}: not enough time steps for configured window")120                if "time_step" not in fields.attrs:121                    raise ValueError(f"{path}: fields.attrs['time_step'] is required by ERA5Datapipe")122                time_step = int(fields.attrs["time_step"])123                if time_step != 6 or (self.time_step_hours is not None and time_step != self.time_step_hours):124                    raise ValueError(f"{path}: expected a consistent 6-hour time_step, got {time_step}")125                self.time_step_hours = time_step126                if "global_means" in handle and "global_stds" in handle:127                    means = np.asarray(handle["global_means"])128                    stds = np.asarray(handle["global_stds"])129                else:130                    candidates = (131                        (self.dataset_dir / "stats" / "global_means.npy",132                         self.dataset_dir / "stats" / "global_stds.npy"),133                        (self.dataset_dir / "mean.npy", self.dataset_dir / "std.npy"),134                        (self.dataset_dir.parent / "mean.npy", self.dataset_dir.parent / "std.npy"),135                    )136                    stats_paths = next(((mean, std) for mean, std in candidates137                                        if mean.is_file() and std.is_file()), None)138                    if stats_paths is None:139                        raise ValueError(f"{path}: embedded or external ERA5 statistics are required")140                    self._external_stats = stats_paths141                    means, stds = (np.load(item) for item in stats_paths)142                expected_shape = (1, len(variables), 1, 1)143                if means.shape != expected_shape or stds.shape != expected_shape:144                    raise ValueError(f"{path}: statistics must have shape {expected_shape}")145                if not np.isfinite(means).all() or not np.isfinite(stds).all() or not (stds > 0).all():146                    raise ValueError(f"{path}: statistics must be finite and standard deviations positive")147                if self.global_means is not None and not np.array_equal(means, self.global_means):148                    raise ValueError(f"{path}: global_means differ between yearly files")149                if self.global_stds is not None and not np.array_equal(stds, self.global_stds):150                    raise ValueError(f"{path}: global_stds differ between yearly files")151                self.global_means, self.global_stds = means, stds152 153    def _onescience_dataset_dir(self) -> Path:154        if self._layout_dir is not None:155            return Path(self._layout_dir.name)156        self._layout_dir = tempfile.TemporaryDirectory(prefix="oneforecast_era5_")157        root = Path(self._layout_dir.name)158        data_dir = root / "data"159        data_dir.mkdir()160        for year in self.years:161            source_path = self._year_path(year)162            target_path = data_dir / f"{year}.h5"163            if self.source_grid == SOURCE_GRID:164                import h5py165 166                with h5py.File(source_path, "r") as source_handle:167                    source_fields = source_handle["fields"]168                    layout = h5py.VirtualLayout(169                        shape=(source_fields.shape[0], source_fields.shape[1], *ONEFORECAST_FILE_GRID),170                        dtype=source_fields.dtype,171                    )172                    virtual_source = h5py.VirtualSource(str(source_path), "fields", shape=source_fields.shape)173                    layout[:] = virtual_source[:, :, ::SPATIAL_STRIDE, ::SPATIAL_STRIDE]174                    with h5py.File(target_path, "w", libver="latest") as target_handle:175                        fields = target_handle.create_virtual_dataset("fields", layout)176                        for name, value in source_fields.attrs.items():177                            fields.attrs[name] = value178            else:179                target_path.symlink_to(source_path)180        if self._external_stats is not None:181            stats_dir = root / "stats"182            stats_dir.mkdir()183            (stats_dir / "global_means.npy").symlink_to(self._external_stats[0])184            (stats_dir / "global_stds.npy").symlink_to(self._external_stats[1])185 186        return root187 188    def get_dataloader(self, mode: str):189        """Delegate loading to OneScience, then align native ERA5 to OneForecast's grid."""190        try:191            from onescience.datapipes.climate.era5 import ERA5Datapipe192        except ImportError as exc:193            raise RuntimeError("OneScience ERA5Datapipe is required for data loading") from exc194        datapipe = ERA5Datapipe(195            dataset_dir=str(self._onescience_dataset_dir()), used_years=self.years,196            used_variables=self.source_variables, distributed=self.distributed,197            input_steps=self.input_steps, output_steps=self.output_steps,198            normalize=self.normalize, batch_size=self.batch_size, num_workers=self.num_workers,199        )200        loader, sampler = datapipe.get_dataloader(mode=mode)201        return _SpatiallyAdaptedLoader(loader, self.source_grid), sampler202 203    def inspect(self) -> dict[str, Any]:204        try:205            import h5py206        except ImportError as exc:207            raise RuntimeError("h5py is required to inspect ERA5 HDF5 files") from exc208        path = self._year_path(self.years[0])209        with h5py.File(path, "r") as handle:210            fields = handle["fields"]211            variables = _decode_variables(fields.attrs["variables"])212            indices = [variables.index(name) for name in self.source_variables]213            return {"path": str(path), "fields_shape": list(fields.shape),214                     "source_grid": list(fields.shape[2:]),215                     "oneforecast_file_grid": list(ONEFORECAST_FILE_GRID),216                     "oneforecast_model_grid": [120, 240],217                     "spatial_transform": "identity" if tuple(fields.shape[2:]) == ONEFORECAST_FILE_GRID else "stride_6",218                     "time_step_hours": int(fields.attrs["time_step"]),219                     "variable_count": len(variables), "official_channel_indices": indices,220                     "source_variables": self.source_variables,221                     "statistics_shape": list(self.global_means.shape),222                     "statistics_shared_across_years": True,223                     "official_variables_match": len(indices) == len(OFFICIAL_VARIABLES)}224 225    def selected_statistics(self) -> tuple[np.ndarray, np.ndarray]:226        """Return normalization statistics in the model's 69-channel order."""227        if self.global_means is None or self.global_stds is None:228            raise RuntimeError("ERA5 statistics have not been validated")229        return self.global_means[:, self.channel_indices], self.global_stds[:, self.channel_indices]230 231 232def _adapt_spatial(value: Any, source_grid: tuple[int, int] | None) -> Any:233    if not hasattr(value, "shape") or len(value.shape) < 2:234        return value235    if tuple(value.shape[-2:]) == ONEFORECAST_FILE_GRID:236        return value237    if tuple(value.shape[-2:]) != SOURCE_GRID or source_grid != SOURCE_GRID:238        return value239    return value[..., ::SPATIAL_STRIDE, ::SPATIAL_STRIDE]240 241 242class _SpatiallyAdaptedLoader:243    """Preserve the DataLoader interface while adapting fields after ERA5Datapipe."""244 245    def __init__(self, loader: Any, source_grid: tuple[int, int] | None) -> None:246        self.loader = loader247        self.source_grid = source_grid248 249    def __len__(self) -> int:250        return len(self.loader)251 252    def __iter__(self):253        for batch in self.loader:254            yield tuple(_adapt_spatial(value, self.source_grid) for value in batch)255