NEXRAD radar scans

from datetime import UTC, datetime
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr
import xradar

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__}; xradar {xradar.__version__}; xarray {xr.__version__}")

print(manifest.read_text())

result = pull(manifest)
(item,) = result.fetched
print("File:", item.path.name)
print("Bytes:", item.provenance.size)
print("Source:", item.provenance.source_url)
print("Retrieved (UTC):", item.provenance.retrieved_at.isoformat(timespec="seconds"))
print("Checksum:", item.provenance.checksum)
print("Volume start:", item.asset.time.start.isoformat())

radar = item.open_nexrad(sweep=0)
sweep = radar["/sweep_0"].to_dataset()
moments = [name for name, data in sweep.data_vars.items() if "range" in data.dims]
print(
    "Radar:",
    radar.attrs["instrument_name"],
    f"at {float(radar['latitude']):.3f} °N,",
    f"{float(radar['longitude']):.3f} °E",
)
print("Elevation cuts in the volume:", radar.attrs["actual_elevation_cuts"])
print(f"Sweep 0: {float(sweep.sweep_fixed_angle):.2f}° elevation,", dict(sweep.sizes))
print(f"Range: {float(sweep.range.min()) / 1000:.1f} to {float(sweep.range.max()) / 1000:.1f} km")
print("Ray times:", str(sweep.time.min().values)[:19], "to", str(sweep.time.max().values)[:19])
pd.DataFrame(
    {
        "long_name": [sweep[name].attrs.get("long_name") for name in moments],
        "units": [sweep[name].attrs.get("units") for name in moments],
        "dtype": [str(sweep[name].dtype) for name in moments],
        "valid_gates": [int(sweep[name].count()) for name in moments],
    },
    index=pd.Index(moments, name="moment"),
)

def edges(centres: np.ndarray) -> np.ndarray:
    """Cell edges halfway between centres, extended half a step at each end."""
    middle = (centres[:-1] + centres[1:]) / 2
    return np.concatenate([[2 * centres[0] - middle[0]], middle, [2 * centres[-1] - middle[-1]]])


ordered = sweep.sortby("azimuth")
azimuth = np.deg2rad(edges(ordered.azimuth.values))[:, None]
ground = edges(ordered.range.values)[None, :] * np.cos(np.deg2rad(float(sweep.sweep_fixed_angle)))
east, north = np.sin(azimuth) * ground / 1000, np.cos(azimuth) * ground / 1000
reflectivity = ordered["DBZH"]
correlation = ordered["RHOHV"].where(reflectivity.notnull())

fig, (left, right) = plt.subplots(1, 2, sharey=True, layout="constrained")
panels = [
    (left, reflectivity, "turbo", (-10, 60), "Reflectivity (dBZ)"),
    (right, correlation, "viridis", (0.5, 1.0), "Correlation coefficient"),
]
for ax, values, cmap, (low, high), label in panels:
    mesh = ax.pcolormesh(east, north, values, cmap=cmap, vmin=low, vmax=high, rasterized=True)
    fig.colorbar(mesh, ax=ax, label=label, orientation="horizontal", shrink=0.9)
    ax.plot(0, 0, marker="+", color="black")
    ax.set_aspect("equal")
    ax.set(xlim=(-350, 150), ylim=(-200, 300), xlabel="East of radar (km)")
left.set_ylabel("North of radar (km)")
fig.suptitle(
    f"KTLX lowest sweep, {item.asset.time.start:%H:%M} UTC 6 May 2024: storms far to the "
    "northwest, clear-air echo overhead"
)
plt.show()

near = ordered.range <= 150_000
far = ordered.range >= 200_000
near_echo = reflectivity.where(near)
near_correlation = correlation.where(near)
strong_far = correlation.where(far & (reflectivity >= 40))
print(f"Echo within 150 km: {int(near_echo.count()):,} gates, max {float(near_echo.max()):.1f} dBZ")
print(f"  below 15 dBZ: {float((near_echo < 15).sum() / near_echo.count()):.0%}")
low = float((near_correlation < 0.9).sum() / near_correlation.count())
print(f"  correlation below 0.9: {low:.0%}")
high = float((strong_far >= 0.95).sum() / strong_far.count())
print(f"Echo of 40 dBZ or more beyond 200 km: {int(strong_far.count()):,} gates")
print(f"  correlation of 0.95 or more: {high:.0%}")
print(f"Strongest echo: {float(reflectivity.max()):.1f} dBZ")

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