Which severe reports came with rotation and lightning?

import gc
import tempfile
import time
from datetime import UTC, datetime, timedelta, timezone
from importlib.metadata import version
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pyproj import Geod

from usdata import build_query, cite_lockfile, get, pull, select_by_time, verify
from usdata.providers import load_adapter

# One figure style for every usdata notebook, so previews look alike.
plt.rcParams.update(
    {
        "figure.figsize": (8, 4.5),
        "figure.dpi": 120,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.grid": True,
        "grid.alpha": 0.3,
        "font.size": 10,
    }
)
manifest = Path("dataset.yaml")
notebook_started = time.perf_counter()
print("Executed (UTC):", datetime.now(UTC).isoformat(timespec="seconds"))
print(
    "; ".join(
        f"{name} {version(name)}"
        for name in ("usdata", "pandas", "xarray", "xradar", "eccodes", "pyproj")
    )
)

print(manifest.read_text().strip())

started = time.perf_counter()
result = pull(manifest)
by_dataset: dict[str, list] = {}
for item in result.fetched:
    by_dataset.setdefault(item.asset.dataset_id, []).append(item)
print("\nPull seconds:", round(time.perf_counter() - started, 2))
print("Total source bytes:", sum(item.provenance.size for item in result.fetched))
summary = pd.DataFrame(
    {
        "dataset": dataset,
        "files": len(items),
        "bytes": sum(item.provenance.size for item in items),
        "first": min(item.asset.id for item in items),
        "retrieved_utc": max(item.provenance.retrieved_at for item in items).isoformat(
            timespec="seconds"
        ),
    }
    for dataset, items in by_dataset.items()
)
assert set(summary.dataset) == {
    "noaa:storm-events",
    "noaa:nexrad-level2",
    "noaa:mrms",
    "noaa:goes-glm",
}
summary

columns = [
    "EVENT_ID",
    "EVENT_TYPE",
    "STATE",
    "CZ_NAME",
    "BEGIN_DATE_TIME",
    "END_DATE_TIME",
    "CZ_TIMEZONE",
    "BEGIN_LAT",
    "BEGIN_LON",
    "TOR_F_SCALE",
    "MAGNITUDE",
]
frame = by_dataset["noaa:storm-events"][0].open_csv(usecols=columns)
reports = frame.loc[
    frame.STATE.eq("OKLAHOMA") & frame.EVENT_TYPE.isin(["Tornado", "Hail", "Thunderstorm Wind"])
].copy()
if set(reports.CZ_TIMEZONE) != {"CST-6"}:
    raise ValueError(f"Review source timezones before matching: {set(reports.CZ_TIMEZONE)}")
local = pd.to_datetime(reports.BEGIN_DATE_TIME, format="%d-%b-%y %H:%M:%S")
reports["begin_utc"] = local.dt.tz_localize(timezone(timedelta(hours=-6))).dt.tz_convert("UTC")
assert (reports.begin_utc == reports.BEGIN_UTC).all(), "the reader derived a different UTC time"
reports["BEGIN_LAT"] = reports.BEGIN_LAT.astype(float)
reports["BEGIN_LON"] = reports.BEGIN_LON.astype(float)
assert reports[["BEGIN_LAT", "BEGIN_LON"]].notna().all().all()

case = reports.loc[reports.EVENT_ID.eq("1184052")].iloc[0]
event_time = case.begin_utc.to_pydatetime()
lat, lon = float(case.BEGIN_LAT), float(case.BEGIN_LON)
print(
    "Case:", case[["EVENT_ID", "EVENT_TYPE", "CZ_NAME", "BEGIN_DATE_TIME", "TOR_F_SCALE"]].to_dict()
)
print("Case start UTC:", event_time.isoformat())

window = reports.loc[
    reports.begin_utc.between("2024-05-07T04:30:00Z", "2024-05-07T04:42:00Z")
].sort_values(["begin_utc", "EVENT_ID"])
selected = window.head(12).reset_index(drop=True)
assert set(selected.EVENT_TYPE) == {"Tornado", "Hail", "Thunderstorm Wind"}
assert case.EVENT_ID in set(selected.EVENT_ID)
print("Reports in the 04:30-04:42 UTC window:", len(window), "· selected:", len(selected))
print(selected.EVENT_TYPE.value_counts().to_dict())
del frame
gc.collect();

