Rand index computation

rand_index.py

# Heavily inspired by the answer by Tom in https://stats.stackexchange.com/questions/89030/rand-index-calculation

def rand_score(labels_true, labels_pred):
    """
    Args:
        labels_true: Array of shape :math:`(N)` denoting the class identities of each element.
        labels_pred: Array of shape :math:`(N)` denoting the cluster identities of each element.
    
    Returns:
        (Unadjusted) Rand index.
    """
    # contingency[c, k] -> number of instances in cluster k labelled with class c
    contingency = contingency_matrix(labels_true, labels_pred, sparse=True)
    tp_plus_fp = sum(comb2(n_class_elems) for n_class_elems in contingency.sum(axis=1).flat)
    tp_plus_fn = sum(comb2(n_cluster_elems) for n_cluster_elems in contingency.sum(axis=0).flat)
    tp = sum(comb2(n) for n in contingency.data)
    n_examples = len(labels_true)
    n = comb2(n_examples)
    fp = tp_plus_fp - tp
    fn = tp_plus_fn - tp
    tn = n - tp - fp - fn
    return (tp + tn) / (tp + tn + fp + fn)