Improved documentation and formatting of Sankey class and its demo

kdavies4 · kdavies4 · commit dec14bd8a28c · 2011-09-16T21:34:28.000-04:00
diff --git a/doc/api/api_changes.rst b/doc/api/api_changes.rst
@@ -14,6 +14,8 @@ For new features that were added to matplotlib, please see
 Changes in 1.1.x
 ================
 
+* Added new :class:`matplotlib.sankey.Sankey` for generating Sankey diagrams.
+
 * In :meth:`~matplotlib.pyplot.imshow`, setting *interpolation* to 'nearest'
   will now always mean that the nearest-neighbor interpolation is performed.
   If you want the no-op interpolation to be performed, choose 'none'.
diff --git a/examples/api/sankey_demo.py b/examples/api/sankey_demo.py
@@ -6,12 +6,6 @@
 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.
@@ -35,8 +29,8 @@
 #   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
+#   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")
@@ -49,20 +43,17 @@
            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()
+           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.
+#   1. Since the sum of the flows is nonzero, the width of the trunk isn't
+#      uniform.  If verbose.level is helpful (in matplotlibrc), a message is
+#      given in the terminal window.
+#   2. The second flow doesn't appear because its value is zero.  Again, if
+#      verbose.level is helpful, a message is given in the terminal window.
 
 # Example 3
 # This demonstrates:
@@ -79,18 +70,17 @@
            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.
+# 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):
+    """Generate a side chain.
+    """
     prior = len(sankey.diagrams)
     colors = cycle(['orange', 'b', 'g', 'r', 'c', 'm', 'y'])
     for i in range(0, 2*n, 2):
@@ -101,14 +91,15 @@ def side(sankey, n=1):
                    patchlabel=str(prior+i+1), facecolor=colors.next(),
                    prior=prior+i, connect=(1, 0), alpha=0.5)
 def corner(sankey):
+    """Generate a corner link.
+    """
     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.)")
+                     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',
@@ -124,26 +115,20 @@ def corner(sankey):
 # 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.
+# 2. The first diagram is rotated 45 deg, so all other diagrams are rotated
+#    accordingly.
 
 # Example 5
-# This demonstrates a practical example -- a Rankine power cycle.
+# This is a practical example of 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,
+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])*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])
+        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],
@@ -199,10 +184,6 @@ def corner(sankey):
     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
diff --git a/lib/matplotlib/sankey.py b/lib/matplotlib/sankey.py
