File size: 4,324 Bytes
2e913c2 | 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 | """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()
|