File size: 3,673 Bytes
0fa8141 | 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 | """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)
# Smooth EnKF-like errors are tied to convection and dry-region mass drift.
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
# Shared synthetic climatology keeps train and validation normalization identical.
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()
|