Allow kwargs in FunctionSampler

Author: glemaitre

imblearn/misc.py

@@ -30,12 +30,10 @@ class FunctionSampler(SamplerMixin):
 
     """
 
-    def __init__(self, func=None, accept_sparse=True, kw_args=None,
-                 random_state=None):
+    def __init__(self, func=None, accept_sparse=True, kw_args=None):
         self.func = func
         self.accept_sparse = accept_sparse
         self.kw_args = kw_args
-        self.random_state = random_state
         self.logger = logging.getLogger(__name__)
 
     def _check_X_y(self, X, y):
@@ -48,15 +46,14 @@ def _check_X_y(self, X, y):
         return X, y
 
     def fit(self, X, y):
-        print(self.accept_sparse)
         X, y = self._check_X_y(X, y)
         self.X_hash_, self.y_hash_ = hash_X_y(X, y)
         # when using a sampler, ratio_ is supposed to exist after fit
         self.ratio_ = 'is_fitted'
 
         return self
 
-    def _sample(self, X, y, func, kw_args):
+    def _sample(self, X, y, func=None, kw_args=None):
         X, y = self._check_X_y(X, y)
         check_is_fitted(self, 'ratio_')
         X_hash, y_hash = hash_X_y(X, y)
@@ -66,7 +63,7 @@ def _sample(self, X, y, func, kw_args):
         if func is None:
             func = _identity
 
-        return func(X, y, **(kw_args if kw_args else {}))
+        return func(X, y, **(kw_args if self.kw_args else {}))
 
     def sample(self, X, y):
         return self._sample(X, y, func=self.func, kw_args=self.kw_args)

imblearn/tests/test_misc.py

@@ -3,8 +3,45 @@
 # Authors: Guillaume Lemaitre <[email protected]>
 # License: MIT
 
+from sklearn.datasets import load_iris
+from sklearn.utils.testing import assert_array_equal
+
+from imblearn.datasets import make_imbalance
 from imblearn.misc import FunctionSampler
+from imblearn.under_sampling import RandomUnderSampler
+
+iris = load_iris()
+X, y = make_imbalance(iris.data, iris.target, ratio={0: 10, 1: 25},
+                      random_state=0)
+
+
+def test_function_sampler_identity():
+    sampler = FunctionSampler()
+    X_res, y_res = sampler.fit_sample(X, y)
+    assert_array_equal(X_res, X)
+    assert_array_equal(y_res, y)
+
+
+def test_function_sampler_func():
+
+    def func(X, y):
+        return X[:10], y[:10]
+
+    sampler = FunctionSampler(func=func)
+    X_res, y_res = sampler.fit_sample(X, y)
+    assert_array_equal(X_res, X[:10])
+    assert_array_equal(y_res, y[:10])
+
+
+def test_function_sampler_func_kwargs():
 
+    def func(X, y, ratio, random_state):
+        rus = RandomUnderSampler(ratio=ratio, random_state=random_state)
+        return rus.fit_sample(X, y)
 
-def function_sampler_identity():
-    sampler = FunctionSampler(1)
+    sampler = FunctionSampler(func=func, kw_args={'ratio': 'auto',
+                                                  'random_state': 0})
+    X_res, y_res = sampler.fit_sample(X, y)
+    X_res_2, y_res_2 = RandomUnderSampler(random_state=0).fit_sample(X, y)
+    assert_array_equal(X_res, X_res_2)
+    assert_array_equal(y_res, y_res_2)