Merge pull request #5061 from grlee77/fast_voronoi_plotting

ENH: spatial: faster 2D Voronoi and Convex Hull plotting

scipy/spatial/_plotutils.py

@@ -94,20 +94,25 @@ def convex_hull_plot_2d(hull, ax=None):
     Requires Matplotlib.
 
     """
+    from matplotlib.collections import LineCollection
+
     if hull.points.shape[1] != 2:
         raise ValueError("Convex hull is not 2-D")
 
     ax.plot(hull.points[:,0], hull.points[:,1], 'o')
+    line_segments = []
     for simplex in hull.simplices:
-        ax.plot(hull.points[simplex,0], hull.points[simplex,1], 'k-')
-
+        line_segments.append([(x, y) for x, y in hull.points[simplex]])
+    ax.add_collection(LineCollection(line_segments,
+                                     colors='k',
+                                     linestyle='solid'))
     _adjust_bounds(ax, hull.points)
 
     return ax.figure
 
 
 @_held_figure
-def voronoi_plot_2d(vor, ax=None):
+def voronoi_plot_2d(vor, ax=None, **kw):
     """
     Plot the given Voronoi diagram in 2-D
 
@@ -117,6 +122,8 @@ def voronoi_plot_2d(vor, ax=None):
         Diagram to plot
     ax : matplotlib.axes.Axes instance, optional
         Axes to plot on
+    show_vertices : bool, optional
+        Add the Voronoi vertices to the plot.
 
     Returns
     -------
@@ -132,19 +139,27 @@ def voronoi_plot_2d(vor, ax=None):
     Requires Matplotlib.
 
     """
+    from matplotlib.collections import LineCollection
+
     if vor.points.shape[1] != 2:
         raise ValueError("Voronoi diagram is not 2-D")
 
     ax.plot(vor.points[:,0], vor.points[:,1], '.')
-    ax.plot(vor.vertices[:,0], vor.vertices[:,1], 'o')
+    if kw.get('show_vertices', True):
+        ax.plot(vor.vertices[:,0], vor.vertices[:,1], 'o')
 
+    line_segments = []
     for simplex in vor.ridge_vertices:
         simplex = np.asarray(simplex)
         if np.all(simplex >= 0):
-            ax.plot(vor.vertices[simplex,0], vor.vertices[simplex,1], 'k-')
+            line_segments.append([(x, y) for x, y in vor.vertices[simplex]])
 
+    ax.add_collection(LineCollection(line_segments,
+                                     colors='k',
+                                     linestyle='solid'))
     ptp_bound = vor.points.ptp(axis=0)
 
+    line_segments = []
     center = vor.points.mean(axis=0)
     for pointidx, simplex in zip(vor.ridge_points, vor.ridge_vertices):
         simplex = np.asarray(simplex)
@@ -159,9 +174,12 @@ def voronoi_plot_2d(vor, ax=None):
             direction = np.sign(np.dot(midpoint - center, n)) * n
             far_point = vor.vertices[i] + direction * ptp_bound.max()
 
-            ax.plot([vor.vertices[i,0], far_point[0]],
-                    [vor.vertices[i,1], far_point[1]], 'k--')
+            line_segments.append([(vor.vertices[i, 0], vor.vertices[i, 1]),
+                                  (far_point[0], far_point[1])])
 
+    ax.add_collection(LineCollection(line_segments,
+                                     colors='k',
+                                     linestyle='dashed'))
     _adjust_bounds(ax, vor.points)
 
     return ax.figure

scipy/spatial/tests/test__plotutils.py

@@ -6,6 +6,7 @@
     import matplotlib
     matplotlib.rcParams['backend'] = 'Agg'
     import matplotlib.pyplot as plt
+    from matplotlib.collections import LineCollection
     has_matplotlib = True
 except:
     has_matplotlib = False
@@ -37,6 +38,7 @@ def test_voronoi(self):
         r = voronoi_plot_2d(obj, ax=fig.gca())
         assert_(r is fig)
         voronoi_plot_2d(obj)
+        voronoi_plot_2d(obj, show_vertices=False)
 
     @dec.skipif(not has_matplotlib, "Matplotlib not available")
     def test_convex_hull(self):