File size: 5,747 Bytes
1616901 | 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 | """ERA5-style HDF5 loader used by training and inference."""
from __future__ import annotations
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Sequence, Tuple
import h5py
import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
def _decode(value: object) -> str:
return value.decode() if isinstance(value, bytes) else str(value)
def resolve_data_dir(path: str | Path, project_root: Path | None = None) -> Path:
root = project_root or Path(__file__).resolve().parents[1]
candidate = Path(path).expanduser()
return candidate if candidate.is_absolute() else (root / candidate).resolve()
def read_metadata(data_dir: str | Path, channels: Sequence[str]) -> Dict[str, np.ndarray | int | List[str]]:
data_dir = Path(data_dir)
files = sorted((data_dir / "data").glob("*.h5"))
if not files:
raise FileNotFoundError(f"No yearly HDF5 files found under {data_dir / 'data'}")
with h5py.File(files[0], "r") as source:
fields = source["fields"]
variables = [_decode(item) for item in fields.attrs["variables"]]
time_step = int(fields.attrs.get("time_step", 6))
means = np.asarray(source["global_means"][:], dtype=np.float32)
stds = np.asarray(source["global_stds"][:], dtype=np.float32)
shape = tuple(int(item) for item in fields.shape)
missing = [name for name in channels if name not in variables]
if missing:
raise ValueError(f"Variables missing from synthetic/ERA5 data: {missing}")
indices = np.asarray([variables.index(name) for name in channels], dtype=np.int64)
return {
"variables": variables,
"indices": indices,
"time_step": time_step,
"means": means[:, indices, :, :],
"stds": np.maximum(stds[:, indices, :, :], 1.0e-6),
"shape": shape,
}
class ERA5WindowDataset(Dataset):
def __init__(
self,
data_dir: str | Path,
years: Sequence[int],
channels: Sequence[str],
input_steps: int = 2,
rollout_steps: int = 1,
normalize: bool = True,
) -> None:
self.data_dir = Path(data_dir)
self.years = [int(year) for year in years]
self.channels = list(channels)
self.input_steps = max(1, int(input_steps))
self.rollout_steps = max(1, int(rollout_steps))
self.normalize = bool(normalize)
metadata = read_metadata(self.data_dir, self.channels)
self.channel_indices = metadata["indices"]
self.time_step = int(metadata["time_step"])
self.means = torch.from_numpy(metadata["means"])
self.stds = torch.from_numpy(metadata["stds"])
self.shape = metadata["shape"]
self.samples: List[Tuple[int, int]] = []
for year in self.years:
path = self.data_dir / "data" / f"{year}.h5"
if not path.exists():
raise FileNotFoundError(f"Missing year file: {path}")
with h5py.File(path, "r") as source:
timesteps = int(source["fields"].shape[0])
count = timesteps - self.input_steps - self.rollout_steps + 1
if count <= 0:
raise ValueError(
f"Year {year} has {timesteps} steps, but input={self.input_steps} "
f"and rollout={self.rollout_steps} require at least {self.input_steps + self.rollout_steps}"
)
self.samples.extend((year, index) for index in range(count))
def __len__(self) -> int:
return len(self.samples)
def _timestamp(self, year: int, index: int) -> str:
value = datetime(year, 1, 1) + timedelta(hours=index * self.time_step)
return value.strftime("%Y%m%d%H")
def __getitem__(self, item: int):
year, start = self.samples[item]
path = self.data_dir / "data" / f"{year}.h5"
with h5py.File(path, "r") as source:
fields = source["fields"]
input_data = np.asarray(
fields[start : start + self.input_steps, self.channel_indices, :, :], dtype=np.float32
)
target_data = np.asarray(
fields[
start + self.input_steps : start + self.input_steps + self.rollout_steps,
self.channel_indices,
:,
:,
],
dtype=np.float32,
)
input_tensor = torch.from_numpy(input_data)
target_tensor = torch.from_numpy(target_data)
if self.normalize:
input_tensor = (input_tensor - self.means) / self.stds
target_tensor = (target_tensor - self.means) / self.stds
timestamp = self._timestamp(year, start + self.input_steps)
return input_tensor, target_tensor, timestamp
def make_dataloader(
data_dir: str | Path,
years: Sequence[int],
channels: Sequence[str],
input_steps: int,
rollout_steps: int,
batch_size: int,
num_workers: int = 0,
distributed: bool = False,
train: bool = False,
pin_memory: bool = False,
):
dataset = ERA5WindowDataset(
data_dir=data_dir,
years=years,
channels=channels,
input_steps=input_steps,
rollout_steps=rollout_steps,
)
sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
loader = DataLoader(
dataset,
batch_size=max(1, int(batch_size)),
shuffle=train and sampler is None,
sampler=sampler,
num_workers=max(0, int(num_workers)),
pin_memory=bool(pin_memory),
drop_last=False,
)
return loader, sampler
|