NEXRAD derived radar products

import struct
from datetime import UTC, date, datetime, timedelta
from pathlib import Path

import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd

import usdata
from usdata import cite_lockfile, pull, verify
from usdata.readers import UnsupportedFormat

# 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__}; pandas {pd.__version__}")

print(manifest.read_text())

result = pull(manifest)
files = pd.DataFrame(
    {
        "product": item.asset.id.split("_")[1],
        "scan_utc": item.asset.time.start,
        "bytes": item.provenance.size,
        "file": item.path.name,
    }
    for item in result.fetched
)
print("Files:", len(files), "· bytes:", f"{files.bytes.sum():,}")
print("Source:", result.fetched[0].provenance.source_url.rsplit("/", 1)[0] + "/")
print(
    "Retrieved (UTC):",
    min(i.provenance.retrieved_at for i in result.fetched).isoformat(timespec="seconds"),
)
print("First checksum:", result.fetched[0].provenance.checksum)
files.groupby("product").agg(
    files=("file", "size"),
    first=("scan_utc", "min"),
    last=("scan_utc", "max"),
    min_bytes=("bytes", "min"),
    max_bytes=("bytes", "max"),
)

try:
    result.fetched[0].open()
except UnsupportedFormat as error:
    print(error)

def level3_header(path: Path) -> dict:
    """Identity fields from a Level III file's WMO header and product description block."""
    data = path.read_bytes()
    wmo, product_line, rest = data.split(b"\r\r\n", 2)
    code, _, _, length = struct.unpack(">hhii", rest[:12])
    fields = struct.unpack(">hiihhhhhhhi", rest[18:46])
    _, lat, lon, height, _, mode, vcp, _, volume, scan_day, scan_seconds = fields
    scan = datetime.combine(date(1969, 12, 31), datetime.min.time(), UTC) + timedelta(
        days=scan_day, seconds=scan_seconds
    )
    return {
        "wmo": wmo.decode(),
        "product_line": product_line.decode(),
        "code": code,
        "message_bytes": length,
        "radar_lat": lat / 1000,
        "radar_lon": lon / 1000,
        "height_ft": height,
        "mode": {1: "clear air", 2: "precipitation"}.get(mode, mode),
        "vcp": vcp,
        "volume_number": volume,
        "volume_start": scan,
    }


headers = pd.DataFrame(level3_header(item.path) for item in result.fetched)
headers["name_matches_header"] = headers.volume_start.eq(files.scan_utc)
headers.drop_duplicates("code")[
    ["wmo", "product_line", "code", "radar_lat", "radar_lon", "mode", "vcp", "volume_start"]
]

print("Every file's name matches its header's scan start:", bool(headers.name_matches_header.all()))
print(
    "Operating modes:", headers["mode"].unique().tolist(), "· VCPs:", headers.vcp.unique().tolist()
)
print("File bytes minus message bytes:", sorted((files.bytes - headers.message_bytes).unique()))

colors = {"EET": "#2563a6", "NMD": "#b45631", "NST": "#4d7c0f"}
fig, ax = plt.subplots(layout="constrained")
for row, (product, group) in enumerate(files.groupby("product")):
    ax.scatter(
        group.scan_utc,
        [row] * len(group),
        s=group.bytes / 12,
        color=colors[product],
        alpha=0.8,
        zorder=3,
    )
    for scan, size in zip(group.scan_utc, group.bytes, strict=True):
        ax.annotate(
            f"{size:,} B",
            (scan, row),
            xytext=(0, 14),
            textcoords="offset points",
            ha="center",
            fontsize=8,
        )
switch = headers.loc[headers["mode"] == "precipitation", "volume_start"].min()
start = pd.Timestamp("2024-05-06 20:30", tz="UTC")
ax.axvspan(start, switch, color="#9ca3af", alpha=0.15)
ax.text(start, 2.65, "  clear-air mode, VCP 35", fontsize=8, va="center")
ax.text(switch, 2.65, "  precipitation mode, VCP 212", fontsize=8, va="center")
ax.set_yticks(range(3), ["EET echo tops", "NMD mesocyclones", "NST storm tracks"])
ax.set_ylim(-0.6, 2.8)
ax.set_xlim(start, pd.Timestamp("2024-05-06 21:01", tz="UTC"))
ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M", tz=UTC))
ax.set(
    xlabel="Volume scan start (UTC, 6 May 2024)",
    title="KTLX Level III files in the archive; marker area by file size",
)
plt.show()

volumes = files.join(headers[["volume_number", "vcp", "mode"]])
print(volumes.loc[volumes["product"] == "EET", ["scan_utc", "volume_number", "vcp", "mode"]])
print("NMD and NST begin at", volumes.loc[volumes["product"] != "EET", "scan_utc"].min())
print("NMD sizes:", sorted(int(b) for b in files.loc[files["product"] == "NMD", "bytes"].unique()))

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