Fix streamplot to work with colorbar.

tonysyu · tonysyu · commit 72abc5d37d16 · 2012-01-16T21:22:16.000-05:00
Streamplot no longer returns an arrow patch and, instead, only returns a dummy LineCollection.
diff --git a/boilerplate.py b/boilerplate.py
@@ -124,7 +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])',
+    'streamplot' :  'sci(%(ret)s)',
     '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)',
diff --git a/examples/pylab_examples/streamplot_demo.py b/examples/pylab_examples/streamplot_demo.py
@@ -6,12 +6,14 @@
 V = 1 + X - Y**2
 speed = np.sqrt(U*U + V*V)
 
+plt.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn)
+plt.colorbar()
+
 f, (ax1, ax2) = plt.subplots(ncols=2)
 ax1.streamplot(X, Y, U, V)
 
 lw = 5*speed/speed.max()
-ax2.streamplot(X, Y, U, V, density=0.6,
-               color=U, cmap=plt.cm.autumn, linewidth=lw)
+ax2.streamplot(X, Y, U, V, density=0.6, color=U, linewidth=lw)
 
 plt.show()
 
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -6307,15 +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()
-        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
+        lines = mstream.streamplot(self, x, y, u, v,
+                                   density=density,
+                                   linewidth=linewidth,
+                                   color=color,
+                                   cmap=cmap,
+                                   arrowsize=arrowsize,
+                                   arrowstyle=arrowstyle,
+                                   minlength=minlength)
+        return lines
     streamplot.__doc__ = mstream.streamplot.__doc__
 
     @docstring.dedent_interpd
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2663,7 +2663,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None, arr
         draw_if_interactive()
     finally:
         ax.hold(washold)
-    sci(ret[0])
+    sci(ret)
     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
@@ -134,10 +134,18 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=1, color='k', cmap=None,
         p = mpp.FancyArrowPatch(arrow_tail, arrow_head, **arrow_kw)
         axes.add_patch(p)
 
+    # Add dummy line collection so that colorbar works correctly.
+    if type(color) == np.ndarray:
+        lc = matplotlib.collections.LineCollection([], **line_kw)
+        lc.set_array(color.ravel())
+        lc.set_cmap(cmap)
+        lc.set_norm(norm)
+        axes.add_collection(lc)
+
     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
+    # TODO: Currently this returns only dummy streamline
+    return lc
 
 
 # Coordinate definitions
