File size: 3,582 Bytes
186a48a | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 | #!/usr/bin/env python3
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)
# Forecast meteorology evolves, while observed flow is frozen at issue time.
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}")
|