diff --git a/metriks/ranking.py b/metriks/ranking.py index b61ca66..4f6204e 100644 --- a/metriks/ranking.py +++ b/metriks/ranking.py @@ -4,10 +4,15 @@ Metrics to use for ranking models. """ +from typing import Tuple + import numpy as np +from numpy.ma import MaskedArray +from pytypes import typechecked -def check_arrays(y_true, y_prob): +@typechecked +def check_arrays(y_true: np.ndarray, y_prob: np.ndarray) -> None: # Make sure that inputs this conforms to our expectations assert isinstance(y_true, np.ndarray), AssertionError( 'Expect y_true to be a {expected}. Got {actual}' @@ -36,7 +41,8 @@ def check_arrays(y_true, y_prob): ) -def check_k(n_items, k): +@typechecked +def check_k(n_items: int, k: int) -> None: # Make sure that inputs conform to our expectations assert isinstance(k, int), AssertionError( 'Expect k to be a {expected}. Got {actual}' @@ -49,7 +55,8 @@ def check_k(n_items, k): ) -def recall_at_k(y_true, y_prob, k): +@typechecked +def recall_at_k(y_true: np.ndarray, y_prob: np.ndarray, k: int) -> float: """ Calculates recall at k for binary classification ranking problems. Recall at k measures the proportion of total relevant items that are found in the @@ -68,7 +75,7 @@ def recall_at_k(y_true, y_prob, k): sorted order by y_prob Returns: - recall (~np.ndarray): The recall at k + recall (float): The recall at k Example: >>> y_true = np.array([ @@ -118,7 +125,8 @@ def recall_at_k(y_true, y_prob, k): return recall -def precision_at_k(y_true, y_prob, k): +@typechecked +def precision_at_k(y_true: np.ndarray, y_prob: np.ndarray, k: int) -> float: """ Calculates precision at k for binary classification ranking problems. Precision at k measures the proportion of items in the top k (in ranked @@ -134,7 +142,7 @@ def precision_at_k(y_true, y_prob, k): sorted order by y_prob Returns: - precision_k (~np.ndarray): The precision at k + precision_k (float): The precision at k Example: >>> y_true = np.array([ @@ -178,7 +186,8 @@ def precision_at_k(y_true, y_prob, k): return precision -def mean_reciprocal_rank(y_true, y_prob): +@typechecked +def mean_reciprocal_rank(y_true: np.ndarray, y_prob: np.ndarray) -> MaskedArray: """ Gets a positional score about how well you did at rank 1, rank 2, etc. The resulting vector is of size (n_items,) but element 0 corresponds to @@ -217,7 +226,8 @@ def mean_reciprocal_rank(y_true, y_prob): return ma.mean(axis=0) -def label_mean_reciprocal_rank(y_true, y_prob): +@typechecked +def label_mean_reciprocal_rank(y_true: np.ndarray, y_prob: np.ndarray) -> MaskedArray: """ Determines the average rank each label was placed across samples. Only labels that are relevant in the true data set are considered in the calculation. @@ -241,7 +251,8 @@ def label_mean_reciprocal_rank(y_true, y_prob): return ma.mean(axis=0) -def ndcg(y_true, y_prob, k=0): +@typechecked +def ndcg(y_true: np.ndarray, y_prob: np.ndarray, k=0) -> np.float64: """ A score for measuring the quality of a set of ranked results. The resulting score is between 0 and 1.0 - results that are relevant and appear earlier in the result set are given a heavier weight, so the @@ -322,7 +333,8 @@ def ndcg(y_true, y_prob, k=0): return np.mean(sample_ndcg, dtype=np.float64) -def generate_y_pred_at_k(y_prob, k): +@typechecked +def generate_y_pred_at_k(y_prob: np.ndarray, k: int) -> np.ndarray: """ Generates a matrix of binary predictions from a matrix of probabilities by evaluating the top k items (in ranked order by y_prob) as true. @@ -362,11 +374,13 @@ def generate_y_pred_at_k(y_prob, k): index_array = np.argsort(y_prob, axis=1) col_idx = np.arange(y_prob.shape[0]).reshape(-1, 1) y_pred = np.zeros(np.shape(y_prob)) - y_pred[col_idx, index_array[:, n_items-k:n_items]] = 1 + y_pred[col_idx, index_array[:, n_items - k:n_items]] = 1 return y_pred -def confusion_matrix_at_k(y_true, y_prob, k): +@typechecked +def confusion_matrix_at_k(y_true: np.ndarray, y_prob: np.ndarray, k: int) -> Tuple[ + np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Generates binary predictions from probabilities by evaluating the top k items (in ranked order by y_prob) as true. Uses these binary predictions along with @@ -433,4 +447,3 @@ def confusion_matrix_at_k(y_true, y_prob, k): fn = np.count_nonzero((y_pred - 1) * y_true, axis=0) return tn, fp, fn, tp - diff --git a/setup.py b/setup.py index 9c70672..61cd427 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,7 @@ # Insert dependencies list here install_requires=[ 'numpy', + 'pytypes' ], setup_requires=["setuptools-scm"], extras_require={ diff --git a/tests/test_label_mean_reciprocal_rank.py b/tests/test_label_mean_reciprocal_rank.py index 374e28d..f8a6b86 100644 --- a/tests/test_label_mean_reciprocal_rank.py +++ b/tests/test_label_mean_reciprocal_rank.py @@ -16,6 +16,7 @@ def test_label_mrr_with_errors(): expected = np.ma.array([1.0, (7/4) / 3, 2/3, 1/3]) np.testing.assert_allclose(result, expected) + print(result) diff --git a/tests/test_ndcg.py b/tests/test_ndcg.py index 7f0144c..dafa738 100644 --- a/tests/test_ndcg.py +++ b/tests/test_ndcg.py @@ -207,6 +207,7 @@ def test_ndcg_0_over_0_error(): np.testing.assert_allclose([actual], [expected]) + if __name__ == '__main__': test_ndcg_perfect() test_ndcg_errors()