Coverage for /usr/lib/python3/dist-packages/matplotlib/stackplot.py: 12%

1""" 

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

3answer: 

4https://stackoverflow.com/q/2225995/ 

5 

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

7""" 

8 

9import itertools 

10 

11import numpy as np 

12 

13from matplotlib import _api 

14 

15__all__ = ['stackplot'] 

16 

17 

18def stackplot(axes, x, *args, 

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

20 **kwargs): 

21 """ 

22 Draw a stacked area plot. 

23 

24 Parameters 

25 ---------- 

26 x : (N,) array-like 

27 

28 y : (M, N) array-like 

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

30 calls is legal:: 

31 

32 stackplot(x, y) # where y has shape (M, N) 

33 stackplot(x, y1, y2, y3) # where y1, y2, y3, y4 have length N 

34 

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

36 Method used to calculate the baseline: 

37 

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

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

40 'ThemeRiver'. 

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

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

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

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

45 

46 labels : list of str, optional 

47 A sequence of labels to assign to each data series. If unspecified, 

48 then no labels will be applied to artists. 

49 

50 colors : list of color, optional 

51 A sequence of colors to be cycled through and used to color the stacked 

52 areas. The sequence need not be exactly the same length as the number 

53 of provided *y*, in which case the colors will repeat from the 

54 beginning. 

55 

56 If not specified, the colors from the Axes property cycle will be used. 

57 

58 data : indexable object, optional 

59 DATA_PARAMETER_PLACEHOLDER 

60 

61 **kwargs 

62 All other keyword arguments are passed to `.Axes.fill_between`. 

63 

64 Returns 

65 ------- 

66 list of `.PolyCollection` 

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

68 stacked area plot. 

69 """ 

70 

71 y = np.row_stack(args) 

72 

73 labels = iter(labels) 

74 if colors is not None: 

75 colors = itertools.cycle(colors) 

76 else: 

77 colors = (axes._get_lines.get_next_color() for _ in y) 

78 

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

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

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

82 

83 _api.check_in_list(['zero', 'sym', 'wiggle', 'weighted_wiggle'], 

84 baseline=baseline) 

85 if baseline == 'zero': 

86 first_line = 0. 

87 

88 elif baseline == 'sym': 

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

90 stack += first_line[None, :] 

91 

92 elif baseline == 'wiggle': 

93 m = y.shape[0] 

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

95 first_line /= -m 

96 stack += first_line 

97 

98 elif baseline == 'weighted_wiggle': 

99 total = np.sum(y, 0) 

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

101 inv_total = np.zeros_like(total) 

102 mask = total > 0 

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

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

105 below_size = total - stack 

106 below_size += 0.5 * y 

107 move_up = below_size * inv_total 

108 move_up[:, 0] = 0.5 

109 center = (move_up - 0.5) * increase 

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

111 first_line = center - 0.5 * total 

112 stack += first_line 

113 

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

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

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

117 **kwargs) 

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

119 r = [coll] 

120 

121 # Color between array i-1 and array i 

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

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

124 facecolor=next(colors), 

125 label=next(labels, None), 

126 **kwargs)) 

127 return r