-
-
Notifications
You must be signed in to change notification settings - Fork 8.2k
Expand file tree
/
Copy pathstackplot.py
More file actions
143 lines (111 loc) · 4.72 KB
/
stackplot.py
File metadata and controls
143 lines (111 loc) · 4.72 KB
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
"""
Stacked area plot for 1D arrays inspired by Douglas Y'barbo's stackoverflow
answer:
https://stackoverflow.com/q/2225995/
(https://stackoverflow.com/users/66549/doug)
"""
import numpy as np
from matplotlib import cbook, collections, _api, _style_helpers
__all__ = ['stackplot']
def stackplot(axes, x, *args,
labels=(), colors=None, baseline='zero',
**kwargs):
"""
Draw a stacked area plot or a streamgraph.
Parameters
----------
x : (N,) array-like
y : (M, N) array-like
The data can be either stacked or unstacked. Each of the following
calls is legal::
stackplot(x, y) # where y has shape (M, N) e.g. y = [y1, y2, y3, y4]
stackplot(x, y1, y2, y3, y4) # where y1, y2, y3, y4 have length N
baseline : {'zero', 'sym', 'wiggle', 'weighted_wiggle'}
Method used to calculate the baseline:
- ``'zero'``: Constant zero baseline, i.e. a simple stacked plot.
- ``'sym'``: Symmetric around zero and is sometimes called
'ThemeRiver'.
- ``'wiggle'``: Minimizes the sum of the squared slopes.
- ``'weighted_wiggle'``: Does the same but weights to account for
size of each layer. It is also called 'Streamgraph'-layout. More
details can be found at http://leebyron.com/streamgraph/.
labels : list of str, optional
A sequence of labels to assign to each data series. If unspecified,
then no labels will be applied to artists.
colors : list of :mpltype:`color`, optional
A sequence of colors to be cycled through and used to color the stacked
areas. The sequence need not be exactly the same length as the number
of provided *y*, in which case the colors will repeat from the
beginning.
If not specified, the colors from the Axes property cycle will be used.
data : indexable object, optional
DATA_PARAMETER_PLACEHOLDER
**kwargs
All other keyword arguments are passed to `.Axes.fill_between`. The
following parameters additionally accept a sequence of values
corresponding to the *y* datasets:
- *hatch*
- *edgecolor*
- *facecolor*
- *linewidth*
- *linestyle*
.. versionadded:: 3.9
Allowing a sequence of strings for *hatch*.
.. versionadded:: 3.11
Allowing sequences of values in above listed `.Axes.fill_between`
parameters.
Returns
-------
list of `.PolyCollection`
A list of `.PolyCollection` instances, one for each element in the
stacked area plot.
"""
y = np.vstack(args)
labels = iter(labels)
if colors is None:
colors = [axes._get_lines.get_next_color() for _ in y]
kwargs = cbook.normalize_kwargs(kwargs, collections.PolyCollection)
kwargs.setdefault('facecolor', colors)
kwargs, style_gen = _style_helpers.style_generator(kwargs)
# Assume data passed has not been 'stacked', so stack it here.
# We'll need a float buffer for the upcoming calculations.
stack = np.cumsum(y, axis=0, dtype=np.promote_types(y.dtype, np.float32))
_api.check_in_list(['zero', 'sym', 'wiggle', 'weighted_wiggle'],
baseline=baseline)
if baseline == 'zero':
first_line = 0.
elif baseline == 'sym':
first_line = -np.sum(y, 0) * 0.5
stack += first_line[None, :]
elif baseline == 'wiggle':
m = y.shape[0]
first_line = (y * (m - 0.5 - np.arange(m)[:, None])).sum(0)
first_line /= -m
stack += first_line
elif baseline == 'weighted_wiggle':
total = np.sum(y, 0)
# multiply by 1/total (or zero) to avoid infinities in the division:
inv_total = np.zeros_like(total)
mask = total > 0
inv_total[mask] = 1.0 / total[mask]
increase = np.hstack((y[:, 0:1], np.diff(y)))
below_size = total - stack
below_size += 0.5 * y
move_up = below_size * inv_total
move_up[:, 0] = 0.5
center = (move_up - 0.5) * increase
center = np.cumsum(center.sum(0))
first_line = center - 0.5 * total
stack += first_line
# Color between x = 0 and the first array.
coll = axes.fill_between(x, first_line, stack[0, :],
label=next(labels, None),
**next(style_gen), **kwargs)
coll.sticky_edges.y[:] = [0]
r = [coll]
# Color between array i-1 and array i
for i in range(len(y) - 1):
r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :],
label=next(labels, None),
**next(style_gen), **kwargs))
return r