| """Pair and split the deterministic spherical sequence for SFNO training.""" | |
| from torch.utils.data import Dataset | |
| class SphericalPairDataset(Dataset): | |
| def __init__(self, fields, start, stop): | |
| self.fields = fields | |
| self.start = start | |
| self.stop = min(stop, fields.shape[0] - 1) | |
| if self.stop <= self.start: | |
| raise ValueError("A spherical split must contain at least one input/target pair") | |
| def __len__(self): | |
| return self.stop - self.start | |
| def __getitem__(self, index): | |
| step = self.start + index | |
| return self.fields[step], self.fields[step + 1] | |