File size: 2,807 Bytes
78dea75 | 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 83 | #!/usr/bin/env python3
"""Aggregate GHCN-Daily station files into weather-heavy JSONL documents."""
from __future__ import annotations
import argparse
import sys
from collections import defaultdict
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from tqdm import tqdm
from weather_llm.config_loader import load_yaml, repo_root_from
from weather_llm.data.jsonl_utils import write_jsonl
from weather_llm.data.noaa import (
aggregate_monthly,
iter_ghcn_daily_csv,
monthly_docs_for_state,
resolve_noaa_states_and_max_per_state,
station_id_to_state_map,
state_rollup_doc,
)
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--config", type=Path, default=None)
args = ap.parse_args()
root = repo_root_from(ROOT)
cfg_path = args.config or (root / "configs/data_default.yaml")
cfg = load_yaml(cfg_path)
noaa = cfg["noaa"]
paths = cfg["paths"]
raw_dir = root / paths["raw_noaa"]
out_path = root / paths["processed"] / "weather.jsonl"
files = sorted(raw_dir.glob("US*.csv.gz")) + sorted(raw_dir.glob("US*.csv"))
if not files:
raise SystemExit(f"No station CSV(.gz) files under {raw_dir}. Run scripts/download_noaa.py first.")
states_cfg, _ = resolve_noaa_states_and_max_per_state(noaa)
fallback_state = states_cfg[0]
inv_path = raw_dir / "ghcnd-stations.txt"
if inv_path.is_file():
sid_to_state = station_id_to_state_map(inv_path.read_text(encoding="utf-8", errors="replace"))
else:
sid_to_state = {}
all_docs: list[dict] = []
seed = int(cfg["merge"]["seed"])
y0, y1 = int(noaa["year_start"]), int(noaa["year_end"])
monthly_global: dict = {}
stations_by_state: dict[str, list[str]] = defaultdict(list)
for fp in tqdm(files, desc="stations"):
sid = fp.name.split(".")[0]
st = sid_to_state.get(sid, fallback_state)
stations_by_state[st].append(sid)
daily = iter_ghcn_daily_csv(fp)
monthly = aggregate_monthly(daily, y0, y1)
for k, v in monthly.items():
monthly_global[k] = v
docs = monthly_docs_for_state(st, [sid], monthly, y0, y1, seed)
all_docs.extend(docs)
# One July rollup per state represented in the download (correct geography in text)
for st in sorted(stations_by_state.keys()):
sids = stations_by_state[st]
sub_seed = seed + sum(ord(c) for c in st) * 1_000_003
rollup = state_rollup_doc(st, sids, monthly_global, y0, y1, month=7, seed=sub_seed)
if rollup:
all_docs.append(rollup)
write_jsonl(out_path, iter(all_docs))
print("Wrote", len(all_docs), "documents to", out_path)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|