| |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| parser = argparse.ArgumentParser(description="Generate deterministic hydrometeorological smoke-test data") |
| parser.add_argument("--config", default="conf/config.yaml") |
| args = parser.parse_args() |
| with open(args.config, encoding="utf-8") as handle: |
| config = json.load(handle) |
|
|
| rng = np.random.default_rng(config["seed"]) |
| data = config["data"] |
| gauges, train_n, val_n = len(data["gauges"]), data["train_samples"], data["validation_samples"] |
| cases, leads, history, variables = data["forecast_cases"], data["forecast_steps"], 28, 23 |
| means = np.array([0.37, 8.24, 20.5, 3.28, 4.47, 49.0, 26.0, 0.14, 1.19, 8.55], dtype=np.float32) |
|
|
|
|
| def forcing(gauge, count, phase_offset=0.0): |
| phase = rng.uniform(0, 2 * np.pi, (count, 1)) + phase_offset |
| t = np.arange(history, dtype=np.float32)[None] / 4.0 |
| x = rng.normal(0, 0.35, (count, history, variables)).astype(np.float32) |
| temperature = 9 + 12 * np.sin(2 * np.pi * (t / 365 + phase / (2 * np.pi))) - gauge * 0.35 |
| precipitation = rng.gamma(1.2, 0.8, (count, history)) * (0.5 + gauge / 18) |
| x[:, :, 2] = temperature |
| x[:, :, 5] = precipitation |
| x[:, :, 6] = np.maximum(0, precipitation * 0.55 + rng.normal(0, 0.1, precipitation.shape)) |
| x[:, :, 10:14] = np.clip(0.45 + np.cumsum(precipitation[:, :, None] * 0.004, axis=1), 0, 1) |
| response = means[gauge] + means[gauge] * (0.08 * precipitation[:, -8:].mean(1) + 0.025 * x[:, -1, 6]) |
| response += rng.normal(0, max(float(means[gauge]) * 0.03, 0.02), count) |
| raw = np.maximum(0, response * (1.12 - gauge * 0.012) + rng.normal(0, max(float(means[gauge]) * 0.05, 0.03), count)) |
| x[:, :, 20], x[:, :, 21], x[:, :, 22] = raw[:, None], (0.92 * raw)[:, None], means[gauge] |
| return x, np.maximum(0, response).astype(np.float32) |
|
|
|
|
| train_x = np.empty((gauges, train_n, history, variables), np.float32) |
| train_y = np.empty((gauges, train_n), np.float32) |
| val_x = np.empty((gauges, val_n, history, variables), np.float32) |
| val_y = np.empty((gauges, val_n), np.float32) |
| forecast_x = np.empty((gauges, cases, leads, history, variables), np.float32) |
| forecast_y = np.empty((gauges, cases, leads), np.float32) |
| persistence = np.empty_like(forecast_y) |
| glofas = np.empty_like(forecast_y) |
| for g in range(gauges): |
| train_x[g], train_y[g] = forcing(g, train_n) |
| val_x[g], val_y[g] = forcing(g, val_n, 0.3) |
| for case in range(cases): |
| base_x, base_y = forcing(g, 1, case / 7) |
| initial = float(base_y[0]) |
| persistence[g, case] = initial |
| state = initial |
| for lead in range(leads): |
| window, target = forcing(g, 1, case / 7 + lead / 80) |
| |
| window[:, :, 20] = initial |
| window[:, :, 21] *= 1.0 + 0.003 * lead |
| state = 0.75 * state + 0.25 * float(target[0]) |
| forecast_x[g, case, lead] = window[0] |
| forecast_y[g, case, lead] = state |
| glofas[g, case, lead] = max(0, state * (1.10 - 0.002 * lead) + rng.normal(0, max(float(means[g]) * 0.06, 0.03))) |
|
|
| path = Path(data["path"]) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(path, train_x=train_x, train_y=train_y, val_x=val_x, val_y=val_y, |
| forecast_x=forecast_x, forecast_y=forecast_y, persistence=persistence, |
| glofas=glofas, gauges=np.array(data["gauges"]), lead_hours=np.arange(1, leads + 1) * 6) |
| print(f"wrote {path}: train={train_x.shape}, forecast={forecast_x.shape}") |
|
|
