OneScience-Group/ML-MODIS

1#!/usr/bin/env python32"""Generate structured synthetic ERA5-MODIS monthly pairs for an executable demo."""3 4from __future__ import annotations5 6import argparse7import sys8from pathlib import Path9 10import numpy as np11import yaml12 13ROOT = Path(__file__).resolve().parents[1]14sys.path.insert(0, str(ROOT / "model"))15from ml_modis import PRESSURE_LEVELS, PRESSURE_VARIABLES, SINGLE_FEATURES, feature_names, validate_multimodal_keys16 17 18def parse_args() -> argparse.Namespace:19    parser = argparse.ArgumentParser()20    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))21    parser.add_argument("--samples", type=int, default=None)22    parser.add_argument("--output", default=None)23    return parser.parse_args()24 25 26def ocean_mask(lat: np.ndarray, lon: np.ndarray) -> np.ndarray:27    """Analytic North Atlantic mask excluding coarse Greenland/Europe land shapes."""28    greenland = (lat > 59) & (lon > -53) & (lon < -20 + 0.55 * (lat - 59))29    europe = (lat > 50) & (lon > -10 + 0.35 * (lat - 50))30    iceland = (lat > 63) & (lat < 67) & (lon > -25) & (lon < -13)31    north_america = (lon < -52 + 0.3 * (lat - 45))32    return ~(greenland | europe | iceland | north_america)33 34 35def main() -> None:36    args = parse_args()37    config = yaml.safe_load(Path(args.config).read_text())38    n = int(args.samples or config["data"]["samples"])39    rng = np.random.default_rng(config["runtime"]["seed"])40    years = np.asarray(config["data"]["years"], dtype=np.int16)41    months = np.asarray(config["data"]["months"], dtype=np.int8)42    platforms = np.asarray(config["data"]["platforms"], dtype="U5")43 44    records = []45    used = set()46    while len(records) < n:47        year = int(rng.choice(years))48        month = int(rng.choice(months))49        platform = str(rng.choice(platforms))50        lat = int(rng.integers(45, 76))51        lon = int(rng.integers(-60, 31))52        key = (year, month, platform, lat, lon)53        if key in used or not ocean_mask(np.array([lat]), np.array([lon]))[0]:54            continue55        used.add(key)56        records.append(key)57    year = np.asarray([r[0] for r in records], dtype=np.int16)58    month = np.asarray([r[1] for r in records], dtype=np.int8)59    platform = np.asarray([r[2] for r in records], dtype="U5")60    lat = np.asarray([r[3] for r in records], dtype=np.float32)61    lon = np.asarray([r[4] for r in records], dtype=np.float32)62    hour = np.where(platform == "Terra", 11.0, 13.0).astype(np.float32)63 64    phase = np.deg2rad(lon + 25) + (month - 9) * 0.3565    maritime = np.cos(np.deg2rad(lat - 58)) * np.cos(np.deg2rad(lon + 25))66    synoptic = np.sin(phase * 1.7 + (year - 2001) * 0.43) + 0.45 * np.cos(np.deg2rad(lat * 3))67    sst = 286.0 - 0.42 * (lat - 45) + 1.1 * np.cos(phase) - 0.35 * (month - 9) + 0.025 * (year - 2001)68    surface_pressure = 101300 + 900 * synoptic - 8 * (lat - 55) + rng.normal(0, 160, n)69    humidity_base = np.clip(0.82 - 0.008 * (lat - 45) + 0.08 * maritime + 0.04 * synoptic, 0.35, 0.98)70    stability = 0.7 * (lat - 55) - 1.8 * synoptic + rng.normal(0, 0.7, n)71    x = np.empty((n, 114), dtype=np.float32)72    column = 073    for variable in PRESSURE_VARIABLES:74        for level in PRESSURE_LEVELS:75            z = (1000 - level) / 50.076            if variable == "temperature": value = sst - 1.7 - 3.15 * z + 0.15 * stability77            elif variable == "specific_humidity": value = 0.010 * humidity_base * np.exp(-0.23 * z)78            elif variable == "relative_humidity": value = np.clip(humidity_base - 0.025 * z + 0.04 * np.sin(phase + z), 0.05, 1.0)79            elif variable == "u_wind": value = 5 + 0.8 * z + 2.2 * np.sin(phase) + 0.12 * (lat - 55)80            elif variable == "v_wind": value = 1.5 + 1.6 * np.cos(phase * 1.3) - 0.25 * z81            elif variable == "omega": value = -0.025 * synoptic * np.exp(-0.08 * z)82            elif variable == "geopotential": value = z * 50 * 9.81 + 4 * synoptic83            elif variable == "cloud_liquid": value = np.maximum(0, 2.2e-4 * (humidity_base - 0.55) * np.exp(-0.18 * z))84            else: value = np.clip((humidity_base - 0.55) * 1.8 * np.exp(-0.12 * z), 0, 1)85            x[:, column] = value + rng.normal(0, max(float(np.std(value)) * 0.035, 1e-6), n)86            column += 187    cos_sza = np.clip(np.cos(np.deg2rad(lat - 20)) * (0.97 - 0.01 * (hour - 11)), 0, 1)88    singles = np.column_stack([89        sst, surface_pressure, surface_pressure + 35, sst - 0.4, sst - 1.1,90        sst - (1 - humidity_base) * 12, x[:, 30], x[:, 40], 190 * cos_sza,91        315 - 2.5 * (sst - 278), 65 + 18 * synoptic, 18 + 8 * stability,92        650 + 120 * humidity_base + 20 * synoptic, 16 + 30 * humidity_base,93        0.08 + 0.18 * np.maximum(synoptic, 0), 80 * np.maximum(synoptic, 0),94        -25 * np.maximum(-synoptic, 0), np.clip(0.25 + 0.45 * humidity_base + 0.05 * synoptic, 0, 1),95        np.clip((lat - 68) / 8, 0, 1), np.maximum(0, 1.8 + 1.5 * synoptic),96        cos_sza, lat, lon, hour,97    ]).astype(np.float32)98    x[:, 90:] = singles99 100    platform_term = np.where(platform == "Aqua", 1.0, -1.0)101    low_cloud = np.clip(0.22 + 0.55 * humidity_base + 0.035 * stability + 0.025 * synoptic, 0.05, 0.9)102    nd = 62 + 48 * humidity_base + 5 * synoptic + 0.32 * (lat - 55) + 1.8 * platform_term103    reff = 18.5 - 0.035 * nd + 0.055 * (sst - 278) - 0.10 * stability104    lwp = 58 + 115 * low_cloud + 10 * synoptic - 2.0 * stability105    cf = np.clip(low_cloud + 0.018 * platform_term, 0.03, 0.95)106 107    plume = np.exp(-((lat - 60) / 10) ** 2 - ((lon + 20) / 25) ** 2)108    eruption = (year == 2014).astype(np.float32) * (0.72 + 0.28 * (month == 10)) * plume109    nd *= 1 + 0.28 * eruption110    reff *= 1 - 0.08 * eruption111    lwp *= 1 + 0.008 * eruption112    cf = np.clip(cf * (1 + 0.11 * eruption), 0.01, 0.99)113    y = np.column_stack([114        nd + rng.normal(0, 3.0, n), reff + rng.normal(0, 0.28, n),115        lwp + rng.normal(0, 5.0, n), cf + rng.normal(0, 0.018, n),116    ]).astype(np.float32)117    y[:, 0:3] = np.maximum(y[:, 0:3], 1e-3)118    y[:, 3] = np.clip(y[:, 3], 0.001, 0.999)119 120    payload = {"X": x, "Y": y, "year": year, "month": month, "platform": platform,121               "platform_hour": hour, "latitude": lat, "longitude": lon,122               "feature_names": np.asarray(feature_names()), "target_names": np.asarray(config["data"]["variables"]["targets"]["names"]),123               "format_version": np.array(config["format_version"]),124               "is_ocean": np.ones(n, dtype=bool), "eruption_strength": eruption.astype(np.float32)}125    validate_multimodal_keys(payload)126    output = ROOT / (args.output or config["data"]["path"])127    output.parent.mkdir(parents=True, exist_ok=True)128    np.savez_compressed(output, **payload)129    print(f"output={output.relative_to(ROOT)} samples={n} shape={list(x.shape)} "130          f"eruption_samples={int((year == 2014).sum())}")131 132 133if __name__ == "__main__":134    main()135