| """Generate normalized precipitation sequences at the paper's 288x288 resolution.""" |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def make_split(path, count, config, seed): |
| rng = np.random.default_rng(seed) |
| data = config["data"] |
| inputs, outputs = int(data["input_frames"]), int(data["output_frames"]) |
| height, width = int(data["height"]), int(data["width"]) |
| y, x = np.mgrid[-1:1:complex(height), -1:1:complex(width)].astype(np.float32) |
| sequences = np.empty((count, inputs + outputs, height, width), np.float32) |
| for sample in range(count): |
| centers = rng.uniform(-0.45, 0.45, (2, 2)) |
| velocity = rng.uniform(-0.025, 0.025, (2, 2)) |
| amplitudes = rng.uniform(0.35, 1.0, 2) |
| widths = rng.uniform(0.10, 0.28, 2) |
| for step in range(inputs + outputs): |
| field = np.zeros((height, width), np.float32) |
| for storm in range(2): |
| cy, cx = centers[storm] + velocity[storm] * step |
| radius = ((x - cx) ** 2 + (y - cy) ** 2) / (2 * widths[storm] ** 2) |
| field += amplitudes[storm] * np.exp(-radius) |
| field += rng.normal(0, 0.01, field.shape) |
| sequences[sample, step] = np.clip(field, 0, 1) |
| np.savez_compressed( |
| path, |
| format_version=np.asarray(data["format_version"]), |
| data_source=np.asarray("synthetic_knmi_like"), |
| inputs=sequences[:, :inputs], |
| targets=sequences[:, inputs:], |
| ) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| output = ROOT / config["data"]["root"] |
| output.mkdir(parents=True, exist_ok=True) |
| for offset, (filename, count) in enumerate(( |
| ("train.npz", int(config["data"]["train_samples"])), |
| ("test.npz", int(config["data"]["test_samples"])), |
| )): |
| target = output / filename |
| if not target.exists(): |
| make_split(target, count, config, int(config["seed"]) + offset) |
| print(f"generated={target.relative_to(ROOT)} samples={count} shape=({config['data']['height']},{config['data']['width']})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
