hotfix bug in feature-cover-injector

.gitignore

@@ -14,6 +14,7 @@ _build
 *.coverage
 *.DS_Store
 .idea/
+tmp
 
 # Images
 

menelaus/injection/feature_manipulation.py

@@ -105,7 +105,8 @@ def __call__(self, data, col, sample_size, random_state=None):
         ret, (col,) = self._preprocess(data, col, return_df=True)
 
         # hide and reorder
-        ret = ret.groupby(col).sample(n=sample_size, random_state=random_state)
+        n = sample_size // len(ret[col].unique())
+        ret = ret.groupby(col).sample(n=n, random_state=random_state)
         ret = ret.drop(columns=[col]).reset_index(drop=True)
 
         # handle type and return

tests/menelaus/injection/test_feature_manipulation.py

@@ -16,9 +16,8 @@ def test_feature_cover_1():
     data = np.array([[0,2], [0,2], [0,2], [1,3], [1,3], [1,3]])
     i = FeatureCoverInjector()
     copy = i(data, col=0, sample_size=2, random_state=0)
-    assert copy.shape == (4,1)
-    assert np.array_equal(np.where(copy[:,0]==2)[0], [0,1])
-    assert np.array_equal(np.where(copy[:,0]==3)[0], [2,3])
+    assert copy.shape == (2,1)
+    assert np.array_equal(copy, [[2], [3]])
 
 def test_feature_shift_1():
     ''' Check correct feature cover behavior '''