Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2Stacked area plot for 1D arrays inspired by Douglas Y'barbo's stackoverflow 

3answer: 

4http://stackoverflow.com/questions/2225995/how-can-i-create-stacked-line-graph-with-matplotlib 

5 

6(http://stackoverflow.com/users/66549/doug) 

7 

8""" 

9import numpy as np 

10 

11import matplotlib.cbook as cbook 

12 

13__all__ = ['stackplot'] 

14 

15 

16def stackplot(axes, x, *args, 

17 labels=(), colors=None, baseline='zero', 

18 **kwargs): 

19 """ 

20 Draw a stacked area plot. 

21 

22 Parameters 

23 ---------- 

24 x : 1d array of dimension N 

25 

26 y : 2d array (dimension MxN), or sequence of 1d arrays (each dimension 1xN) 

27 

28 The data is assumed to be unstacked. Each of the following 

29 calls is legal:: 

30 

31 stackplot(x, y) # where y is MxN 

32 stackplot(x, y1, y2, y3, y4) # where y1, y2, y3, y4, are all 1xNm 

33 

34 baseline : {'zero', 'sym', 'wiggle', 'weighted_wiggle'} 

35 Method used to calculate the baseline: 

36 

37 - ``'zero'``: Constant zero baseline, i.e. a simple stacked plot. 

38 - ``'sym'``: Symmetric around zero and is sometimes called 

39 'ThemeRiver'. 

40 - ``'wiggle'``: Minimizes the sum of the squared slopes. 

41 - ``'weighted_wiggle'``: Does the same but weights to account for 

42 size of each layer. It is also called 'Streamgraph'-layout. More 

43 details can be found at http://leebyron.com/streamgraph/. 

44 

45 labels : Length N sequence of strings 

46 Labels to assign to each data series. 

47 

48 colors : Length N sequence of colors 

49 A list or tuple of colors. These will be cycled through and used to 

50 colour the stacked areas. 

51 

52 **kwargs 

53 All other keyword arguments are passed to `Axes.fill_between()`. 

54 

55 

56 Returns 

57 ------- 

58 list : list of `.PolyCollection` 

59 A list of `.PolyCollection` instances, one for each element in the 

60 stacked area plot. 

61 """ 

62 

63 y = np.row_stack(args) 

64 

65 labels = iter(labels) 

66 if colors is not None: 

67 axes.set_prop_cycle(color=colors) 

68 

69 # Assume data passed has not been 'stacked', so stack it here. 

70 # We'll need a float buffer for the upcoming calculations. 

71 stack = np.cumsum(y, axis=0, dtype=np.promote_types(y.dtype, np.float32)) 

72 

73 cbook._check_in_list(['zero', 'sym', 'wiggle', 'weighted_wiggle'], 

74 baseline=baseline) 

75 if baseline == 'zero': 

76 first_line = 0. 

77 

78 elif baseline == 'sym': 

79 first_line = -np.sum(y, 0) * 0.5 

80 stack += first_line[None, :] 

81 

82 elif baseline == 'wiggle': 

83 m = y.shape[0] 

84 first_line = (y * (m - 0.5 - np.arange(m)[:, None])).sum(0) 

85 first_line /= -m 

86 stack += first_line 

87 

88 elif baseline == 'weighted_wiggle': 

89 total = np.sum(y, 0) 

90 # multiply by 1/total (or zero) to avoid infinities in the division: 

91 inv_total = np.zeros_like(total) 

92 mask = total > 0 

93 inv_total[mask] = 1.0 / total[mask] 

94 increase = np.hstack((y[:, 0:1], np.diff(y))) 

95 below_size = total - stack 

96 below_size += 0.5 * y 

97 move_up = below_size * inv_total 

98 move_up[:, 0] = 0.5 

99 center = (move_up - 0.5) * increase 

100 center = np.cumsum(center.sum(0)) 

101 first_line = center - 0.5 * total 

102 stack += first_line 

103 

104 # Color between x = 0 and the first array. 

105 color = axes._get_lines.get_next_color() 

106 coll = axes.fill_between(x, first_line, stack[0, :], 

107 facecolor=color, label=next(labels, None), 

108 **kwargs) 

109 coll.sticky_edges.y[:] = [0] 

110 r = [coll] 

111 

112 # Color between array i-1 and array i 

113 for i in range(len(y) - 1): 

114 color = axes._get_lines.get_next_color() 

115 r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :], 

116 facecolor=color, label=next(labels, None), 

117 **kwargs)) 

118 return r