File size: 2,144 Bytes
380b161 | 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 | from pathlib import Path
import sys
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.ace2 import (
PRECIPITATION,
Q_INDICES,
RADIATION_INDICES,
SURFACE_PRESSURE,
forcing_for_hours,
load_config,
seed_all,
)
ROOT = Path(__file__).resolve().parents[1]
def main():
cfg = load_config(ROOT)
seed_all(cfg["seed"])
d = cfg["data"]
shape = (d["samples"], d["time_steps"], d["channels"], d["height"], d["width"])
rng = np.random.default_rng(cfg["seed"])
lat = np.deg2rad(np.linspace(-89.5, 89.5, d["height"], dtype=np.float32))[None, None, None, :, None]
lon = np.deg2rad(np.linspace(0.5, 359.5, d["width"], dtype=np.float32))[None, None, None, None, :]
time = np.arange(d["time_steps"], dtype=np.float32)[None, :, None, None, None]
channel = np.arange(d["channels"], dtype=np.float32)[None, None, :, None, None]
state = (0.2 * np.sin(lat * (1 + channel % 3)) + 0.1 * np.cos(lon + time / 4)
+ 0.002 * channel + 0.003 * time).astype(np.float32)
state = np.broadcast_to(state, shape).copy()
state += rng.normal(0, 0.005, (d["samples"], 1, d["channels"], 1, 1)).astype(np.float32)
state[:, :, Q_INDICES] = np.maximum(state[:, :, Q_INDICES] * 0.01 + 0.003, 0)
state[:, :, SURFACE_PRESSURE] = 1.0 + 0.01 * np.cos(lat[:, :, 0])
state[:, :, PRECIPITATION] = np.maximum(0, 0.001 * (np.sin(lon[:, :, 0] + time[:, :, 0]) + 1))
state[:, :, RADIATION_INDICES] = np.maximum(state[:, :, RADIATION_INDICES] + 0.5, 0)
hours = (np.arange(d["samples"])[:, None] * d["time_steps"] + np.arange(d["time_steps"])[None]) * d["dt_hours"]
forcing = forcing_for_hours(hours.reshape(-1), d["height"], d["width"]).reshape(
d["samples"], d["time_steps"], 4, d["height"], d["width"]
)
path = ROOT / d["path"]
path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(path, state=state.astype(np.float16), forcing=forcing.astype(np.float16), hours=hours)
print(f"created {path}: state={state.shape}, forcing={forcing.shape}, dt={d['dt_hours']}h")
if __name__ == "__main__":
main()
|