File size: 2,183 Bytes
8b64ae3 | 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 | """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()
|