ENH dataframe in/out for all samplers

Author: glemaitre

doc/whats_new/v0.6.rst

@@ -54,7 +54,11 @@ Enhancement
 - :class:`imblearn.under_sampling.RandomUnderSampler`,
   :class:`imblearn.over_sampling.RandomOverSampler` can resample when non
   finite values are present in ``X``.
-  :pr:`643` by `Guillaume Lemaitre <glemaitre>`.
+  :pr:`643` by :user:`Guillaume Lemaitre <glemaitre>`.
+
+- All samplers will output a Pandas DataFrame if a Pandas DataFrame was given
+  as an input.
+  :pr:`644` by :user:`Guillaume Lemaitre <glemaitre>`.
 
 Deprecation
 ...........

imblearn/base.py

@@ -78,12 +78,20 @@ def fit_resample(self, X, y):
 
         output = self._fit_resample(X, y)
 
+        if self._columns is not None:
+            import pandas as pd
+            X_ = pd.DataFrame(output[0], columns=self._columns)
+        else:
+            X_ = output[0]
+
         if binarize_y:
             y_sampled = label_binarize(output[1], np.unique(y))
             if len(output) == 2:
-                return output[0], y_sampled
-            return output[0], y_sampled, output[2]
-        return output
+                return X_, y_sampled
+            return X_, y_sampled, output[2]
+        if len(output) == 2:
+            return X_, output[1]
+        return X_, output[1], output[2]
 
     #  define an alias for back-compatibility
     fit_sample = fit_resample
@@ -124,8 +132,9 @@ class BaseSampler(SamplerMixin):
     def __init__(self, sampling_strategy="auto"):
         self.sampling_strategy = sampling_strategy
 
-    @staticmethod
-    def _check_X_y(X, y, accept_sparse=None):
+    def _check_X_y(self, X, y, accept_sparse=None):
+        # store the columns name to reconstruct a dataframe
+        self._columns = X.columns if hasattr(X, "loc") else None
         if accept_sparse is None:
             accept_sparse = ["csr", "csc"]
         y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
@@ -238,6 +247,8 @@ def fit_resample(self, X, y):
         y_resampled : array-like of shape (n_samples_new,)
             The corresponding label of `X_resampled`.
         """
+        # store the columns name to reconstruct a dataframe
+        self._columns = X.columns if hasattr(X, "loc") else None
         if self.validate:
             check_classification_targets(y)
             X, y, binarize_y = self._check_X_y(
@@ -250,12 +261,20 @@ def fit_resample(self, X, y):
 
         output = self._fit_resample(X, y)
 
+        if self._columns is not None:
+            import pandas as pd
+            X_ = pd.DataFrame(output[0], columns=self._columns)
+        else:
+            X_ = output[0]
+
         if self.validate and binarize_y:
             y_sampled = label_binarize(output[1], np.unique(y))
             if len(output) == 2:
-                return output[0], y_sampled
-            return output[0], y_sampled, output[2]
-        return output
+                return X_, y_sampled
+            return X_, y_sampled, output[2]
+        if len(output) == 2:
+            return X_, output[1]
+        return X_, output[1], output[2]
 
     def _fit_resample(self, X, y):
         func = _identity if self.func is None else self.func

imblearn/over_sampling/_random_over_sampler.py

@@ -74,13 +74,12 @@ def __init__(self, sampling_strategy="auto", random_state=None):
         super().__init__(sampling_strategy=sampling_strategy)
         self.random_state = random_state
 
-    @staticmethod
-    def _check_X_y(X, y):
+    def _check_X_y(self, X, y):
+        # store the columns name to reconstruct a dataframe
+        self._columns = X.columns if hasattr(X, "loc") else None
         y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
-        if not hasattr(X, "loc"):
-            # Do not convert dataframe
-            X = check_array(X, accept_sparse=["csr", "csc"], dtype=None,
-                            force_all_finite=False)
+        X = check_array(X, accept_sparse=["csr", "csc"], dtype=None,
+                        force_all_finite=False)
         y = check_array(
             y, accept_sparse=["csr", "csc"], dtype=None, ensure_2d=False
         )

imblearn/under_sampling/_prototype_selection/_random_under_sampler.py

@@ -80,13 +80,12 @@ def __init__(
         self.random_state = random_state
         self.replacement = replacement
 
-    @staticmethod
-    def _check_X_y(X, y):
+    def _check_X_y(self, X, y):
+        # store the columns name to reconstruct a dataframe
+        self._columns = X.columns if hasattr(X, "loc") else None
         y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
-        if not hasattr(X, "loc"):
-            # Do not convert dataframe
-            X = check_array(X, accept_sparse=["csr", "csc"], dtype=None,
-                            force_all_finite=False)
+        X = check_array(X, accept_sparse=["csr", "csc"], dtype=None,
+                        force_all_finite=False)
         y = check_array(
             y, accept_sparse=["csr", "csc"], dtype=None, ensure_2d=False
         )

imblearn/utils/estimator_checks.py

@@ -262,7 +262,7 @@ def check_samplers_pandas(name, Sampler):
         weights=[0.2, 0.3, 0.5],
         random_state=0,
     )
-    X_pd = pd.DataFrame(X)
+    X_pd = pd.DataFrame(X, columns=[str(i) for i in range(X.shape[1])])
     sampler = Sampler()
     if isinstance(Sampler(), NearMiss):
         samplers = [Sampler(version=version) for version in (1, 2, 3)]
@@ -274,7 +274,11 @@ def check_samplers_pandas(name, Sampler):
         set_random_state(sampler)
         X_res_pd, y_res_pd = sampler.fit_resample(X_pd, y)
         X_res, y_res = sampler.fit_resample(X, y)
-        assert_allclose(X_res_pd, X_res)
+
+        # check that we return a pandas dataframe if a dataframe was given in
+        assert isinstance(X_res_pd, pd.DataFrame)
+        assert X_pd.columns.to_list() == X_res_pd.columns.to_list()
+        assert_allclose(X_res_pd.to_numpy(), X_res)
         assert_allclose(y_res_pd, y_res)