Map_Postprocess.py

"""Local post-processing of the wall-to-wall VM inference outputs (see Wall_to_Wall_Maps.ipynb
Steps 4-5): mosaic the predicted-class / confidence tiles that come back from GCS into one
raster per year, and visualize the class map.

Kept separate from Map_Export.py on purpose: that module is the Earth Engine side (AOI +
embedding export) and imports `ee`; these are local rasterio operations on the downloaded
tiles, so they carry different (and heavier) dependencies. Nothing here touches Earth Engine.
"""
import os

import geopandas as gpd
import numpy as np
import rasterio
import rasterio.warp
from rasterio.enums import Resampling
from rasterio.windows import Window

# NODATA value written by vm_predict.py (outside both the 1-7 class range and 0-100 confidence
# range), and the Level-1 class palette. The tuned model predicts [1, 3, 4, 5, 6, 7]; there is
# no Ice/Snow (id 2), and glaciers are masked out, so id 2 should not appear in a valid map.
NODATA = 255
CLASS_COLORS = {
    1: '#3a7ca5',  # Water
    2: '#d9f0f7',  # Ice/Snow (not predicted; here for completeness)
    3: '#ab2b2b',  # Developed
    4: '#d2b48c',  # Barren/Sparse
    5: '#1b7837',  # Trees
    6: '#c2a55c',  # Shrubs
    7: '#a6d96a',  # Herbaceous
}


def mosaic_tiles(tifs, out_path, nodata=NODATA):
    """Mosaic grid-aligned, non-overlapping GeoTIFF tiles into one file, memory-safe.

    Streams each tile into its window of the output, one at a time, so a state-sized mosaic
    never loads fully into RAM (rasterio.merge would build the whole array). The AlphaEarth
    export tiles share one aligned grid and do not overlap, so windowed writes reconstruct the
    mosaic exactly. The output is first filled with `nodata` in row strips so any gap not
    covered by a tile reads as nodata rather than the default 0 (0 is a valid confidence value,
    so this matters for the confidence mosaic).

    Args:
        tifs: List of tile paths (all sharing one CRS, resolution, and aligned grid).
        out_path: Output GeoTIFF path.
        nodata: NODATA value for gaps and for the output profile.

    Returns:
        (out_path, (height, width)) of the written mosaic.
    """
    with rasterio.open(tifs[0]) as s0:
        res_x, res_y = s0.res
        crs, dtype, count = s0.crs, s0.dtypes[0], s0.count
    bounds = []
    for f in tifs:                               # metadata only -- cheap
        with rasterio.open(f) as s:
            bounds.append(s.bounds)
    left, top = min(b.left for b in bounds), max(b.top for b in bounds)
    right, bottom = max(b.right for b in bounds), min(b.bottom for b in bounds)
    width = int(round((right - left) / res_x))
    height = int(round((top - bottom) / res_y))
    transform = rasterio.transform.from_origin(left, top, res_x, res_y)
    profile = dict(driver='GTiff', height=height, width=width, count=count, dtype=dtype, crs=crs,
                   transform=transform, nodata=nodata, compress='deflate', predictor=2,
                   tiled=True, bigtiff='if_safer')
    with rasterio.open(out_path, 'w', **profile) as dst:
        for r in range(0, height, 4096):         # initialise the grid to nodata in row strips
            h = min(4096, height - r)
            dst.write(np.full((count, h, width), nodata, dtype=dtype), window=Window(0, r, width, h))
        for f in tifs:                           # paint each tile into its aligned window
            with rasterio.open(f) as s:
                col = int(round((s.bounds.left - left) / res_x))
                row = int(round((top - s.bounds.top) / res_y))
                dst.write(s.read(), window=Window(col, row, s.width, s.height))
    return out_path, (height, width)


