| """Generate correlated 113-feature WoFS ensemble storm-track examples.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def correlated_block(rng, latent, count, noise, offset): |
| projection_rng = np.random.default_rng(1000 + int(offset * 100)) |
| weights = projection_rng.normal(0, 0.35, (latent.shape[1], count)) |
| for index in range(min(latent.shape[1], count)): |
| weights[index, index::latent.shape[1]] += 0.8 |
| values = latent @ weights + rng.normal(0, noise, (len(latent), count)) |
| values += 0.08 * np.sin(np.arange(count)[None, :] * 0.21 + offset) |
| return values.astype(np.float32) |
|
|
|
|
| def make_split(path, samples, config, seed): |
| rng = np.random.default_rng(seed) |
| lead_start = rng.choice(np.arange(0, 125, 5), samples) |
| lead_group = (lead_start > 60).astype(np.int64) |
| organization = rng.normal(size=samples) |
| rotation = 0.65 * organization + rng.normal(0, 0.75, samples) |
| instability = rng.normal(size=samples) |
| shear = 0.35 * instability + rng.normal(0, 0.9, samples) |
| cold_pool = 0.45 * organization + rng.normal(0, 0.85, samples) |
| spread = np.maximum(0.15, rng.lognormal(-0.2 + 0.2 * lead_group, 0.35, samples)) |
| latent = np.column_stack((organization, rotation, instability, shear, cold_pool, spread)) |
| amplitude = correlated_block(rng, latent, 30, 0.32, 0.0) |
| spatial = correlated_block(rng, latent, 76, 0.48 + 0.08 * lead_group[:, None], 0.7) |
| area = np.exp(0.45 * organization + rng.normal(5.8, 0.35, samples)) |
| eccentricity = 1 / (1 + np.exp(-(0.5 * shear + rng.normal(0, 0.6, samples)))) |
| major = np.sqrt(area) * (1.2 + eccentricity) |
| minor = area / np.maximum(major, 1) |
| orientation = np.arctan2(spatial[:, 4], spatial[:, 3]) / np.pi |
| extent = np.clip(0.72 - 0.15 * spread + rng.normal(0, 0.08, samples), 0.15, 1) |
| initialization_time = rng.uniform(0, 1, samples) |
| object_properties = np.column_stack((np.log1p(area), eccentricity, orientation, major / 100, |
| minor / 100, extent, initialization_time)).astype(np.float32) |
| features = np.concatenate((amplitude, spatial, object_properties), axis=1).astype(np.float32) |
| risk = np.column_stack((1.25 * rotation + 0.70 * shear + 0.35 * organization, |
| 1.15 * instability + 0.55 * organization + 0.45 * amplitude[:, 5], |
| 1.00 * cold_pool + 0.65 * shear + 0.35 * spatial[:, 12])) |
| risk -= lead_group[:, None] * np.array((0.35, 0.22, 0.18)) |
| |
| intercept = np.array((-2.35, -1.75, -1.45)) |
| probabilities = 1 / (1 + np.exp(-(risk + intercept))) |
| targets = (rng.random((samples, 3)) < probabilities).astype(np.float32) |
| np.savez_compressed(path, features=features, targets=targets, lead_group=lead_group, |
| lead_start_minutes=lead_start.astype(np.int16), |
| lead_end_minutes=(lead_start + 30).astype(np.int16), |
| format_version=np.asarray(config["data"]["format_version"]), |
| feature_group_sizes=np.asarray((30, 76, 7), dtype=np.int16), |
| hazards=np.asarray(config["data"]["hazards"]), |
| lead_group_names=np.asarray(config["data"]["lead_groups"]), |
| ensemble_members=np.asarray(18), grid_spacing_km=np.asarray(3), |
| forecast_window_minutes=np.asarray(30), forecast_interval_minutes=np.asarray(5), |
| data_source=np.asarray("structured_synthetic_elevated_event_rates"), |
| paper_event_rates=np.asarray((0.012, 0.025, 0.040), dtype=np.float32)) |
| return targets.mean(0) |
|
|
|
|
| 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) |
| for offset, (name, count) in enumerate((("train.npz", config["data"]["train_samples"]), |
| ("test.npz", config["data"]["test_samples"]))): |
| path = output / name |
| if args.force or not path.exists(): |
| rates = make_split(path, int(count), config, int(config["seed"]) + offset) |
| else: |
| rates = np.load(path)["targets"].mean(0) |
| print(f"generated={path.relative_to(ROOT)} shape=({count},113) event_rates={rates.round(3).tolist()}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
