PythonでMAPとMRRを計算する方法

Mean Average Precision (@k) や Mean Reciprocal Rank (@k) という、レコメンドのオフライン評価で使われる指標があるのですが、scikit-learn では計算できなさげだったので、ChatGPT先生に書いてもらいました。

ちなみにMAPの説明はここが分かりやすかった。

mAP (mean average precision) for Recommender systems and Object detection algorithms

Before we start, I would like to point out this blog is my understanding of the other blogs and source code that I have mentioned under the…

martian1231-py.medium.com

MAPのコード

import numpy as np


def map_at_k(y_true, y_scores, k=None):
    """
    Computes Mean Average Precision at K (MAP@K)
    
    Parameters:
    y_true: np.ndarray of shape (n_samples, n_labels)
        Binary indicator matrix of relevant items (1 if relevant, 0 otherwise)
    y_scores: np.ndarray of shape (n_samples, n_labels)
        Predicted scores for each label
    k: int or None, default=None
        Number of top elements to consider. If None, considers all elements.
    
    Returns:
    float: MAP@K score
    """
    n_samples, n_labels = y_true.shape
    avg_precisions = []
    
    for i in range(n_samples):
        # Sort indices by score in descending order
        sorted_indices = np.argsort(y_scores[i])[::-1]
        
        if k is not None:
            sorted_indices = sorted_indices[:k]
        
        relevant = y_true[i][sorted_indices]
        
        # Compute AP@K
        num_hits = 0.0
        score = 0.0
        for j, rel in enumerate(relevant):
            if rel == 1:
                num_hits += 1.0
                score += num_hits / (j + 1.0)
        
        # Avoid division by zero
        relevant_count = np.sum(y_true[i])
        if relevant_count == 0:
            avg_precisions.append(0.0)  # No relevant items, AP is 0
        else:
            avg_precisions.append(score / min(relevant_count, k if k is not None else n_labels))
    
    return np.mean(avg_precisions)


# Example usage
y_true = np.array([[0, 0, 1, 1], [0, 0, 0, 0]])  # 2nd sample has no relevant items
y_scores = np.array([[0.1, 0.4, 0.35, 0.8], [0.2, 0.3, 0.5, 0.1]])
k = None
print("MAP@{}:".format(k), map_at_k(y_true, y_scores, k))

MRRのコード

import numpy as np


def mean_reciprocal_rank_at_k(y_true, y_scores, k=None):
    """
    Computes Mean Reciprocal Rank at K (MRR@K)
    
    Parameters:
    y_true: np.ndarray of shape (n_samples, n_labels)
        Binary indicator matrix of relevant items (1 if relevant, 0 otherwise)
    y_scores: np.ndarray of shape (n_samples, n_labels)
        Predicted scores for each label
    k: int or None, default=None
        Number of top elements to consider. If None, considers all elements.
    
    Returns:
    float: MRR@K score
    """
    n_samples, n_labels = y_true.shape
    reciprocal_ranks = []
    
    for i in range(n_samples):
        sorted_indices = np.argsort(y_scores[i])[::-1]  # Sort by predicted scores
        
        if k is not None:
            sorted_indices = sorted_indices[:k]
        
        relevant = y_true[i][sorted_indices]
        
        # Find the rank of the first relevant item
        rank = np.where(relevant == 1)[0]
        
        if rank.size > 0:
            reciprocal_ranks.append(1.0 / (rank[0] + 1))  # Convert to 1-based index
        else:
            reciprocal_ranks.append(0.0)  # No relevant items found
    
    return np.mean(reciprocal_ranks)


# Example usage
y_true = np.array([[0, 0, 1, 1], [0, 0, 0, 1]])
y_scores = np.array([[0.1, 0.4, 0.35, 0.8], [0.2, 0.3, 0.5, 0.9]])
k = 3
print("MRR@{}:".format(k), mean_reciprocal_rank_at_k(y_true, y_scores, k))