ENH: Minor change to parallel_coordinates (pandas-dev#15908)

Add option to sort class lables, add to test

Author: stangirala

pandas/tests/plotting/test_misc.py

@@ -241,6 +241,22 @@ def test_parallel_coordinates(self):
         with tm.assert_produces_warning(FutureWarning):
             parallel_coordinates(df, 'Name', colors=colors)
 
+        df = DataFrame({"feat": [i for i in range(30)],
+                        "class": [2 for _ in range(10)] +
+                        [3 for _ in range(10)] +
+                        [1 for _ in range(10)]})
+        ax = parallel_coordinates(df, 'class', sort_labels=True)
+        polylines, lables = ax.get_legend_handles_labels()
+        color_label_tuples = \
+            zip([polyline.get_color() for polyline in polylines], labels)
+        ordered_color_label_tuples = sorted(color_label_tuples,
+                                            key=lambda x: x[1])
+        prev_next_tupels = zip([i for i in ordered_color_label_tuples[0:-1]],
+                               [i for i in ordered_color_label_tuples[1:]])
+        for prev, nxt in prev_next_tupels:
+            # lables and colors are ordered strictly increasing
+            assert prev[1] < nxt[1] and prev[0] < nxt[0]
+
     @slow
     def test_radviz(self):
         from pandas.tools.plotting import radviz

pandas/tools/plotting.py

@@ -694,7 +694,8 @@ def bootstrap_plot(series, fig=None, size=50, samples=500, **kwds):
 @deprecate_kwarg(old_arg_name='data', new_arg_name='frame', stacklevel=3)
 def parallel_coordinates(frame, class_column, cols=None, ax=None, color=None,
                          use_columns=False, xticks=None, colormap=None,
-                         axvlines=True, axvlines_kwds=None, **kwds):
+                         axvlines=True, axvlines_kwds=None, sort_labels=False,
+                         **kwds):
     """Parallel coordinates plotting.
 
     Parameters
@@ -718,6 +719,8 @@ def parallel_coordinates(frame, class_column, cols=None, ax=None, color=None,
         If true, vertical lines will be added at each xtick
     axvlines_kwds: keywords, optional
         Options to be passed to axvline method for vertical lines
+    sort_labels: bool, optional
+        Sort class_column labels, useful when assigning colours
     kwds: keywords
         Options to pass to matplotlib plotting method
 
@@ -774,6 +777,8 @@ def parallel_coordinates(frame, class_column, cols=None, ax=None, color=None,
                                         colormap=colormap, color_type='random',
                                         color=color)
 
+    if sort_labels is True:
+        classes = sorted(classes)
     colors = dict(zip(classes, color_values))
 
     for i in range(n):