| """Create spatially coherent daily anomalies with weather and warming signals.""" |
|
|
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.precipdd import DATA_FORMAT_VERSION, load_config |
|
|
|
|
| def agmt_for_year(year, member, rng): |
| forced = -0.45 + 0.0042 * (year - 1850) + 0.000030 * max(year - 1970, 0) ** 2 |
| return forced + 0.08 * np.sin((year - 1850) / 8.0 + member) + rng.normal(0, 0.045) |
|
|
|
|
| def main(): |
| config = load_config(ROOT / "conf/config.yaml") |
| settings = config["data"] |
| rng = np.random.default_rng(config["project"]["seed"]) |
| train_n, val_n = settings["train_samples"], settings["validation_samples"] |
| first, last = settings["test_years"] |
| test_year = np.repeat(np.arange(first, last + 1), settings["test_days_per_year"]) |
| train_year = rng.integers(1850, 2101, train_n) |
| val_year = rng.integers(1850, 2101, val_n) |
| year = np.concatenate((train_year, val_year, test_year)).astype(np.int16) |
| split = np.concatenate((np.zeros(train_n, np.int8), np.ones(val_n, np.int8), np.full(len(test_year), 2, np.int8))) |
| day = np.concatenate((rng.integers(1, 366, train_n + val_n), np.tile(np.linspace(15, 350, settings["test_days_per_year"], dtype=int), last - first + 1))).astype(np.int16) |
| member = rng.integers(0, settings["synthetic_members"], len(year), dtype=np.int16) |
| agmt = np.array([agmt_for_year(int(y), int(m), rng) for y, m in zip(year, member)], dtype=np.float32) |
|
|
| latitude = np.arange(-60.0, 77.5, 2.5, dtype=np.float32) |
| base_longitude = np.arange(0.0, 360.0, 2.5, dtype=np.float32) |
| longitude = np.arange(0.0, 400.0, 2.5, dtype=np.float32) |
| lat2d, lon2d = np.meshgrid(latitude, base_longitude, indexing="ij") |
| east_pacific = np.exp(-((lat2d / 16) ** 2 + ((lon2d - 245) / 35) ** 2)) |
| storm_tracks = np.exp(-((np.abs(lat2d) - 45) / 12) ** 2) * (0.65 + 0.35 * np.cos(np.deg2rad(lon2d * 2))) |
| fingerprint = (east_pacific + storm_tracks).astype(np.float32) |
| fingerprint /= fingerprint.max() |
| coarse = torch.from_numpy(rng.normal(size=(len(year), 1, 14, 36)).astype(np.float32)) |
| weather = F.interpolate(coarse, size=(55, 144), mode="bilinear", align_corners=False).numpy() |
| synoptic = np.sin(2 * np.pi * day[:, None, None] / 7.0 + np.deg2rad(lon2d)[None]) |
| noise = rng.normal(0, 0.28, size=weather.shape).astype(np.float32) |
| amplitude = 1.0 + 0.34 * np.maximum(agmt, -0.5)[:, None, None] * fingerprint[None] |
| fields = weather[:, 0] * amplitude + 0.32 * synoptic * (0.3 + fingerprint[None]) + noise[:, 0] |
| fields += 0.22 * agmt[:, None, None] * (fingerprint[None] - 0.35) |
| fields = fields.astype(np.float32) |
| fields = np.concatenate((fields, fields[:, :, :16]), axis=2) |
| fingerprint = np.concatenate((fingerprint, fingerprint[:, :16]), axis=1) |
| fields -= fields.mean(axis=2, keepdims=True) |
| zonal_std = fields.std(axis=2, keepdims=True).mean(axis=0, keepdims=True) |
| fields /= np.maximum(zonal_std, 1e-5) |
|
|
| output = ROOT / config["paths"]["data"] |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, format_version=np.array(DATA_FORMAT_VERSION), precipitation=fields[:, None], agmt=agmt, |
| split=split, year=year, day_of_year=day, member=member, latitude=latitude, |
| longitude=longitude, synthetic=np.array(True), fingerprint=fingerprint) |
| metadata = {"synthetic": True, "samples": len(year), "splits": {"train": train_n, "validation": val_n, "test": len(test_year)}, |
| "precipitation_shape": [len(year), 1, 55, 160], "target_shape": [len(year)], |
| "science_note": "Structured engineering data with daily weather variability and an AGMT-dependent spatial variance signal; not CESM2 LE."} |
| (output.parent / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8") |
| print(f"data={output.relative_to(ROOT)} shape={fields[:, None].shape} target={agmt.shape}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
