|
| 1 | +""" |
| 2 | +========== |
| 3 | +Streamplot |
| 4 | +========== |
| 5 | +A streamplot, or streamline plot, is used to display 2D vector fields. This |
| 6 | +example shows a few features of the stream plot function: |
| 7 | +
|
| 8 | + * Varying the color along a streamline. |
| 9 | + * Varying the density of streamlines. |
| 10 | + * Varying the line width along a stream line. |
| 11 | + * Controlling the start points of streamlines. |
| 12 | + * Streamlines skipping masked regions and NaN values. |
| 13 | +""" |
| 14 | +import numpy as np |
| 15 | +import matplotlib.pyplot as plt |
| 16 | + |
| 17 | +w = 3 |
| 18 | +Y, X = np.mgrid[-w:w:100j, -w:w:100j] |
| 19 | +U = -1 - X**2 + Y |
| 20 | +V = 1 + X - Y**2 |
| 21 | +speed = np.sqrt(U*U + V*V) |
| 22 | + |
| 23 | +fig = plt.figure() |
| 24 | + |
| 25 | +# Varying density along a streamline |
| 26 | +ax0 = fig.add_subplot(321) |
| 27 | +ax0.streamplot(X, Y, U, V, density=[0.5, 1]) |
| 28 | +ax0.set_title('Varying Density') |
| 29 | + |
| 30 | +# Varying color along a streamline |
| 31 | +ax1 = fig.add_subplot(322) |
| 32 | +strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn) |
| 33 | +fig.colorbar(strm.lines) |
| 34 | +ax1.set_title('Varying color') |
| 35 | + |
| 36 | +# Varying line width along a streamline |
| 37 | +ax2 = fig.add_subplot(323) |
| 38 | +lw = 5*speed / speed.max() |
| 39 | +ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw) |
| 40 | +ax2.set_title('Varying Line Width') |
| 41 | + |
| 42 | +# Controlling the starting points of the streamlines |
| 43 | +X, Y = (np.linspace(-3, 3, 100), |
| 44 | + np.linspace(-3, 3, 100)) |
| 45 | +U, V = np.mgrid[-3:3:100j, 0:0:100j] |
| 46 | +seed_points = np.array([[-2, -1, 0, 1, 2, -1], [-2, -1, 0, 1, 2, 2]]) |
| 47 | + |
| 48 | +ax3 = fig.add_subplot(324) |
| 49 | +strm = ax3.streamplot(X, Y, U, V, color=U, linewidth=2, |
| 50 | + cmap=plt.cm.autumn, start_points=seed_points.T) |
| 51 | +fig.colorbar(strm.lines) |
| 52 | +ax3.set_title('Controlling Starting Points') |
| 53 | + |
| 54 | +# Displaying the starting points with red symbols. |
| 55 | +ax3.plot(seed_points[0], seed_points[1], 'bo') |
| 56 | + |
| 57 | +ax3.axis((-3, 3, -3, 3)) |
| 58 | + |
| 59 | +# Create a mask |
| 60 | +w = 3 |
| 61 | +Y, X = np.mgrid[-w:w:100j, -w:w:100j] |
| 62 | +U = -1 - X**2 + Y |
| 63 | +V = 1 + X - Y**2 |
| 64 | +speed = np.sqrt(U*U + V*V) |
| 65 | + |
| 66 | +mask = np.zeros(U.shape, dtype=bool) |
| 67 | +mask[40:60, 40:60] = True |
| 68 | +U[:20, :20] = np.nan |
| 69 | +U = np.ma.array(U, mask=mask) |
| 70 | + |
| 71 | +ax4 = fig.add_subplot(325) |
| 72 | +ax4.streamplot(X, Y, U, V, color='r') |
| 73 | +ax4.set_title('Streamline with Masking') |
| 74 | + |
| 75 | +ax4.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, |
| 76 | + interpolation='nearest', cmap=plt.cm.gray, aspect='auto') |
| 77 | + |
| 78 | +plt.tight_layout() |
| 79 | +plt.show() |
0 commit comments