Merge branch 'streamplot' of https://github.com/tonysyu/matplotlib into tonysyu-streamplot

WeatherGod · WeatherGod · commit b2ca027a363e · 2012-02-27T10:33:28.000-06:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,6 @@
+2012-01-08 Add axes.streamplot to plot streamlines of a velocity field.
+		   Adapted from Tom Flannaghan streamplot implementation. -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
 
diff --git a/boilerplate.py b/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)',
     '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
@@ -0,0 +1,19 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+Y, X = np.mgrid[-3:3:100j, -3:3:100j]
+U = -1 - X**2 + Y
+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='k', linewidth=lw)
+
+plt.show()
+
diff --git a/examples/pylab_examples/streamplot_with_mask.py b/examples/pylab_examples/streamplot_with_mask.py
@@ -0,0 +1,23 @@
+"""
+Demonstrate the use of the `streamplot` function using a masked array
+and NaN values.
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+
+w = 3
+Y, X = np.mgrid[-w:w:100j, -w:w:100j]
+U = -1 - X**2 + Y
+V = 1 + X - Y**2
+speed = np.sqrt(U*U + V*V)
+
+mask = np.zeros(U.shape, dtype=bool)
+mask[40:60, 40:60] = 1
+U = np.ma.array(U, mask=mask)
+U[:20, :20] = np.nan
+
+plt.streamplot(X, Y, U, V, color='r')
+plt.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, interpolation='nearest')
+
+plt.show()
+
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -29,6 +29,7 @@
 import matplotlib.spines as mspines
 import matplotlib.quiver as mquiver
 import matplotlib.scale as mscale
+import matplotlib.streamplot as mstream
 import matplotlib.table as mtable
 import matplotlib.text as mtext
 import matplotlib.ticker as mticker
@@ -6303,6 +6304,20 @@ def quiver(self, *args, **kw):
         return q
     quiver.__doc__ = mquiver.Quiver.quiver_doc
 
+    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 = 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
     def barbs(self, *args, **kw):
         """
diff --git a/lib/matplotlib/pyplot.py b/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)
+    return ret
+
 # This function was autogenerated by boilerplate.py.  Do not edit as
 # changes will be lost
 @autogen_docstring(Axes.tricontour)
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
