DOC improve make_imbalance example

Author: glemaitre

examples/datasets/plot_make_imbalance.py

@@ -1,9 +1,11 @@
 """
-===========================
-make_imbalance function
-===========================
+============================
+Create an imbalanced dataset
+============================
 
-An illustration of the make_imbalance function
+An illustration of the :func:`imblearn.datasets.make_imbalance` function to
+create an imbalanced dataset from a balanced dataset. We show the ability of
+:func:`imblearn.datasets.make_imbalance` of dealing with Pandas DataFrame.
 
 """
 
@@ -14,36 +16,29 @@
 
 from collections import Counter
 
+import pandas as pd
 import matplotlib.pyplot as plt
+
 from sklearn.datasets import make_moons
 
 from imblearn.datasets import make_imbalance
 
 print(__doc__)
 
-
-def plot_decoration(ax):
-    ax.spines['top'].set_visible(False)
-    ax.spines['right'].set_visible(False)
-    ax.get_xaxis().tick_bottom()
-    ax.get_yaxis().tick_left()
-    ax.spines['left'].set_position(('outward', 10))
-    ax.spines['bottom'].set_position(('outward', 10))
-    ax.set_xlim([-4, 4])
-
-
 # Generate the dataset
 X, y = make_moons(n_samples=200, shuffle=True, noise=0.5, random_state=10)
+X = pd.DataFrame(X, columns=["feature 1", "feature 2"])
 
 # Two subplots, unpack the axes array immediately
 f, axs = plt.subplots(2, 3)
 
 axs = [a for ax in axs for a in ax]
 
-axs[0].scatter(X[y == 0, 0], X[y == 0, 1], label="Class #0", alpha=0.5)
-axs[0].scatter(X[y == 1, 0], X[y == 1, 1], label="Class #1", alpha=0.5)
+X.plot.scatter(
+    x='feature 1', y='feature 2', c=y, ax=axs[0], colormap='viridis',
+    colorbar=False
+)
 axs[0].set_title('Original set')
-plot_decoration(axs[0])
 
 
 def ratio_func(y, multiplier, minority_class):
@@ -58,10 +53,11 @@ def ratio_func(y, multiplier, minority_class):
     X_, y_ = make_imbalance(X, y, sampling_strategy=ratio_func,
                             **{"multiplier": multiplier,
                                "minority_class": 1})
-    ax.scatter(X_[y_ == 0, 0], X_[y_ == 0, 1], label="Class #0", alpha=0.5)
-    ax.scatter(X_[y_ == 1, 0], X_[y_ == 1, 1], label="Class #1", alpha=0.5)
-    ax.set_title('sampling_strategy = {}'.format(multiplier))
-    plot_decoration(ax)
+    X_.plot.scatter(
+        x='feature 1', y='feature 2', c=y_, ax=ax, colormap='viridis',
+        colorbar=False
+    )
+    ax.set_title('Sampling ratio = {}'.format(multiplier))
 
 plt.tight_layout()
 plt.show()