File size: 3,622 Bytes
702a3cf | 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 | from pathlib import Path
import sys
import numpy as np
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.pbl_emulator import INPUT_NAMES, OUTPUT_NAMES, load_yaml
def main() -> None:
config = load_yaml(ROOT / "conf/config.yaml")
count, seed = int(config["data"]["samples"]), int(config["data"]["seed"])
interval_hours = int(config["data"]["interval_hours"])
rng = np.random.default_rng(seed)
hours = (np.arange(count) * interval_hours) % 24
days = (np.arange(count) * interval_hours / 24) % 365
daylight = np.maximum(0.0, np.sin(np.pi * (hours - 6) / 12))
season = np.sin(2 * np.pi * (days - 172) / 365)
swdown = (760 + 110 * season) * daylight
t2 = 284 + 10 * season + 5 * np.sin(2 * np.pi * (hours - 14) / 24) + rng.normal(0, 0.8, count)
q2 = np.clip(0.0065 + 0.004 * (season + 1) / 2 + 0.0012 * daylight + rng.normal(0, 0.0003, count), 0.001, None)
ug = 5.5 + 1.8 * np.sin(2 * np.pi * days / 9) + rng.normal(0, 0.5, count)
vg = 2.0 + 1.5 * np.cos(2 * np.pi * days / 13) + rng.normal(0, 0.5, count)
u10 = 0.62 * ug + rng.normal(0, 0.45, count); v10 = 0.62 * vg + rng.normal(0, 0.45, count)
hfx = 12 + 150 * daylight * (0.75 + 0.25 * season) + rng.normal(0, 5, count)
lh = 25 + 110 * daylight * (0.8 + 0.2 * season) + rng.normal(0, 5, count)
pblh = np.clip(170 + 1200 * daylight * (0.75 + 0.25 * season) + rng.normal(0, 45, count), 80, 1800)
ust = np.clip(0.12 + 0.0014 * hfx + 0.022 * np.hypot(u10, v10), 0.05, 1.2)
tsk = t2 + 1.5 + 4.0 * daylight; tslb = 283 + 7.0 * season
smois = np.clip(0.25 + 0.04 * np.sin(2 * np.pi * days / 40) - 0.025 * season + rng.normal(0, 0.008, count), 0.08, 0.45)
grdflx = 0.12 * hfx - 8 * (1 - daylight); glw = 285 + 2.2 * (t2 - 273.15) + 450 * q2
x = np.column_stack((q2, t2, u10, v10, grdflx, swdown, glw, lh, hfx, pblh, ust, tsk, tslb, smois, ug, vg)).astype(np.float32)
heights = np.array([10, 35, 70, 110, 160, 220, 290, 370, 460, 560, 680, 820, 980, 1160, 1370, 1600, 1880], dtype=np.float32)
z = heights[None, :]; mix = np.exp(-z / pblh[:, None])
u = ug[:, None] + (u10 - ug)[:, None] * mix + 0.18 * np.sin(z / 260)
v = vg[:, None] + (v10 - vg)[:, None] * mix + 0.14 * np.cos(z / 300)
w = 0.20 * daylight[:, None] * np.sin(np.pi * np.minimum(z / pblh[:, None], 1)) - 0.015 * (1 - daylight[:, None]) + rng.normal(0, 0.008, (count, 17))
tk = t2[:, None] - 0.0062 * z + 0.9 * mix * daylight[:, None]
qv = np.clip(q2[:, None] * np.exp(-z / 2100) * (0.88 + 0.12 * mix) + 1e-4 * smois[:, None], 1e-5, None)
y = np.stack((u, v, w, tk, qv), axis=-1).astype(np.float32)
y += rng.normal(size=y.shape).astype(np.float32) * np.array([0.035, 0.035, 0.003, 0.04, 0.000015], dtype=np.float32)
n_train, n_val = int(0.70 * count), int(0.15 * count)
slices = {"train": slice(0, n_train), "val": slice(n_train, n_train + n_val), "test": slice(n_train + n_val, count)}
output = ROOT / config["paths"]["data"]; output.parent.mkdir(parents=True, exist_ok=True)
payload = {"heights_m": heights, "input_names": np.asarray(INPUT_NAMES), "output_names": np.asarray(OUTPUT_NAMES)}
for name, section in slices.items():
payload.update({f"x_{name}": x[section], f"y_{name}": y[section], f"time_{name}": (np.arange(count)[section] * interval_hours).astype(np.int64), f"pblh_{name}": pblh[section].astype(np.float32)})
np.savez_compressed(output, **payload)
print(f"saved {output}: x={x.shape}, y={y.shape}, split={[len(range(*s.indices(count))) for s in slices.values()]}")
if __name__ == "__main__":
main()
|