Add streamplot axes method to plot streamlines

tonysyu · tonysyu · commit 687af189381b · 2012-01-08T17:16:23.000-05:00
diff --git a/examples/pylab_examples/streamplot_demo.py b/examples/pylab_examples/streamplot_demo.py
@@ -0,0 +1,17 @@
+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)
+
+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)
+
+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,21 @@ 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,
+                   integrator='RK4'):
+        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,
+                           integrator=integrator)
+    streamplot.__doc__ = mstream.streamplot.__doc__
+
     @docstring.dedent_interpd
     def barbs(self, *args, **kw):
         """
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
