Add Game.compute function and related test to the Game class

shapiq/games/base.py

@@ -419,6 +419,30 @@ def precompute(self, coalitions: np.ndarray | None = None) -> None:
         self.coalition_lookup = coalitions_dict
         self.precompute_flag = True
 
+    def compute(
+        self, coalitions: np.ndarray | None = None, *, return_normalization: bool = False
+    ) -> tuple:
+        """Compute the game values for all or a given set of coalitions.
+
+        Args:
+            coalitions: The coalitions to evaluate.
+            return_normalization: Whether to return the normalization value. Defaults to ``False``.
+
+        Returns:
+            A tuple containing:
+            - np.ndarray: The values of the coalitions.
+            - dict[tuple[int, ...], int]: The lookup of the coalitions
+            - float: The normalization value (optional, if return_normalization is 'True')
+
+        """
+        coalitions: np.ndarray = self._check_coalitions(coalitions)
+        values = self.value_function(coalitions)
+        game_values = values - self.normalization_value
+
+        if return_normalization:
+            return (game_values, self.coalition_lookup, self.normalization_value)
+        return game_values, self.coalition_lookup
+
     def save_values(self, path: Path | str) -> None:
         """Saves the game values to the given path.
 

tests/tests_games/test_base_game.py

@@ -326,3 +326,28 @@ def test_exact_computer_call():
     sv = game.exact_values(index=index, order=order)
     assert sv.index == index
     assert sv.max_order == order
+
+
+def test_compute_with_and_without_normalization():
+    """Tests the compute function with and without returned normalization."""
+    normalization_value = 1.0  # not zero
+
+    n_players = 3
+    game = DummyGame(n=n_players, interaction=(0, 1))
+
+    coalitions = np.array([[1, 0, 0], [0, 1, 1]])
+
+    # Make sure normalization value is added
+    game.normalization_value = normalization_value
+    assert game.normalize
+
+    # Test without returned normalization
+    result = game.compute(coalitions=coalitions, return_normalization=False)
+    assert len(result[0]) == len(coalitions)
+    assert len(result) == 2  # game_values and coalition_lookup
+
+    # Test with returned normalization
+    result = game.compute(coalitions=coalitions, return_normalization=True)
+    assert len(result[0]) == len(coalitions)
+    assert result[2] == 1.0  # normalization_value
+    assert len(result) == 3  # game_values, normalization_value and coalition_lookup