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