stackplot.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. """
  2. Stacked area plot for 1D arrays inspired by Douglas Y'barbo's stackoverflow
  3. answer:
  4. http://stackoverflow.com/questions/2225995/how-can-i-create-stacked-line-graph-with-matplotlib
  5. (http://stackoverflow.com/users/66549/doug)
  6. """
  7. import numpy as np
  8. import matplotlib.cbook as cbook
  9. __all__ = ['stackplot']
  10. def stackplot(axes, x, *args,
  11. labels=(), colors=None, baseline='zero',
  12. **kwargs):
  13. """
  14. Draw a stacked area plot.
  15. Parameters
  16. ----------
  17. x : 1d array of dimension N
  18. y : 2d array (dimension MxN), or sequence of 1d arrays (each dimension 1xN)
  19. The data is assumed to be unstacked. Each of the following
  20. calls is legal::
  21. stackplot(x, y) # where y is MxN
  22. stackplot(x, y1, y2, y3, y4) # where y1, y2, y3, y4, are all 1xNm
  23. baseline : {'zero', 'sym', 'wiggle', 'weighted_wiggle'}
  24. Method used to calculate the baseline:
  25. - ``'zero'``: Constant zero baseline, i.e. a simple stacked plot.
  26. - ``'sym'``: Symmetric around zero and is sometimes called
  27. 'ThemeRiver'.
  28. - ``'wiggle'``: Minimizes the sum of the squared slopes.
  29. - ``'weighted_wiggle'``: Does the same but weights to account for
  30. size of each layer. It is also called 'Streamgraph'-layout. More
  31. details can be found at http://leebyron.com/streamgraph/.
  32. labels : Length N sequence of strings
  33. Labels to assign to each data series.
  34. colors : Length N sequence of colors
  35. A list or tuple of colors. These will be cycled through and used to
  36. colour the stacked areas.
  37. **kwargs
  38. All other keyword arguments are passed to `.Axes.fill_between`.
  39. Returns
  40. -------
  41. list of `.PolyCollection`
  42. A list of `.PolyCollection` instances, one for each element in the
  43. stacked area plot.
  44. """
  45. y = np.row_stack(args)
  46. labels = iter(labels)
  47. if colors is not None:
  48. axes.set_prop_cycle(color=colors)
  49. # Assume data passed has not been 'stacked', so stack it here.
  50. # We'll need a float buffer for the upcoming calculations.
  51. stack = np.cumsum(y, axis=0, dtype=np.promote_types(y.dtype, np.float32))
  52. cbook._check_in_list(['zero', 'sym', 'wiggle', 'weighted_wiggle'],
  53. baseline=baseline)
  54. if baseline == 'zero':
  55. first_line = 0.
  56. elif baseline == 'sym':
  57. first_line = -np.sum(y, 0) * 0.5
  58. stack += first_line[None, :]
  59. elif baseline == 'wiggle':
  60. m = y.shape[0]
  61. first_line = (y * (m - 0.5 - np.arange(m)[:, None])).sum(0)
  62. first_line /= -m
  63. stack += first_line
  64. elif baseline == 'weighted_wiggle':
  65. total = np.sum(y, 0)
  66. # multiply by 1/total (or zero) to avoid infinities in the division:
  67. inv_total = np.zeros_like(total)
  68. mask = total > 0
  69. inv_total[mask] = 1.0 / total[mask]
  70. increase = np.hstack((y[:, 0:1], np.diff(y)))
  71. below_size = total - stack
  72. below_size += 0.5 * y
  73. move_up = below_size * inv_total
  74. move_up[:, 0] = 0.5
  75. center = (move_up - 0.5) * increase
  76. center = np.cumsum(center.sum(0))
  77. first_line = center - 0.5 * total
  78. stack += first_line
  79. # Color between x = 0 and the first array.
  80. color = axes._get_lines.get_next_color()
  81. coll = axes.fill_between(x, first_line, stack[0, :],
  82. facecolor=color, label=next(labels, None),
  83. **kwargs)
  84. coll.sticky_edges.y[:] = [0]
  85. r = [coll]
  86. # Color between array i-1 and array i
  87. for i in range(len(y) - 1):
  88. color = axes._get_lines.get_next_color()
  89. r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :],
  90. facecolor=color, label=next(labels, None),
  91. **kwargs))
  92. return r