File size: 3,570 Bytes
7f71cfd | 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 71 | 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()
|