def find_map(map_dir, year):
    """Path to the {year} class map in map_dir, or None.

    Prefers the clipped LULC_Class1_forest_{year}.tif (Step 4b), falling back to the
    _mosaic.tif from Step 4a.
    """
    for name in (f'LULC_Class1_forest_{year}.tif', f'LULC_Class1_forest_{year}_mosaic.tif'):
        p = os.path.join(map_dir, name)
        if os.path.exists(p):
            return p
    return None


def find_confidence(map_dir, year):
    """Path to the {year} confidence map in map_dir, or None.

    Prefers the clipped confidence_{year}.tif, falling back to the confidence_{year}_mosaic.tif
    from Step 4a.
    """
    for name in (f'confidence_{year}.tif', f'confidence_{year}_mosaic.tif'):
        p = os.path.join(map_dir, name)
        if os.path.exists(p):
            return p
    return None


def confidence_by_class(class_path, conf_path, nodata=NODATA):
    """Per-class histogram of confidence values (0-100), computed block-by-block.

    Returns {class_id: np.ndarray(101)}, where element k is the count of pixels of that class
    whose confidence equals k. Memory-safe (one block at a time). The class and confidence
    rasters must share the grid -- vm_predict.py writes both per tile, so they do. Pixels where
    either raster is nodata are skipped.
    """
    hist = {}
    with rasterio.open(class_path) as cs, rasterio.open(conf_path) as fs:
        if (cs.width, cs.height) != (fs.width, fs.height):
            raise ValueError(f'class and confidence rasters differ in size: '
                             f'{(cs.width, cs.height)} vs {(fs.width, fs.height)}.')
        for _, window in cs.block_windows(1):
            cls = cs.read(1, window=window).ravel()
            conf = fs.read(1, window=window).ravel()
            valid = (cls != nodata) & (conf != nodata)
            if not valid.any():
                continue
            cv, fv = cls[valid], conf[valid]
            for c in np.unique(cv):
                h = np.bincount(fv[cv == c], minlength=101)[:101]
                hist[c] = hist.get(c, np.zeros(101, dtype=np.int64)) + h
    return hist


def class_areas(path, nodata=NODATA, strip_rows=1024):
    """Class histogram of a class raster -> {class_id: hectares}, computed in row strips.

    Never loads the whole raster into RAM. Assumes an equal-area CRS (EPSG:5070), where every
    pixel is a true fixed area, so a plain pixel count times pixel-area is unbiased.

    Reads horizontal strips rather than the file's native 256x256 blocks: a state-sized 10 m
    raster holds roughly 59,000 such blocks and the per-read overhead dominates, turning a
    seconds-long histogram into a multi-minute one. At the default a full state is about 50
    reads costing ~70 MB each. The result is identical either way; lower strip_rows if memory
    is tight.
    """
    counts = np.zeros(256, dtype=np.int64)
    with rasterio.open(path) as src:
        for row in range(0, src.height, strip_rows):
            h = min(strip_rows, src.height - row)
            window = Window(0, row, src.width, h)
            counts += np.bincount(src.read(1, window=window).ravel(), minlength=256)
        px_ha = abs(src.transform.a * src.transform.e) / 1e4  # EPSG:5070 -> 0.01 ha/pixel at 10 m
    return {c: counts[c] * px_ha for c in range(256) if counts[c] and c != nodata}


def hex_to_rgb(hexc):
    """'#rrggbb' -> [r, g, b] floats in [0, 1] for matplotlib."""
    return [int(hexc[i:i + 2], 16) / 255 for i in (1, 3, 5)]


