test: aggregators

hakeem/core/aggregation/aggregators/mallows.py

@@ -5,27 +5,28 @@
 
 
 class StandardApprovalAggregator(WeightedAggregator):
-	@staticmethod
-	def compute_weights(annotations: pd.DataFrame) -> pd.Series:
-		return pd.Series(1, index=annotations.index)
+    @staticmethod
+    def compute_weights(annotations: pd.DataFrame) -> pd.Series:
+        return pd.Series(1, index=annotations.index)
 
 
 class EuclidAggregator(WeightedAggregator):
-	@staticmethod
-	def compute_weights(annotations: pd.DataFrame) -> pd.Series:
-		vote_size = annotations.sum(axis=1)
-		return np.sqrt(vote_size + 1) - np.sqrt(vote_size - 1)
+    @staticmethod
+    def compute_weights(annotations: pd.DataFrame) -> pd.Series:
+        vote_size = annotations.sum(axis=1)
+        return np.sqrt(vote_size + 1) - np.sqrt(vote_size - 1)
 
 
 class JaccardAggregator(WeightedAggregator):
-	@staticmethod
-	def compute_weights(annotations: pd.DataFrame) -> pd.Series:
-		vote_size = annotations.sum(axis=1)
-		return 1 / vote_size
+    @staticmethod
+    def compute_weights(annotations: pd.DataFrame) -> pd.Series:
+        vote_size = annotations.sum(axis=1)
+        assert np.all(vote_size > 0), "Jaccard weights are not defined for empty votes"
+        return 1 / vote_size
 
 
 class DiceAggregator(WeightedAggregator):
-	@staticmethod
-	def compute_weights(annotations: pd.DataFrame) -> pd.Series:
-		vote_size = annotations.sum(axis=1)
-		return 2 / (vote_size + 1)
+    @staticmethod
+    def compute_weights(annotations: pd.DataFrame) -> pd.Series:
+        vote_size = annotations.sum(axis=1)
+        return 2 / (vote_size + 1)

tests/core/aggregation/aggregators.py/test_mallows.py

@@ -0,0 +1,67 @@
+import pandas as pd
+import pytest
+
+
+@pytest.mark.ut
+@pytest.mark.parametrize(
+    ["annotations", "expected_result"],
+    [
+        (
+            pd.DataFrame(
+                {"task": ["q1", "q1"], "worker": ["v1", "v2"], "a": [1, 0], "b": [0, 1]}
+            ).set_index(["task", "worker"]),
+            pd.Series(
+                [1, 1],
+                index=pd.MultiIndex.from_tuples(
+                    [("q1", "v1"), ("q1", "v2")], names=["task", "worker"]
+                ),
+            ),
+        ),
+    ],
+)
+def test_StandardApprovalAggregator_compute_weights(
+    annotations: pd.DataFrame, expected_result: pd.Series
+) -> None:
+    # Given
+    from hakeem.core.aggregation.aggregators.mallows import StandardApprovalAggregator
+
+    # When
+    result = StandardApprovalAggregator().compute_weights(annotations)
+
+    # Then
+    pd.testing.assert_series_equal(expected_result, result)
+
+
+@pytest.mark.ut
+@pytest.mark.parametrize(
+    ["annotations", "expected_result"],
+    [
+        (
+            pd.DataFrame(
+                {
+                    "task": ["q1", "q1", "q2"],
+                    "worker": ["v1", "v2", "v1"],
+                    "a": [1, 1, 1],
+                    "b": [0, 1, 1],
+                }
+            ).set_index(["task", "worker"]),
+            pd.Series(
+                [1, 0.5, 0.5],
+                index=pd.MultiIndex.from_tuples(
+                    [("q1", "v1"), ("q1", "v2"), ("q2", "v1")], names=["task", "worker"]
+                ),
+            ),
+        ),
+    ],
+)
+def test_JaccardAggregator_compute_weights(
+    annotations: pd.DataFrame, expected_result: pd.Series
+) -> None:
+    # Given
+    from hakeem.core.aggregation.aggregators.mallows import JaccardAggregator
+
+    # When
+    result = JaccardAggregator().compute_weights(annotations)
+
+    # Then
+    pd.testing.assert_series_equal(expected_result, result)