OneScience-Group/CRAI-ClimateExtremes

1"""Create structurally realistic full-grid samples with irregular HadEX-style masks."""2 3from pathlib import Path4import json5import numpy as np6 7 8INDICES = np.array(["TX90p", "TN90p", "TX10p", "TN10p"])9H, W = 144, 19210 11 12def europe_mask(lat, lon):13    yy, xx = np.meshgrid(lat, lon, indexing="ij")14    broad = (yy >= 30) & (yy <= 72) & (xx >= -25) & (xx <= 45)15    # A coarse geographic silhouette keeps the scientific global-grid contract.16    atlantic_cut = (xx < -10) & (yy < 44)17    southeast_cut = (xx > 30) & (yy < 40)18    north_cut = (yy > 68) & ((xx < 5) | (xx > 30))19    return (broad & ~atlantic_cut & ~southeast_cut & ~north_cut).astype(np.float32)20 21 22def main():23    rng = np.random.default_rng(42)24    root = Path(__file__).resolve().parents[1]25    output = root / "data"26    output.mkdir(exist_ok=True)27    lat = np.linspace(-89.375, 89.375, H, dtype=np.float32)28    lon = np.linspace(-179.0625, 179.0625, W, dtype=np.float32)29    land = europe_mask(lat, lon)30    yy, xx = np.meshgrid(lat, lon, indexing="ij")31    n = 832    target = np.zeros((n, 1, H, W), dtype=np.float32)33    valid = np.zeros_like(target)34    index_ids = np.arange(n, dtype=np.int64) % 435    for sample in range(n):36        phase = 0.55 * sample37        field = 50 + 21 * np.sin(np.deg2rad(2.3 * xx) + phase)38        field += 16 * np.cos(np.deg2rad(3.2 * yy) - 0.4 * phase)39        field += 5 * np.sin(np.deg2rad(xx + yy) * 4 + phase)40        field += rng.normal(0, 1.2, (H, W))41        if index_ids[sample] >= 2:42            field = 100 - field43        target[sample, 0] = np.clip(field, 0, 100) * land44        observed = land.copy()45        observed[rng.random((H, W)) < (0.35 + 0.04 * (sample % 3))] = 046        for _ in range(5):47            cy, cx = rng.integers(45, 99), rng.integers(78, 121)48            ry, rx = rng.integers(3, 11), rng.integers(4, 15)49            hole = ((np.arange(H)[:, None] - cy) / ry) ** 250            hole = hole + ((np.arange(W)[None, :] - cx) / rx) ** 251            observed[hole < 1] = 052        valid[sample, 0] = observed53    observed_values = target * valid54    np.savez_compressed(55        output / "crai_fake.npz", target=target, observed=observed_values,56        valid_mask=valid, europe_mask=land, index_ids=index_ids,57        index_names=INDICES, latitude=lat, longitude=lon,58    )59    metadata = {60        "kind": "structured_synthetic",61        "shape": [n, 1, H, W],62        "grid_resolution": {"longitude_degrees": 1.875, "latitude_degrees": 1.25},63        "indices": INDICES.tolist(),64        "mask": "global grid with coarse Europe land support and irregular missing regions",65    }66    (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")67    print(f"wrote {output / 'crai_fake.npz'} with shape {target.shape}")68 69 70if __name__ == "__main__":71    main()72