Retrieving and Gridding PRISM Precipitation in Python

import logging
from pathlib import Path

import geopandas as gpd
import numpy as np
import pandas as pd
import rasterio
from rasterio.mask import mask as rio_mask

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


def basin_mean_series(
    grid_dir: str,
    catchment_path: str,
    pattern: str = "*.tif",
    date_from_name=None,
    stamp: str = "end",
    tz_shift_hours: float = 0.0,
) -> pd.Series:
    """
    Area-weighted basin-mean precipitation from a directory of grid files.

    Parameters
    ----------
    grid_dir       : Directory of single-timestep precipitation rasters.
    catchment_path : Catchment polygon, any CRS.
    pattern        : Glob for the raster files.
    date_from_name : Callable mapping a filename to a timestamp.
    stamp          : "end" if a file's timestamp marks the END of its
                     accumulation interval, "start" otherwise.
    tz_shift_hours : Hours to add to convert the product's clock to the
                     model's clock.

    Returns
    -------
    Series of basin-mean depth indexed by the START of each interval.
    """
    files = sorted(Path(grid_dir).glob(pattern))
    if not files:
        raise FileNotFoundError(f"no rasters matching {pattern} in {grid_dir}")

    catchment = gpd.read_file(catchment_path)
    records = []
    step_hours = None

    for i, path in enumerate(files):
        with rasterio.open(path) as src:
            geoms = catchment.to_crs(src.crs).geometry
            # --- all_touched=True brings in every partially covered cell;
            # the weights below then scale each by its true overlap, which is
            # what a centre-in-polygon rule cannot do. ---
            data, _ = rio_mask(src, geoms, crop=True, all_touched=True,
                               filled=True, nodata=np.nan)
            weights, _ = rio_mask(
                src, geoms, crop=True, all_touched=True, filled=True,
                nodata=0, indexes=1,
            )
            arr = data[0].astype("float64")

            # Weight is the fraction of each cell inside the polygon. Computing
            # it exactly needs a coverage rasteriser; a high-resolution
            # supersample is a good approximation and needs no extra dependency.
            frac = _coverage_fraction(src, geoms, arr.shape, supersample=8)
            valid = np.isfinite(arr) & (frac > 0)
            if not valid.any():
                log.warning("%s: no overlapping cells", path.name)
                continue
            mean = float(np.average(arr[valid], weights=frac[valid]))
            naive = float(np.nanmean(arr[valid]))

        ts = date_from_name(path.name) if date_from_name else pd.Timestamp(i, unit="h")
        records.append({"stamp": ts, "depth_mm": mean, "naive_mm": naive})

    df = pd.DataFrame(records).sort_values("stamp").reset_index(drop=True)
    if len(df) > 1:
        step_hours = (df["stamp"].iloc[1] - df["stamp"].iloc[0]).total_seconds() / 3600.0

    # --- Normalise to interval-START stamping. Reading an end-stamped product
    # as start-stamped shifts the whole hyetograph by one step, which shows up
    # in calibration as a timing error nobody attributes to the input. ---
    shift = pd.Timedelta(hours=tz_shift_hours)
    if stamp == "end" and step_hours:
        shift -= pd.Timedelta(hours=step_hours)
        log.info("Product is end-stamped at a %.1f h step — shifting the index "
                 "back to interval starts", step_hours)
    df["start"] = df["stamp"] + shift

    bias = 100.0 * (df["naive_mm"].sum() - df["depth_mm"].sum()) / max(
        df["depth_mm"].sum(), 1e-9)
    log.info("%d intervals, total %.1f mm basin-mean; a cell-centre average "
             "would differ by %+.1f %%", len(df), df["depth_mm"].sum(), bias)

    out = df.set_index("start")["depth_mm"]
    out.index.name = "interval_start"
    return out


def _coverage_fraction(src, geoms, shape, supersample: int = 8) -> np.ndarray:
    """Approximate the fraction of each cell inside the polygon."""
    from rasterio.features import rasterize
    from rasterio.transform import from_origin

    # Rasterize at N× resolution inside the same window, then block-average.
    win_transform = src.window_transform(
        rasterio.windows.from_bounds(*src.bounds, transform=src.transform))
    fine_transform = from_origin(
        win_transform.c, win_transform.f,
        abs(win_transform.a) / supersample, abs(win_transform.e) / supersample)
    fine = rasterize(
        [(g, 1) for g in geoms], fill=0, dtype="uint8",
        out_shape=(shape[0] * supersample, shape[1] * supersample),
        transform=fine_transform,
    )
    return fine.reshape(shape[0], supersample, shape[1], supersample).mean(axis=(1, 3))


# --- Example usage ---
# import re
# def name_to_ts(name):
#     m = re.search(r"(\d{8})", name)
#     return pd.Timestamp(m.group(1))
# series = basin_mean_series("prism_daily/", "catchment.gpkg",
#                            date_from_name=name_to_ts, stamp="end")
# print(series.head())