File size: 13,779 Bytes
04fe5fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 | """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)
|