Improve stackplot example

timhoffm · timhoffm · commit 316684340bfa · 2020-05-20T01:41:46.000+02:00
diff --git a/examples/lines_bars_and_markers/stackplot_demo.py b/examples/lines_bars_and_markers/stackplot_demo.py
@@ -1,54 +1,59 @@
 """
-==============
-Stackplot Demo
-==============
-
-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.
+===========================
+Stackplots and streamgraphs
+===========================
 """
-import numpy as np
-import matplotlib.pyplot as plt
 
+##############################################################################
+# Stackplots
+# ----------
+#
+# Stackplots draw multiple datasets as vertically stacked areas. This is
+# useful when the individual data values and addtionally their cumulative
+# value are of interest.
 
-# Fixing random state for reproducibility
-np.random.seed(19680801)
 
-x = [1, 2, 3, 4, 5]
-y1 = [1, 1, 2, 3, 5]
-y2 = [0, 4, 2, 6, 8]
-y3 = [1, 3, 5, 7, 9]
+import numpy as np
+import matplotlib.pyplot as plt
 
-y = np.vstack([y1, y2, y3])
+year = [1950, 1960, 1970, 1980, 1990, 2000, 2010, 2018]
+population_by_contient = {
+    '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],
+}
 
-labels = ["Fibonacci ", "Evens", "Odds"]
+plt.stackplot(year, population_by_contient.values(),
+              labels=population_by_contient.keys())
+plt.legend(loc='upper left')
+plt.title('World population')
+plt.xlabel('Year')
+plt.ylabel('Number of people (millions)')
 
-fig, ax = plt.subplots()
-ax.stackplot(x, y1, y2, y3, labels=labels)
-ax.legend(loc='upper left')
 plt.show()
 
-fig, ax = plt.subplots()
-ax.stackplot(x, y)
-plt.show()
+##############################################################################
+# Streamgraphs
+# ------------
+#
+# Using the *baseline* parameter, you can turn an ordinary stacked area plot
+# with baseline 0 into a stream graph.
 
-###############################################################################
-# Here we show an example of making a streamgraph using stackplot
+
+# Fixing random state for reproducibility
+np.random.seed(19680801)
 
 
 def layers(n, m):
-    """
-    Return *n* random Gaussian mixtures, each of length *m*.
-    """
+    """Return *n* random Gaussian mixtures, each of length *m*."""
     def bump(a):
-        x = 1 / (.1 + np.random.random())
+        amplitude = 1 / (.1 + np.random.random())
+        x = np.linspace(0, 1, m)
         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 += amplitude * np.exp(-((x - y) * z)**2)
     a = np.zeros((m, n))
     for i in range(n):
         for j in range(5):
