| """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] |
|
|
