| from pathlib import Path |
| import math |
| import yaml |
| import numpy as np |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def main(): |
| cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| rng = np.random.default_rng(cfg["seed"]) |
| ncol = cfg["data"]["columns"] |
| nwin = cfg["data"]["windows_per_column"] |
| steps = cfg["data"]["window_steps"] |
| levels = cfg["data"]["levels"] |
| dt = cfg["data"]["step_hours"] * 3600.0 |
| total_steps = nwin * steps + 64 |
| sigma = np.linspace(0.02, 0.995, levels, dtype=np.float32) |
| pressure = 100000.0 * sigma |
| interfaces = np.linspace(0.0, 100000.0, levels + 1, dtype=np.float32) |
| layer_mass = np.diff(interfaces) / 9.80665 |
| time = np.arange(total_steps + 1, dtype=np.float32) |
| states = np.empty((ncol, total_steps + 1, 68), dtype=np.float32) |
| surface = np.empty((ncol, total_steps, 3), dtype=np.float32) |
| advection = np.empty((ncol, total_steps + 1, 68), dtype=np.float32) |
| for c in range(ncol): |
| phase = 2.0 * math.pi * c / ncol |
| sl0 = 2.82e5 + 1.25e5 * (1.0 - sigma) + 1400.0 * np.sin(math.pi * sigma + phase) |
| qt0 = 0.0175 * sigma ** 2.6 + 0.00012 + 0.0004 * np.sin(2 * math.pi * sigma + phase) |
| states[c, 0] = np.concatenate((sl0, np.maximum(qt0, 2e-5))) |
| solar = np.maximum(0.0, np.sin(2.0 * math.pi * (time[:-1] % 8.0) / 8.0 + phase)) |
| surface[c, :, 0] = 30.0 + 18.0 * np.sin(2 * math.pi * time[:-1] / 8.0 + phase) |
| surface[c, :, 1] = 92.0 + 34.0 * np.sin(2 * math.pi * time[:-1] / 8.0 + phase - 0.5) |
| surface[c, :, 2] = 760.0 * solar |
| vertical = np.sin(math.pi * sigma) |
| for t in range(total_steps + 1): |
| wave = math.sin(2 * math.pi * t / 16.0 + phase) |
| advection[c, t, :levels] = (0.010 * wave * vertical + rng.normal(0, 0.0003, levels)) |
| advection[c, t, levels:] = (1.4e-9 * wave * vertical + rng.normal(0, 4e-11, levels)) |
| for t in range(total_steps): |
| adv = 0.5 * (advection[c, t] + advection[c, t + 1]) |
| sl, qt = np.split(states[c, t], 2) |
| shf, lhf, solin = surface[c, t] |
| heat = 2.0e-7 * (300000.0 - sl) + 1.0e-5 * shf * sigma + 1.0e-7 * solin * vertical |
| moist = 1.4e-7 * (0.012 * sigma ** 2.4 - qt) + 2.5e-11 * lhf * sigma |
| next_state = states[c, t] + dt * (adv + np.concatenate((heat, moist))) |
| next_state[levels:] = np.maximum(next_state[levels:], 1e-6) |
| states[c, t + 1] = next_state |
| initial, targets, surf_windows, adv_windows, source = [], [], [], [], [] |
| for c in range(ncol): |
| for w in range(nwin): |
| start = w * steps |
| initial.append(states[c, start]) |
| targets.append(states[c, start:start + steps + 1]) |
| surf_windows.append(surface[c, start:start + steps]) |
| adv_windows.append(advection[c, start:start + steps + 1]) |
| source.append(c) |
| out = ROOT / cfg["data"]["file"] |
| out.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(out, initial=np.array(initial), target=np.array(targets), |
| surface=np.array(surf_windows), advection=np.array(adv_windows), |
| layer_mass=layer_mass, pressure=pressure, source=np.array(source), |
| long_state=states, long_surface=surface, long_advection=advection, |
| dt_seconds=np.float32(dt), format_version=cfg["data"]["format_version"]) |
| print(f"saved {out}: windows={len(initial)}, T={steps}, levels={levels}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
