removes for loop improving runtime of baseline imputer (#498)

Author: mmschlk

src/shapiq/imputer/baseline_imputer.py

@@ -105,10 +105,7 @@ def value_function(self, coalitions: np.ndarray) -> np.ndarray:
                ``(n_subsets, n_outputs)``.
 
         """
-        n_coalitions = coalitions.shape[0]
-        data = np.tile(np.copy(self.x), (n_coalitions, 1))
-        for i in range(n_coalitions):
-            data[i, ~coalitions[i]] = self.baseline_values[0, ~coalitions[i]]
+        data = np.where(coalitions, self.x, self.baseline_values)
         return self.predict(data)
 
     def init_background(self, data: np.ndarray) -> BaselineImputer:

tests/shapiq/tests_unit/tests_imputer/test_baseline_imputer.py

@@ -93,6 +93,27 @@ def model_cat(x: np.ndarray) -> np.ndarray:
     assert imputer.baseline_values[0, 2] == np.mean(data[:, 2])
 
 
+def test_baseline_value_function_imputes_correctly():
+    """Test that value_function substitutes baseline values for absent features."""
+    captured = []
+
+    def model(x: np.ndarray) -> np.ndarray:
+        captured.append(x.copy())
+        return np.zeros(x.shape[0])
+
+    x = np.array([[1, 2, 3]])
+    baseline = np.array([[10, 20, 30]])
+    imputer = BaselineImputer(model=model, data=baseline, x=x)
+
+    coalitions = np.array([[True, False, True], [False, True, False]])
+    imputer(coalitions)
+
+    assert len(captured) == 1
+    result = captured[0]
+    np.testing.assert_array_equal(result[0], [1, 20, 3])
+    np.testing.assert_array_equal(result[1], [10, 2, 30])
+
+
 def test_baseline_imputer_init():
     """Test the initialization of the marginal imputer."""