| """Generate structured annual forcing windows and climate responses.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def make_split(path: Path, count: int, config: dict, seed: int, test: bool) -> None: |
| rng = np.random.default_rng(seed) |
| data = config["data"] |
| time, height, width = int(data["time_steps"]), int(data["height"]), int(data["width"]) |
| lat = np.linspace(-90, 90, height, dtype=np.float32) |
| lon = np.linspace(0, 360, width, endpoint=False, dtype=np.float32) |
| lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij") |
| windows, responses, years = [], [], [] |
| for sample in range(count): |
| end_year = 2100 if test else 2060 + 10 * sample |
| annual_years = np.arange(end_year - time + 1, end_year + 1, dtype=np.int32) |
| progress = (annual_years - 1850).astype(np.float32) |
| co2 = 0.90 * progress + 0.0018 * progress ** 2 |
| ch4 = 0.42 * progress + 12.0 * np.sin(progress / 38.0) |
| industrial = np.exp(-((lat2d - 35) / 22) ** 2) * (0.65 + 0.35 * np.cos(np.deg2rad(lon2d - 90))) |
| tropical = np.exp(-(lat2d / 20) ** 2) * (0.65 + 0.35 * np.sin(np.deg2rad(2 * lon2d))) |
| forcing = np.empty((time, 4, height, width), dtype=np.float32) |
| for index, year_progress in enumerate(progress): |
| forcing[index, 0] = co2[index] |
| forcing[index, 1] = ch4[index] |
| forcing[index, 2] = industrial * (32.0 * np.exp(-((year_progress - 125) / 75) ** 2)) |
| forcing[index, 3] = tropical * (5.0 + 2.0 * np.sin(year_progress / 27.0)) |
| forcing += rng.normal(0, 0.01, forcing.shape).astype(np.float32) |
| co2_level = forcing[-1, 0].mean() / 300.0 |
| methane = forcing[-3:, 1].mean() / 120.0 |
| aerosol = forcing[-3:, 2:].mean(axis=(0, 1)) / 25.0 |
| arctic = 1.0 + 1.3 * (np.abs(lat2d) / 90.0) ** 2 |
| land_pattern = np.cos(np.deg2rad(2 * lon2d)) * np.cos(np.deg2rad(lat2d)) |
| itcz = np.sin(np.deg2rad(lon2d)) * np.exp(-(lat2d / 17) ** 2) |
| tas = 1.35 * co2_level * arctic + 0.22 * methane - 0.30 * aerosol + 0.05 * land_pattern |
| dtr = 0.10 * co2_level + 0.23 * aerosol * land_pattern - 0.04 * methane |
| pr = 0.16 * co2_level * np.cos(np.deg2rad(lat2d)) + 0.20 * itcz - 0.07 * aerosol |
| pr90 = 1.35 * pr + 0.08 * co2_level * np.exp(-(lat2d / 28) ** 2) |
| response = np.stack([tas, dtr, pr, pr90]).astype(np.float32) |
| response += rng.normal(0, 0.003, response.shape).astype(np.float32) |
| windows.append(forcing) |
| responses.append(response) |
| years.append(annual_years) |
| evaluation = config["evaluation"] |
| np.savez_compressed( |
| path, inputs=np.stack(windows).astype(np.float32), targets=np.stack(responses).astype(np.float32), |
| years=np.stack(years), latitude=lat, longitude=lon, |
| channel_names=np.asarray(data["channels"]), target_names=np.asarray(data["targets"]), |
| format_version=np.asarray(data["format_version"]), storage_layout=np.asarray("NTCHW"), |
| temporal_resolution=np.asarray("annual"), data_source=np.asarray("structured_synthetic"), |
| scenario=np.asarray(evaluation["scenario"] if test else "synthetic_training_scenarios"), |
| target_aggregation=np.asarray("2080-2100 climatological mean" if test else "annual response"), |
| evaluation_start_year=np.asarray(evaluation["start_year"], dtype=np.int32), |
| evaluation_end_year=np.asarray(evaluation["end_year"], dtype=np.int32), |
| ) |
|
|
|
|
| def main() -> None: |
| 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, test) in enumerate((("train.npz", config["data"]["train_samples"], False), |
| ("test.npz", config["data"]["test_samples"], True))): |
| path = output / name |
| if args.force or not path.exists(): |
| make_split(path, int(count), config, int(config["seed"]) + offset, test) |
| print(f"generated={path.relative_to(ROOT)} samples={count} layout=NTCHW annual=true") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
