iter

Author: glemaitre

imblearn/base.py

@@ -18,7 +18,10 @@
     _deprecate_positional_args,
     get_classes_counts,
 )
-from .utils.wrapper import check_classification_targets
+from .utils.wrapper import (
+    check_classification_targets,
+    label_binarize,
+)
 
 
 class SamplerMixin(BaseEstimator, metaclass=ABCMeta):
@@ -49,7 +52,11 @@ def fit(self, X, y):
         self : object
             Return the instance itself.
         """
+        arrays_transformer = ArraysTransformer(X, y)
         dask_collection = any([is_dask_collection(arr) for arr in (X, y)])
+        if dask_collection:
+            X, y = arrays_transformer.to_dask_array(X, y)
+
         if (not dask_collection or
                 (dask_collection and self.validate_if_dask_collection)):
             X, y, _ = self._check_X_y(X, y)
@@ -83,12 +90,14 @@ def fit_resample(self, X, y):
         """
         arrays_transformer = ArraysTransformer(X, y)
         dask_collection = any([is_dask_collection(arr) for arr in (X, y)])
+        if dask_collection:
+            X, y = arrays_transformer.to_dask_array(X, y)
+
         if (not dask_collection or
                 (dask_collection and self.validate_if_dask_collection)):
             check_classification_targets(y)
             X, y, binarize_y = self._check_X_y(X, y)
         else:
-            X, y = arrays_transformer.to_dask_array(X, y)
             binarize_y = False
 
         self._classes_counts = get_classes_counts(y)
@@ -98,9 +107,10 @@ def fit_resample(self, X, y):
 
         output = self._fit_resample(X, y)
 
-        # TODO: label binarize is not implemented with dask
-        y_ = (label_binarize(output[1], np.unique(y))
-              if binarize_y else output[1])
+        if binarize_y:
+            y_ = label_binarize(output[1], classes=np.unique(y))
+        else:
+            y_ = output[1]
 
         X_, y_ = arrays_transformer.transform(output[0], y_)
         return (X_, y_) if len(output) == 2 else (X_, y_, output[2])
@@ -281,9 +291,10 @@ def fit_resample(self, X, y):
         output = self._fit_resample(X, y)
 
         if self.validate:
-
-            y_ = (label_binarize(output[1], np.unique(y))
-                  if binarize_y else output[1])
+            if binarize_y:
+                y_ = label_binarize(output[1], classes=np.unique(y))
+            else:
+                y_ = output[1]
             X_, y_ = arrays_transformer.transform(output[0], y_)
             return (X_, y_) if len(output) == 2 else (X_, y_, output[2])
 

imblearn/dask/preprocessing.py

@@ -0,0 +1,10 @@
+import numpy as np
+
+
+def label_binarize(y, *, classes):
+    import pandas as pd
+    from dask import dataframe
+
+    cat_dtype = pd.CategoricalDtype(categories=classes)
+    y = dataframe.from_array(y).astype(cat_dtype)
+    return dataframe.get_dummies(y).to_dask_array()

imblearn/utils/estimator_checks.py

@@ -295,9 +295,10 @@ def check_samplers_pandas(name, sampler):
     assert_allclose(y_res_s.to_numpy(), y_res)
 
 
-def check_samplers_dask_array(name, sampler):
+def check_samplers_dask_array(name, sampler_orig):
     pytest.importorskip("dask")
     from dask import array
+    sampler = clone(sampler_orig)
     # Check that the samplers handle dask array
     X, y = make_classification(
         n_samples=1000,
@@ -309,20 +310,25 @@ def check_samplers_dask_array(name, sampler):
     X_dask = array.from_array(X, chunks=100)
     y_dask = array.from_array(y, chunks=100)
 
-    X_res_dask, y_res_dask = sampler.fit_resample(X_dask, y_dask)
-    X_res, y_res = sampler.fit_resample(X, y)
+    for validate_if_dask_collection in (True, False):
+        sampler.set_params(
+            validate_if_dask_collection=validate_if_dask_collection
+        )
+        X_res_dask, y_res_dask = sampler.fit_resample(X_dask, y_dask)
+        X_res, y_res = sampler.fit_resample(X, y)
 
-    # check that we return the same type for dataframes or series types
-    assert isinstance(X_res_dask, array.Array)
-    assert isinstance(y_res_dask, array.Array)
+        # check that we return the same type for dataframes or series types
+        assert isinstance(X_res_dask, array.Array)
+        assert isinstance(y_res_dask, array.Array)
 
-    assert_allclose(X_res_dask, X_res)
-    assert_allclose(y_res_dask, y_res)
+        assert_allclose(X_res_dask, X_res)
+        assert_allclose(y_res_dask, y_res)
 
 
-def check_samplers_dask_dataframe(name, sampler):
+def check_samplers_dask_dataframe(name, sampler_orig):
     pytest.importorskip("dask")
     from dask import dataframe
+    sampler = clone(sampler_orig)
     # Check that the samplers handle dask dataframe and dask series
     X, y = make_classification(
         n_samples=1000,
@@ -337,18 +343,22 @@ def check_samplers_dask_dataframe(name, sampler):
     y_s = dataframe.from_array(y)
     y_s = y_s.rename("target")
 
-    X_res_df, y_res_s = sampler.fit_resample(X_df, y_s)
-    X_res, y_res = sampler.fit_resample(X, y)
+    for validate_if_dask_collection in (True, False):
+        sampler.set_params(
+            validate_if_dask_collection=validate_if_dask_collection
+        )
+        X_res_df, y_res_s = sampler.fit_resample(X_df, y_s)
+        X_res, y_res = sampler.fit_resample(X, y)
 
-    # check that we return the same type for dataframes or series types
-    assert isinstance(X_res_df, dataframe.DataFrame)
-    assert isinstance(y_res_s, dataframe.Series)
+        # check that we return the same type for dataframes or series types
+        assert isinstance(X_res_df, dataframe.DataFrame)
+        assert isinstance(y_res_s, dataframe.Series)
 
-    assert X_df.columns.to_list() == X_res_df.columns.to_list()
-    assert y_s.name == y_res_s.name
+        assert X_df.columns.to_list() == X_res_df.columns.to_list()
+        assert y_s.name == y_res_s.name
 
-    assert_allclose(np.array(X_res_df), X_res)
-    assert_allclose(np.array(y_res_s), y_res)
+        assert_allclose(np.array(X_res_df), X_res)
+        assert_allclose(np.array(y_res_s), y_res)
 
 
 def check_samplers_list(name, sampler):

imblearn/utils/wrapper.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from sklearn.preprocessing import label_binarize as sklearn_label_binarize
 from sklearn.utils.multiclass import check_classification_targets as \
     sklearn_check_classification_targets
 from sklearn.utils.multiclass import type_of_target as sklearn_type_of_target
@@ -49,3 +50,11 @@ def check_classification_targets(y):
 
         return dask_check_classification_targets(y)
     return sklearn_check_classification_targets(y)
+
+
+def label_binarize(y, *, classes):
+    if is_dask_collection(y):
+        from ..dask.preprocessing import label_binarize as dask_label_binarize
+
+        return dask_label_binarize(y, classes=classes)
+    return sklearn_label_binarize(y, classes=classes)