SPC convective and fire weather outlook areas

from datetime import UTC, datetime
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd
import shapefile
from matplotlib.patches import PathPatch
from matplotlib.path import Path as Outline

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

print(manifest.read_text())

result = pull(manifest)
for name in ("layers", "cutters"):
    item = result.one(name)
    print(f"{name}: {item.path.name}, {item.provenance.size:,} bytes")
print("Source:", result.one("layers").provenance.source_url)
print("Transformations:", *result.one("layers").provenance.transformations, sep="\n  ")

def read_zip(path: Path, times: list[str]) -> pd.DataFrame:
    """One row per shape: the DBF attributes, the named UTC times parsed, and the shape."""
    with shapefile.Reader(str(path)) as reader:
        names = [field[0] for field in reader.fields[1:]]
        frame = pd.DataFrame([list(record) for record in reader.records()], columns=names)
        frame["shape"] = reader.shapes()
    for column in times:
        frame[column] = pd.to_datetime(
            frame[column], format="%Y%m%d%H%M", utc=True, errors="coerce"
        )
    return frame


def patch(shape, **style) -> PathPatch | None:
    """A matplotlib patch for a polygon shape, its parts as rings so holes stay holes."""
    if not shape.points:
        return None
    vertices, codes = [], []
    for start, stop in zip(shape.parts, [*shape.parts[1:], len(shape.points)], strict=True):
        ring = shape.points[start:stop]
        vertices += ring
        codes += [Outline.MOVETO] + [Outline.LINETO] * (len(ring) - 1)
    return PathPatch(Outline(vertices, codes), **style)


TIMES = ["ISSUE", "EXPIRE", "PRODISS"]
layers = read_zip(result.one("layers").path, TIMES)
cutters = read_zip(result.one("cutters").path, TIMES)
issuances = layers.groupby("PRODISS").agg(
    cycle=("CYCLE", "first"),
    valid_from=("ISSUE", "first"),
    valid_to=("EXPIRE", "first"),
    areas=("CATEGORY", "size"),
)
issuances

moment = pd.Timestamp("2024-05-07T02:12Z")
known = layers[layers["PRODISS"] <= moment]
latest = known[known["PRODISS"] == known["PRODISS"].max()]
print(
    "Latest Day 1 outlook issued before", f"{moment:%Y-%m-%d %H:%M} UTC:", latest["PRODISS"].iloc[0]
)
latest.groupby("CATEGORY")["THRESHOLD"].agg(lambda values: ", ".join(values))

COLORS = {
    "TSTM": "#c1e9c1",
    "MRGL": "#66a366",
    "SLGT": "#f6f67f",
    "ENH": "#e6c27f",
    "MDT": "#e67f7f",
    "HIGH": "#ff80ff",
}
categorical = cutters[(cutters["CATEGORY"] == "CATEGORICAL") & (cutters["CYCLE"] >= 0)]
cycles = sorted(categorical["CYCLE"].unique())
fig, axes = plt.subplots(1, len(cycles), figsize=(11, 2.8), layout="constrained")
for ax, cycle in zip(axes, cycles, strict=True):
    outlook = categorical[categorical["CYCLE"] == cycle]
    for threshold, color in COLORS.items():
        for shape in outlook.loc[outlook["THRESHOLD"] == threshold, "shape"]:
            if (area := patch(shape, facecolor=color, edgecolor="0.4", linewidth=0.3)) is not None:
                ax.add_patch(area)
    ax.set_xlim(-125, -66)
    ax.set_ylim(24, 50)
    ax.set_aspect(1 / 0.77)
    ax.set_xticks([])
    ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set(visible=True, color="0.8")
    ax.set_title(f"{outlook['PRODISS'].iloc[0]:%H:%M} UTC", fontsize=9)
handles = [plt.Rectangle((0, 0), 1, 1, facecolor=color) for color in COLORS.values()]
fig.legend(handles, COLORS, loc="outside lower center", ncols=len(COLORS), frameon=False)
fig.suptitle("SPC Day 1 categorical outlooks issued on 6 May 2024")
plt.show()

tornado = layers[layers["CATEGORY"] == "TORNADO"]
numeric = tornado[tornado["THRESHOLD"] != "SIGN"].assign(
    probability=lambda frame: frame["THRESHOLD"].astype(float)
)
summary = numeric.groupby("PRODISS").agg(cycle=("CYCLE", "first"), highest=("probability", "max"))
summary["significant"] = tornado.groupby("PRODISS")["THRESHOLD"].apply(lambda t: "SIGN" in set(t))
summary

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