| 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() |
|
|
