OneScience-Group/Streamflow-LSTM

1#!/usr/bin/env python32import argparse3import json4from pathlib import Path5 6import numpy as np7 8parser = argparse.ArgumentParser(description="Generate deterministic hydrometeorological smoke-test data")9parser.add_argument("--config", default="conf/config.yaml")10args = parser.parse_args()11with open(args.config, encoding="utf-8") as handle:12    config = json.load(handle)13 14rng = np.random.default_rng(config["seed"])15data = config["data"]16gauges, train_n, val_n = len(data["gauges"]), data["train_samples"], data["validation_samples"]17cases, leads, history, variables = data["forecast_cases"], data["forecast_steps"], 28, 2318means = 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)19 20 21def forcing(gauge, count, phase_offset=0.0):22    phase = rng.uniform(0, 2 * np.pi, (count, 1)) + phase_offset23    t = np.arange(history, dtype=np.float32)[None] / 4.024    x = rng.normal(0, 0.35, (count, history, variables)).astype(np.float32)25    temperature = 9 + 12 * np.sin(2 * np.pi * (t / 365 + phase / (2 * np.pi))) - gauge * 0.3526    precipitation = rng.gamma(1.2, 0.8, (count, history)) * (0.5 + gauge / 18)27    x[:, :, 2] = temperature28    x[:, :, 5] = precipitation29    x[:, :, 6] = np.maximum(0, precipitation * 0.55 + rng.normal(0, 0.1, precipitation.shape))30    x[:, :, 10:14] = np.clip(0.45 + np.cumsum(precipitation[:, :, None] * 0.004, axis=1), 0, 1)31    response = means[gauge] + means[gauge] * (0.08 * precipitation[:, -8:].mean(1) + 0.025 * x[:, -1, 6])32    response += rng.normal(0, max(float(means[gauge]) * 0.03, 0.02), count)33    raw = np.maximum(0, response * (1.12 - gauge * 0.012) + rng.normal(0, max(float(means[gauge]) * 0.05, 0.03), count))34    x[:, :, 20], x[:, :, 21], x[:, :, 22] = raw[:, None], (0.92 * raw)[:, None], means[gauge]35    return x, np.maximum(0, response).astype(np.float32)36 37 38train_x = np.empty((gauges, train_n, history, variables), np.float32)39train_y = np.empty((gauges, train_n), np.float32)40val_x = np.empty((gauges, val_n, history, variables), np.float32)41val_y = np.empty((gauges, val_n), np.float32)42forecast_x = np.empty((gauges, cases, leads, history, variables), np.float32)43forecast_y = np.empty((gauges, cases, leads), np.float32)44persistence = np.empty_like(forecast_y)45glofas = np.empty_like(forecast_y)46for g in range(gauges):47    train_x[g], train_y[g] = forcing(g, train_n)48    val_x[g], val_y[g] = forcing(g, val_n, 0.3)49    for case in range(cases):50        base_x, base_y = forcing(g, 1, case / 7)51        initial = float(base_y[0])52        persistence[g, case] = initial53        state = initial54        for lead in range(leads):55            window, target = forcing(g, 1, case / 7 + lead / 80)56            # Forecast meteorology evolves, while observed flow is frozen at issue time.57            window[:, :, 20] = initial58            window[:, :, 21] *= 1.0 + 0.003 * lead59            state = 0.75 * state + 0.25 * float(target[0])60            forecast_x[g, case, lead] = window[0]61            forecast_y[g, case, lead] = state62            glofas[g, case, lead] = max(0, state * (1.10 - 0.002 * lead) + rng.normal(0, max(float(means[g]) * 0.06, 0.03)))63 64path = Path(data["path"])65path.parent.mkdir(parents=True, exist_ok=True)66np.savez_compressed(path, train_x=train_x, train_y=train_y, val_x=val_x, val_y=val_y,67                    forecast_x=forecast_x, forecast_y=forecast_y, persistence=persistence,68                    glofas=glofas, gauges=np.array(data["gauges"]), lead_hours=np.arange(1, leads + 1) * 6)69print(f"wrote {path}: train={train_x.shape}, forecast={forecast_x.shape}")70