| """Generate seasonal directed VAR networks and precipitation fields for CME.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def stable_network(rng, nodes, season_index): |
| coefficients = np.zeros((nodes, nodes, 10), dtype=np.float32) |
| coefficients[np.arange(nodes), np.arange(nodes), 0] = 0.48 + 0.04 * season_index |
| for source in range(nodes): |
| target = (source + 3 + season_index) % nodes |
| lag = (source + 2 * season_index) % 4 |
| coefficients[source, target, lag] = (0.13 + 0.03 * (source % 3)) * (-1 if source % 5 == 0 else 1) |
| if source % 4 == 0: |
| coefficients[source, (source + 11) % nodes, (lag + 1) % 6] = -0.11 |
| return coefficients |
|
|
|
|
| def simulate(rng, coefficients, samples, time_steps): |
| nodes, max_lag = coefficients.shape[0], coefficients.shape[2] |
| output = np.zeros((samples, time_steps, nodes), dtype=np.float32) |
| for sample in range(samples): |
| series = rng.normal(0, 0.35, (time_steps + max_lag, nodes)).astype(np.float32) |
| for time in range(max_lag, time_steps + max_lag): |
| forcing = np.zeros(nodes, dtype=np.float32) |
| for lag in range(1, max_lag + 1): |
| forcing += series[time - lag] @ coefficients[:, :, lag - 1] |
| series[time] = forcing + rng.normal(0, 0.38, nodes) |
| output[sample] = series[max_lag:] |
| return output |
|
|
|
|
| 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()) |
| data, seed = config["data"], int(config["seed"]) |
| path = ROOT / config["paths"]["dataset"] |
| path.parent.mkdir(parents=True, exist_ok=True) |
| if path.exists() and not args.force: |
| print(f"exists={path.relative_to(ROOT)} use --force to regenerate") |
| return |
| rng = np.random.default_rng(seed) |
| seasons, models = len(data["seasons"]), int(data["model_count"]) |
| nodes, samples, steps = int(data["nodes"]), int(data["samples"]), int(data["time_steps"]) |
| reference_coefficients = np.stack([stable_network(rng, nodes, season) for season in range(seasons)]) |
| reference_series = np.stack([simulate(rng, reference_coefficients[s], samples, steps) for s in range(seasons)]) |
| model_series = np.empty((models, seasons, samples, steps, nodes), dtype=np.float32) |
| model_coefficients = np.empty((models, seasons, nodes, nodes, 10), dtype=np.float32) |
| quality = np.linspace(0.92, 0.35, models).astype(np.float32) |
| for model in range(models): |
| for season in range(seasons): |
| coefficients = reference_coefficients[season].copy() |
| cross = ~np.eye(nodes, dtype=bool) |
| coefficients[cross] *= quality[model] |
| mutation_count = 3 + 3 * model |
| for _ in range(mutation_count): |
| source, target = rng.integers(0, nodes, 2) |
| if source != target: |
| coefficients[source, target, rng.integers(0, 7)] = rng.choice([-1, 1]) * rng.uniform(0.08, 0.16) |
| model_coefficients[model, season] = coefficients |
| model_series[model, season] = simulate(rng, coefficients, samples, steps) |
| lat_count, lon_count = map(int, data["grid_shape"]) |
| latitude = np.linspace(-90.0, 90.0, lat_count, dtype=np.float32) |
| longitude = np.linspace(0, 360, lon_count, endpoint=False, dtype=np.float32) |
| lat2d, lon2d = np.meshgrid(latitude, longitude, indexing="ij") |
| reference_precip = (3.0 + 2.1 * np.cos(np.deg2rad(lat2d)) ** 2 + |
| 0.45 * np.sin(np.deg2rad(2 * lon2d))).astype(np.float32) |
| precip_fields, delta = [], [] |
| for model, q in enumerate(quality): |
| bias = (1 - q) * (0.8 * np.sin(np.deg2rad(lat2d)) + 0.35 * np.cos(np.deg2rad(lon2d))) |
| precip_fields.append(reference_precip * (0.88 + 0.12 * q) + bias + rng.normal(0, 0.06 + 0.08 * (1 - q), reference_precip.shape)) |
| delta.append(0.18 + 0.95 * (q - 0.62) ** 2 + rng.normal(0, 0.025)) |
| np.savez_compressed(path, format_version=np.asarray(data["format_version"]), seasons=np.asarray(data["seasons"]), |
| reference_series=reference_series, model_series=model_series, |
| reference_coefficients=reference_coefficients, model_coefficients=model_coefficients, |
| reference_precipitation=reference_precip, model_precipitation=np.asarray(precip_fields, np.float32), |
| delta_precipitation=np.asarray(delta, np.float32), latitude_degrees=latitude, |
| longitude_degrees=longitude, model_quality=quality, time_step_days=np.asarray(data["time_step_days"])) |
| print(f"generated={path.relative_to(ROOT)} reference={reference_series.shape} models={model_series.shape}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
