File size: 5,442 Bytes
80cf062 | 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 | from __future__ import annotations
from pathlib import Path
from typing import Any, Sequence
import h5py
EXPECTED_CHANNELS = 20
SOURCE_SIZE = (721, 1440)
MODEL_SIZE = (720, 1440)
def _decode_variables(values: Sequence[Any]) -> list[str]:
return [value.decode() if isinstance(value, bytes) else str(value) for value in values]
def validate_variables(variables: Sequence[str], expected_channels: int = EXPECTED_CHANNELS) -> list[str]:
names = list(variables)
if len(names) != expected_channels:
raise ValueError(
f"W-MAE requires exactly {expected_channels} explicitly ordered variables; "
f"received {len(names)}. The official repository does not publish a complete "
"channel-name mapping, so no default mapping is assumed."
)
if len(set(names)) != len(names):
raise ValueError("W-MAE variable names must be unique and explicitly ordered.")
return names
def inspect_era5_contract(
dataset_dir: str | Path,
years: Sequence[int],
variables: Sequence[str],
expected_channels: int = EXPECTED_CHANNELS,
) -> dict[str, Any]:
"""Validate the HDF5 metadata before importing the torch-based datapipe."""
dataset_dir = Path(dataset_dir)
names = validate_variables(variables, expected_channels)
missing_years = [year for year in years if not (dataset_dir / "data" / f"{year}.h5").is_file()]
if missing_years:
raise ValueError(f"ERA5 year files are missing: {missing_years}")
first_file = dataset_dir / "data" / f"{years[0]}.h5"
with h5py.File(first_file, "r") as handle:
if "fields" not in handle:
raise ValueError(f"{first_file} does not contain a 'fields' dataset.")
fields = handle["fields"]
if fields.ndim != 4:
raise ValueError(f"fields must have shape [T,C,H,W], got {fields.shape}.")
if tuple(fields.shape[-2:]) != SOURCE_SIZE:
raise ValueError(f"W-MAE expects source grid {SOURCE_SIZE}, got {fields.shape[-2:]}.")
if "variables" not in fields.attrs or "time_step" not in fields.attrs:
raise ValueError("fields must define 'variables' and 'time_step' attributes.")
available = _decode_variables(fields.attrs["variables"])
fields_shape = tuple(fields.shape)
time_step = int(fields.attrs["time_step"])
missing_variables = [name for name in names if name not in available]
if missing_variables:
raise ValueError(f"Configured variables are absent from HDF5 metadata: {missing_variables}")
if "global_means" not in handle or "global_stds" not in handle:
stats_dir = dataset_dir / "stats"
if not (stats_dir / "global_means.npy").is_file() or not (stats_dir / "global_stds.npy").is_file():
raise ValueError("Normalization statistics are missing from HDF5 and dataset_dir/stats.")
return {
"file": str(first_file),
"fields_shape": fields_shape,
"time_step": time_step,
"selected_variables": names,
"channel_indices": [available.index(name) for name in names],
"crop": "fields[..., :720, :]",
"model_size": MODEL_SIZE,
}
def _crop_last_latitude(value: Any) -> Any:
if tuple(value.shape[-2:]) == MODEL_SIZE:
return value
if tuple(value.shape[-2:]) != SOURCE_SIZE:
raise ValueError(f"Expected trailing spatial shape {SOURCE_SIZE}, got {tuple(value.shape[-2:])}.")
return value[..., : MODEL_SIZE[0], :]
class WMAEERA5Dataset:
"""Thin W-MAE adapter around OneScience ERA5Dataset.
The official W-MAE loader removes the final latitude row from a
721x1440 ERA5 field. This wrapper preserves OneScience loading and
normalization while applying that exact spatial convention.
"""
def __init__(
self,
dataset_dir: str | Path,
years: Sequence[int],
variables: Sequence[str],
task: str = "forecast",
input_steps: int = 1,
output_steps: int = 1,
normalize: bool = True,
) -> None:
if task not in {"pretrain", "forecast"}:
raise ValueError("task must be either 'pretrain' or 'forecast'.")
names = validate_variables(variables)
inspect_era5_contract(dataset_dir, years, names)
try:
from onescience.datapipes.climate.era5 import ERA5Dataset
except (ImportError, OSError) as error:
raise RuntimeError(
"OneScience ERA5Dataset could not be imported. Verify the active "
"OneScience/PyTorch runtime before constructing WMAEERA5Dataset."
) from error
self.task = task
self.dataset = ERA5Dataset(
dataset_dir=str(dataset_dir),
used_years=list(years),
used_variables=names,
input_steps=input_steps,
output_steps=output_steps,
normalize=normalize,
)
def __len__(self) -> int:
return len(self.dataset)
def __getitem__(self, index: int) -> tuple[Any, Any, Any, int, list[str]]:
invar, outvar, cos_zenith, step_idx, time_index = self.dataset[index]
invar = _crop_last_latitude(invar)
outvar = _crop_last_latitude(outvar)
cos_zenith = _crop_last_latitude(cos_zenith)
target = invar if self.task == "pretrain" else outvar
return invar, target, cos_zenith, step_idx, time_index
|