model/fake_data.py

"""Indexed compact MetNet-3 fake samples."""

import torch
from torch.utils.data import Dataset


def make_fake(config, batch_size=2, height=8, width=8, seed=7):
    torch.manual_seed(seed)
    def field(t, c, h=height, w=width): return torch.randn(batch_size, t, c, h, w)
    batch = {"mrms_high": field(config.high_frames, config.high_channels), "mrms_low": field(2, config.low_channels, height // 2, width // 2), "omo": field(config.omo_frames, config.omo_channels), "hrrr": field(1, config.hrrr_channels), "goes": field(1, config.goes_channels), "elevation": field(1, 1, height, width)[:, 0], "coordinates": field(1, 2, height, width)[:, 0], "topography_embedding": field(1, 1, height, width)[:, 0], "current_time": torch.rand(batch_size, 1), "lead_time": torch.rand(batch_size, 1), "omo_input_mask": torch.rand(batch_size, 1, height, width) > .25}
    targets = {"precipitation": torch.randint(config.precipitation_bins, (batch_size, 1, height, width)), "ground": torch.randint(config.ground_bins, (batch_size, config.ground_targets, height, width)), "hrrr": torch.randn(batch_size, config.hrrr_channels, height, width), "precipitation_mask": torch.ones(batch_size, 1, height, width), "ground_mask": torch.ones(batch_size, config.ground_targets, height, width), "hrrr_mask": torch.ones(batch_size, config.hrrr_channels, height, width)}
    return batch, targets


class FakeMetNetDataset(Dataset):
    def __init__(self, config, samples, seed): self.config, self.samples, self.seed = config, samples, seed
    def __len__(self): return self.samples
    def __getitem__(self, index):
        batch, targets = make_fake(self.config, batch_size=1, seed=self.seed + index)
        return {key: value.squeeze(0) for key, value in batch.items()}, {key: value.squeeze(0) for key, value in targets.items()}