| """Training pairs and evaluation inputs of a dataset package (``trajectories_fine.npz``, ``split.json``, |
| ``normalization.npz``). |
| |
| * ``get_dataloaders(cfg, seed)`` -> ``(train_loader, val_loader)``: segment pairs ``(x_tk, x_tk1, t_k, t_k1)`` for |
| Cascaded SI (``method.data_mode: fine_segment_pairs``) or endpoint pairs ``(x_0, x_T)`` for Direct SI and EDM |
| (``fine_endpoint_pairs``), from the pair manifest named by ``method.fine_slice_manifest``. |
| * ``get_split_set(cfg, split, seed, x0_bank)`` -> ``(x0_per_gid, truth_clouds, gids)`` at the forecast horizon |
| ``method.fine_horizon_idx``, the inputs of sampling and of validation. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Tuple |
|
|
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from .fine import (EndpointPairDataset, SegmentPairDataset, ValEndpointPairDataset, |
| ValSegmentPairDataset, get_fine_split_set) |
|
|
| DATA_MODES = ("fine_segment_pairs", "fine_endpoint_pairs") |
|
|
|
|
| def get_dataloaders(cfg: dict, seed: int = 0) -> Tuple[DataLoader, DataLoader]: |
| packaged_dir = cfg["data"]["packaged_dataset_path"] |
| batch_size = int(cfg["training"]["batch_size"]) |
| method_cfg = cfg.get("method", {}) |
| data_mode = method_cfg.get("data_mode") |
| num_workers = int(cfg["training"].get("num_workers", 0)) |
| if data_mode == "fine_segment_pairs": |
| train_ds = SegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed, |
| expected_t_values=method_cfg.get("sampling", {}).get("t_values")) |
| val_ds = ValSegmentPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], |
| n_members=int(method_cfg.get("fine_val_members", 20))) |
| elif data_mode == "fine_endpoint_pairs": |
| train_ds = EndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], seed=seed, |
| expected_t_values=method_cfg.get("sampling", {}).get("t_values")) |
| val_ds = ValEndpointPairDataset(packaged_dir, method_cfg["fine_slice_manifest"], |
| n_members=int(method_cfg.get("fine_val_members", 20))) |
| else: |
| raise ValueError(f"method.data_mode must be one of {DATA_MODES}, got {data_mode!r}") |
| g = torch.Generator() |
| g.manual_seed(seed) |
| train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, generator=g, |
| pin_memory=num_workers > 0) |
| val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, |
| pin_memory=num_workers > 0) |
| return train_loader, val_loader |
|
|
|
|
| def get_split_set(cfg: dict, split_name: str, seed: int = 0, x0_bank: str | None = None): |
| return get_fine_split_set(cfg, split_name, seed=seed, x0_bank=x0_bank) |
|
|
