File size: 7,894 Bytes
5c365c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Array-backed ACE dataset and NPZ loader."""
from __future__ import annotations
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset
from ACE.model.variables import (
DIAGNOSTIC_CHANNELS,
FORCING_CHANNELS,
INPUT_CHANNELS,
OUTPUT_CHANNELS,
PROGNOSTIC_CHANNELS,
INPUT_UNITS,
OUTPUT_UNITS,
)
def gaussian_latitudes(height: int) -> np.ndarray:
"""Return south-to-north Gaussian latitudes in radians."""
nodes, _ = np.polynomial.legendre.leggauss(height)
return np.arcsin(nodes).astype(np.float32)
class ArrayPairDataset(Dataset):
def __init__(self, inputs: np.ndarray, targets: np.ndarray) -> None:
self.inputs = np.asarray(inputs, dtype=np.float32)
self.targets = np.asarray(targets, dtype=np.float32)
if self.inputs.ndim != 4 or self.targets.ndim != 4:
raise ValueError("inputs and targets must be [N,C,H,W]")
if self.inputs.shape[0] != self.targets.shape[0]:
raise ValueError("inputs and targets must have equal sample counts")
if self.inputs.shape[1] != len(INPUT_CHANNELS):
raise ValueError(f"inputs require {len(INPUT_CHANNELS)} channels")
if self.targets.shape[1] != len(OUTPUT_CHANNELS):
raise ValueError(f"targets require {len(OUTPUT_CHANNELS)} channels")
if self.inputs.shape[-2:] != self.targets.shape[-2:]:
raise ValueError("inputs and targets must share spatial shape")
def __len__(self) -> int:
return self.inputs.shape[0]
def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
return torch.from_numpy(self.inputs[index]), torch.from_numpy(self.targets[index])
def load_npz(path: str | Path) -> ArrayPairDataset:
data = np.load(path)
if "inputs" not in data or "targets" not in data:
raise KeyError("NPZ must contain inputs and targets arrays")
return ArrayPairDataset(data["inputs"], data["targets"])
def make_fake_pairs(
num_samples: int = 8,
height: int = 180,
width: int = 360,
seed: int = 0,
) -> tuple[np.ndarray, np.ndarray]:
"""Create ACE-shaped smooth global fields with the paper's variable semantics.
The official training samples are 6-hour pairs on a 1-degree Gaussian grid:
``inputs[N,40,180,360]`` and ``targets[N,44,180,360]``. This generator keeps
that contract while using analytic spherical modes instead of white noise.
``height`` and ``width`` remain configurable for small CPU smoke tests.
"""
if min(num_samples, height, width) <= 0:
raise ValueError("num_samples, height and width must be positive")
rng = np.random.default_rng(seed)
lat = gaussian_latitudes(height)
lon = np.arange(width, dtype=np.float32) * (2 * np.pi / width)
latitude, longitude = np.meshgrid(lat, lon, indexing="ij")
cos_lat = np.cos(latitude)
sin_lat = np.sin(latitude)
topography = 1200.0 * (
0.55 * np.sin(2.0 * latitude) ** 2 * np.cos(2.0 * longitude)
+ 0.25 * np.cos(3.0 * latitude + longitude)
)
land_fraction = np.clip(
0.5 + 0.35 * np.sin(2.0 * latitude) * np.cos(longitude)
+ 0.15 * np.cos(3.0 * longitude),
0.0,
1.0,
)
sea_ice_fraction = np.clip((np.abs(latitude) - np.deg2rad(55.0)) / np.deg2rad(30.0), 0.0, 0.9)
ocean_fraction = np.clip(1.0 - land_fraction - sea_ice_fraction, 0.0, 1.0)
inputs = np.empty((num_samples, len(INPUT_CHANNELS), height, width), dtype=np.float32)
targets = np.empty((num_samples, len(OUTPUT_CHANNELS), height, width), dtype=np.float32)
for sample in range(num_samples):
phase = float(rng.uniform(0.0, 2.0 * np.pi))
seasonal = float(rng.normal(0.0, 0.08))
wave = (
cos_lat * np.cos(longitude + phase)
+ 0.35 * np.sin(2.0 * latitude - 0.13 * sample) * np.sin(2.0 * longitude - phase)
+ 0.15 * np.cos(3.0 * latitude + 0.07 * sample)
).astype(np.float32)
wave_next = (
cos_lat * np.cos(longitude + phase + 0.08)
+ 0.35 * np.sin(2.0 * latitude - 0.13 * sample - 0.03) * np.sin(2.0 * longitude - phase + 0.05)
+ 0.15 * np.cos(3.0 * latitude + 0.07 * sample + 0.02)
).astype(np.float32)
prognostic = np.empty((len(PROGNOSTIC_CHANNELS), height, width), dtype=np.float32)
prognostic_next = np.empty_like(prognostic)
for level in range(8):
height_factor = 1.0 - 0.06 * level
humidity_factor = np.exp(-0.28 * level)
offset = 4 * level
prognostic[offset] = 215.0 + 48.0 * height_factor + 7.0 * wave
prognostic[offset + 1] = np.maximum(1.0e-6, 0.00015 + 0.012 * humidity_factor * (1.0 + 0.25 * wave))
prognostic[offset + 2] = 18.0 * np.sin(2.0 * latitude) * np.cos(longitude + phase) * height_factor
prognostic[offset + 3] = 12.0 * np.sin(longitude - phase) * cos_lat * height_factor
prognostic_next[offset] = prognostic[offset] + 0.35 * (wave_next - wave)
prognostic_next[offset + 1] = np.maximum(1.0e-6, prognostic[offset + 1] * (1.0 + 0.015 * (wave_next - wave)))
prognostic_next[offset + 2] = prognostic[offset + 2] + 0.08 * (wave_next - wave)
prognostic_next[offset + 3] = prognostic[offset + 3] - 0.05 * (wave_next - wave)
prognostic[32] = 276.0 + 11.0 * wave + 2.0 * land_fraction
prognostic[33] = 101325.0 + 1800.0 * wave - 0.18 * topography
prognostic_next[32] = prognostic[32] + 0.4 * (wave_next - wave)
prognostic_next[33] = prognostic[33] + 30.0 * (wave_next - wave)
forcing = np.stack(
[
340.0 * np.maximum(cos_lat, 0.0) * (1.0 + seasonal),
288.0 + 7.0 * wave + 0.5 * seasonal,
topography,
land_fraction,
ocean_fraction,
sea_ice_fraction,
]
).astype(np.float32)
diagnostics = np.stack(
[
70.0 * np.maximum(cos_lat, 0.0) * (1.0 + 0.05 * wave_next),
235.0 + 10.0 * wave_next,
35.0 * np.maximum(cos_lat, 0.0) * (1.0 - land_fraction),
310.0 + 8.0 * wave_next,
250.0 * np.maximum(cos_lat, 0.0),
280.0 + 5.0 * wave_next,
np.maximum(0.0, 1.5e-5 * (1.0 + wave_next) * (1.0 - 0.4 * sea_ice_fraction)),
2.0e-6 * (wave_next - wave),
70.0 * (1.0 - sea_ice_fraction) * (1.0 + 0.1 * wave_next),
12.0 * land_fraction * (1.0 + 0.1 * wave_next),
]
).astype(np.float32)
inputs[sample] = np.concatenate([prognostic, forcing], axis=0)
targets[sample] = np.concatenate([prognostic_next, diagnostics], axis=0)
return inputs, targets
def save_fake_pairs(path: str | Path, **kwargs: int) -> Path:
"""Save :func:`make_fake_pairs` with grid and variable metadata."""
output = Path(path)
output.parent.mkdir(parents=True, exist_ok=True)
inputs, targets = make_fake_pairs(**kwargs)
height, width = inputs.shape[-2:]
np.savez(
output,
inputs=inputs,
targets=targets,
lat=np.rad2deg(gaussian_latitudes(height)),
lon=np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32),
input_channels=np.asarray(INPUT_CHANNELS),
output_channels=np.asarray(OUTPUT_CHANNELS),
input_units=np.asarray(INPUT_UNITS),
output_units=np.asarray(OUTPUT_UNITS),
prognostic_channels=np.asarray(PROGNOSTIC_CHANNELS),
forcing_channels=np.asarray(FORCING_CHANNELS),
diagnostic_channels=np.asarray(DIAGNOSTIC_CHANNELS),
time_step_hours=np.asarray(6, dtype=np.int32),
source=np.asarray("synthetic FV3GFS-compatible analytic fixture"),
)
return output
|