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