Thanks to visit codestin.com
Credit goes to github.com

Skip to content

add sankey module and new demo #465

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 14, 2011
Merged

add sankey module and new demo #465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
399 changes: 211 additions & 188 deletions examples/api/sankey_demo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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()
Loading