cxcscmu/skilllearnbench/greedy-matching-evaluation

import numpy as np
from scipy.spatial.distance import cdist
def evaluate_clustering(centroids, experts, max_dist=100):
    if len(centroids) == 0:
        return 0.0, np.nan
    if len(experts) == 0:
        return 0.0, np.nan
    distances = cdist(centroids, experts, metric='euclidean')
    
    # Greedy matching
    indices = np.where(distances <= max_dist)
    pairs = sorted(zip(indices[0], indices[1]), key=lambda x: distances[x[0], x[1]])
    
    matched_centroids = set()
    matched_experts = set()
    match_distances = []
    
    for c_idx, e_idx in pairs:
        if c_idx not in matched_centroids and e_idx not in matched_experts:
            matched_centroids.add(c_idx)
            matched_experts.add(e_idx)
            match_distances.append(distances[c_idx, e_idx])
            
    tp = len(matched_centroids)
    fp = len(centroids) - tp
    fn = len(experts) - tp
    
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
    
    avg_delta = np.mean(match_distances) if len(match_distances) > 0 else np.nan
    
    return f1, avg_delta