Merge pull request matplotlib#1671 from dmcdougall/feature_stack_base

Feature stack base

Author: dmcdougall

CHANGELOG

@@ -4,6 +4,8 @@
 
 2013-01-07 Add framealpha keyword argument to legend - PO
 
+2013-01-16 Till Stensitzki added a baseline feature to stackplot
+
 2012-12-22 Added classes for interpolation within triangular grids
            (LinearTriInterpolator) and to find the triangles in which points
            lie (TrapezoidMapTriFinder) to matplotlib.tri module. - IMT

doc/users/whats_new.rst

@@ -22,6 +22,13 @@ revision, see the :ref:`github-stats`.
 new in matplotlib-1.3
 =====================
 
+Baselines for stackplot
+-----------------------
+Till Stensitzki added non-zero baselines to :func:`~matplotlib.pyplot.stackplot`.
+They may be symmetric or weighted.
+
+.. plot:: mpl_examples/pylab_examples/stackplot_demo2.py
+
 Initialize a rotated rectangle
 ------------------------------
 Damon McDougall extended the :class:`~matplotlib.patches.Rectangle` constructor

examples/pylab_examples/stackplot_demo2.py

@@ -0,0 +1,26 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+np.random.seed(0)
+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 / float(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])
+    return a
+
+d = layers(3, 100)
+
+plt.subplots()
+plt.stackplot(range(100), d.T, baseline='wiggle')
+plt.show()

lib/matplotlib/stackplot.py

@@ -25,6 +25,16 @@ def stackplot(axes, x, *args, **kwargs):
 
     Keyword arguments:
 
+    *baseline* : ['zero', 'sym', 'wiggle', 'weighted_wiggle']
+                Method used to calculate the baseline. 'zero' is just a
+                simple stacked plot. 'sym' is 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://www.leebyron.com/else/streamgraph/.
+
+
     *colors* : A list or tuple of colors. These will be cycled through and
                used to colour the stacked areas.
                All other keyword arguments are passed to
@@ -44,17 +54,51 @@ def stackplot(axes, x, *args, **kwargs):
     if colors is not None:
         axes.set_color_cycle(colors)
 
+    baseline = kwargs.pop('baseline', 'zero')
     # Assume data passed has not been 'stacked', so stack it here.
-    y_stack = np.cumsum(y, axis=0)
+    stack = np.cumsum(y, axis=0)
 
     r = []
+    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(0, m)[:, None])).sum(0)
+        first_line /= -m
+        stack += first_line
+
+    elif baseline == 'weighted_wiggle':
+        m, n = y.shape
+        center = np.zeros(n)
+        total = np.sum(y, 0)
+        increase = np.hstack((y[:, 0:1], np.diff(y)))
+        below_size = total - stack
+        below_size += 0.5 * y
+        move_up = below_size / 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
+    else:
+        errstr = "Baseline method %s not recognised. " % baseline
+        errstr += "Expected 'zero', 'sym', 'wiggle' or 'weighted_wiggle'"
+        raise ValueError(errstr)
 
     # Color between x = 0 and the first array.
-    r.append(axes.fill_between(x, 0, y_stack[0, :],
-             facecolor=axes._get_lines.color_cycle.next(), **kwargs))
+    r.append(axes.fill_between(x, first_line, stack[0, :],
+                               facecolor=axes._get_lines.color_cycle.next(),
+                               **kwargs))
 
     # Color between array i-1 and array i
     for i in xrange(len(y) - 1):
-        r.append(axes.fill_between(x, y_stack[i, :], y_stack[i + 1, :],
-                 facecolor=axes._get_lines.color_cycle.next(), **kwargs))
+        color = axes._get_lines.color_cycle.next()
+        r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :],
+                                   facecolor= color,
+                                   **kwargs))
     return r