test_streamplot.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. import sys
  2. import numpy as np
  3. from numpy.testing import assert_array_almost_equal
  4. import matplotlib.pyplot as plt
  5. from matplotlib.testing.decorators import image_comparison
  6. import matplotlib.transforms as mtransforms
  7. on_win = (sys.platform == 'win32')
  8. on_mac = (sys.platform == 'darwin')
  9. def velocity_field():
  10. Y, X = np.mgrid[-3:3:100j, -3:3:100j]
  11. U = -1 - X**2 + Y
  12. V = 1 + X - Y**2
  13. return X, Y, U, V
  14. def swirl_velocity_field():
  15. x = np.linspace(-3., 3., 100)
  16. y = np.linspace(-3., 3., 100)
  17. X, Y = np.meshgrid(x, y)
  18. a = 0.1
  19. U = np.cos(a) * (-Y) - np.sin(a) * X
  20. V = np.sin(a) * (-Y) + np.cos(a) * X
  21. return x, y, U, V
  22. @image_comparison(['streamplot_startpoints'], remove_text=True, style='mpl20')
  23. def test_startpoints():
  24. X, Y, U, V = velocity_field()
  25. start_x = np.linspace(X.min(), X.max(), 10)
  26. start_y = np.linspace(Y.min(), Y.max(), 10)
  27. start_points = np.column_stack([start_x, start_y])
  28. plt.streamplot(X, Y, U, V, start_points=start_points)
  29. plt.plot(start_x, start_y, 'ok')
  30. @image_comparison(['streamplot_colormap'],
  31. tol=.04, remove_text=True, style='mpl20')
  32. def test_colormap():
  33. X, Y, U, V = velocity_field()
  34. plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2,
  35. cmap=plt.cm.autumn)
  36. plt.colorbar()
  37. @image_comparison(['streamplot_linewidth'], remove_text=True, style='mpl20')
  38. def test_linewidth():
  39. X, Y, U, V = velocity_field()
  40. speed = np.hypot(U, V)
  41. lw = 5 * speed / speed.max()
  42. # Compatibility for old test image
  43. df = 25 / 30
  44. ax = plt.figure().subplots()
  45. ax.set(xlim=(-3.0, 2.9999999999999947),
  46. ylim=(-3.0000000000000004, 2.9999999999999947))
  47. ax.streamplot(X, Y, U, V, density=[0.5 * df, 1. * df], color='k',
  48. linewidth=lw)
  49. @image_comparison(['streamplot_masks_and_nans'],
  50. remove_text=True, style='mpl20', tol=0.04 if on_win else 0)
  51. def test_masks_and_nans():
  52. X, Y, U, V = velocity_field()
  53. mask = np.zeros(U.shape, dtype=bool)
  54. mask[40:60, 40:60] = 1
  55. U[:20, :20] = np.nan
  56. U = np.ma.array(U, mask=mask)
  57. # Compatibility for old test image
  58. ax = plt.figure().subplots()
  59. ax.set(xlim=(-3.0, 2.9999999999999947),
  60. ylim=(-3.0000000000000004, 2.9999999999999947))
  61. with np.errstate(invalid='ignore'):
  62. ax.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)
  63. @image_comparison(['streamplot_maxlength.png'],
  64. remove_text=True, style='mpl20',
  65. tol=0.002 if on_mac else 0)
  66. def test_maxlength():
  67. x, y, U, V = swirl_velocity_field()
  68. ax = plt.figure().subplots()
  69. ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
  70. linewidth=2, density=2)
  71. assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3
  72. # Compatibility for old test image
  73. ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413))
  74. @image_comparison(['streamplot_direction.png'],
  75. remove_text=True, style='mpl20')
  76. def test_direction():
  77. x, y, U, V = swirl_velocity_field()
  78. plt.streamplot(x, y, U, V, integration_direction='backward',
  79. maxlength=1.5, start_points=[[1.5, 0.]],
  80. linewidth=2, density=2)
  81. def test_streamplot_limits():
  82. ax = plt.axes()
  83. x = np.linspace(-5, 10, 20)
  84. y = np.linspace(-2, 4, 10)
  85. y, x = np.meshgrid(y, x)
  86. trans = mtransforms.Affine2D().translate(25, 32) + ax.transData
  87. plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)
  88. # The calculated bounds are approximately the bounds of the original data,
  89. # this is because the entire path is taken into account when updating the
  90. # datalim.
  91. assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),
  92. decimal=1)