Use registry pattern for ClassificationMetricSet

rexmex/metricset.py

@@ -1,18 +1,6 @@
 from typing import Collection, List, Tuple
 
-from rexmex.metrics.classification import (
-    accuracy_score,
-    average_precision_score,
-    balanced_accuracy_score,
-    f1_score,
-    fowlkes_mallows_index,
-    matthews_correlation_coefficient,
-    pr_auc_score,
-    precision_score,
-    recall_score,
-    roc_auc_score,
-    specificity,
-)
+from rexmex.metrics.classification import classifications
 from rexmex.metrics.rating import (
     mean_absolute_error,
     mean_absolute_percentage_error,
@@ -104,17 +92,14 @@ class ClassificationMetricSet(MetricSet):
     """
 
     def __init__(self):
-        self["roc_auc"] = roc_auc_score
-        self["pr_auc"] = pr_auc_score
-        self["average_precision"] = average_precision_score
-        self["f1_score"] = binarize(f1_score)
-        self["matthews_correlation_coefficent"] = binarize(matthews_correlation_coefficient)
-        self["fowlkes_mallows_index"] = binarize(fowlkes_mallows_index)
-        self["precision"] = binarize(precision_score)
-        self["recall"] = binarize(recall_score)
-        self["specificity"] = binarize(specificity)
-        self["accuracy"] = binarize(accuracy_score)
-        self["balanced_accuracy"] = binarize(balanced_accuracy_score)
+        super().__init__()
+        for func in classifications:
+            name = func.__name__
+            if name.endswith("_score"):
+                name = name[: -len("_score")]
+            if func.binarize:
+                func = binarize(func)
+            self[name] = func
 
     def __repr__(self):
         """

rexmex/utils.py

@@ -61,7 +61,13 @@ def metric_wrapper(*args, **kwargs):
 
 
 class Annotator:
-    """A class to wrap annotations to make the registry pattern easier later."""
+    """A class to wrap annotations that generates a registry."""
+
+    def __init__(self):
+        self.funcs = {}
+
+    def __iter__(self):
+        return iter(self.funcs.values())
 
     def annotate(
         self,
@@ -79,6 +85,7 @@ def annotate(
         """Annotate a classification function."""
 
         def _wrapper(func):
+            self.funcs[func.__name__] = func
             func.name = name or func.__name__.replace("_", " ").title()
             func.lower = lower
             func.lower_inclusive = lower_inclusive