1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
| import numpy as np import matplotlib.pyplot as plt import matplotlib.gridspec as gridspec
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)
fig = plt.figure(figsize=(7, 9)) gs = gridspec.GridSpec(nrows=3, ncols=2, height_ratios=[1, 1, 2])
ax0 = fig.add_subplot(gs[0, 0]) ax0.streamplot(X, Y, U, V, density=[0.5, 1]) ax0.set_title('Varying Density')
ax1 = fig.add_subplot(gs[0, 1]) strm = ax1.streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn') fig.colorbar(strm.lines) ax1.set_title('Varying Color')
ax2 = fig.add_subplot(gs[1, 0]) lw = 5*speed / speed.max() ax2.streamplot(X, Y, U, V, density=0.6, color='k', linewidth=lw) ax2.set_title('Varying Line Width')
seed_points = np.array([[-2, -1, 0, 1, 2, -1], [-2, -1, 0, 1, 2, 2]])
ax3 = fig.add_subplot(gs[1, 1]) strm = ax3.streamplot(X, Y, U, V, color=U, linewidth=2, cmap='autumn', start_points=seed_points.T) fig.colorbar(strm.lines) ax3.set_title('Controlling Starting Points')
ax3.plot(seed_points[0], seed_points[1], 'bo') ax3.axis((-w, w, -w, w))
mask = np.zeros(U.shape, dtype=bool) mask[40:60, 40:60] = True U[:20, :20] = np.nan U = np.ma.array(U, mask=mask)
ax4 = fig.add_subplot(gs[2:, :]) ax4.streamplot(X, Y, U, V, color='r') ax4.set_title('Streamplot with Masking')
ax4.imshow(~mask, extent=(-w, w, -w, w), alpha=0.5, interpolation='nearest', cmap='gray', aspect='auto') ax4.set_aspect('equal')
plt.tight_layout() plt.show()
|