File size: 1,254 Bytes
1aeffbb | 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 | """Synthetic GLORYS12-shaped data for smoke tests and local development."""
import glob
import torch
from torch.utils.data import Dataset
try:
import h5py
except ImportError:
h5py = None
class SyntheticOceanDataset(Dataset):
def __init__(self, samples, channels, grid, input_steps=2, output_steps=1, seed=42, data_dir=None):
generator = torch.Generator().manual_seed(seed)
files = sorted(glob.glob(f"{data_dir}/data/*.h5")) if data_dir else []
if files and h5py is not None:
with h5py.File(files[0], "r") as handle:
fields = torch.from_numpy(handle["fields"][:]).float()
total = min(samples, fields.shape[0] - input_steps - output_steps + 1)
self.x = torch.stack([fields[i:i + input_steps] for i in range(total)])
self.y = torch.stack([fields[i + input_steps:i + input_steps + output_steps] for i in range(total)])
else:
self.x = torch.randn(samples, input_steps, channels, *grid, generator=generator)
self.y = torch.randn(samples, output_steps, channels, *grid, generator=generator)
def __len__(self):
return self.x.shape[0]
def __getitem__(self, index):
return self.x[index], self.y[index]
|