Clustering trajectories with HDBSCAN on a distance matrix

import numpy as np
import pandas as pd
import hdbscan


def cluster_trajectories(
    dist: np.ndarray,
    trip_ids: pd.Index,
    min_cluster_size: int = 5,
    min_samples: int | None = None,
    selection: str = "eom",
) -> pd.DataFrame:
    """
    Cluster trajectories from a precomputed distance matrix.

    Parameters
    ----------
    dist : np.ndarray
        Square, symmetric, zero-diagonal distance matrix in METRES (or
        whatever unit the measure returns — it must be consistent).
    min_cluster_size : int
        The smallest group you would be willing to call a route. This is a
        domain decision, not a tuning knob: 5 means "five trips make a route".
    min_samples : int | None
        Conservativeness. Lower values assign more points to clusters;
        defaults to min_cluster_size, which is usually too conservative for
        trajectory data — try min_cluster_size // 2.
    selection : str
        'eom' (excess of mass) prefers a few large stable clusters; 'leaf'
        prefers many fine ones. Route networks usually want 'leaf'.

    Returns
    -------
    pd.DataFrame
        Per trip: label, membership probability, and outlier score.

    Raises
    ------
    ValueError
        If the matrix is not square, symmetric or non-negative.
    """
    if dist.ndim != 2 or dist.shape[0] != dist.shape[1]:
        raise ValueError(f"Distance matrix must be square; got {dist.shape}.")
    if len(trip_ids) != dist.shape[0]:
        raise ValueError("trip_ids length does not match the matrix.")
    if not np.allclose(dist, dist.T, atol=1e-6, equal_nan=True):
        raise ValueError(
            "Matrix is not symmetric. DTW and Fréchet are symmetric by "
            "definition, so an asymmetric matrix means the fill loop only "
            "wrote the upper triangle."
        )
    if np.nanmin(dist) < 0:
        raise ValueError("Negative distances — a similarity was passed instead.")
    if not np.allclose(np.diag(dist), 0.0, atol=1e-9):
        raise ValueError("Non-zero diagonal — a trip is not distance 0 from itself.")

    # HDBSCAN cannot handle NaN. Blocked-out pairs must be a large finite
    # value, not NaN: infinity poisons the mutual-reachability computation.
    finite = np.nan_to_num(dist, nan=np.nanmax(dist) * 10.0, posinf=np.nanmax(dist) * 10.0)

    clusterer = hdbscan.HDBSCAN(
        metric="precomputed",
        min_cluster_size=int(min_cluster_size),
        min_samples=int(min_samples) if min_samples else None,
        cluster_selection_method=selection,
        allow_single_cluster=False,
    )
    labels = clusterer.fit_predict(finite.astype(np.float64))

    return pd.DataFrame(
        {
            "label": labels,
            "probability": clusterer.probabilities_,
            "outlier_score": clusterer.outlier_scores_,
        },
        index=trip_ids,
    )

def validate_clustering(res: pd.DataFrame, dist: np.ndarray,
                        max_noise_frac: float = 0.4) -> dict:
    """Three checks that separate a real clustering from a degenerate one."""
    noise = (res["label"] == -1)
    n_clusters = res.loc[~noise, "label"].nunique()

    # 1. Not everything is noise, and not everything is one cluster.
    assert n_clusters >= 2, (
        f"{n_clusters} cluster(s) — try cluster_selection_method='leaf' or "
        "a lower min_samples before concluding the routes are homogeneous"
    )
    assert noise.mean() < max_noise_frac, (
        f"{noise.mean():.0%} noise — the distance measure may be separating "
        "trips that a human would call the same route"
    )

    # 2. Clusters should be tighter internally than the overall spread,
    #    otherwise the labels carry no information.
    overall = np.nanmedian(dist[np.triu_indices_from(dist, k=1)])
    for lab, idx in res.loc[~noise].groupby("label").groups.items():
        pos = res.index.get_indexer(idx)
        sub = dist[np.ix_(pos, pos)]
        within = np.nanmedian(sub[np.triu_indices_from(sub, k=1)]) if len(pos) > 1 else 0.0
        assert within < 0.6 * overall, (
            f"cluster {lab} is as spread out as the whole dataset"
        )

    # 3. Membership probabilities should not be uniformly near zero.
    assert res.loc[~noise, "probability"].median() > 0.4, (
        "median membership probability is very low — clusters are marginal"
    )
    return {"clusters": int(n_clusters), "noise_frac": float(noise.mean())}