File size: 6,120 Bytes
c059069 | 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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Generate one complete structured coarse-grid snapshot at every paper scale."""
import sys
from pathlib import Path
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.rf_climparam import (DIFF_INPUT_NAMES, DIFF_OUTPUT_NAMES,
FORMAT_VERSION, TEND_INPUT_NAMES,
TEND_OUTPUT_NAMES)
def make_columns(count, factor, seed, times=None, latitudes=None, longitudes=None):
rng = np.random.default_rng(seed)
if times is None:
times = np.arange(count, dtype=np.float32) * 3.0
if latitudes is None:
latitudes = rng.uniform(-90.0, 90.0, count).astype(np.float32)
if longitudes is None:
longitudes = rng.uniform(0.0, 360.0, count).astype(np.float32)
times, latitudes, longitudes = map(lambda x: np.asarray(x, np.float32), (times, latitudes, longitudes))
z = np.linspace(0.0, 1.0, 48, dtype=np.float32)[None, :]
lat = np.deg2rad(latitudes)[:, None]
lon = np.deg2rad(longitudes)[:, None]
phase = (2.0 * np.pi * times / (24.0 * 45.0))[:, None]
noise = 0.8 / np.sqrt(factor)
temperature = 300.0 - 78.0 * z - 13.0 * np.sin(lat) ** 2 + 1.5 * np.sin(lon + phase) * np.exp(-2 * z)
temperature += rng.normal(0, noise, (count, 48))
humidity = 0.018 * np.exp(-5.0 * z) * (0.35 + 0.65 * np.cos(lat) ** 2)
humidity *= 1.0 + 0.12 * np.sin(lon - phase)
humidity += rng.normal(0, noise * 2e-4, (count, 48))
humidity = np.maximum(humidity, 2e-6)
condensate = np.maximum(humidity - (0.009 * np.exp(-3.8 * z)), 0.0)
condensate += np.maximum(rng.normal(0, noise * 2e-5, (count, 48)), 0.0)
abs_y = np.abs(latitudes)[:, None] / 90.0
instability = np.maximum(temperature[:, :1] - temperature[:, 20:21] - 30.0, 0.0)
convective = (1.0 + 0.02 * instability) * humidity * np.exp(-1.5 * z)
h_tend = (-1.8e-5 * (temperature - temperature.mean(1, keepdims=True)) +
3e-4 * convective + rng.normal(0, noise * 2e-6, (count, 48)))
qt_tend = (-1.4e-4 * convective + 2e-6 * np.cos(lon + 3 * z) +
rng.normal(0, noise * 2e-7, (count, 48)))
qp_tend = np.maximum(8e-5 * convective + 5e-5 * condensate, 0.0)
tend_inputs = np.concatenate([temperature, humidity, condensate, abs_y], axis=1).astype(np.float32)
tend_targets = np.concatenate([h_tend, qt_tend, qp_tend], axis=1).astype(np.float32)
low_z = z[:, :15]
u = 7.0 * np.cos(lat) * (1.0 + low_z) + 2.0 * np.sin(lon + phase + low_z)
v = 3.0 * np.sin(2 * lat) * (1.0 - low_z) + 1.5 * np.cos(lon - phase + low_z)
v_nh = v * np.where(latitudes[:, None] < 0, -1.0, 1.0)
windsurf = np.sqrt(u[:, :1] ** 2 + v[:, :1] ** 2)
dbar = np.maximum((0.35 + 0.08 * windsurf) * np.exp(-3 * low_z) +
rng.normal(0, noise * 0.015, (count, 15)), 0.0)
h_flux = (12.0 + 0.8 * windsurf[:, 0] + 5.0 * np.cos(lat[:, 0]) ** 2).reshape(-1, 1)
q_flux = (2e-5 + 2e-6 * windsurf[:, 0] + 8e-6 * np.cos(lat[:, 0]) ** 2).reshape(-1, 1)
diff_inputs = np.concatenate([temperature[:, :15], humidity[:, :15], u, v_nh,
windsurf, abs_y], axis=1).astype(np.float32)
diff_targets = np.concatenate([dbar, h_flux, q_flux], axis=1).astype(np.float32)
return tend_inputs, tend_targets, diff_inputs, diff_targets, times, latitudes, longitudes
def validate_flattening(rows, columns, ny, nx):
expected = np.arange(ny * nx)
restored = rows.astype(np.int64) * nx + columns.astype(np.int64)
if not np.array_equal(restored, expected):
raise ValueError(f"grid [{ny},{nx}] is not reversibly flattened in C order")
def make_snapshot(ny, nx, factor, seed, time_hours=0.0):
rows = np.repeat(np.arange(ny, dtype=np.int32), nx)
columns = np.tile(np.arange(nx, dtype=np.int32), ny)
latitude = np.repeat(np.linspace(-90.0, 90.0, ny, dtype=np.float32), nx)
longitude = np.tile(np.linspace(0.0, 360.0, nx, endpoint=False, dtype=np.float32), ny)
time = np.full(ny * nx, time_hours, dtype=np.float32)
validate_flattening(rows, columns, ny, nx)
return make_columns(ny * nx, factor, seed, time, latitude, longitude), rows, columns
def save(path, scale, arrays, grid_shape):
arrays, rows, columns = arrays
tend_x, tend_y, diff_x, diff_y, time, lat, lon = arrays
np.savez_compressed(path, format_version=np.array(FORMAT_VERSION), scale=np.array(scale),
grid_shape=np.asarray(grid_shape, dtype=np.int32),
snapshot_count=np.array(1, dtype=np.int32),
tend_inputs=tend_x, tend_targets=tend_y,
diff_inputs=diff_x, diff_targets=diff_y,
time_hours=time, latitude_deg=lat, longitude_deg=lon,
grid_row=rows, grid_column=columns,
tend_input_names=np.asarray(TEND_INPUT_NAMES),
tend_output_names=np.asarray(TEND_OUTPUT_NAMES),
diff_input_names=np.asarray(DIFF_INPUT_NAMES),
diff_output_names=np.asarray(DIFF_OUTPUT_NAMES))
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
output = ROOT / config["data"]["root"]
output.mkdir(parents=True, exist_ok=True)
snapshots, seed = int(config["data"]["snapshots_per_scale"]), int(config["seed"])
if snapshots != 1:
raise ValueError("compact synthetic protocol requires exactly one snapshot per scale")
for index, scale in enumerate(config["data"]["scales"]):
grid = config["data"]["coarse_grids"][scale]
factor, (ny, nx) = int(grid["factor"]), map(int, grid["grid"])
save(output / f"{scale}.npz", scale,
make_snapshot(ny, nx, factor, seed + index), (ny, nx))
ny, nx = map(int, config["evaluation"]["online_proxy_grid"])
save(output / "x32_online_proxy.npz", "x32",
make_snapshot(ny, nx, 32, seed + 100), (ny, nx))
print(f"generated=4 complete coarse snapshots online_native=1x{ny}x{nx}")
if __name__ == "__main__":
main()
|