| """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: |
| channels[index] = 0.45 + 0.02 * rain + 0.12 * terrain - 0.025 * level + noise |
| elif variable == 4: |
| 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() |
|
|
