| """Generate compact structured daily precipitation data on the real target grid.""" |
|
|
| from datetime import date, timedelta |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| data_config = config["data"] |
| samples = int(data_config["samples"]) |
| height, width = map(int, data_config["target_grid"]) |
| coarse_height, coarse_width = map(int, data_config["coarse_grid"]) |
| if (height, width) != (216, 488) or samples < 3: |
| raise ValueError("target grid must be 216x488 and at least three days are required") |
| rng = np.random.default_rng(int(config["seed"])) |
| y, x = np.mgrid[-1:1:complex(height), -1:1:complex(width)] |
| elevation = (2600 * np.exp(-((x + 0.2) ** 2 / 0.12 + (y - 0.05) ** 2 / 0.35)) + |
| 900 * np.exp(-((x - 0.55) ** 2 / 0.035 + (y + 0.35) ** 2 / 0.09))) |
| elevation += 180 * (np.sin(8 * x) * np.cos(5 * y) + 1) |
| elevation = np.maximum(elevation, 0).astype(np.float32) |
| target = np.empty((samples, 1, height, width), dtype=np.float32) |
| for day in range(samples): |
| phase = 2 * np.pi * day / max(samples, 4) |
| cx, cy = -0.75 + 1.5 * day / max(samples - 1, 1), 0.35 * np.sin(phase) |
| moving_system = 18 * np.exp(-((x - cx) ** 2 / 0.10 + (y - cy) ** 2 / 0.16)) |
| front = 7 * np.exp(-((y - 0.28 * np.sin(2 * x + phase)) ** 2) / 0.025) |
| terrain_enhancement = 5.5 * (elevation / max(float(elevation.max()), 1.0)) ** 1.4 |
| extreme = np.zeros_like(x) |
| if day % 4 == 2: |
| extreme = 42 * np.exp(-((x - 0.35) ** 2 + (y + 0.18) ** 2) / 0.008) |
| noise = rng.gamma(1.2, 0.35, (height, width)) |
| target[day, 0] = np.maximum(moving_system + front + terrain_enhancement + extreme + noise - 3.0, 0) |
| |
| row_edges = np.linspace(0, height, coarse_height + 1, dtype=int) |
| col_edges = np.linspace(0, width, coarse_width + 1, dtype=int) |
| coarse = np.empty((samples, 1, coarse_height, coarse_width), dtype=np.float32) |
| for row in range(coarse_height): |
| for column in range(coarse_width): |
| block = target[:, :, row_edges[row]:row_edges[row + 1], col_edges[column]:col_edges[column + 1]] |
| coarse[:, :, row, column] = block.mean(axis=(2, 3)) |
| start = date.fromisoformat(data_config["start_date"]) |
| timestamps = np.asarray([(start + timedelta(days=i)).isoformat() for i in range(samples)]) |
| years = np.asarray([(start + timedelta(days=i)).year for i in range(samples)], dtype=np.int32) |
| output = ROOT / data_config["root"] |
| output.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output / "daily_precipitation.npz", |
| format_version=np.array(data_config["format_version"]), |
| coarse_precipitation=coarse, elevation=elevation[None, None], |
| target_precipitation=target, timestamps=timestamps, years=years, |
| units=np.array("mm/day"), target_grid=np.array([height, width], np.int32)) |
| print(f"generated={samples} target_shape={target.shape} coarse_shape={coarse.shape}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
