ConvLSTM / scripts /fake_data.py
zhangrenchao's picture
Add engineering reproduction package
9be39c5 verified
Raw
History Blame Contribute Delete
2.17 kB
"""Generate 5-to-15 synthetic radar sequences at the paper's 100x100 size."""
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"]
total = 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, total, 1, height, width), np.float32)
for sample in range(count):
centers = rng.uniform(-0.55, 0.55, (3, 2))
velocities = rng.uniform(-0.035, 0.035, (3, 2))
amplitudes = rng.uniform(0.25, 0.95, 3)
scales = rng.uniform(0.10, 0.28, 3)
for step in range(total):
field = np.zeros((height, width), np.float32)
for storm in range(3):
cy, cx = centers[storm] + velocities[storm] * step
distance = ((x - cx) ** 2 + (y - cy) ** 2) / (2 * scales[storm] ** 2)
field += amplitudes[storm] * np.exp(-distance)
sequences[sample, step, 0] = np.clip(field + rng.normal(0, 0.01, field.shape), 0, 1)
split = int(data["input_frames"])
np.savez_compressed(path, format_version=np.asarray(data["format_version"]),
data_source=np.asarray("synthetic_hko_radar_like"),
inputs=sequences[:, :split], targets=sequences[:, split:])
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", config["data"]["train_samples"]),
("test.npz", config["data"]["test_samples"]))):
target = output / filename
if not target.exists():
make_split(target, int(count), config, int(config["seed"]) + offset)
print(f"generated={target.relative_to(ROOT)} input=5x1x100x100 target=15x1x100x100")
if __name__ == "__main__":
main()