tolerance = timedelta(minutes=5)
locked = {dataset: {item.asset.id for item in items} for dataset, items in by_dataset.items()}
checks = []
for dataset_id, params in [
    ("noaa:nexrad-level2", {"site": "KTLX"}),
    ("noaa:mrms", {"product": "RotationTrackML30min_00.50"}),
    ("noaa:goes-glm", {"satellite": 16}),
]:
    query = build_query(start=event_time - tolerance, end=event_time + tolerance, **params)
    with load_adapter(get(dataset_id)) as provider:
        candidates = provider.list_assets(query)
    if dataset_id == "noaa:goes-glm":
        missing = {asset.id for asset in candidates} - locked[dataset_id]
        assert not missing, f"GLM files listed live but not locked: {sorted(missing)[:3]}"
        checks.append({"dataset": dataset_id, "candidates": len(candidates), "locked_match": "all"})
        continue
    selection = select_by_time(
        candidates, target=event_time, tolerance=tolerance, direction="nearest"
    )
    assert selection.asset is not None, f"No start within tolerance: {dataset_id}"
    assert selection.asset.id in locked[dataset_id], "Revisit the manifest selection"
    checks.append(
        {
            "dataset": dataset_id,
            "candidates": selection.candidate_count,
            "locked_match": selection.asset.id,
            "start_offset_seconds": selection.offset_seconds,
        }
    )
pd.DataFrame(checks)

radar_item = by_dataset["noaa:nexrad-level2"][0]
radar = radar_item.open_nexrad(sweep=0)
sweep = radar["sweep_0"].to_dataset()
assert sweep.DBZH.attrs["units"] == "dBZ"
radar_lat, radar_lon = float(radar.latitude), float(radar.longitude)
geod = Geod(ellps="WGS84")
azimuth = np.deg2rad(sweep.azimuth.values)[:, None]
slant = sweep.range.values[None, :]
gate_x, gate_y = slant * np.sin(azimuth), slant * np.cos(azimuth)
reflectivity = sweep.DBZH.values


def reflectivity_max(point_lat: float, point_lon: float, radius_km: float = 25.0) -> float:
    """Maximum first-sweep DBZH within a flat-plane ground radius of the point."""
    bearing, _, distance = geod.inv(radar_lon, radar_lat, point_lon, point_lat)
    px, py = distance * np.sin(np.deg2rad(bearing)), distance * np.cos(np.deg2rad(bearing))
    inside = np.hypot(gate_x - px, gate_y - py) <= radius_km * 1000
    assert inside.any(), "report lies outside the sweep"
    return float(np.nanmax(np.where(inside, reflectivity, np.nan)))


started = time.perf_counter()
rotation_crops: dict[datetime, object] = {}
for item in by_dataset["noaa:mrms"]:
    grid = item.open()
    (name,) = grid.data_vars
    crop = grid[name].sel(latitude=slice(37.2, 33.4), longitude=slice(257.0, 266.0)).load()
    rotation_crops[item.asset.time.start] = crop.where(crop >= 0)
    del grid, crop
    gc.collect()
print(
    "MRMS grids decoded and cropped:",
    len(rotation_crops),
    "in",
    round(time.perf_counter() - started, 1),
    "s",
)


def rotation_max(
    when: datetime, point_lat: float, point_lon: float, half: float = 0.1
) -> tuple[datetime, float]:
    """Nearest grid stamp and the maximum rotation track inside a ±half-degree box."""
    stamp = min(rotation_crops, key=lambda start: (abs((start - when).total_seconds()), start))
    box = rotation_crops[stamp].sel(
        latitude=slice(point_lat + half, point_lat - half),
        longitude=slice(point_lon + 360 - half, point_lon + 360 + half),
    )
    return stamp, float(box.max()) if int(box.count()) else float("nan")


flashes: dict[datetime, tuple[np.ndarray, np.ndarray]] = {}
for item in by_dataset["noaa:goes-glm"]:
    detections = item.open()
    flashes[item.asset.time.start] = (
        detections.flash_lat.values.copy(),
        detections.flash_lon.values.copy(),
    )
    del detections
print(
    "GLM files:",
    len(flashes),
    "· flashes in files:",
    sum(len(pair[0]) for pair in flashes.values()),
)


def flash_count(when: datetime, point_lat: float, point_lon: float, half: float = 0.25) -> int:
    """GLM flashes within ±half degrees in files starting within the tolerance of when."""
    count = 0
    for start, (flash_lat, flash_lon) in flashes.items():
        if abs((start - when).total_seconds()) <= tolerance.total_seconds():
            near = (np.abs(flash_lat - point_lat) <= half) & (np.abs(flash_lon - point_lon) <= half)
            count += int(near.sum())
    return count


stamp, case_rotation = rotation_max(event_time, lat, lon)
case_row = {
    "event_id": case.EVENT_ID,
    "utc_time": event_time.isoformat(),
    "lat": lat,
    "lon": lon,
    "label": "tornado",
    "ef": case.TOR_F_SCALE,
    "refl_max_dbz": reflectivity_max(lat, lon),
    "rotation_max": case_rotation,
    "rotation_grid_utc": stamp.strftime("%H:%M"),
    "flash_count": flash_count(event_time, lat, lon),
}
print(
    "Report ground distance from KTLX (km):",
    round(geod.inv(radar_lon, radar_lat, lon, lat)[2] / 1000, 1),
)
del radar, sweep
gc.collect()
pd.Series(case_row)

