GOES lightning detections

from datetime import UTC, datetime
from pathlib import Path

import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr
from matplotlib.patches import Rectangle

import usdata
from usdata import cite_lockfile, pull, verify

# 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")
print("Executed (UTC):", datetime.now(UTC).isoformat(timespec="seconds"))
print(f"usdata {usdata.__version__}; xarray {xr.__version__}; pandas {pd.__version__}")

print(manifest.read_text())

result = pull(manifest)
items = result.fetched
sizes = [item.provenance.size for item in items]
print("Files:", len(items))
print(f"Bytes: {sum(sizes):,} ({min(sizes):,} to {max(sizes):,} per file)")
print("First:", items[0].path.name)
print("Last:", items[-1].path.name)
print("Covers:", items[0].asset.time.start.isoformat(), "to", items[-1].asset.time.end.isoformat())
print("Source:", items[0].provenance.source_url.rsplit("/", 1)[0] + "/")
retrieved = [item.provenance.retrieved_at for item in items]
print("Retrieved (UTC):", min(retrieved).isoformat(timespec="seconds"))
print("Checksums:", sum(item.provenance.checksum.startswith("sha256:") for item in items), "sha256")

first = items[0].open()
print({name: size for name, size in first.sizes.items() if name.startswith("number_of")})
print("Flash coordinates:", sorted(name for name in first.coords if name.startswith("flash")))
print("Flash variables:", sorted(name for name in first.data_vars if name.startswith("flash_")))
print(
    "Energy:", first["flash_energy"].attrs["units"], "· area:", first["flash_area"].attrs["units"]
)
first.close()

columns = ["flash_time_offset_of_first_event", "flash_lat", "flash_lon", "flash_quality_flag"]
tables = []
for item in items:
    detections = item.open()
    tables.append(detections[columns].reset_coords()[columns].to_dataframe())
    detections.close()
flashes = pd.concat(tables, ignore_index=True).rename(
    columns={"flash_time_offset_of_first_event": "first_event_utc"}
)
print("Flashes in the hour, whole field of view:", len(flashes))
print("First-event times:", flashes.first_event_utc.min(), "to", flashes.first_event_utc.max())
print("Quality flags:", flashes.flash_quality_flag.value_counts().sort_index().to_dict())
print(
    f"Latitude {flashes.flash_lat.min():.1f} to {flashes.flash_lat.max():.1f} °N; "
    f"longitude {flashes.flash_lon.min():.1f} to {flashes.flash_lon.max():.1f} °E"
)
flashes.head(3)

REPORT_UTC = pd.Timestamp("2024-05-07 04:39")
REPORT_LAT, REPORT_LON = 35.378, -97.543
HALF = 1.0
region = flashes[flashes.flash_lat.between(28, 44) & flashes.flash_lon.between(-110, -80)]
fig, ax = plt.subplots(layout="constrained")
bins = ax.hexbin(
    region.flash_lon,
    region.flash_lat,
    gridsize=(90, 30),
    bins="log",
    mincnt=1,
    cmap="YlOrRd",
    extent=(-110, -80, 28, 44),
)
ax.add_patch(
    Rectangle(
        (REPORT_LON - HALF, REPORT_LAT - HALF),
        2 * HALF,
        2 * HALF,
        fill=False,
        edgecolor="#1f2937",
        linewidth=1.2,
    )
)
ax.plot(REPORT_LON, REPORT_LAT, marker="x", color="#1f2937", markersize=6)
ax.annotate(
    "tornado report",
    (REPORT_LON + HALF, REPORT_LAT - HALF),
    xytext=(4, -12),
    textcoords="offset points",
    fontsize=9,
)
ax.set_aspect(1 / np.cos(np.radians(36)))
ax.set(
    xlim=(-110, -80),
    ylim=(28, 44),
    xlabel="Longitude (°E)",
    ylabel="Latitude (°N)",
    title="GOES-16 GLM flashes, 03:40 to 04:39 UTC, 7 May 2024",
)
fig.colorbar(bins, ax=ax, label="Flashes per bin")
plt.show()

north_south = flashes.flash_lat.between(REPORT_LAT - HALF, REPORT_LAT + HALF)
east_west = flashes.flash_lon.between(REPORT_LON - HALF, REPORT_LON + HALF)
storm = flashes[north_south & east_west]
cells = flashes.groupby(
    [
        np.floor((flashes.flash_lat - (REPORT_LAT - HALF)) / 2),
        np.floor((flashes.flash_lon - (REPORT_LON - HALF)) / 2),
    ]
).size()
rank = int((cells > cells.loc[(0.0, 0.0)]).sum()) + 1
print(f"Flashes in the box: {len(storm):,} of {len(flashes):,} ({len(storm) / len(flashes):.1%})")
print(f"The box ranks {rank} of {len(cells)} two-degree cells that hold any flash")
print(f"Busiest cell: {cells.max():,} flashes ({cells.max() / len(flashes):.1%})")

window = pd.date_range("2024-05-07 03:40", "2024-05-07 04:39", freq="min")


def per_minute(frame: pd.DataFrame) -> pd.Series:
    """Flash counts for each whole minute of the window, zero-filled."""
    minute = frame.first_event_utc.dt.floor("min")
    return frame.groupby(minute).size().reindex(window, fill_value=0)


counts = pd.DataFrame({"box": per_minute(storm), "field_of_view": per_minute(flashes)})
fig, ax = plt.subplots(figsize=(8, 3), layout="constrained")
ax.plot(counts.index, counts.box, color="#b45631")
ax.axvline(REPORT_UTC, color="#444444", linewidth=0.9, linestyle="--")
ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M"))
ax.set(
    ylim=(0, None),
    xlabel="UTC, 7 May 2024",
    ylabel="Flashes per minute",
    title="Flashes in the box; the dashed line is the report",
)
plt.show()
early, late = counts.head(10).mean(), counts.tail(10).mean()
change = 100 * (late - early) / early
print(f"Box: {early.box:.0f} per minute in the first ten minutes, {late.box:.0f} in the last")
print(f"Change: box {change.box:.0f}%, whole field of view {change.field_of_view:.0f}%")
print(
    "Flashes whose first event precedes the window:", len(flashes) - int(counts.field_of_view.sum())
)

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