File size: 4,895 Bytes
20cdc88 | 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 86 87 88 89 90 91 92 | """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()
|