File size: 3,962 Bytes
30f7852 | 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 | """Generate grouped synthetic rain days at the real 39x56x56 input shape."""
from pathlib import Path
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
def smooth(field: np.ndarray, rounds: int = 2) -> np.ndarray:
for _ in range(rounds):
field = sum(np.roll(np.roll(field, y, -2), x, -1) for y, x in ((0, 0), (1, 0), (-1, 0), (0, 1), (0, -1))) / 5
return field
def make_day(rng: np.random.Generator, steps: int, height: int, width: int):
yy, xx = np.mgrid[-1:1:complex(height), -1:1:complex(width)].astype(np.float32)
terrain = np.clip(0.55 + 0.25 * yy + 0.16 * np.sin(2.5 * xx) - 0.12 * np.cos(3 * yy), 0, 1)
center = rng.uniform(-0.45, 0.45, 2)
velocity = rng.uniform(-0.10, 0.10, 2)
width0 = rng.uniform(0.13, 0.27)
amplitude = rng.uniform(12, 28)
inputs = np.empty((steps, 39, height, width), np.float32)
targets = np.empty((steps, height, width), np.float32)
for step in range(steps):
cy, cx = center + velocity * step
rain = amplitude * np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2 * width0**2))
rain += 0.35 * amplitude * np.exp(-((xx + cx * 0.5) ** 2 + (yy - cy * 0.6) ** 2) / (3 * width0**2))
rain = np.clip(rain * (0.8 + 0.3 * terrain) + rng.normal(0, 0.12, rain.shape), 0, None)
forecast = np.clip(np.roll(rain, shift=(1, -1), axis=(0, 1)) * (1.12 + 0.08 * terrain) - 0.25, 0, None)
channels = np.empty((39, height, width), np.float32)
for variable in range(8):
for level in range(4):
index = variable * 4 + level
noise = smooth(rng.normal(0, 0.05, rain.shape).astype(np.float32))
if variable == 1: # humidity follows rain and terrain
channels[index] = 0.45 + 0.02 * rain + 0.12 * terrain - 0.025 * level + noise
elif variable == 4: # vertical velocity is strongest near moving rain
channels[index] = -0.08 * rain / (level + 1) + 0.04 * (xx * velocity[1] + yy * velocity[0]) + noise
else:
channels[index] = (variable + 1) * 0.12 + level * 0.04 + 0.08 * xx - 0.05 * yy + noise
for index in range(32, 37):
channels[index] = 0.1 * (index - 31) + 0.03 * rain + smooth(rng.normal(0, 0.04, rain.shape).astype(np.float32))
channels[37] = forecast
channels[38] = terrain
inputs[step], targets[step] = channels, rain
return inputs, targets, terrain
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
data = config["data"]
expected = (8, 39, 56, 56)
actual = (int(data["time_steps"]), int(data["channels"]), int(data["height"]), int(data["width"]))
if actual != expected or len(data["lead_hours"]) != expected[0]:
raise ValueError(f"fixed data contract is {expected} with eight lead hours, got {actual}")
counts = data["rain_days"]
output = ROOT / data["root"]
output.mkdir(parents=True, exist_ok=True)
seed = int(config["seed"])
group_id = 0
for split_index, split in enumerate(("train", "val", "test")):
arrays, targets, terrains, groups = [], [], [], []
for day in range(int(counts[split])):
x, y, terrain = make_day(np.random.default_rng(seed + group_id), int(data["time_steps"]), int(data["height"]), int(data["width"]))
arrays.append(x); targets.append(y); terrains.append(terrain); groups.append(group_id)
group_id += 1
np.savez_compressed(
output / f"{split}.npz", inputs=np.stack(arrays), targets=np.stack(targets),
terrain=np.stack(terrains), group_id=np.asarray(groups), split=np.asarray(split),
format_version=np.asarray(data["format_version"]), lead_hours=np.asarray(data["lead_hours"]),
)
print(f"split={split} rain_days={counts[split]} inputs={np.stack(arrays).shape}")
if __name__ == "__main__":
main()
|