Merge pull request #1574 from kdavies4/sankey6

WeatherGod · WeatherGod · commit 3e6e735b0f9c · 2012-12-07T08:52:48.000-08:00
Improvements to Sankey class
diff --git a/examples/api/sankey_demo_rankine.py b/examples/api/sankey_demo_rankine.py
@@ -5,7 +5,7 @@
 
 from matplotlib.sankey import Sankey
 
-fig = plt.figure(figsize=(8, 12))
+fig = plt.figure(figsize=(8, 9))
 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")
diff --git a/lib/matplotlib/sankey.py b/lib/matplotlib/sankey.py
@@ -2,6 +2,7 @@
 """
 Module for creating Sankey diagrams using matplotlib
 """
+from __future__ import print_function
 __author__ = "Kevin L. Davies"
 __credits__ = ["Yannick Copin"]
 __license__ = "BSD"
@@ -12,8 +13,8 @@
 #     http://matplotlib.org/examples/api/sankey_demo_old.html
 # Modifications by Kevin Davies (kld@alumni.carnegiemellon.edu) 6/3/2011:
 #   --Used arcs for the curves (so that the widths of the paths are uniform)
-#   --Converted the function to a class and created methods to join
-#     multiple simple Sankey diagrams
+#   --Converted the function to a class and created methods to join multiple
+#     simple Sankey diagrams
 #   --Provided handling for cases where the total of the inputs isn't 100
 #     Now, the default layout is based on the assumption that the inputs sum to
 #     1.  A scaling parameter can be used in other cases.
@@ -29,7 +30,7 @@
 #     inputs/outputs are now specified via an orientation of 0, and there may
 #     be several of each.
 #   --Added assertions to catch common calling errors
-#    -Added the physical unit as a string argument to be used in the labels, so
+#   --Added the physical unit as a string argument to be used in the labels, so
 #     that the values of the flows can usually be applied automatically
 #   --Added an argument for a minimum magnitude below which flows are not shown
 #   --Added a tapered trunk in the case that the flows do not sum to 0
@@ -395,10 +396,10 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
                             *orientations* == 0, inputs will break in from the
                             left and outputs will break away to the right.
           *labels*          list of specifications of the labels for the flows
-                            Each value may be None (no labels), '' (just label
-                            the quantities), or a labeling string.  If a single
-                            value is provided, it will be applied to all flows.
-                            If an entry is a non-empty string, then the
+                            Each value may be *None* (no labels), '' (just
+                            label the quantities), or a labeling string.  If a
+                            single value is provided, it will be applied to all
+                            flows.  If an entry is a non-empty string, then the
                             quantity for the corresponding flow will be shown
                             below the string.  However, if the *unit* of the
                             main diagram is None, then quantities are never
@@ -417,7 +418,7 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
           *connect*         a (prior, this) tuple indexing the flow of the
                             prior diagram and the flow of this diagram which
                             should be connected
-                            If this is the first diagram or *prior* is None,
+                            If this is the first diagram or *prior* is *None*,
                             *connect* will be ignored.
           *rotation*        angle of rotation of the diagram [deg]
                             *rotation* is ignored if this diagram is connected
@@ -460,7 +461,7 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
         if rotation is None:
             rotation = 0
         else:
-            # In the code below, angles are expressed in deg/90
+            # In the code below, angles are expressed in deg/90.
             rotation /= 90.0
         if orientations is None:
             orientations = [0, 0]
@@ -741,10 +742,10 @@ def _get_angle(a, r):
             kwds = dict(s=patchlabel, ha='center', va='center')
             text = self.ax.text(*offset, **kwds)
         if False:  # Debug
-            print "llpath\n", llpath
-            print "ulpath\n", self._revert(ulpath)
-            print "urpath\n", urpath
-            print "lrpath\n", self._revert(lrpath)
+            print("llpath\n", llpath)
+            print("ulpath\n", self._revert(ulpath))
+            print("urpath\n", urpath)
+            print("lrpath\n", self._revert(lrpath))
             xs, ys = zip(*vertices)
             self.ax.plot(xs, ys, 'go-')
         patch = PathPatch(Path(vertices, codes),
@@ -755,22 +756,18 @@ def _get_angle(a, r):
         self.ax.add_patch(patch)
 
         # Add the path labels.
-        for i, (number, angle) in enumerate(zip(flows, angles)):
-            if labels[i] is None or angle is None:
-                labels[i] = ''
+        texts = []
+        for number, angle, label, location in zip(flows, angles, labels,
+                                                  label_locations):
+            if label is None or angle is None:
+                label = ''
             elif self.unit is not None:
                 quantity = self.format % abs(number) + self.unit
-                if labels[i] != '':
-                    labels[i] += "\n"
-                labels[i] += quantity
-        texts = []
-        for i, (label, location) in enumerate(zip(labels, label_locations)):
-            if label:
-                s = label
-            else:
-                s = ''
+                if label != '':
+                    label += "\n"
+                label += quantity
             texts.append(self.ax.text(x=location[0], y=location[1],
-                                      s=s,
+                                      s=label,
                                       ha='center', va='center'))
         # Text objects are placed even they are empty (as long as the magnitude
         # of the corresponding flow is larger than the tolerance) in case the
@@ -820,7 +817,7 @@ def finish(self):
                             (DOWN), and an output from the top side will have
                             an angle of 1 (UP).  If a flow has been skipped
                             (because its magnitude is less than *tolerance*),
-                            then its angle will be None.
+                            then its angle will be *None*.
           *tips*            array in which each row is an [x, y] pair
                             indicating the positions of the tips (or "dips") of
                             the flow paths
@@ -844,5 +841,3 @@ def finish(self):
                       self.extent[3] + self.margin])
         self.ax.set_aspect('equal', adjustable='datalim')
         return self.diagrams
-
-
