Merge pull request #803 from tonysyu/streamplot-return-arrows

efiring · efiring · commit 7a61238c930d · 2012-09-02T10:24:17.000-07:00
Return arrow collection as 2nd argument of streamplot.
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -6601,17 +6601,17 @@ def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
                    cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
                    minlength=0.1, transform=None):
         if not self._hold: self.cla()
-        lines = mstream.streamplot(self, x, y, u, v,
-                                   density=density,
-                                   linewidth=linewidth,
-                                   color=color,
-                                   cmap=cmap,
-                                   norm=norm,
-                                   arrowsize=arrowsize,
-                                   arrowstyle=arrowstyle,
-                                   minlength=minlength,
-                                   transform=transform)
-        return lines
+        stream_container = mstream.streamplot(self, x, y, u, v,
+                                              density=density,
+                                              linewidth=linewidth,
+                                              color=color,
+                                              cmap=cmap,
+                                              norm=norm,
+                                              arrowsize=arrowsize,
+                                              arrowstyle=arrowstyle,
+                                              minlength=minlength,
+                                              transform=transform)
+        return stream_container
     streamplot.__doc__ = mstream.streamplot.__doc__
 
     @docstring.dedent_interpd
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.lines)
     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
@@ -9,6 +9,7 @@
 import matplotlib.collections as mcollections
 import matplotlib.patches as patches
 
+
 __all__ = ['streamplot']
 
 
@@ -46,11 +47,16 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
     *minlength* : float
         Minimum length of streamline in axes coordinates.
 
-    Returns *streamlines* : :class:`~matplotlib.collections.LineCollection`
-        Line collection with all streamlines as a series of line segments.
-        Currently, there is no way to differentiate between line segments
-        on different streamlines (other than manually checking that segments
-        are connected).
+    Returns
+    -------
+    *stream_container* : StreamplotSet
+        Container object with attributes
+            lines : `matplotlib.collections.LineCollection` of streamlines
+            arrows : collection of `matplotlib.patches.FancyArrowPatch` objects
+                repesenting arrows half-way along stream lines.
+        This container will probably change in the future to allow changes to
+        the colormap, alpha, etc. for both lines and arrows, but these changes
+        should be backward compatible.
     """
     grid = Grid(x, y)
     mask = StreamMask(density)
@@ -108,6 +114,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 +146,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 +159,17 @@ 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
+
+    ac = matplotlib.collections.PatchCollection(arrows)
+    stream_container = StreamplotSet(lc, ac)
+    return stream_container
+
+
+class StreamplotSet(object):
+
+    def __init__(self, lines, arrows, **kwargs):
+        self.lines = lines
+        self.arrows = arrows
 
 
 # Coordinate definitions
