Merge pull request #17456 from timhoffm/doc-stackplot

QuLogic · web-flow · commit 02194b6379e4 · 2020-06-09T03:03:26.000-04:00
Improve stackplot example
diff --git a/examples/lines_bars_and_markers/stackplot_demo.py b/examples/lines_bars_and_markers/stackplot_demo.py
@@ -1,63 +1,71 @@
 """
-==============
-Stackplot Demo
-==============
+===========================
+Stackplots and streamgraphs
+===========================
+"""
+
+##############################################################################
+# Stackplots
+# ----------
+#
+# Stackplots draw multiple datasets as vertically stacked areas. This is
+# useful when the individual data values and additionally their cumulative
+# value are of interest.
 
-How to create stackplots with Matplotlib.
 
-Stackplots are generated by plotting different datasets vertically on
-top of one another rather than overlapping with one another. Below we
-show some examples to accomplish this with Matplotlib.
-"""
 import numpy as np
 import matplotlib.pyplot as plt
 
+# data from United Nations World Population Prospects (Revision 2019)
+# https://population.un.org/wpp/, license: CC BY 3.0 IGO
+year = [1950, 1960, 1970, 1980, 1990, 2000, 2010, 2018]
+population_by_continent = {
+    'africa': [228, 284, 365, 477, 631, 814, 1044, 1275],
+    'americas': [340, 425, 519, 619, 727, 840, 943, 1006],
+    'asia': [1394, 1686, 2120, 2625, 3202, 3714, 4169, 4560],
+    'europe': [220, 253, 276, 295, 310, 303, 294, 293],
+    'oceania': [12, 15, 19, 22, 26, 31, 36, 39],
+}
 
-# Fixing random state for reproducibility
-np.random.seed(19680801)
+fig, ax = plt.subplots()
+ax.stackplot(year, population_by_continent.values(),
+             labels=population_by_continent.keys())
+ax.legend(loc='upper left')
+ax.set_title('World population')
+ax.set_xlabel('Year')
+ax.set_ylabel('Number of people (millions)')
 
-x = [1, 2, 3, 4, 5]
-y1 = [1, 1, 2, 3, 5]
-y2 = [0, 4, 2, 6, 8]
-y3 = [1, 3, 5, 7, 9]
+plt.show()
 
-y = np.vstack([y1, y2, y3])
+##############################################################################
+# Streamgraphs
+# ------------
+#
+# Using the *baseline* parameter, you can turn an ordinary stacked area plot
+# with baseline 0 into a stream graph.
 
-labels = ["Fibonacci ", "Evens", "Odds"]
 
-fig, ax = plt.subplots()
-ax.stackplot(x, y1, y2, y3, labels=labels)
-ax.legend(loc='upper left')
-plt.show()
+# Fixing random state for reproducibility
+np.random.seed(19680801)
 
-fig, ax = plt.subplots()
-ax.stackplot(x, y)
-plt.show()
 
-###############################################################################
-# Here we show an example of making a streamgraph using stackplot
-
-
-def layers(n, m):
-    """
-    Return *n* random Gaussian mixtures, each of length *m*.
-    """
-    def bump(a):
-        x = 1 / (.1 + np.random.random())
-        y = 2 * np.random.random() - .5
-        z = 10 / (.1 + np.random.random())
-        for i in range(m):
-            w = (i / m - y) * z
-            a[i] += x * np.exp(-w * w)
-    a = np.zeros((m, n))
-    for i in range(n):
-        for j in range(5):
-            bump(a[:, i])
+def gaussian_mixture(x, n=5):
+    """Return a random mixture of *n* Gaussians, evaluated at positions *x*."""
+    def add_random_gaussian(a):
+        amplitude = 1 / (.1 + np.random.random())
+        dx = x[-1] - x[0]
+        x0 = (2 * np.random.random() - .5) * dx
+        z = 10 / (.1 + np.random.random()) / dx
+        a += amplitude * np.exp(-(z * (x - x0))**2)
+    a = np.zeros_like(x)
+    for j in range(n):
+        add_random_gaussian(a)
     return a
 
 
-d = layers(3, 100)
+x = np.linspace(0, 100, 101)
+ys = [gaussian_mixture(x) for _ in range(3)]
 
 fig, ax = plt.subplots()
-ax.stackplot(range(100), d.T, baseline='wiggle')
+ax.stackplot(x, ys, baseline='wiggle')
 plt.show()
