MRMS gridded radar products

import gc
import sys
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 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
files = pd.DataFrame(
    {
        "stamp_utc": item.asset.time.start.strftime("%H:%M"),
        "file": item.path.name,
        "bytes": item.provenance.size,
        "checksum": item.provenance.checksum[:19] + "…",
    }
    for item in items
)
print("Files:", len(files), "· bytes:", f"{files.bytes.sum():,}")
print("Source:", items[0].provenance.source_url.rsplit("/", 1)[0] + "/")
print("Retrieved (UTC):", items[0].provenance.retrieved_at.isoformat(timespec="seconds"))
print("Asset time of the first file:", items[0].asset.time.start, "to", items[0].asset.time.end)
files.head(3)

def peak_memory_gb() -> float:
    """High-water resident memory of this process, or NaN where it cannot be read."""
    try:
        import resource
    except ImportError:  # Windows
        return float("nan")
    scale = 1e9 if sys.platform == "darwin" else 1e6  # bytes on macOS, kilobytes on Linux
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / scale


before = peak_memory_gb()
grid = items[0].open()
(name,) = grid.data_vars
field = grid[name]
values = field.values
print("Variable:", name, "· units attribute:", field.attrs.get("units"))
print("Shape:", field.shape, "·", field.dtype, f"· {field.nbytes / 1e9:.2f} GB decoded")
print(f"Peak memory: {before:.2f} GB before the open, {peak_memory_gb():.2f} GB after")
print(f"Latitude {float(field.latitude[0]):.4f} to {float(field.latitude[-1]):.4f}")
print(f"Longitude {float(field.longitude[0]):.4f} to {float(field.longitude[-1]):.4f}")
print("NaN:", int(np.isnan(values).sum()), "· negative:", int((values < 0).sum()))
print(f"Above zero: {int((values > 0).sum()):,} of {values.size:,} points")
print("Whole numbers only:", bool(np.array_equal(values, np.round(values))))
del grid, field, values
_ = gc.collect()

REPORT_LAT, REPORT_LON = 35.378, -97.543
crops = []
for item in items:
    grid = item.open()
    (name,) = grid.data_vars
    crop = grid[name].sel(
        latitude=slice(REPORT_LAT + 1, REPORT_LAT - 1),
        longitude=slice(REPORT_LON + 360 - 2, REPORT_LON + 360 + 2),
    )
    crops.append(crop.load().expand_dims(stamp=[item.asset.time.start.replace(tzinfo=None)]))
    del grid
    gc.collect()
tracks = xr.concat(crops, dim="stamp")
tracks = tracks.assign_coords(longitude=tracks.longitude - 360)
print("Cropped:", dict(tracks.sizes))
print(f"Peak memory after eleven decodes: {peak_memory_gb():.2f} GB")

swath = tracks.max("stamp")
fig, ax = plt.subplots(layout="constrained")
mesh = ax.pcolormesh(
    swath.longitude,
    swath.latitude,
    swath.where(swath > 0),
    cmap="YlOrRd",
    vmin=0,
    vmax=22,
    shading="nearest",
    rasterized=True,
)
ax.plot(REPORT_LON, REPORT_LAT, marker="*", markersize=13, color="#2563a6", linestyle="none")
ax.annotate(
    "tornado report 1184052, 04:39",
    (REPORT_LON, REPORT_LAT),
    xytext=(8, -14),
    textcoords="offset points",
    fontsize=9,
    color="#2563a6",
)
ax.set_aspect(1 / np.cos(np.radians(REPORT_LAT)))
ax.set(
    xlabel="Longitude (°E)",
    ylabel="Latitude (°N)",
    title="MRMS mid-level rotation tracks, 04:00 to 04:50 UTC, 7 May 2024",
)
fig.colorbar(mesh, ax=ax, label="Azimuthal shear, 3 to 6 km (0.001 s⁻¹)")
plt.show()

def distance_km(lat, lon):
    """Great-circle distance from the report, on a sphere of radius 6371 km."""
    lat1, lat2 = np.radians(REPORT_LAT), np.radians(lat)
    dlat, dlon = lat2 - lat1, np.radians(lon - REPORT_LON)
    h = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
    return 2 * 6371 * np.arcsin(np.sqrt(h))


def peak(box: xr.DataArray) -> dict:
    """The box maximum, how many points tie at it, and their mean position."""
    rows, cols = np.nonzero(box.values >= float(box.max()))
    lat = float(box.latitude.values[rows].mean())
    lon = float(box.longitude.values[cols].mean())
    return {"max": float(box.max()), "tied": len(rows), "km": round(distance_km(lat, lon), 1)}


box = tracks.sel(
    latitude=slice(REPORT_LAT + 0.25, REPORT_LAT - 0.25),
    longitude=slice(REPORT_LON - 0.25, REPORT_LON + 0.25),
)
near = pd.DataFrame(
    [peak(box.isel(stamp=i)) for i in range(box.sizes["stamp"])],
    index=pd.DatetimeIndex(box.stamp.values).strftime("%H:%M").rename("stamp_utc"),
)
strongest = swath.where(swath == swath.max(), drop=True)
print(
    f"Strongest in the crop: {float(swath.max()):.0f} at "
    f"{distance_km(float(strongest.latitude.mean()), float(strongest.longitude.mean())):.0f} km "
    "from the report"
)
near.T

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