| """Generate small, structured modified-shallow-water analysis pairs.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def periodic_gaussian(x, center, width): |
| distance = np.minimum(np.abs(x - center), 1.0 - np.abs(x - center)) |
| return np.exp(-0.5 * (distance / width) ** 2) |
|
|
|
|
| def make_split(path, count, config, seed): |
| rng = np.random.default_rng(seed) |
| n = int(config["data"]["grid_points"]) |
| x = np.arange(n, dtype=np.float32) / n |
| xa = np.empty((count, 3, n), dtype=np.float32) |
| target = np.empty_like(xa) |
| radar = np.empty((count, 1, n), dtype=np.float32) |
| for sample in range(count): |
| phase = rng.uniform(0.0, 1.0) |
| wave = np.sin(2 * np.pi * (x - phase)) |
| harmonic = np.sin(4 * np.pi * (x - 0.6 * phase)) |
| convective = periodic_gaussian(x, (phase + 0.23) % 1.0, 0.045) |
| secondary = periodic_gaussian(x, (phase + 0.66) % 1.0, 0.07) |
| u_true = 0.75 * wave + 0.22 * harmonic - 0.28 * np.gradient(convective) |
| h_true = 10.0 + 0.35 * np.cos(2 * np.pi * (x - phase)) + 0.5 * convective |
| convergence = np.maximum(-np.gradient(u_true), 0.0) |
| r_true = np.maximum(0.0, 0.7 * convective + 0.28 * convergence - 0.09) |
| rain_mask = (r_true > 0.08).astype(np.float32) |
|
|
| |
| dry = 1.0 - rain_mask |
| u_error = 0.11 * secondary - 0.07 * convective + 0.025 * harmonic |
| h_error = 0.16 * dry + 0.08 * secondary - 0.05 * convective |
| r_error = 0.13 * secondary * dry - 0.06 * convective |
| xa[sample, 0] = u_true + u_error |
| xa[sample, 1] = h_true + h_error |
| xa[sample, 2] = np.maximum(0.0, r_true + r_error) |
| target[sample] = np.stack((u_true, h_true, r_true)) |
| radar[sample, 0] = rain_mask |
|
|
| |
| means = np.asarray([0.0, 10.0], dtype=np.float32) |
| stds = np.asarray([0.6, 0.4, 0.3], dtype=np.float32) |
| normalized_x = xa.copy() |
| normalized_y = target.copy() |
| normalized_x[:, :2] = (xa[:, :2] - means[None, :, None]) / stds[None, :2, None] |
| normalized_y[:, :2] = (target[:, :2] - means[None, :, None]) / stds[None, :2, None] |
| normalized_x[:, 2] = xa[:, 2] / stds[2] |
| normalized_y[:, 2] = target[:, 2] / stds[2] |
| inputs = np.concatenate((normalized_x, radar), axis=1).astype(np.float32) |
| np.savez_compressed( |
| path, inputs=inputs, targets=normalized_y.astype(np.float32), xa=xa, |
| targets_physical=target, radar=radar, climate_mean_uh=means, |
| climate_std_uhr=stds, format_version=np.asarray(config["data"]["format_version"]), |
| variable_order=np.asarray(["u", "h", "r"]), input_layout=np.asarray("BCX"), |
| data_source=np.asarray("structured_synthetic_msw"), |
| ) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--force", action="store_true") |
| args = parser.parse_args() |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| output = ROOT / config["data"]["root"] |
| output.mkdir(parents=True, exist_ok=True) |
| splits = (("train.npz", int(config["data"]["train_samples"])), |
| ("validation.npz", int(config["data"]["validation_samples"]))) |
| for offset, (name, count) in enumerate(splits): |
| path = output / name |
| if args.force or not path.exists(): |
| make_split(path, count, config, int(config["seed"]) + offset) |
| print(f"generated={path.relative_to(ROOT)} samples={count} shape=({count},4,250)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