rows = []
for report in selected.itertuples(index=False):
    when = report.begin_utc.to_pydatetime()
    stamp, rotation = rotation_max(when, report.BEGIN_LAT, report.BEGIN_LON)
    rows.append(
        {
            "event_id": report.EVENT_ID,
            "utc_time": when.strftime("%H:%M"),
            "county": report.CZ_NAME,
            "label": report.EVENT_TYPE,
            "ef_or_magnitude": report.TOR_F_SCALE
            if report.EVENT_TYPE == "Tornado"
            else report.MAGNITUDE,
            "lat": round(report.BEGIN_LAT, 3),
            "lon": round(report.BEGIN_LON, 3),
            "rotation_max": rotation,
            "rotation_grid_utc": stamp.strftime("%H:%M"),
            "flash_count": flash_count(when, report.BEGIN_LAT, report.BEGIN_LON),
        }
    )
labeled = pd.DataFrame(rows)
assert len(labeled) == 12 and labeled.rotation_max.notna().all()
print(
    labeled.groupby("label")[["rotation_max", "flash_count"]].agg(["min", "max", "mean"]).round(1)
)
labeled

colors = {"Tornado": "#d62728", "Hail": "#1f77b4", "Thunderstorm Wind": "#7f7f7f"}
case_stamp, _ = rotation_max(event_time, lat, lon)
margin = 0.2
west, east = selected.BEGIN_LON.min() - margin, selected.BEGIN_LON.max() + margin
south, north = selected.BEGIN_LAT.min() - margin, selected.BEGIN_LAT.max() + margin
crop = rotation_crops[case_stamp].sel(
    latitude=slice(north, south), longitude=slice(west + 360, east + 360)
)
fig, (ax_map, ax_features) = plt.subplots(
    1, 2, layout="constrained", gridspec_kw={"width_ratios": [1.5, 1]}
)
mesh = ax_map.pcolormesh(
    crop.longitude.values - 360,
    crop.latitude.values,
    crop.values,
    cmap="magma_r",
    shading="auto",
    vmin=0,
    vmax=25,
)
ax_map.scatter(radar_lon, radar_lat, marker="^", s=70, c="white", edgecolors="black", zorder=3)
ax_map.annotate("KTLX", (radar_lon, radar_lat), xytext=(5, -10), textcoords="offset points")
for label, color in colors.items():
    group = selected.loc[selected.EVENT_TYPE.eq(label)]
    ax_map.scatter(
        group.BEGIN_LON,
        group.BEGIN_LAT,
        marker="*" if label == "Tornado" else "o",
        s=180 if label == "Tornado" else 30,
        c=color,
        edgecolors="black",
        zorder=4,
    )
ax_map.grid(False)
ax_map.set(
    xlim=(west, east),
    ylim=(south, north),
    xlabel="Longitude (°E)",
    ylabel="Latitude (°N)",
    title=f"MRMS rotation track, {case_stamp:%H:%M} UTC",
    aspect=1 / np.cos(np.deg2rad(lat)),
)
fig.colorbar(
    mesh, ax=ax_map, label="30-min azimuthal shear (0.001 s$^{-1}$)", shrink=0.7, location="bottom"
)
for label, color in colors.items():
    group = labeled.loc[labeled.label.eq(label)]
    ax_features.scatter(
        group.flash_count,
        group.rotation_max,
        marker="*" if label == "Tornado" else "o",
        s=180 if label == "Tornado" else 45,
        c=color,
        edgecolors="black",
        label=label,
        zorder=3,
    )
ax_features.set(
    xlabel="GLM flashes (±0.25°, ±5 min)",
    ylabel="Rotation max (0.001 s$^{-1}$)",
    title="Twelve reports, two features",
)
ax_features.legend(fontsize=8, loc="upper center")
fig.suptitle("Oklahoma severe reports, 04:30 to 04:42 UTC on 7 May 2024")
plt.show()
plt.close(fig)
del crop
gc.collect();

assert verify(manifest) == []
for citation in cite_lockfile(manifest):
    print(citation.as_text())
    print()

again = pull(manifest)
assert again.from_lockfile and all(item.from_cache for item in again.fetched)
assert verify(manifest) == []
with tempfile.TemporaryDirectory(prefix="usdata-tornado-restore-") as directory:
    started = time.perf_counter()
    restored = pull(manifest, root=Path(directory))
    assert restored.from_lockfile and not any(item.from_cache for item in restored.fetched)
    assert verify(manifest, root=Path(directory)) == []
    assert sorted(item.provenance.checksum for item in restored.fetched) == sorted(
        item.provenance.checksum for item in result.fetched
    )
    print("Empty-cache restoration seconds:", round(time.perf_counter() - started, 2))
print(
    f"All {len(result.fetched)} inputs verified; cached and empty-cache locked restoration passed."
)
print("Notebook seconds:", round(time.perf_counter() - notebook_started, 1))