Add streamplot function to pyplot.

As a temporary fix, streamplot returns the last plotted streamline and arrow for use with colorbars and legends. Instead, all lines should instead be grouped into a single LineCollection, and all FancyArrowPatches should be grouped into a patch collection.

Author: tonysyu

CHANGELOG

@@ -1,3 +1,5 @@
+2012-01-08 Add axes.streamplot to plot streamlines of a velocity field. -TSY
+
 2011-12-29 ps and pdf markers are now stroked only if the line width
            is nonzero for consistency with agg, fixes issue #621. - JKS
 

boilerplate.py

@@ -89,6 +89,7 @@ def %(func)s(%(argspec)s):
     #'spy',
     'stem',
     'step',
+    'streamplot',
     'tricontour',
     'tricontourf',
     'tripcolor',
@@ -123,6 +124,7 @@ def %(func)s(%(argspec)s):
     #'spy'    :  'sci(%(ret)s)',  ### may return image or Line2D
     'quiver' :  'sci(%(ret)s)',
     'specgram'  : 'sci(%(ret)s[-1])',
+    'streamplot' :  'sci(%(ret)s[0])',
     'tricontour' : 'if %(ret)s._A is not None: sci(%(ret)s)',
     'tricontourf': 'if %(ret)s._A is not None: sci(%(ret)s)',
     'tripcolor'  : 'sci(%(ret)s)',

lib/matplotlib/axes.py

@@ -6307,14 +6307,15 @@ def quiver(self, *args, **kw):
     def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
                    cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.1):
         if not self._hold: self.cla()
-        mstream.streamplot(self, x, y, u, v,
-                           density=density,
-                           linewidth=linewidth,
-                           color=color,
-                           cmap=cmap,
-                           arrowsize=arrowsize,
-                           arrowstyle=arrowstyle,
-                           minlength=minlength)
+        lines, arrows = mstream.streamplot(self, x, y, u, v,
+                                           density=density,
+                                           linewidth=linewidth,
+                                           color=color,
+                                           cmap=cmap,
+                                           arrowsize=arrowsize,
+                                           arrowstyle=arrowstyle,
+                                           minlength=minlength)
+        return lines, arrows
     streamplot.__doc__ = mstream.streamplot.__doc__
 
     @docstring.dedent_interpd

lib/matplotlib/pyplot.py

@@ -2615,15 +2615,15 @@ def specgram(x, NFFT=256, Fs=2, Fc=0, detrend=mlab.detrend_none, window=mlab.win
 # This function was autogenerated by boilerplate.py.  Do not edit as
 # changes will be lost
 @autogen_docstring(Axes.stem)
-def stem(x, y, linefmt='b-', markerfmt='bo', basefmt='r-', hold=None):
+def stem(x, y, linefmt='b-', markerfmt='bo', basefmt='r-', bottom=None, label=None, hold=None):
     ax = gca()
     # allow callers to override the hold state by passing hold=True|False
     washold = ax.ishold()
 
     if hold is not None:
         ax.hold(hold)
     try:
-        ret = ax.stem(x, y, linefmt, markerfmt, basefmt)
+        ret = ax.stem(x, y, linefmt, markerfmt, basefmt, bottom, label)
         draw_if_interactive()
     finally:
         ax.hold(washold)
@@ -2648,6 +2648,24 @@ def step(x, y, *args, **kwargs):
 
     return ret
 
+# This function was autogenerated by boilerplate.py.  Do not edit as
+# changes will be lost
+@autogen_docstring(Axes.streamplot)
+def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, arrowsize=1, arrowstyle='-|>', minlength=0.10000000000000001, hold=None):
+    ax = gca()
+    # allow callers to override the hold state by passing hold=True|False
+    washold = ax.ishold()
+
+    if hold is not None:
+        ax.hold(hold)
+    try:
+        ret = ax.streamplot(x, y, u, v, density, linewidth, color, cmap, arrowsize, arrowstyle, minlength)
+        draw_if_interactive()
+    finally:
+        ax.hold(washold)
+    sci(ret[0])
+    return ret
+
 # This function was autogenerated by boilerplate.py.  Do not edit as
 # changes will be lost
 @autogen_docstring(Axes.tricontour)

lib/matplotlib/streamplot.py

@@ -136,6 +136,8 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=1, color='k', cmap=None,
 
     axes.update_datalim(((x.min(), y.min()), (x.max(), y.max())))
     axes.autoscale_view(tight=True)
+    # TODO: Currently this returns only the last streamline and arrow patch.
+    return lc, p
 
 
 # Coordinate definitions