Feature engineering for transport mode classifiers

import numpy as np
import pandas as pd
import geopandas as gpd


def add_context_features(
    seg_feats: pd.DataFrame,
    seg_lines: gpd.GeoDataFrame,
    rail: gpd.GeoDataFrame,
    cycle: gpd.GeoDataFrame,
    bus_routes: gpd.GeoDataFrame,
    stop_dists: pd.Series | None = None,
    max_search_m: float = 400.0,
) -> pd.DataFrame:
    """
    Attach network-proximity and stop-regularity features to segment features.

    Parameters
    ----------
    seg_feats : pd.DataFrame
        Kinematic features, indexed by segment id.
    seg_lines : gpd.GeoDataFrame
        One LineString per segment, same index, in a PROJECTED METRIC CRS.
        Distances computed in degrees are latitude-dependent and useless here.
    rail, cycle, bus_routes : gpd.GeoDataFrame
        Network layers in the SAME CRS as seg_lines.
    stop_dists : pd.Series | None
        Optional Series of arrays: distances between consecutive stops within
        each segment, in metres. Missing entries give a null regularity score.

    Raises
    ------
    ValueError
        On index mismatch, missing CRS, or a CRS that is not projected.
    """
    if not seg_lines.index.equals(seg_feats.index):
        raise ValueError("seg_lines and seg_feats must share an index.")
    if seg_lines.crs is None or seg_lines.crs.is_geographic:
        raise ValueError(
            "seg_lines must be in a PROJECTED metric CRS; got "
            f"{seg_lines.crs}. Reproject before computing distances."
        )
    for name, layer in (("rail", rail), ("cycle", cycle), ("bus", bus_routes)):
        if layer.crs != seg_lines.crs:
            raise ValueError(f"{name} layer CRS {layer.crs} != segments {seg_lines.crs}")

    out = seg_feats.copy()

    # ── Network proximity via sjoin_nearest ───────────────────────────
    # max_distance caps the search; segments with nothing within it get the
    # cap rather than NaN, which keeps the feature monotone for the tree.
    for name, layer in (("rail", rail), ("cycle", cycle), ("bus", bus_routes)):
        joined = gpd.sjoin_nearest(
            seg_lines[["geometry"]], layer[["geometry"]],
            how="left", max_distance=max_search_m, distance_col=f"d_{name}",
        )
        # sjoin_nearest can emit ties; keep the closest per segment.
        d = joined.groupby(level=0)[f"d_{name}"].min()
        out[f"dist_{name}_m"] = d.reindex(out.index).fillna(max_search_m)

    # ── Stop regularity ───────────────────────────────────────────────
    # A bus stops at roughly even spacing; a car stops at junctions and
    # queues, which is highly irregular. Coefficient of variation captures
    # exactly that, with no timetable required.
    def _cv(arr) -> float:
        a = np.asarray(arr, dtype=float)
        a = a[np.isfinite(a) & (a > 0)]
        if a.size < 3:
            return np.nan          # too few stops to say anything
        m = a.mean()
        return float(a.std() / m) if m > 0 else np.nan

    if stop_dists is not None:
        out["stop_spacing_cv"] = stop_dists.reindex(out.index).map(_cv)
    else:
        out["stop_spacing_cv"] = np.nan

    # A ratio feature carries more signal than either distance alone: a bus
    # is close to a bus route AND far from a rail line.
    out["rail_over_bus"] = (
        out["dist_rail_m"] / out["dist_bus_m"].clip(lower=1.0)
    )
    return out


def add_neighbour_context(
    seg_feats: pd.DataFrame,
    oof_proba: pd.DataFrame,
    trip_col: str = "trip_id",
) -> pd.DataFrame:
    """
    Add the previous and next segment's OUT-OF-FOLD probabilities.

    oof_proba must come from cross-validated predictions on the training
    data. Using in-fold predictions leaks the label and produces a model
    that validates beautifully and fails in production.
    """
    if trip_col not in seg_feats.columns:
        raise ValueError(f"seg_feats needs a {trip_col} column to order segments.")
    if not oof_proba.index.equals(seg_feats.index):
        raise ValueError("oof_proba must share seg_feats' index.")

    out = seg_feats.copy()
    grouped = oof_proba.groupby(seg_feats[trip_col], sort=False)
    for cls in oof_proba.columns:
        out[f"prev_p_{cls}"] = grouped[cls].shift(1).fillna(0.0)
        out[f"next_p_{cls}"] = grouped[cls].shift(-1).fillna(0.0)
    return out

def validate_context_features(df: pd.DataFrame, max_search_m: float = 400.0) -> None:
    """Catch the three ways these joins go wrong without raising."""
    # 1. A distance column that is entirely at the cap means the join found
    #    nothing — almost always a CRS mismatch rather than a rural fleet.
    for c in ("dist_rail_m", "dist_cycle_m", "dist_bus_m"):
        at_cap = (df[c] >= max_search_m - 1e-6).mean()
        assert at_cap < 0.95, f"{c}: {at_cap:.0%} at the search cap — check the CRS"

    # 2. Distances must be non-negative and finite.
    assert (df[["dist_rail_m", "dist_cycle_m", "dist_bus_m"]] >= 0).all().all()

    # 3. Regularity is a ratio, so a value above ~3 means the stop detector
    #    is emitting spurious stops rather than the bus being irregular.
    cv = df["stop_spacing_cv"].dropna()
    assert (cv < 3.0).mean() > 0.98, "stop_spacing_cv tail suggests spurious stops"
    print(f"OK — {len(df)} segments, {df['stop_spacing_cv'].notna().mean():.0%} with a regularity score")