Rendering KDE surfaces to GeoTIFF and web tiles

import numpy as np
import rasterio
from rasterio.transform import from_origin
from rasterio.enums import Resampling
from rasterio.warp import calculate_default_transform, reproject


def write_density_geotiff(
    density: np.ndarray,
    x_min: float,
    y_max: float,
    cell_m: float,
    crs: str,
    path: str,
    nodata: float = -9999.0,
) -> dict:
    """
    Write a density array as a correctly georeferenced GeoTIFF.

    Parameters
    ----------
    density : np.ndarray
        (rows, cols) array, row 0 at the TOP. If it was built with a
        bottom-up y axis — as meshgrid output often is — flip it first, or
        the map is mirrored and every hotspot moves to the wrong side.
    x_min, y_max : float
        Coordinates of the top-left CORNER of the top-left cell, not its
        centre. Passing the centre shifts the whole layer half a cell.
    nodata : float
        A distinct value for "not estimated". Zero density is a measurement;
        outside the study area is not, and a viewer cannot tell them apart
        unless they are different values.

    Raises
    ------
    ValueError
        On a non-2D array, a non-positive cell size, or a missing CRS.
    """
    if density.ndim != 2:
        raise ValueError(f"density must be 2-D; got shape {density.shape}.")
    if cell_m <= 0:
        raise ValueError("cell_m must be positive.")
    if not crs:
        raise ValueError("A CRS is required; an unreferenced raster is not a map.")

    transform = from_origin(x_min, y_max, cell_m, cell_m)   # corner-anchored
    data = np.where(np.isfinite(density), density, nodata).astype("float32")

    profile = {
        "driver": "GTiff", "height": data.shape[0], "width": data.shape[1],
        "count": 1, "dtype": "float32", "crs": crs, "transform": transform,
        "nodata": nodata, "compress": "deflate", "predictor": 3,
        "tiled": True, "blockxsize": 512, "blockysize": 512,
    }
    with rasterio.open(path, "w", **profile) as dst:
        dst.write(data, 1)
        # Overviews are what make a large surface usable in a viewer; without
        # them every pan reads the full-resolution band.
        dst.build_overviews([2, 4, 8, 16], Resampling.average)
        dst.update_tags(units="events per km2", cell_size_m=cell_m)

    return {"path": path, "shape": data.shape, "transform": transform}


def reproject_for_tiles(src_path: str, dst_path: str,
                        dst_crs: str = "EPSG:3857") -> dict:
    """
    Reproject to Web Mercator for tiling.

    Resampling.average is correct for a RATE surface being downsampled — it
    preserves the mean. Use Resampling.sum only for count rasters, and never
    use nearest for either: it drops cells entirely at reduced resolution.
    """
    with rasterio.open(src_path) as src:
        transform, width, height = calculate_default_transform(
            src.crs, dst_crs, src.width, src.height, *src.bounds
        )
        profile = src.profile.copy()
        profile.update(crs=dst_crs, transform=transform, width=width, height=height)

        with rasterio.open(dst_path, "w", **profile) as dst:
            reproject(
                source=rasterio.band(src, 1),
                destination=rasterio.band(dst, 1),
                src_transform=src.transform, src_crs=src.crs,
                dst_transform=transform, dst_crs=dst_crs,
                resampling=Resampling.average,
                src_nodata=src.nodata, dst_nodata=src.nodata,
            )
    return {"path": dst_path, "crs": dst_crs, "shape": (height, width)}


def global_breaks(density: np.ndarray, n_classes: int = 6,
                  method: str = "quantile") -> list[float]:
    """
    Class breaks computed ONCE over the whole surface.

    Recomputing per tile makes adjacent tiles use different scales, so a
    hotspot appears or disappears at a tile boundary — the single most
    common visual bug in a published density tile set.
    """
    v = density[np.isfinite(density) & (density > 0)]
    if v.size == 0:
        raise ValueError("No positive density values to classify.")
    if method == "quantile":
        qs = np.linspace(0, 100, n_classes + 1)[1:-1]
        return [float(x) for x in np.percentile(v, qs)]
    if method == "linear":
        return [float(x) for x in np.linspace(v.min(), v.max(), n_classes + 1)[1:-1]]
    raise ValueError(f"Unknown method {method!r}; expected quantile or linear.")

def validate_raster(path: str, expected_total: float | None = None,
                    reference_point: tuple | None = None) -> dict:
    """Three checks that catch georeferencing and resampling errors."""
    with rasterio.open(path) as ds:
        assert ds.crs is not None, "raster has no CRS — it is an image, not a map"
        assert ds.nodata is not None, "no nodata value: zero and missing are conflated"
        band = ds.read(1, masked=True)

        # 1. Orientation. A north-up raster has a NEGATIVE y pixel size.
        assert ds.transform.e < 0, (
            "positive y pixel size — the array is bottom-up and the map is "
            "vertically mirrored"
        )

        # 2. Mass conservation across the write, if the caller knows the total.
        if expected_total is not None:
            cell_area_km2 = (ds.transform.a * -ds.transform.e) / 1e6
            total = float(band.sum() * cell_area_km2)
            err = abs(total - expected_total) / expected_total
            assert err < 0.02, f"total density off by {err:.1%} after writing"

        # 3. A known point lands where it should. This is the only check that
        #    catches a half-cell offset.
        if reference_point is not None:
            x, y, expected_row, expected_col = reference_point
            row, col = ds.index(x, y)
            assert (row, col) == (expected_row, expected_col), (
                f"reference point maps to ({row}, {col}) not "
                f"({expected_row}, {expected_col}) — check corner vs centre anchoring"
            )
    return {"crs": str(ds.crs), "shape": band.shape}