File size: 866 Bytes
989c6ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | """Deterministic fake cubed-sphere fields; not a substitute for ERA5."""
import torch
def make_fake_batch(batch=2, channels=2, faces=6, height=8, width=8, seed=7):
if faces != 6:
raise ValueError("cubed sphere requires six faces")
g = torch.Generator().manual_seed(seed)
yy, xx = torch.meshgrid(torch.linspace(-1, 1, height), torch.linspace(-1, 1, width), indexing="ij")
continuous = (xx + 0.5 * yy).expand(batch, 1, faces, height, width)
face_id = torch.arange(faces).view(1, 1, faces, 1, 1).float().expand(batch, 1, faces, height, width)
bank = torch.cat((continuous, face_id, torch.ones_like(continuous)), 1)
if channels > bank.shape[1]:
bank = torch.cat((bank, torch.zeros(batch, channels - bank.shape[1], faces, height, width)), 1)
x = bank[:, :channels]
return x + 0.01 * torch.randn(x.shape, generator=g)
|