Merge pull request matplotlib#465 from jdh2358/sankey2

add sankey module and new demo

Author: jdh2358

examples/api/sankey_demo.py

@@ -1,188 +1,211 @@
-#!/usr/bin/env python
-
-__author__ = "Yannick Copin <[email protected]>"
-__version__ = "Time-stamp: <10/02/2010 16:49 [email protected]>"
-
-import numpy as N
-
-def sankey(ax,
-           outputs=[100.], outlabels=None,
-           inputs=[100.], inlabels='',
-           dx=40, dy=10, outangle=45, w=3, inangle=30, offset=2, **kwargs):
-    """Draw a Sankey diagram.
-
-    outputs: array of outputs, should sum up to 100%
-    outlabels: output labels (same length as outputs),
-               or None (use default labels) or '' (no labels)
-    inputs and inlabels: similar for inputs
-    dx: horizontal elongation
-    dy: vertical elongation
-    outangle: output arrow angle [deg]
-    w: output arrow shoulder
-    inangle: input dip angle
-    offset: text offset
-    **kwargs: propagated to Patch (e.g. fill=False)
-
-    Return (patch,[intexts,outtexts])."""
-
-    import matplotlib.patches as mpatches
-    from matplotlib.path import Path
-
-    outs = N.absolute(outputs)
-    outsigns = N.sign(outputs)
-    outsigns[-1] = 0                       # Last output
-
-    ins = N.absolute(inputs)
-    insigns = N.sign(inputs)
-    insigns[0] = 0                         # First input
-
-    assert sum(outs)==100, "Outputs don't sum up to 100%"
-    assert sum(ins)==100, "Inputs don't sum up to 100%"
-
-    def add_output(path, loss, sign=1):
-        h = (loss/2+w)*N.tan(outangle/180.*N.pi) # Arrow tip height
-        move,(x,y) = path[-1]           # Use last point as reference
-        if sign==0:                     # Final loss (horizontal)
-            path.extend([(Path.LINETO,[x+dx,y]),
-                         (Path.LINETO,[x+dx,y+w]),
-                         (Path.LINETO,[x+dx+h,y-loss/2]), # Tip
-                         (Path.LINETO,[x+dx,y-loss-w]),
-                         (Path.LINETO,[x+dx,y-loss])])
-            outtips.append((sign,path[-3][1]))
-        else:                           # Intermediate loss (vertical)
-            path.extend([(Path.CURVE4,[x+dx/2,y]),
-                         (Path.CURVE4,[x+dx,y]),
-                         (Path.CURVE4,[x+dx,y+sign*dy]),
-                         (Path.LINETO,[x+dx-w,y+sign*dy]),
-                         (Path.LINETO,[x+dx+loss/2,y+sign*(dy+h)]), # Tip
-                         (Path.LINETO,[x+dx+loss+w,y+sign*dy]),
-                         (Path.LINETO,[x+dx+loss,y+sign*dy]),
-                         (Path.CURVE3,[x+dx+loss,y-sign*loss]),
-                         (Path.CURVE3,[x+dx/2+loss,y-sign*loss])])
-            outtips.append((sign,path[-5][1]))
-
-    def add_input(path, gain, sign=1):
-        h = (gain/2)*N.tan(inangle/180.*N.pi) # Dip depth
-        move,(x,y) = path[-1]           # Use last point as reference
-        if sign==0:                     # First gain (horizontal)
-            path.extend([(Path.LINETO,[x-dx,y]),
-                         (Path.LINETO,[x-dx+h,y+gain/2]), # Dip
-                         (Path.LINETO,[x-dx,y+gain])])
-            xd,yd = path[-2][1]         # Dip position
-            indips.append((sign,[xd-h,yd]))
-        else:                           # Intermediate gain (vertical)
-            path.extend([(Path.CURVE4,[x-dx/2,y]),
-                         (Path.CURVE4,[x-dx,y]),
-                         (Path.CURVE4,[x-dx,y+sign*dy]),
-                         (Path.LINETO,[x-dx-gain/2,y+sign*(dy-h)]), # Dip
-                         (Path.LINETO,[x-dx-gain,y+sign*dy]),
-                         (Path.CURVE3,[x-dx-gain,y-sign*gain]),
-                         (Path.CURVE3,[x-dx/2-gain,y-sign*gain])])
-            xd,yd = path[-4][1]         # Dip position
-            indips.append((sign,[xd,yd+sign*h]))
-
-    outtips = []                        # Output arrow tip dir. and positions
-    urpath = [(Path.MOVETO,[0,100])]    # 1st point of upper right path
-    lrpath = [(Path.LINETO,[0,0])]      # 1st point of lower right path
-    for loss,sign in zip(outs,outsigns):
-        add_output(sign>=0 and urpath or lrpath, loss, sign=sign)
-
-    indips = []                         # Input arrow tip dir. and positions
-    llpath = [(Path.LINETO,[0,0])]      # 1st point of lower left path
-    ulpath = [(Path.MOVETO,[0,100])]    # 1st point of upper left path
-    for gain,sign in zip(ins,insigns)[::-1]:
-        add_input(sign<=0 and llpath or ulpath, gain, sign=sign)
-
-    def revert(path):
-        """A path is not just revertable by path[::-1] because of Bezier
-        curves."""
-        rpath = []
-        nextmove = Path.LINETO
-        for move,pos in path[::-1]:
-            rpath.append((nextmove,pos))
-            nextmove = move
-        return rpath
-
-    # Concatenate subpathes in correct order
-    path = urpath + revert(lrpath) + llpath + revert(ulpath)
-    
-    codes,verts = zip(*path)
-    verts = N.array(verts)
-
-    # Path patch
-    path = Path(verts,codes)
-    patch = mpatches.PathPatch(path, **kwargs)
-    ax.add_patch(patch)
-
-    if False:                           # DEBUG
-        print "urpath", urpath
-        print "lrpath", revert(lrpath)
-        print "llpath", llpath
-        print "ulpath", revert(ulpath)
-
-        xs,ys = zip(*verts)
-        ax.plot(xs,ys,'go-')
-
-    # Labels
-
-    def set_labels(labels,values):
-        """Set or check labels according to values."""
-        if labels=='':                   # No labels
-            return labels
-        elif labels is None:             # Default labels
-            return [ '%2d%%' % val for val in values ]
-        else:
-            assert len(labels)==len(values)
-            return labels
-
-    def put_labels(labels,positions,output=True):
-        """Put labels to positions."""
-        texts = []
-        lbls = output and labels or labels[::-1]
-        for i,label in enumerate(lbls):
-            s,(x,y) = positions[i]      # Label direction and position
-            if s==0:
-                t = ax.text(x+offset,y,label,
-                            ha=output and 'left' or 'right', va='center')
-            elif s>0:
-                t = ax.text(x,y+offset,label, ha='center', va='bottom')
-            else:
-                t = ax.text(x,y-offset,label, ha='center', va='top')
-            texts.append(t)
-        return texts
-
-    outlabels = set_labels(outlabels, outs)
-    outtexts = put_labels(outlabels, outtips, output=True)
-
-    inlabels = set_labels(inlabels, ins)
-    intexts = put_labels(inlabels, indips, output=False)
-
-    # Axes management
-    ax.set_xlim(verts[:,0].min()-dx, verts[:,0].max()+dx)
-    ax.set_ylim(verts[:,1].min()-dy, verts[:,1].max()+dy)
-    ax.set_aspect('equal', adjustable='datalim')
-
-    return patch,[intexts,outtexts]
-
-if __name__=='__main__':
-
-    import matplotlib.pyplot as P
-
-    outputs = [10.,-20.,5.,15.,-10.,40.]
-    outlabels = ['First','Second','Third','Fourth','Fifth','Hurray!']
-    outlabels = [ s+'\n%d%%' % abs(l) for l,s in zip(outputs,outlabels) ]
-
-    inputs = [60.,-25.,15.]
-
-    fig = P.figure()
-    ax = fig.add_subplot(1,1,1, xticks=[],yticks=[],
-                         title="Sankey diagram"
-                         )
-
-    patch,(intexts,outtexts) = sankey(ax, outputs=outputs, outlabels=outlabels,
-                                      inputs=inputs, inlabels=None,
-                                      fc='g', alpha=0.2)
-    outtexts[1].set_color('r')
-    outtexts[-1].set_fontweight('bold')
-
-    P.show()
+"""Demonstrate the Sankey class.
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+
+from matplotlib.sankey import Sankey
+from itertools import cycle
+
+
+"""Demonstrate the Sankey class.
+"""
+import matplotlib.pyplot as plt
+from itertools import cycle
+
+# Example 1 -- Mostly defaults
+# This demonstrates how to create a simple diagram by implicitly calling the
+# Sankey.add() method and by appending finish() to the call to the class.
+Sankey(flows=[0.25, 0.15, 0.60, -0.20, -0.15, -0.05, -0.50, -0.10],
+       labels=['', '', '', 'First', 'Second', 'Third', 'Fourth', 'Fifth'],
+       orientations=[-1, 1, 0, 1, 1, 1, 0, -1]).finish()
+plt.title("The default settings produce a diagram like this.")
+# Notice:
+#   1. Axes weren't provided when Sankey() was instantiated, so they were
+#      created automatically.
+#   2. The scale argument wasn't necessary since the data was already
+#      normalized.
+#   3. By default, the lengths of the paths are justified.
+
+# Example 2
+# This demonstrates:
+#   1. Setting one path longer than the others
+#   2. Placing a label in the middle of the diagram
+#   3. Using the the scale argument to normalize the flows
+#   4. Implicitly passing keyword arguments to PathPatch()
+#   5. Changing the angle of the arrow heads
+#   6. Changing the offset between the tips of the paths and their labels
+#   7. Formatting the numbers in the path labels and the associated unit
+#   8. Changing the appearance of the patch and the labels after the figure
+#      is created
+fig = plt.figure()
+ax = fig.add_subplot(1, 1, 1, xticks=[], yticks=[],
+                     title="Flow Diagram of a Widget")
+sankey = Sankey(ax=ax, scale=0.01, offset=0.2, head_angle=180,
+                format='%.0f', unit='%')
+sankey.add(flows=[25, 0, 60, -10, -20, -5, -15, -10, -40],
+           labels = ['', '', '', 'First', 'Second', 'Third', 'Fourth',
+                     'Fifth', 'Hurray!'],
+           orientations=[-1, 1, 0, 1, 1, 1, -1, -1, 0],
+           pathlengths = [0.25, 0.25, 0.25, 0.25, 0.25, 0.6, 0.25, 0.25,
+                          0.25],
+           patchlabel="Widget\nA",
+           alpha=0.2, lw=2.0)  # Arguments to matplotlib.patches.PathPatch()
+diagrams = sankey.finish()
+diagrams[0].patch.set_facecolor('#37c959')
+diagrams[0].texts[-1].set_color('r')
+diagrams[0].text.set_fontweight('bold')
+# Without namedtuple:
+#diagrams[0][0].set_facecolor('#37c959')
+#diagrams[0][5][-1].set_color('r')
+#diagrams[0][4].set_fontweight('bold')
+# Notice:
+#   1. Since the sum of the flows isn't zero, the width of the trunk isn't
+#      uniform.  A message is given in the terminal window.
+#   2. The second flow doesn't appear because its value is zero.  A messsage
+#      is given in the terminal window.
+
+# Example 3
+# This demonstrates:
+#   1. Connecting two systems
+#   2. Turning off the labels of the quantities
+#   3. Adding a legend
+fig = plt.figure()
+ax = fig.add_subplot(1, 1, 1, xticks=[], yticks=[], title="Two Systems")
+flows = [0.25, 0.15, 0.60, -0.10, -0.05, -0.25, -0.15, -0.10, -0.35]
+sankey = Sankey(ax=ax, unit=None)
+sankey.add(flows=flows, label='one',
+           orientations=[-1, 1, 0, 1, 1, 1, -1, -1, 0])
+sankey.add(flows=[-0.25, 0.15, 0.1], fc='#37c959', label='two',
+           orientations=[-1, -1, -1], prior=0, connect=(0, 0))
+diagrams = sankey.finish()
+diagrams[-1].patch.set_hatch('/')
+# Without namedtuple:
+#diagrams[-1][0].set_hatch('/')
+
+plt.legend(loc='best')
+# Notice that only one connection is specified, but the systems form a
+# circuit since: (1) the lengths of the paths are justified and (2) the
+# orientation and ordering of the flows is mirrored.
+
+# Example 4
+# This tests a long chain of connections.
+links_per_side = 6
+def side(sankey, n=1):
+    prior = len(sankey.diagrams)
+    colors = cycle(['orange', 'b', 'g', 'r', 'c', 'm', 'y'])
+    for i in range(0, 2*n, 2):
+        sankey.add(flows=[1, -1], orientations=[-1, -1],
+                   patchlabel=str(prior+i), facecolor=colors.next(),
+                   prior=prior+i-1, connect=(1, 0), alpha=0.5)
+        sankey.add(flows=[1, -1], orientations=[1, 1],
+                   patchlabel=str(prior+i+1), facecolor=colors.next(),
+                   prior=prior+i, connect=(1, 0), alpha=0.5)
+def corner(sankey):
+    prior = len(sankey.diagrams)
+    sankey.add(flows=[1, -1], orientations=[0, 1],
+               patchlabel=str(prior), facecolor='k',
+               prior=prior-1, connect=(1, 0), alpha=0.5)
+fig = plt.figure()
+ax = fig.add_subplot(1, 1, 1, xticks=[], yticks=[],
+                     title="Why would you want to do this?" \
+                           "\n(But you could.)")
+sankey = Sankey(ax=ax, unit=None)
+sankey.add(flows=[1, -1], orientations=[0, 1],
+           patchlabel="0", facecolor='k',
+           rotation=45)
+side(sankey, n=links_per_side)
+corner(sankey)
+side(sankey, n=links_per_side)
+corner(sankey)
+side(sankey, n=links_per_side)
+corner(sankey)
+side(sankey, n=links_per_side)
+sankey.finish()
+# Notice:
+# 1. The alignment doesn't drift significantly (if at all; with 16007
+#    subdiagrams there is still closure).
+# 2. The first diagram is rotated 45 degrees, so all other diagrams are
+#    rotated accordingly.
+
+# Example 5
+# This demonstrates a practical example -- a Rankine power cycle.
+fig = plt.figure(figsize=(8, 12))
+ax = fig.add_subplot(1, 1, 1, xticks=[], yticks=[],
+                     title="Rankine Power Cycle: Example 8.6 from Moran and Shapiro\n"
+                           + "\x22Fundamentals of Engineering Thermodynamics\x22, 6th ed., 2008")
+Hdot = np.array([260.431, 35.078, 180.794, 221.115, 22.700,
+        142.361, 10.193, 10.210, 43.670, 44.312,
+        68.631, 10.758, 10.758, 0.017, 0.642,
+        232.121, 44.559, 100.613, 132.168])*1.0e6 # W
+sankey = Sankey(ax=ax, format='%.3G', unit='W', gap=0.5, scale=1.0/Hdot[0])
+# Shared copy:
+#Hdot = [260.431, 35.078, 180.794, 221.115, 22.700,
+#        142.361, 10.193, 10.210, 43.670, 44.312,
+#        68.631, 10.758, 10.758, 0.017, 0.642,
+#        232.121, 44.559, 100.613, 132.168] # MW
+#sankey = Sankey(ax=ax, format='%.3G', unit=' MW', gap=0.5, scale=1.0/Hdot[0])
+sankey.add(patchlabel='\n\nPump 1', rotation=90, facecolor='#37c959',
+           flows=[Hdot[13], Hdot[6], -Hdot[7]],
+           labels=['Shaft power', '', None],
+           pathlengths=[0.4, 0.883, 0.25],
+           orientations=[1, -1, 0])
+sankey.add(patchlabel='\n\nOpen\nheater', facecolor='#37c959',
+           flows=[Hdot[11], Hdot[7], Hdot[4], -Hdot[8]],
+           labels=[None, '', None, None],
+           pathlengths=[0.25, 0.25, 1.93, 0.25],
+           orientations=[1, 0, -1, 0], prior=0, connect=(2, 1))
+sankey.add(patchlabel='\n\nPump 2', facecolor='#37c959',
+           flows=[Hdot[14], Hdot[8], -Hdot[9]],
+           labels=['Shaft power', '', None],
+           pathlengths=[0.4, 0.25, 0.25],
+           orientations=[1, 0, 0], prior=1, connect=(3, 1))
+sankey.add(patchlabel='Closed\nheater', trunklength=2.914, fc='#37c959',
+           flows=[Hdot[9], Hdot[1], -Hdot[11], -Hdot[10]],
+           pathlengths=[0.25, 1.543, 0.25, 0.25],
+           labels=['', '', None, None],
+           orientations=[0, -1, 1, -1], prior=2, connect=(2, 0))
+sankey.add(patchlabel='Trap', facecolor='#37c959', trunklength=5.102,
+           flows=[Hdot[11], -Hdot[12]],
+           labels=['\n', None],
+           pathlengths=[1.0, 1.01],
+           orientations=[1, 1], prior=3, connect=(2, 0))
+sankey.add(patchlabel='Steam\ngenerator', facecolor='#ff5555',
+           flows=[Hdot[15], Hdot[10], Hdot[2], -Hdot[3], -Hdot[0]],
+           labels=['Heat rate', '', '', None, None],
+           pathlengths=0.25,
+           orientations=[1, 0, -1, -1, -1], prior=3, connect=(3, 1))
+sankey.add(patchlabel='\n\n\nTurbine 1', facecolor='#37c959',
+           flows=[Hdot[0], -Hdot[16], -Hdot[1], -Hdot[2]],
+           labels=['', None, None, None],
+           pathlengths=[0.25, 0.153, 1.543, 0.25],
+           orientations=[0, 1, -1, -1], prior=5, connect=(4, 0))
+sankey.add(patchlabel='\n\n\nReheat', facecolor='#37c959',
+           flows=[Hdot[2], -Hdot[2]],
+           labels=[None, None],
+           pathlengths=[0.725, 0.25],
+           orientations=[-1, 0], prior=6, connect=(3, 0))
+sankey.add(patchlabel='Turbine 2', trunklength=3.212, facecolor='#37c959',
+           flows=[Hdot[3], Hdot[16], -Hdot[5], -Hdot[4], -Hdot[17]],
+           labels=[None, 'Shaft power', None, '', 'Shaft power'],
+           pathlengths=[0.751, 0.15, 0.25, 1.93, 0.25],
+           orientations=[0, -1, 0, -1, 1], prior=6, connect=(1, 1))
+sankey.add(patchlabel='Condenser', facecolor='#58b1fa', trunklength=1.764,
+           flows=[Hdot[5], -Hdot[18], -Hdot[6]],
+           labels=['', 'Heat rate', None],
+           pathlengths=[0.45, 0.25, 0.883],
+           orientations=[-1, 1, 0], prior=8, connect=(2, 0))
+diagrams = sankey.finish()
+for diagram in diagrams:
+    diagram.text.set_fontweight('bold')
+    diagram.text.set_fontsize('10')
+    for text in diagram.texts:
+    # Without namedtuple:
+    #diagram[4].set_fontweight('bold')
+    #diagram[4].set_fontsize('10')
+    #for text in diagram[5]:
+        text.set_fontsize('10')
+# Notice that the explicit connections are handled automatically, but the
+# implicit ones currently are not.  The lengths of the paths and the trunks
+# must be adjusted manually, and that is a bit tricky.
+
+plt.show()