Segmenting trips by dwell time and gap thresholds

import numpy as np
import pandas as pd


def segment_trips(
    df: pd.DataFrame,
    time_col: str = "t",
    x_col: str = "x",
    y_col: str = "y",
    speed_col: str = "speed_ms",
    entity_col: str = "entity_id",
    stop_speed_ms: float = 0.8,
    stop_radius_m: float = 30.0,
    min_dwell_s: float = 180.0,
    gap_s: float = 600.0,
) -> pd.DataFrame:
    """
    Cut each entity's track into trips at confirmed dwells and at data gaps.

    Parameters
    ----------
    min_dwell_s : float
        How long the vehicle must remain within stop_radius_m below
        stop_speed_ms before the pause counts as a trip boundary. Below about
        120 s, traffic signals start creating trips.
    gap_s : float
        Reporting gap that ends a trip regardless of what the vehicle did.

    Returns
    -------
    pd.DataFrame
        Input plus 'trip_id' and 'boundary_kind' in {start, dwell, gap}.

    Raises
    ------
    ValueError
        On missing columns, an empty frame, or a non-positive dwell.
    """
    required = {time_col, x_col, y_col, speed_col, entity_col}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"Missing required columns: {missing}")
    if df.empty:
        raise ValueError("Input DataFrame is empty.")
    if min_dwell_s <= 0:
        raise ValueError("min_dwell_s must be positive.")

    d = df.sort_values([entity_col, time_col]).reset_index(drop=True).copy()
    trip_id = np.zeros(len(d), dtype=np.int64)
    kind = np.full(len(d), "", dtype=object)
    tid = 0

    for _, idx in d.groupby(entity_col, sort=False).groups.items():
        pos = d.index.get_indexer(idx)
        t = d[time_col].to_numpy()[pos]
        x = d[x_col].to_numpy()[pos]
        y = d[y_col].to_numpy()[pos]
        v = d[speed_col].to_numpy()[pos]

        tid += 1
        trip_id[pos[0]] = tid
        kind[pos[0]] = "start"

        anchor = 0            # first fix of the current candidate stop
        in_candidate = False

        for i in range(1, len(pos)):
            dt = (t[i] - t[i - 1]) / np.timedelta64(1, "s")

            # ── gap: ends the trip, claims nothing about what happened ──
            if dt > gap_s or dt < 0:
                tid += 1
                kind[pos[i]] = "gap"
                in_candidate = False
                trip_id[pos[i]] = tid
                continue

            slow = v[i] < stop_speed_ms
            if slow and not in_candidate:
                in_candidate, anchor = True, i
            elif slow and in_candidate:
                moved = np.hypot(x[i] - x[anchor], y[i] - y[anchor])
                dwell = (t[i] - t[anchor]) / np.timedelta64(1, "s")
                if moved > stop_radius_m:
                    # Drifted out of the radius: it was crawling, not stopping.
                    in_candidate, anchor = True, i
                elif dwell >= min_dwell_s:
                    # Confirmed. The boundary is placed at the START of the
                    # dwell, so the arrival belongs to the trip that arrived.
                    tid += 1
                    kind[pos[anchor]] = "dwell"
                    trip_id[pos[anchor]:pos[i] + 1] = tid
                    in_candidate = False
            else:
                in_candidate = False

            if trip_id[pos[i]] == 0:
                trip_id[pos[i]] = tid

    d["trip_id"] = trip_id
    d["boundary_kind"] = kind
    return d

def validate_trips(d: pd.DataFrame, labelled_trip_count: int | None = None) -> dict:
    """Checks that catch over- and under-segmentation before they reach a report."""
    trips = d.groupby("trip_id")
    durations = trips["t"].agg(lambda s: (s.max() - s.min()).total_seconds())
    n = len(durations)

    # 1. No zero-length trips. These come from a boundary placed on a single fix.
    assert (durations > 0).all(), f"{(durations == 0).sum()} zero-duration trips"

    # 2. Very short trips are a symptom of over-segmentation, not of driving.
    tiny = (durations < 60).mean()
    assert tiny < 0.1, (
        f"{tiny:.0%} of trips are under a minute — min_dwell_s is too low and "
        "traffic signals are creating trips"
    )

    # 3. Gap boundaries must not be counted as dwells anywhere downstream.
    gaps = (d["boundary_kind"] == "gap").sum()
    dwells = (d["boundary_kind"] == "dwell").sum()

    # 4. Against ground truth, if available. This is the only real calibration.
    out = {"trips": n, "gap_boundaries": int(gaps), "dwell_boundaries": int(dwells),
           "median_duration_s": float(durations.median())}
    if labelled_trip_count:
        err = abs(n - labelled_trip_count) / labelled_trip_count
        assert err < 0.15, (
            f"{n} trips against {labelled_trip_count} labelled ({err:.0%} off) "
            "— tune min_dwell_s before shipping this"
        )
        out["count_error"] = err
    return out