Reducing Spatial Leakage in Model Training

import logging
import geopandas as gpd
import numpy as np
from shapely.geometry import box
from sklearn.model_selection import GroupKFold, cross_val_score
from sklearn.ensemble import RandomForestRegressor

logger = logging.getLogger(__name__)


def assign_spatial_blocks(
    gdf: gpd.GeoDataFrame,
    block_size: float = 5000.0,
    id_col: str = "block_id",
) -> gpd.GeoDataFrame:
    """Assign a spatial grid block ID to every observation.

    Parameters
    ----------
    gdf:
        Input GeoDataFrame with point or polygon geometry.
        Must be in a projected (metric) CRS — raises ValueError otherwise.
    block_size:
        Grid cell dimension in CRS units (default 5000 m = 5 km).
        Set this to at least 2× the empirical spatial correlation range
        of your target variable.
    id_col:
        Name of the output block-ID column.

    Returns
    -------
    GeoDataFrame with ``id_col`` appended. Observations that fall
    outside the computed grid receive block_id = -1 and are logged.
    """
    if gdf.crs is None:
        raise ValueError("GeoDataFrame has no CRS. Set one before blocking.")
    if gdf.crs.is_geographic:
        raise ValueError(
            f"CRS {gdf.crs.to_string()} is geographic. "
            "Reproject to a metric CRS (e.g. EPSG:3857 or a local UTM) first."
        )

    x_min, y_min, x_max, y_max = gdf.total_bounds

    x_bins = np.arange(x_min, x_max + block_size, block_size)
    y_bins = np.arange(y_min, y_max + block_size, block_size)

    grid_cells = [
        box(x, y, x + block_size, y + block_size)
        for x in x_bins[:-1]
        for y in y_bins[:-1]
    ]
    grid_gdf = gpd.GeoDataFrame(
        {"_cell_id": range(len(grid_cells))},
        geometry=grid_cells,
        crs=gdf.crs,
    )

    joined = gpd.sjoin(gdf, grid_gdf, how="left", predicate="intersects")
    joined[id_col] = joined["_cell_id"].fillna(-1).astype(int)
    joined = joined.drop(columns=["index_right", "_cell_id"], errors="ignore")

    n_unassigned = (joined[id_col] == -1).sum()
    if n_unassigned > 0:
        logger.warning(
            "%d observation(s) fell outside the grid and received block_id=-1. "
            "Exclude or merge them with a neighbouring block before CV.",
            n_unassigned,
        )

    logger.info(
        "Assigned %d observations to %d spatial blocks (block_size=%.0f m).",
        len(joined),
        joined[id_col].nunique(),
        block_size,
    )
    return joined


def run_spatial_cv(
    gdf_blocked: gpd.GeoDataFrame,
    feature_cols: list[str],
    target_col: str,
    block_id_col: str = "block_id",
    n_splits: int = 5,
) -> dict[str, float]:
    """Run spatial GroupKFold CV and return mean ± std of R².

    Parameters
    ----------
    gdf_blocked:
        GeoDataFrame produced by ``assign_spatial_blocks``.
    feature_cols:
        List of numeric feature column names.
    target_col:
        Name of the regression target column.
    block_id_col:
        Column containing spatial block IDs (used as CV groups).
    n_splits:
        Number of CV folds.

    Returns
    -------
    dict with keys ``mean_r2`` and ``std_r2``.
    """
    valid = gdf_blocked[gdf_blocked[block_id_col] != -1].copy()
    X = valid[feature_cols].values
    y = valid[target_col].values
    groups = valid[block_id_col].values

    gkf = GroupKFold(n_splits=n_splits)
    model = RandomForestRegressor(n_estimators=200, random_state=42)

    scores = cross_val_score(
        model, X, y, cv=gkf.split(X, y, groups=groups), scoring="r2"
    )
    logger.info(
        "Spatial CV R²: %.3f ± %.3f over %d folds.",
        scores.mean(), scores.std(), n_splits,
    )
    return {"mean_r2": float(scores.mean()), "std_r2": float(scores.std())}

gdf = gpd.read_file("observations.geojson")
# Always inspect the CRS before any spatial operation
print(gdf.crs)

# Reproject to Web Mercator (global) or a local UTM zone for better accuracy
gdf_m = gdf.to_crs(epsg=3857)

assert gdf_m.crs.is_projected, "CRS must be projected before blocking"
assert not gdf_m.crs.is_geographic

import libpysal
from esda.moran import Moran

# Build a distance-band weights matrix at candidate ranges
for radius_m in [1000, 2000, 5000, 10000]:
    w = libpysal.weights.DistanceBand.from_dataframe(
        gdf_m, threshold=radius_m, silence_warnings=True
    )
    mi = Moran(gdf_m["target"].values, w)
    print(f"Range {radius_m} m → Moran's I = {mi.I:.3f}  p = {mi.p_sim:.4f}")

gdf_blocked = assign_spatial_blocks(gdf_m, block_size=6000)

# Sanity check: how many blocks have enough observations for CV?
block_counts = gdf_blocked["block_id"].value_counts()
print(block_counts.describe())

# Drop or merge sparse blocks (fewer than 5 observations) to avoid tiny folds
min_obs = 5
valid_blocks = block_counts[block_counts >= min_obs].index
gdf_blocked = gdf_blocked[gdf_blocked["block_id"].isin(valid_blocks)]

feature_cols = ["elevation", "slope", "ndvi", "distance_to_road"]
result = run_spatial_cv(
    gdf_blocked,
    feature_cols=feature_cols,
    target_col="target",
    n_splits=5,
)
print(f"Spatial CV  R² = {result['mean_r2']:.3f} ± {result['std_r2']:.3f}")

from sklearn.model_selection import cross_val_score, KFold

X_all = gdf_blocked[feature_cols].values
y_all = gdf_blocked["target"].values
random_scores = cross_val_score(
    RandomForestRegressor(n_estimators=200, random_state=42),
    X_all, y_all,
    cv=KFold(n_splits=5, shuffle=True, random_state=42),
    scoring="r2",
)
print(f"Random CV   R² = {random_scores.mean():.3f} ± {random_scores.std():.3f}")

LEAKAGE_THRESHOLD = 0.85  # reject if spatial R² < 85% of random R²

ratio = result["mean_r2"] / random_scores.mean()
if ratio < LEAKAGE_THRESHOLD:
    raise RuntimeError(
        f"Spatial CV / random CV ratio = {ratio:.2f} < {LEAKAGE_THRESHOLD}. "
        "Model is over-reliant on spatial proximity. Review block size and features."
    )

# 1. Verify no training observation is within block_size of any validation observation
from shapely.ops import unary_union

for fold_idx, (train_idx, val_idx) in enumerate(
    GroupKFold(n_splits=5).split(
        gdf_blocked, groups=gdf_blocked["block_id"].values
    )
):
    train_geom = unary_union(gdf_blocked.iloc[train_idx].geometry.buffer(0))
    val_pts = gdf_blocked.iloc[val_idx].geometry
    # No validation point should touch the unioned training geometry
    # (for strict buffer exclusion, substitute buffer(block_size/2) above)
    overlaps = val_pts.intersects(train_geom).sum()
    assert overlaps == 0, f"Fold {fold_idx}: {overlaps} validation pts overlap training geometry"
    print(f"Fold {fold_idx}: OK — no spatial overlap")

# 2. Log the spatial CV vs random CV gap
print(f"Leakage gap: {random_scores.mean() - result['mean_r2']:.3f} R² units")