FilesExpand file tree

 master

/

test_streamplot.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

65 lines (51 loc) · 1.98 KB

 master

/

test_streamplot.py

Copy path

Top

File metadata and controls

• Code
• Blame

65 lines (51 loc) · 1.98 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

60

61

62

63

64

65

from __future__ import (absolute_import, division, print_function,

unicode_literals)

import six

import numpy as np

from numpy.testing import assert_array_almost_equal

import matplotlib.pyplot as plt

from matplotlib.testing.decorators import image_comparison, cleanup

import matplotlib.transforms as mtransforms

def velocity_field():

Y, X = np.mgrid[-3:3:100j, -3:3:100j]

U = -1 - X**2 + Y

V = 1 + X - Y**2

return X, Y, U, V

@image_comparison(baseline_images=['streamplot_colormap_test_image'])

def test_colormap():

X, Y, U, V = velocity_field()

plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2,

cmap=plt.cm.autumn)

plt.colorbar()

@image_comparison(baseline_images=['streamplot_linewidth_test_image'])

def test_linewidth():

X, Y, U, V = velocity_field()

speed = np.sqrt(U*U + V*V)

lw = 5*speed/speed.max()

df = 25. / 30. # Compatibility factor for old test image

plt.streamplot(X, Y, U, V, density=[0.5 * df, 1. * df], color='k',

linewidth=lw)

@image_comparison(baseline_images=['streamplot_masks_and_nans_test_image'])

def test_masks_and_nans():

X, Y, U, V = velocity_field()

mask = np.zeros(U.shape, dtype=bool)

mask[40:60, 40:60] = 1

U = np.ma.array(U, mask=mask)

U[:20, :20] = np.nan

plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)

@cleanup

def test_streamplot_limits():

ax = plt.axes()

x = np.linspace(-5, 10, 20)

y = np.linspace(-2, 4, 10)

y, x = np.meshgrid(y, x)

trans = mtransforms.Affine2D().translate(25, 32) + ax.transData

plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)

# The calculated bounds are approximately the bounds of the original data,

# this is because the entire path is taken into account when updating the

# datalim.

assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),

decimal=1)

if __name__=='__main__':

import nose

nose.runmodule()