File size: 3,848 Bytes
7a7efc9 | 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 | """Write a lazy 569-station protocol manifest; fields are generated per chunk."""
from pathlib import Path
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
def station_metadata(count: int):
index = np.arange(count, dtype=np.float32)
latitude = -43.5 + 33.0 * ((index * 0.61803398875) % 1)
longitude = 113.0 + 40.0 * ((index * 0.41421356237) % 1)
elevation = 20 + 1450 * ((index * 0.2718281828) % 1)
return latitude, longitude, elevation.astype(np.float32)
def make_split(path: Path, config: dict):
data = config["data"]
dates = np.asarray(data["valid_dates"], dtype="U10")
lat, lon, elevation = station_metadata(data["station_count"])
history_dates = np.empty((len(dates), data["history_days"]), dtype="U10")
for i, date in enumerate(dates.astype("datetime64[D]")):
history_dates[i] = (date - np.arange(data["history_days"], 0, -1)).astype("U10")
np.savez(path, format_version=np.asarray(data["format_version"]), valid_dates=dates,
history_dates=history_dates, initialization_utc=np.asarray("1200"), lead_hours=np.arange(241),
station_id=np.asarray([f"JIVE-{i:04d}" for i in range(data["station_count"])]),
station_latitude=lat, station_longitude=lon, station_elevation_m=elevation,
variables=np.asarray(data["variables"]), units=np.asarray(data["units"]), sources=np.asarray(data["sources"]),
neighborhood_shape=np.asarray([3, 3]), target_grid_projection=np.asarray("Albers"),
target_grid_nx=np.asarray(None), target_grid_ny=np.asarray(None),
representation=np.asarray(data["target_representation"]))
def generate_chunk(meta, date_index: int, station_selector, history: bool, seed: int, include_patch: bool = True):
"""Emulate lazy reads after authoritative station patches have been extracted."""
station = np.arange(len(meta["station_id"]))[station_selector]
days = 30 if history else 1
lead = np.arange(241, dtype=np.float32)[None, :, None, None]
day = np.arange(days, dtype=np.float32)[:, None, None, None]
lat = meta["station_latitude"][station][None, None, None, :]
lon = meta["station_longitude"][station][None, None, None, :]
phase = date_index * 0.7 + day * 0.11
temp = 20 - 0.35 * (lat + 25) + 4 * np.sin(2 * np.pi * (lead + 12) / 24 + phase)
dew = temp - 5 - 1.5 * np.cos(np.deg2rad(lon) + lead / 48)
wind = 5 + 1.2 * np.abs(np.sin(np.deg2rad(lon) + lead / 18 + phase))
truth = np.concatenate((temp, dew, wind), axis=2).astype(np.float32)
source = np.arange(3, dtype=np.float32)[None, None, None, :, None]
forecast = truth[:, :, :, None] + (source - 0.6) * np.asarray([1.1, 0.8, 0.5], np.float32)[None, None, :, None, None]
forecast += (lead[..., None] / 240) * np.asarray([0.8, -0.5, 0.6], np.float32)[None, None, :, None, None]
rng = np.random.default_rng(seed + date_index * 1000 + int(station[0]))
forecast += rng.normal(0, 0.08, forecast.shape).astype(np.float32)
offsets = np.asarray([[-0.18, -0.10, -0.04], [-0.08, 0.0, 0.09], [0.03, 0.12, 0.20]], np.float32)
elevation_delta = meta["station_elevation_m"][station] - np.mean(meta["station_elevation_m"])
if not include_patch:
return forecast.astype(np.float32), truth, elevation_delta.astype(np.float32)
patches = forecast[..., None, None] + offsets
return patches.astype(np.float32), truth, elevation_delta.astype(np.float32)
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
output = ROOT / config["data"]["root"]
output.mkdir(parents=True, exist_ok=True)
make_split(output / "protocol.npz", config)
print("generated=data/protocol.npz dates=2 history_days=30 leads=241 stations=569 patch=3x3 nx=unknown ny=unknown")
if __name__ == "__main__":
main()
|