def read_overview(path, max_dim=2000):
    """Decimated (nearest-neighbour) read of band 1 -> 2D array, for display of a big raster.

    Nearest-neighbour so categorical classes are not blurred; decimated so a state-sized
    raster stays light in memory. Two rasters on the same grid read with the same max_dim
    decimate to the same shape, so their overviews can be compared pixel-for-pixel.
    """
    with rasterio.open(path) as src:
        scale = max(1, round(max(src.width, src.height) / max_dim))
        return src.read(1, out_shape=(src.height // scale, src.width // scale),
                        resampling=Resampling.nearest)


def overview_rgb(path, max_dim=2000, class_colors=None):
    """Decimated class raster rendered to an (H, W, 3) RGB float image for imshow.

    Pixels with no colour (nodata, unmapped ids) render white.
    """
    class_colors = class_colors or CLASS_COLORS
    arr = read_overview(path, max_dim)
    rgb = np.ones((*arr.shape, 3))               # white where nodata / no colour
    for cid, hexc in class_colors.items():
        m = arr == cid
        if m.any():
            rgb[m] = hex_to_rgb(hexc)
    return rgb


def clip_to_geometry(src_path, out_path, geoms, nodata=NODATA):
    """Set pixels outside `geoms` to nodata, block-by-block (memory-safe).

    `geoms` is an iterable of shapely geometries already in the raster's CRS. Used to trim a
    mosaic to a boundary (e.g. clip an Oregon run's mosaic to the Oregon polygon) without
    loading the whole state-sized raster into RAM. The output keeps the input's grid/profile.
    """
    from rasterio.features import geometry_mask
    with rasterio.open(src_path) as src:
        profile = src.profile
        profile.update(nodata=nodata)
        with rasterio.open(out_path, 'w', **profile) as dst:
            for _, window in src.block_windows(1):
                data = src.read(window=window)
                outside = geometry_mask(geoms, out_shape=(window.height, window.width),
                                        transform=src.window_transform(window))  # True outside geoms
                data[:, outside] = nodata
                dst.write(data, window=window)
    return out_path


def _stratum_groups(stratum_ids):
    """(unique stratum ids, list of position-arrays per stratum, sizes) for 1-D integer labels."""
    stratum_ids = np.asarray(stratum_ids)
    uniq, inv = np.unique(stratum_ids, return_inverse=True)
    order = np.argsort(inv, kind='stable')
    bounds = np.searchsorted(inv[order], np.arange(len(uniq) + 1))
    groups = [order[bounds[k]:bounds[k + 1]] for k in range(len(uniq))]
    return uniq, groups, np.array([len(g) for g in groups])


def _allocate(sizes, n, floor):
    """Per-stratum counts summing to n: a `floor` each, then the remainder by sqrt(size)."""
    n = int(n)
    alloc = np.minimum(floor, sizes).astype(int)          # floor (or the whole stratum if smaller)
    if alloc.sum() < n:                                   # distribute the remainder by sqrt(size)
        cap = sizes - alloc
        w = np.sqrt(sizes) * (cap > 0)
        if w.sum() > 0:
            alloc += np.minimum(np.floor((n - alloc.sum()) * w / w.sum()).astype(int), cap)
        for k in np.argsort(-np.sqrt(sizes)):             # hand out leftover, largest strata first
            while alloc.sum() < n and alloc[k] < sizes[k]:
                alloc[k] += 1
            if alloc.sum() >= n:
                break
    elif alloc.sum() > n:                                 # floors alone exceed n: trim smallest strata
        for k in np.argsort(np.sqrt(sizes)):
            while alloc.sum() > n and alloc[k] > 0:
                alloc[k] -= 1
            if alloc.sum() <= n:
                break
    return alloc


def nested_choice(n_total, k, seed=0):
    """First `k` positions of a seeded permutation of range(n_total) -- a *nested* random draw.

    Unlike `rng.choice(n_total, size=k, replace=False)`, the result for a smaller k is always a
    prefix of the result for a larger k, so growing a sample only ADDS points. (Measured: choice()
    with size=400 vs 600 from the same seed shares *zero* points, which would silently invalidate
    every cache keyed on the previous draw.) Returns a sorted int array."""
    k = int(min(k, n_total))
    return np.sort(np.random.default_rng(seed).permutation(int(n_total))[:k])


def stratified_sample(stratum_ids, n, seed=0, floor=1):
    """Select ~n positions from `stratum_ids` (1-D integer stratum labels), reproducibly.

    Allocation is sqrt(size) across strata with a per-stratum `floor`, then uniform random
    sampling within each stratum. sqrt is the middle ground between proportional (dominated by
    the biggest strata) and equal (over-weights tiny noise strata), so every present category is
    represented while larger strata still get more points. Deterministic for a given `seed`.
    Returns a sorted int array of selected positions into `stratum_ids` (length min(n, total)).

    **The draw is nested in `n`**: each stratum is ordered by a permutation seeded on
    `[seed, stratum_id]` -- independent of `n` and of the other strata -- and the allocation takes
    a prefix of it, so increasing `n` only appends points. This is what lets the reference sample
    grow (3,000 -> 4,000) while every cache keyed on the earlier point ids stays valid. The older
    `rng.choice` form was not nested: growing 3,000 -> 4,000 retained only ~5% of the original
    points, which would have discarded the whole COLD extraction.
    """
    uniq, groups, sizes = _stratum_groups(stratum_ids)
    alloc = _allocate(sizes, min(n, sizes.sum()), floor)

    # per-stratum permutation seeded on [seed, stratum_id]: ordering is independent of `n` and of
    # every other stratum, so a larger `n` takes a longer prefix of the same order (nested draw)
    picks = [np.random.default_rng([seed, int(uniq[k])]).permutation(groups[k])[:int(alloc[k])]
             for k in range(len(uniq)) if alloc[k] > 0]
    return np.sort(np.concatenate(picks)) if picks else np.array([], dtype=int)


def stratified_topup(stratum_ids, existing_positions, n_target, seed=0, floor=1):
    """Positions to ADD so a sample reaches `n_target` while keeping `existing_positions` intact.

    Needed because the reference sample already on disk was drawn with the older (non-nested)
    sampler, so a fresh `stratified_sample(..., n_target)` would not contain it. This instead
    treats the existing points as fixed and draws only the shortfall, per stratum, from cells not
    already used -- so previously extracted points (and every cache keyed on their ids) survive.

    Per-stratum target is the `n_target` allocation; a stratum already over its target keeps all
    its existing points (no removals), and the remaining budget is redistributed over the strata
    that are short. The result is close to, but not identical to, a single stratified draw of
    `n_target` -- worth noting if per-stratum inclusion probabilities are ever needed for
    design-based (Olofsson) area estimation.

    Returns a sorted int array of NEW positions, disjoint from `existing_positions`.
    """
    stratum_ids = np.asarray(stratum_ids)
    existing = np.asarray(list(existing_positions), dtype=np.int64)
    uniq, groups, sizes = _stratum_groups(stratum_ids)
    n_add = int(min(n_target - len(existing), sizes.sum() - len(existing)))
    if n_add <= 0:
        return np.array([], dtype=int)

    used = np.zeros(len(stratum_ids), dtype=bool)
    used[existing] = True
    have = np.array([int(used[g].sum()) for g in groups])          # existing points per stratum
    free = sizes - have                                            # still-available cells
    want = np.maximum(_allocate(sizes, min(n_target, sizes.sum()), floor) - have, 0)
    want = np.minimum(want, free)

    # reconcile to exactly n_add: trim from / add to the strata with the most headroom
    while want.sum() > n_add:
        k = int(np.argmax(want))
        want[k] -= 1
    for k in np.argsort(-np.sqrt(sizes)):
        while want.sum() < n_add and want[k] < free[k]:
            want[k] += 1
        if want.sum() >= n_add:
            break

    picks = []
    for k in range(len(uniq)):
        if want[k] > 0:
            perm = np.random.default_rng([seed, int(uniq[k])]).permutation(groups[k])
            picks.append(perm[~used[perm]][:int(want[k])])
    return np.sort(np.concatenate(picks)) if picks else np.array([], dtype=int)


def sample_rasters_at_points(points_gdf, path_by_name):
    """Sample raster value(s) at each point. Returns a DataFrame indexed like `points_gdf`.

    `path_by_name` is {output_column: raster_path}. Points are reprojected to each raster's CRS
    before sampling (so a lon/lat GeoDataFrame samples an EPSG:5070 raster correctly). Missing
    paths yield an all-NaN column.
    """
    import pandas as pd
    out = pd.DataFrame(index=points_gdf.index)
    for name, path in path_by_name.items():
        if not path or not os.path.exists(path):
            out[name] = np.nan
            continue
        with rasterio.open(path) as ds:
            pts = points_gdf.to_crs(ds.crs)
            vals = [v[0] for v in ds.sample([(g.x, g.y) for g in pts.geometry])]
        out[name] = np.asarray(vals, dtype=float)
    return out


def transition_confidence_bands(from_path, to_path, conf_from_path, conf_to_path, edges,
                                nodata=NODATA):
    """Per from->to transition, area (ha) in each confidence band, computed block-by-block.

    The confidence of a change is the JOINT min(confidence_from, confidence_to): a change is
    only as trustworthy as the less-confident of its two years. `edges` are the band
    boundaries, e.g. [0, 50, 70, 80, 90, 101] gives bands [<50, 50-70, 70-80, 80-90, >=90].
    Returns {(from_id, to_id): np.ndarray(len(edges)-1)} of hectares per band. Memory-safe; the
    four rasters must share one grid (the clipped class + confidence maps do).
    """
    edges = np.asarray(edges)
    nbands = len(edges) - 1
    out = {}
    with rasterio.open(from_path) as sa, rasterio.open(to_path) as sb, \
            rasterio.open(conf_from_path) as ca, rasterio.open(conf_to_path) as cb:
        sizes = {(s.width, s.height) for s in (sa, sb, ca, cb)}
        if len(sizes) != 1:
            raise ValueError(f'class and confidence rasters are not the same grid: {sizes}.')
        px_ha = abs(sa.transform.a * sa.transform.e) / 1e4
        for _, w in sa.block_windows(1):
            a, b = sa.read(1, window=w), sb.read(1, window=w)
            cf, ct = ca.read(1, window=w), cb.read(1, window=w)
            ch = (a != b) & (a != nodata) & (b != nodata) & (cf != nodata) & (ct != nodata)
            if not ch.any():
                continue
            band = np.clip(np.digitize(np.minimum(cf, ct), edges) - 1, 0, nbands - 1)
            fa, ta, ba = a[ch], b[ch], band[ch]
            for fc, tc in set(zip(fa.tolist(), ta.tolist())):
                sel = (fa == fc) & (ta == tc)
                key = (int(fc), int(tc))
                out.setdefault(key, np.zeros(nbands, dtype=np.int64))
                out[key] += np.bincount(ba[sel], minlength=nbands)[:nbands]
    return {k: v * px_ha for k, v in out.items()}


def transition_in_geometry(from_path, to_path, geoms, from_class=5, nodata=NODATA):
    """Area of a from_class -> (not from_class) transition, and how much falls inside `geoms`.

    Block-by-block (memory-safe). `geoms` are shapely polygons already in the rasters' CRS
    (e.g. fire perimeters). Used to check how much detected deforestation (Trees -> non-Trees)
    coincides with mapped fire footprints. Returns (transition_ha, transition_inside_geoms_ha).
    Passing a single unioned (Multi)Polygon in `geoms` is much faster than many small ones.
    """
    from rasterio.features import geometry_mask
    trans_px = inside_px = 0
    with rasterio.open(from_path) as sa, rasterio.open(to_path) as sb:
        px_ha = abs(sa.transform.a * sa.transform.e) / 1e4
        for _, window in sa.block_windows(1):
            a = sa.read(1, window=window)
            b = sb.read(1, window=window)
            trans = (a == from_class) & (b != from_class) & (a != nodata) & (b != nodata)
            if not trans.any():
                continue
            trans_px += int(trans.sum())
            inside = ~geometry_mask(geoms, out_shape=a.shape,
                                    transform=sa.window_transform(window))  # True inside geoms
            inside_px += int((trans & inside).sum())
    return trans_px * px_ha, inside_px * px_ha


def show_map(ax, rgb, raster_path, boundary=None, title=None):
    """Draw a decimated RGB overview on `ax` in its raster's real projected extent.

    Placing the overview at the raster's projected bounds (rather than raw pixel coordinates)
    lets an overlaid boundary line up. If `boundary` is given (a GeoDataFrame or GeoSeries, e.g.
    the Oregon outline), it is reprojected to the raster CRS, outlined, and the axes are framed
    to its full extent, so the whole state shows even where the raster only covers part of it
    (the uncovered area reads as the blank axes background inside the outline).

    boundary carries its own geopandas methods, so this stays free of a geopandas import.
    """
    with rasterio.open(raster_path) as src:
        b, crs = src.bounds, src.crs
    ax.imshow(rgb, extent=(b.left, b.right, b.bottom, b.top))
    if boundary is not None:
        gb = boundary.to_crs(crs)
        gb.boundary.plot(ax=ax, color='black', linewidth=0.8)
        minx, miny, maxx, maxy = gb.total_bounds
        ax.set_xlim(minx, maxx)
        ax.set_ylim(miny, maxy)
    else:
        ax.set_xlim(b.left, b.right)
        ax.set_ylim(b.bottom, b.top)
    ax.set_aspect('equal')
    ax.axis('off')
    if title:
        ax.set_title(title, fontsize=12)


def transition_areas(path_from, path_to, nodata=NODATA):
    """From->to class-transition areas between two aligned class rasters, block-by-block.

    Returns {(from_id, to_id): hectares}, over pixels that are valid (not nodata) in both
    rasters. Includes the no-change pairs (from == to); filter those out for a change table.
    Memory-safe: reads one block window from each raster at a time, so state-sized rasters
    never load fully into RAM.

    Requires the two rasters to share the same grid (same width/height and transform), which
    holds when both years were exported with the same AOI, CRS, scale, and tiling. Assumes an
    equal-area CRS (EPSG:5070), so a pixel count times pixel-area is an unbiased area.
    """
    counts = {}
    with rasterio.open(path_from) as sa, rasterio.open(path_to) as sb:
        if (sa.width, sa.height) != (sb.width, sb.height):
            raise ValueError(f'Rasters are not on the same grid: '
                             f'{(sa.width, sa.height)} vs {(sb.width, sb.height)}.')
        px_ha = abs(sa.transform.a * sa.transform.e) / 1e4
        for _, window in sa.block_windows(1):
            a = sa.read(1, window=window).ravel().astype(np.int32)
            b = sb.read(1, window=window).ravel().astype(np.int32)
            valid = (a != nodata) & (b != nodata)
            if not valid.any():
                continue
            codes = a[valid] * 256 + b[valid]    # encode (from, to) as one integer
            u, c = np.unique(codes, return_counts=True)
            for code, cnt in zip(u.tolist(), c.tolist()):
                key = (code // 256, code % 256)
                counts[key] = counts.get(key, 0) + cnt
    return {k: v * px_ha for k, v in counts.items()}


def read_decimated_grid(paths, target_pixels=9000):
    """Read several same-extent rasters' first band at one shared decimated resolution.

    The scale factor is chosen from `paths[0]` so max(width, height) is about `target_pixels`;
    every path is then read at that shared (H, W) with nearest-neighbour resampling, so a
    handful of coarse full-state arrays fit comfortably in memory for point sampling (see
    Wall_to_Wall_Maps.ipynb Step 8).

    Returns:
        (arrays, transform, crs) -- `arrays` is a list of 2-D arrays aligned to `paths`, and
        `transform`/`crs` describe the shared decimated grid.
    """
    with rasterio.open(paths[0]) as s:
        scale = max(1, round(max(s.width, s.height) / target_pixels))
        H, W = s.height // scale, s.width // scale
        transform = s.transform * rasterio.Affine.scale(s.width / W, s.height / H)
        crs = s.crs
        arrays = [s.read(1, out_shape=(H, W), resampling=Resampling.nearest)]
    for p in paths[1:]:
        with rasterio.open(p) as s:
            arrays.append(s.read(1, out_shape=(H, W), resampling=Resampling.nearest))
    return arrays, transform, crs


def sample_rows_from_grid(pos, kind, W, transform, crs, class_a, class_b, confidence,
                          conf_band_idx, burned, class_dict, band_labels):
    """Reference-sample rows for flat grid positions `pos` on a decimated raster grid.

    Builds the per-point attribute columns (class_2019/2025, transition, confidence, conf_band,
    burned, sample_type) plus point geometry in EPSG:4326, from the decimated arrays
    `read_decimated_grid` produced. `pos` indexes the raveled (H, W) grid; `kind` is the
    sample_type label ('candidate' or 'stable').
    """
    r, c = np.divmod(pos, W)
    xs, ys = rasterio.transform.xy(transform, r, c)
    lons, lats = rasterio.warp.transform(crs, 'EPSG:4326', xs, ys)
    c19, c25 = class_a.ravel()[pos], class_b.ravel()[pos]
    return gpd.GeoDataFrame({
        'class_2019': [class_dict.get(int(v), int(v)) for v in c19],
        'class_2025': [class_dict.get(int(v), int(v)) for v in c25],
        'transition': [f'{class_dict.get(int(u), u)} to {class_dict.get(int(v), v)}'
                      for u, v in zip(c19, c25)],
        'confidence': confidence.ravel()[pos].astype(int),
        'conf_band': [band_labels[i] for i in conf_band_idx.ravel()[pos]],
        'burned': burned.ravel()[pos],
        'sample_type': [kind] * len(pos),
    }, geometry=gpd.points_from_xy(lons, lats), crs='EPSG:4326')


def grid_positions_of(gdf, transform, crs, class_a, class_b, class_dict):
    """Recover flat grid positions for an existing point sample, verifying the mapping.

    Reprojects the sample's EPSG:4326 points into the raster's CRS, looks up row/col via the
    affine transform, and asserts the recovered classes match the sample's stored classes --
    catching a silently wrong transform/CRS before it corrupts a sample-growing step (see
    stratified_topup usage in Wall_to_Wall_Maps.ipynb Step 8).
    """
    xs, ys = rasterio.warp.transform('EPSG:4326', crs, gdf.geometry.x.values, gdf.geometry.y.values)
    r, c = rasterio.transform.rowcol(transform, xs, ys)
    W = class_a.shape[1]
    pos = np.asarray(r, dtype=np.int64) * W + np.asarray(c, dtype=np.int64)
    got19 = [class_dict.get(int(v), int(v)) for v in class_a.ravel()[pos]]
    got25 = [class_dict.get(int(v), int(v)) for v in class_b.ravel()[pos]]
    assert list(gdf.class_2019) == got19 and list(gdf.class_2025) == got25, \
        'recovered grid positions do not reproduce the cached sample - refusing to grow'
    return pos


def class_color(name, class_dict, colors=CLASS_COLORS):
    """Hex color for a class name, looked up via its numeric id in `class_dict` (e.g. class1_dict).

    Falls back to a neutral gray for a name with no matching id -- keeps a plot from erroring on
    an unexpected class label rather than needing every caller to guard against it."""
    for cid, nm in class_dict.items():
        if nm == name:
            return colors.get(cid, '#888')
    return '#888'