Return arrow collection as 2nd argument of streamplot.

tonysyu · Sep 2, 2012 · 391c24a · 391c24a
1 parent cadd152
commit 391c24a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -3055,7 +3055,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
         draw_if_interactive()
     finally:
         ax.hold(washold)
-    sci(ret)
+    sci(ret[0])
     return ret
 
 # This function was autogenerated by boilerplate.py.  Do not edit as

diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -108,6 +108,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
             cmap = cm.get_cmap(cmap)
 
     streamlines = []
+    arrows = []
     for t in trajectories:
         tgx = np.array(t[0])
         tgy = np.array(t[1])
@@ -139,6 +140,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
                                     transform=transform, 
                                     **arrow_kw)
         axes.add_patch(p)
+        arrows.append(p)
 
     lc = mcollections.LineCollection(streamlines, 
                                      transform=transform, 
@@ -151,7 +153,9 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
 
     axes.update_datalim(((x.min(), y.min()), (x.max(), y.max())))
     axes.autoscale_view(tight=True)
-    return lc
+
+    arrow_collection = matplotlib.collections.PatchCollection(arrows)
+    return lc, arrow_collection
 
 
 # Coordinate definitions