OneScience-Group/MassConservingCNN

1"""Generate small, structured modified-shallow-water analysis pairs."""2 3import argparse4from pathlib import Path5 6import numpy as np7import yaml8 9 10ROOT = Path(__file__).resolve().parents[1]11 12 13def periodic_gaussian(x, center, width):14    distance = np.minimum(np.abs(x - center), 1.0 - np.abs(x - center))15    return np.exp(-0.5 * (distance / width) ** 2)16 17 18def make_split(path, count, config, seed):19    rng = np.random.default_rng(seed)20    n = int(config["data"]["grid_points"])21    x = np.arange(n, dtype=np.float32) / n22    xa = np.empty((count, 3, n), dtype=np.float32)23    target = np.empty_like(xa)24    radar = np.empty((count, 1, n), dtype=np.float32)25    for sample in range(count):26        phase = rng.uniform(0.0, 1.0)27        wave = np.sin(2 * np.pi * (x - phase))28        harmonic = np.sin(4 * np.pi * (x - 0.6 * phase))29        convective = periodic_gaussian(x, (phase + 0.23) % 1.0, 0.045)30        secondary = periodic_gaussian(x, (phase + 0.66) % 1.0, 0.07)31        u_true = 0.75 * wave + 0.22 * harmonic - 0.28 * np.gradient(convective)32        h_true = 10.0 + 0.35 * np.cos(2 * np.pi * (x - phase)) + 0.5 * convective33        convergence = np.maximum(-np.gradient(u_true), 0.0)34        r_true = np.maximum(0.0, 0.7 * convective + 0.28 * convergence - 0.09)35        rain_mask = (r_true > 0.08).astype(np.float32)36 37        # Smooth EnKF-like errors are tied to convection and dry-region mass drift.38        dry = 1.0 - rain_mask39        u_error = 0.11 * secondary - 0.07 * convective + 0.025 * harmonic40        h_error = 0.16 * dry + 0.08 * secondary - 0.05 * convective41        r_error = 0.13 * secondary * dry - 0.06 * convective42        xa[sample, 0] = u_true + u_error43        xa[sample, 1] = h_true + h_error44        xa[sample, 2] = np.maximum(0.0, r_true + r_error)45        target[sample] = np.stack((u_true, h_true, r_true))46        radar[sample, 0] = rain_mask47 48    # Shared synthetic climatology keeps train and validation normalization identical.49    means = np.asarray([0.0, 10.0], dtype=np.float32)50    stds = np.asarray([0.6, 0.4, 0.3], dtype=np.float32)51    normalized_x = xa.copy()52    normalized_y = target.copy()53    normalized_x[:, :2] = (xa[:, :2] - means[None, :, None]) / stds[None, :2, None]54    normalized_y[:, :2] = (target[:, :2] - means[None, :, None]) / stds[None, :2, None]55    normalized_x[:, 2] = xa[:, 2] / stds[2]56    normalized_y[:, 2] = target[:, 2] / stds[2]57    inputs = np.concatenate((normalized_x, radar), axis=1).astype(np.float32)58    np.savez_compressed(59        path, inputs=inputs, targets=normalized_y.astype(np.float32), xa=xa,60        targets_physical=target, radar=radar, climate_mean_uh=means,61        climate_std_uhr=stds, format_version=np.asarray(config["data"]["format_version"]),62        variable_order=np.asarray(["u", "h", "r"]), input_layout=np.asarray("BCX"),63        data_source=np.asarray("structured_synthetic_msw"),64    )65 66 67def main():68    parser = argparse.ArgumentParser()69    parser.add_argument("--force", action="store_true")70    args = parser.parse_args()71    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())72    output = ROOT / config["data"]["root"]73    output.mkdir(parents=True, exist_ok=True)74    splits = (("train.npz", int(config["data"]["train_samples"])),75              ("validation.npz", int(config["data"]["validation_samples"])))76    for offset, (name, count) in enumerate(splits):77        path = output / name78        if args.force or not path.exists():79            make_split(path, count, config, int(config["seed"]) + offset)80        print(f"generated={path.relative_to(ROOT)} samples={count} shape=({count},4,250)")81 82 83if __name__ == "__main__":84    main()85