Time-based vs distance-based trajectory segmentation

import numpy as np
import pandas as pd


def segment_by_rule(
    df: pd.DataFrame,
    rule: str = "time",
    time_col: str = "t",
    x_col: str = "x",
    y_col: str = "y",
    entity_col: str = "entity_id",
    time_threshold_s: float = 300.0,
    distance_threshold_m: float = 2000.0,
    gap_threshold_s: float = 60.0,
) -> pd.DataFrame:
    """
    Cut a track into segments by elapsed time or by travelled distance.

    Parameters
    ----------
    rule : {'time', 'distance', 'both'}
        'both' emits two independent columns rather than intersecting them —
        intersecting produces segments shorter than either rule intended,
        which is almost never what a caller means by "use both".
    gap_threshold_s : float
        A data outage always breaks a segment, under either rule. This is a
        different concept from the time rule: a gap means the position is
        unknown, whereas the time rule is just a chosen granularity.

    Returns
    -------
    pd.DataFrame
        Input plus 'seg_time' and/or 'seg_distance', and always 'seg_gap'.

    Raises
    ------
    ValueError
        On missing columns, an empty frame, or an unknown rule.
    """
    if rule not in {"time", "distance", "both"}:
        raise ValueError(f"Unknown rule {rule!r}; expected time, distance or both.")
    required = {time_col, x_col, y_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.")

    d = df.sort_values([entity_col, time_col]).reset_index(drop=True).copy()
    g = d.groupby(entity_col, sort=False)

    dt = g[time_col].diff().dt.total_seconds().fillna(0.0).to_numpy()
    step = np.hypot(g[x_col].diff().fillna(0.0).to_numpy(),
                    g[y_col].diff().fillna(0.0).to_numpy())

    # Gaps always cut. A negative dt is a clock fault and also cuts, because
    # elapsed time across it is unknown.
    gap_cut = (dt > gap_threshold_s) | (dt < 0)
    d["seg_gap"] = gap_cut.cumsum()

    if rule in ("time", "both"):
        # Accumulate elapsed time, resetting at each cut. A vectorised cumsum
        # cannot express "reset on threshold", so this is an explicit loop —
        # at fix-level volumes, do it per entity with numpy inside.
        d["seg_time"] = _accumulate_cuts(dt, time_threshold_s, gap_cut)

    if rule in ("distance", "both"):
        d["seg_distance"] = _accumulate_cuts(step, distance_threshold_m, gap_cut)

    return d


def _accumulate_cuts(delta: np.ndarray, threshold: float, force_cut: np.ndarray) -> np.ndarray:
    """Running accumulator that emits a new segment id whenever it exceeds threshold."""
    seg = np.empty(delta.size, dtype=np.int64)
    acc, sid = 0.0, 0
    for i in range(delta.size):
        if force_cut[i]:
            sid += 1
            acc = 0.0
        elif acc + delta[i] > threshold:
            sid += 1
            acc = 0.0
        acc += max(delta[i], 0.0)
        seg[i] = sid
    return seg

def validate_segments(d: pd.DataFrame, rule: str,
                      time_threshold_s: float = 300.0,
                      distance_threshold_m: float = 2000.0) -> dict:
    """Check the segments obey the rule that produced them."""
    out = {}
    if rule in ("time", "both"):
        dur = d.groupby("seg_time")["t"].agg(lambda s: (s.max() - s.min()).total_seconds())
        # Allow one sampling interval of overshoot; more means the accumulator
        # is not resetting.
        assert dur.max() <= time_threshold_s * 1.2, (
            f"a time segment ran to {dur.max():.0f}s against a "
            f"{time_threshold_s:.0f}s threshold"
        )
        out["time_segments"] = int(d["seg_time"].nunique())

    if rule in ("distance", "both"):
        length = d.groupby("seg_distance").apply(
            lambda g: float(np.hypot(g["x"].diff(), g["y"].diff()).sum()),
            include_groups=False,
        )
        assert length.max() <= distance_threshold_m * 1.2, (
            f"a distance segment ran to {length.max():.0f}m"
        )
        out["distance_segments"] = int(d["seg_distance"].nunique())

    # Every gap must break every segmentation, or a segment spans an outage.
    for col in [c for c in ("seg_time", "seg_distance") if c in d]:
        spans = d.groupby(col)["seg_gap"].nunique()
        assert (spans == 1).all(), f"{col} has segments spanning a data gap"
    return out