Add tv metric

ethicml/metrics/tv.py

@@ -0,0 +1,45 @@
+"""
+Implementation of the variant of toal variation used as the unfairness metric in
+https://arxiv.org/abs/1905.13662
+"""
+import numpy as np
+import pandas as pd
+
+from ethicml.algorithms.utils import DataTuple
+from ethicml.metrics import Metric
+from ethicml.utility.heaviside import Heaviside
+
+
+class TV(Metric):
+    """
+    Requires that predictions be probabilities of the class
+    """
+
+    def score(self, prediction: pd.DataFrame, actual: DataTuple) -> float:
+
+        # Only consider one s column (for now)
+        s_col = actual.s.columns[0]
+        unique_s = actual.s[s_col].unique()
+
+        heavi = Heaviside()
+        class_labels = heavi.apply(prediction.values)
+        unique_y = np.unique(class_labels)
+        class_labels = pd.DataFrame(class_labels, columns=["labels"])
+
+        sum_dist = 0
+
+        for y in unique_y:
+            for s in unique_s:
+                to_check = prediction.loc[actual.s[s_col] == s]
+                to_check = to_check.loc[class_labels["labels"] == y]
+                to_check = to_check.values.mean()
+
+                avg_pred = prediction[class_labels["labels"] == y].values.mean()
+
+                sum_dist += abs(avg_pred - to_check)
+
+        return sum_dist / len(unique_s)
+
+    @property
+    def name(self) -> str:
+        return "TV"

tests/metric_test.py

@@ -22,6 +22,7 @@
 from ethicml.metrics.confusion_matrix import LabelOutOfBounds
 from ethicml.metrics.hsic import Hsic
 from ethicml.metrics.theil import Theil
+from ethicml.metrics.tv import TV
 from ethicml.preprocessing.train_test_split import train_test_split
 from tests.run_algorithm_test import get_train_test
 
@@ -38,6 +39,16 @@ def test_get_acc_of_predictions():
     assert score == 0.89
 
 
+def test_get_tv_of_predictions():
+    train, test = get_train_test()
+    model: InAlgorithm = LRProb()
+    predictions: pd.DataFrame = model.run(train, test)
+    tv: Metric = TV()
+    assert tv.name == "TV"
+    score = tv.score(predictions, test)
+    assert score == pytest.approx(0.048, abs=0.001)
+
+
 def test_mni_preds_and_s():
     train, test = get_train_test()
     model: InAlgorithm = SVM()