From 5c69a4148d1fc090df769a1484bb2df2f3877d16 Mon Sep 17 00:00:00 2001 From: Till Stensitzki Date: Thu, 10 Jan 2013 10:52:37 +0100 Subject: [PATCH 1/2] [ENH] added baseline feature to stacked graph. --- lib/matplotlib/stackplot.py | 61 ++++++++++++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 5 deletions(-) diff --git a/lib/matplotlib/stackplot.py b/lib/matplotlib/stackplot.py index 534e33dfa359..f507ba670772 100644 --- a/lib/matplotlib/stackplot.py +++ b/lib/matplotlib/stackplot.py @@ -25,6 +25,16 @@ def stackplot(axes, x, *args, **kwargs): Keyword arguments: + *baseline* : ['zero', 'sym', 'wiggle', 'weighted_wiggle'] + Method use to calculate to 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 in 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,58 @@ 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) r = [] + if baseline == 'zero': + first_line = 0. + + elif baseline == 'sym': + first_line = -np.sum(y, 0) * 0.5 + y_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 + y_stack += first_line + + elif baseline == 'weighted_wiggle': + #TODO: Vectorize this stuff. + m, n = y.shape + center = np.zeros(n) + total = np.sum(y, 0) + for i in range(n): + if i > 0: + center[i] = center[i - 1] + for j in range(m): + if i == 0: + increase = y[j, i] + moveUp = 0.5 + else: + belowSize = 0.5 * y[j, i] + for k in range(j + 1, m): + belowSize += y[k, i] + increase = y[j, i] - y[j, i - 1] + moveUp = belowSize / total[i] + center[i] += (moveUp - 0.5) * increase + first_line = center - 0.5 * total + y_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, y_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)) + 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)) return r From a766d236b15f9bf37ab675ce978e89daabc7cac3 Mon Sep 17 00:00:00 2001 From: Till Stensitzki Date: Thu, 10 Jan 2013 10:52:37 +0100 Subject: [PATCH 2/2] pep8 compliance --- lib/matplotlib/stackplot.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/lib/matplotlib/stackplot.py b/lib/matplotlib/stackplot.py index f507ba670772..7f66c7d0bc3a 100644 --- a/lib/matplotlib/stackplot.py +++ b/lib/matplotlib/stackplot.py @@ -56,21 +56,21 @@ def stackplot(axes, x, *args, **kwargs): 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 - y_stack += first_line[None, :] + 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 - y_stack += first_line + stack += first_line elif baseline == 'weighted_wiggle': #TODO: Vectorize this stuff. @@ -83,29 +83,30 @@ def stackplot(axes, x, *args, **kwargs): for j in range(m): if i == 0: increase = y[j, i] - moveUp = 0.5 + move_up = 0.5 else: - belowSize = 0.5 * y[j, i] + below_size = 0.5 * y[j, i] for k in range(j + 1, m): - belowSize += y[k, i] + below_size += y[k, i] increase = y[j, i] - y[j, i - 1] - moveUp = belowSize / total[i] - center[i] += (moveUp - 0.5) * increase + move_up = below_size / total[i] + center[i] += (move_up - 0.5) * increase first_line = center - 0.5 * total - y_stack += first_line + 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, first_line, y_stack[0,:], + 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(), + for i in xrange(len(y) - 1): + color = axes._get_lines.color_cycle.next() + r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :], + facecolor= color, **kwargs)) return r