|
|
- """
- Module for creating Sankey diagrams using matplotlib
- """
-
- import logging
- from types import SimpleNamespace
-
- import numpy as np
-
- from matplotlib.cbook import iterable
- from matplotlib.path import Path
- from matplotlib.patches import PathPatch
- from matplotlib.transforms import Affine2D
- from matplotlib import docstring
- from matplotlib import rcParams
-
- _log = logging.getLogger(__name__)
-
- __author__ = "Kevin L. Davies"
- __credits__ = ["Yannick Copin"]
- __license__ = "BSD"
- __version__ = "2011/09/16"
-
- # Angles [deg/90]
- RIGHT = 0
- UP = 1
- # LEFT = 2
- DOWN = 3
-
-
- class Sankey(object):
- """
- Sankey diagram in matplotlib
-
- Sankey diagrams are a specific type of flow diagram, in which
- the width of the arrows is shown proportionally to the flow
- quantity. They are typically used to visualize energy or
- material or cost transfers between processes.
- `Wikipedia (6/1/2011) <https://en.wikipedia.org/wiki/Sankey_diagram>`_
-
- """
-
- def __init__(self, ax=None, scale=1.0, unit='', format='%G', gap=0.25,
- radius=0.1, shoulder=0.03, offset=0.15, head_angle=100,
- margin=0.4, tolerance=1e-6, **kwargs):
- """
- Create a new Sankey instance.
-
- Optional keyword arguments:
-
- =============== ===================================================
- Field Description
- =============== ===================================================
- *ax* axes onto which the data should be plotted
- If *ax* isn't provided, new axes will be created.
- *scale* scaling factor for the flows
- *scale* sizes the width of the paths in order to
- maintain proper layout. The same scale is applied
- to all subdiagrams. The value should be chosen
- such that the product of the scale and the sum of
- the inputs is approximately 1.0 (and the product of
- the scale and the sum of the outputs is
- approximately -1.0).
- *unit* string representing the physical unit associated
- with the flow quantities
- If *unit* is None, then none of the quantities are
- labeled.
- *format* a Python number formatting string to be used in
- labeling the flow as a quantity (i.e., a number
- times a unit, where the unit is given)
- *gap* space between paths that break in/break away
- to/from the top or bottom
- *radius* inner radius of the vertical paths
- *shoulder* size of the shoulders of output arrowS
- *offset* text offset (from the dip or tip of the arrow)
- *head_angle* angle of the arrow heads (and negative of the angle
- of the tails) [deg]
- *margin* minimum space between Sankey outlines and the edge
- of the plot area
- *tolerance* acceptable maximum of the magnitude of the sum of
- flows
- The magnitude of the sum of connected flows cannot
- be greater than *tolerance*.
- =============== ===================================================
-
- The optional arguments listed above are applied to all subdiagrams so
- that there is consistent alignment and formatting.
-
- If :class:`Sankey` is instantiated with any keyword arguments other
- than those explicitly listed above (``**kwargs``), they will be passed
- to :meth:`add`, which will create the first subdiagram.
-
- In order to draw a complex Sankey diagram, create an instance of
- :class:`Sankey` by calling it without any kwargs::
-
- sankey = Sankey()
-
- Then add simple Sankey sub-diagrams::
-
- sankey.add() # 1
- sankey.add() # 2
- #...
- sankey.add() # n
-
- Finally, create the full diagram::
-
- sankey.finish()
-
- Or, instead, simply daisy-chain those calls::
-
- Sankey().add().add... .add().finish()
-
- .. seealso::
-
- :meth:`add`
- :meth:`finish`
-
-
- **Examples:**
-
- .. plot:: gallery/specialty_plots/sankey_basics.py
- """
- # Check the arguments.
- if gap < 0:
- raise ValueError(
- "The gap is negative.\nThis isn't allowed because it "
- "would cause the paths to overlap.")
- if radius > gap:
- raise ValueError(
- "The inner radius is greater than the path spacing.\n"
- "This isn't allowed because it would cause the paths to overlap.")
- if head_angle < 0:
- raise ValueError(
- "The angle is negative.\nThis isn't allowed "
- "because it would cause inputs to look like "
- "outputs and vice versa.")
- if tolerance < 0:
- raise ValueError(
- "The tolerance is negative.\nIt must be a magnitude.")
-
- # Create axes if necessary.
- if ax is None:
- import matplotlib.pyplot as plt
- fig = plt.figure()
- ax = fig.add_subplot(1, 1, 1, xticks=[], yticks=[])
-
- self.diagrams = []
-
- # Store the inputs.
- self.ax = ax
- self.unit = unit
- self.format = format
- self.scale = scale
- self.gap = gap
- self.radius = radius
- self.shoulder = shoulder
- self.offset = offset
- self.margin = margin
- self.pitch = np.tan(np.pi * (1 - head_angle / 180.0) / 2.0)
- self.tolerance = tolerance
-
- # Initialize the vertices of tight box around the diagram(s).
- self.extent = np.array((np.inf, -np.inf, np.inf, -np.inf))
-
- # If there are any kwargs, create the first subdiagram.
- if len(kwargs):
- self.add(**kwargs)
-
- def _arc(self, quadrant=0, cw=True, radius=1, center=(0, 0)):
- """
- Return the codes and vertices for a rotated, scaled, and translated
- 90 degree arc.
-
- Optional keyword arguments:
-
- =============== ==========================================
- Keyword Description
- =============== ==========================================
- *quadrant* uses 0-based indexing (0, 1, 2, or 3)
- *cw* if True, clockwise
- *center* (x, y) tuple of the arc's center
- =============== ==========================================
- """
- # Note: It would be possible to use matplotlib's transforms to rotate,
- # scale, and translate the arc, but since the angles are discrete,
- # it's just as easy and maybe more efficient to do it here.
- ARC_CODES = [Path.LINETO,
- Path.CURVE4,
- Path.CURVE4,
- Path.CURVE4,
- Path.CURVE4,
- Path.CURVE4,
- Path.CURVE4]
- # Vertices of a cubic Bezier curve approximating a 90 deg arc
- # These can be determined by Path.arc(0,90).
- ARC_VERTICES = np.array([[1.00000000e+00, 0.00000000e+00],
- [1.00000000e+00, 2.65114773e-01],
- [8.94571235e-01, 5.19642327e-01],
- [7.07106781e-01, 7.07106781e-01],
- [5.19642327e-01, 8.94571235e-01],
- [2.65114773e-01, 1.00000000e+00],
- # Insignificant
- # [6.12303177e-17, 1.00000000e+00]])
- [0.00000000e+00, 1.00000000e+00]])
- if quadrant == 0 or quadrant == 2:
- if cw:
- vertices = ARC_VERTICES
- else:
- vertices = ARC_VERTICES[:, ::-1] # Swap x and y.
- elif quadrant == 1 or quadrant == 3:
- # Negate x.
- if cw:
- # Swap x and y.
- vertices = np.column_stack((-ARC_VERTICES[:, 1],
- ARC_VERTICES[:, 0]))
- else:
- vertices = np.column_stack((-ARC_VERTICES[:, 0],
- ARC_VERTICES[:, 1]))
- if quadrant > 1:
- radius = -radius # Rotate 180 deg.
- return list(zip(ARC_CODES, radius * vertices +
- np.tile(center, (ARC_VERTICES.shape[0], 1))))
-
- def _add_input(self, path, angle, flow, length):
- """
- Add an input to a path and return its tip and label locations.
- """
- if angle is None:
- return [0, 0], [0, 0]
- else:
- x, y = path[-1][1] # Use the last point as a reference.
- dipdepth = (flow / 2) * self.pitch
- if angle == RIGHT:
- x -= length
- dip = [x + dipdepth, y + flow / 2.0]
- path.extend([(Path.LINETO, [x, y]),
- (Path.LINETO, dip),
- (Path.LINETO, [x, y + flow]),
- (Path.LINETO, [x + self.gap, y + flow])])
- label_location = [dip[0] - self.offset, dip[1]]
- else: # Vertical
- x -= self.gap
- if angle == UP:
- sign = 1
- else:
- sign = -1
-
- dip = [x - flow / 2, y - sign * (length - dipdepth)]
- if angle == DOWN:
- quadrant = 2
- else:
- quadrant = 1
-
- # Inner arc isn't needed if inner radius is zero
- if self.radius:
- path.extend(self._arc(quadrant=quadrant,
- cw=angle == UP,
- radius=self.radius,
- center=(x + self.radius,
- y - sign * self.radius)))
- else:
- path.append((Path.LINETO, [x, y]))
- path.extend([(Path.LINETO, [x, y - sign * length]),
- (Path.LINETO, dip),
- (Path.LINETO, [x - flow, y - sign * length])])
- path.extend(self._arc(quadrant=quadrant,
- cw=angle == DOWN,
- radius=flow + self.radius,
- center=(x + self.radius,
- y - sign * self.radius)))
- path.append((Path.LINETO, [x - flow, y + sign * flow]))
- label_location = [dip[0], dip[1] - sign * self.offset]
-
- return dip, label_location
-
- def _add_output(self, path, angle, flow, length):
- """
- Append an output to a path and return its tip and label locations.
-
- .. note:: *flow* is negative for an output.
- """
- if angle is None:
- return [0, 0], [0, 0]
- else:
- x, y = path[-1][1] # Use the last point as a reference.
- tipheight = (self.shoulder - flow / 2) * self.pitch
- if angle == RIGHT:
- x += length
- tip = [x + tipheight, y + flow / 2.0]
- path.extend([(Path.LINETO, [x, y]),
- (Path.LINETO, [x, y + self.shoulder]),
- (Path.LINETO, tip),
- (Path.LINETO, [x, y - self.shoulder + flow]),
- (Path.LINETO, [x, y + flow]),
- (Path.LINETO, [x - self.gap, y + flow])])
- label_location = [tip[0] + self.offset, tip[1]]
- else: # Vertical
- x += self.gap
- if angle == UP:
- sign = 1
- else:
- sign = -1
-
- tip = [x - flow / 2.0, y + sign * (length + tipheight)]
- if angle == UP:
- quadrant = 3
- else:
- quadrant = 0
- # Inner arc isn't needed if inner radius is zero
- if self.radius:
- path.extend(self._arc(quadrant=quadrant,
- cw=angle == UP,
- radius=self.radius,
- center=(x - self.radius,
- y + sign * self.radius)))
- else:
- path.append((Path.LINETO, [x, y]))
- path.extend([(Path.LINETO, [x, y + sign * length]),
- (Path.LINETO, [x - self.shoulder,
- y + sign * length]),
- (Path.LINETO, tip),
- (Path.LINETO, [x + self.shoulder - flow,
- y + sign * length]),
- (Path.LINETO, [x - flow, y + sign * length])])
- path.extend(self._arc(quadrant=quadrant,
- cw=angle == DOWN,
- radius=self.radius - flow,
- center=(x - self.radius,
- y + sign * self.radius)))
- path.append((Path.LINETO, [x - flow, y + sign * flow]))
- label_location = [tip[0], tip[1] + sign * self.offset]
- return tip, label_location
-
- def _revert(self, path, first_action=Path.LINETO):
- """
- A path is not simply revertable by path[::-1] since the code
- specifies an action to take from the **previous** point.
- """
- reverse_path = []
- next_code = first_action
- for code, position in path[::-1]:
- reverse_path.append((next_code, position))
- next_code = code
- return reverse_path
- # This might be more efficient, but it fails because 'tuple' object
- # doesn't support item assignment:
- # path[1] = path[1][-1:0:-1]
- # path[1][0] = first_action
- # path[2] = path[2][::-1]
- # return path
-
- @docstring.dedent_interpd
- def add(self, patchlabel='', flows=None, orientations=None, labels='',
- trunklength=1.0, pathlengths=0.25, prior=None, connect=(0, 0),
- rotation=0, **kwargs):
- """
- Add a simple Sankey diagram with flows at the same hierarchical level.
-
- Return value is the instance of :class:`Sankey`.
-
- Optional keyword arguments:
-
- =============== ===================================================
- Keyword Description
- =============== ===================================================
- *patchlabel* label to be placed at the center of the diagram
- Note: *label* (not *patchlabel*) will be passed to
- the patch through ``**kwargs`` and can be used to
- create an entry in the legend.
- *flows* array of flow values
- By convention, inputs are positive and outputs are
- negative.
- *orientations* list of orientations of the paths
- Valid values are 1 (from/to the top), 0 (from/to
- the left or right), or -1 (from/to the bottom). If
- *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
- 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
- shown, regardless of the value of this argument.
- *trunklength* length between the bases of the input and output
- groups
- *pathlengths* list of lengths of the arrows before break-in or
- after break-away
- If a single value is given, then it will be applied
- to the first (inside) paths on the top and bottom,
- and the length of all other arrows will be
- justified accordingly. The *pathlengths* are not
- applied to the horizontal inputs and outputs.
- *prior* index of the prior diagram to which this diagram
- should be connected
- *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*,
- *connect* will be ignored.
- *rotation* angle of rotation of the diagram [deg]
- *rotation* is ignored if this diagram is connected
- to an existing one (using *prior* and *connect*).
- The interpretation of the *orientations* argument
- will be rotated accordingly (e.g., if *rotation*
- == 90, an *orientations* entry of 1 means to/from
- the left).
- =============== ===================================================
-
- Valid kwargs are :meth:`matplotlib.patches.PathPatch` arguments:
-
- %(Patch)s
-
- As examples, ``fill=False`` and ``label='A legend entry'``.
- By default, ``facecolor='#bfd1d4'`` (light blue) and
- ``linewidth=0.5``.
-
- The indexing parameters (*prior* and *connect*) are zero-based.
-
- The flows are placed along the top of the diagram from the inside out
- in order of their index within the *flows* list or array. They are
- placed along the sides of the diagram from the top down and along the
- bottom from the outside in.
-
- If the sum of the inputs and outputs is nonzero, the discrepancy
- will appear as a cubic Bezier curve along the top and bottom edges of
- the trunk.
-
- .. seealso::
-
- :meth:`finish`
- """
- # Check and preprocess the arguments.
- if flows is None:
- flows = np.array([1.0, -1.0])
- else:
- flows = np.array(flows)
- n = flows.shape[0] # Number of flows
- if rotation is None:
- rotation = 0
- else:
- # In the code below, angles are expressed in deg/90.
- rotation /= 90.0
- if orientations is None:
- orientations = [0, 0]
- if len(orientations) != n:
- raise ValueError(
- "orientations and flows must have the same length.\n"
- "orientations has length %d, but flows has length %d."
- % (len(orientations), n))
- if labels != '' and getattr(labels, '__iter__', False):
- # iterable() isn't used because it would give True if labels is a
- # string
- if len(labels) != n:
- raise ValueError(
- "If labels is a list, then labels and flows must have the "
- "same length.\nlabels has length %d, but flows has length %d."
- % (len(labels), n))
- else:
- labels = [labels] * n
- if trunklength < 0:
- raise ValueError(
- "trunklength is negative.\nThis isn't allowed, because it would "
- "cause poor layout.")
- if np.abs(np.sum(flows)) > self.tolerance:
- _log.info("The sum of the flows is nonzero (%f).\nIs the "
- "system not at steady state?", np.sum(flows))
- scaled_flows = self.scale * flows
- gain = sum(max(flow, 0) for flow in scaled_flows)
- loss = sum(min(flow, 0) for flow in scaled_flows)
- if not (0.5 <= gain <= 2.0):
- _log.info(
- "The scaled sum of the inputs is %f.\nThis may "
- "cause poor layout.\nConsider changing the scale so"
- " that the scaled sum is approximately 1.0.", gain)
- if not (-2.0 <= loss <= -0.5):
- _log.info(
- "The scaled sum of the outputs is %f.\nThis may "
- "cause poor layout.\nConsider changing the scale so"
- " that the scaled sum is approximately 1.0.", gain)
- if prior is not None:
- if prior < 0:
- raise ValueError("The index of the prior diagram is negative.")
- if min(connect) < 0:
- raise ValueError(
- "At least one of the connection indices is negative.")
- if prior >= len(self.diagrams):
- raise ValueError(
- "The index of the prior diagram is %d, but there are "
- "only %d other diagrams.\nThe index is zero-based."
- % (prior, len(self.diagrams)))
- if connect[0] >= len(self.diagrams[prior].flows):
- raise ValueError(
- "The connection index to the source diagram is %d, but "
- "that diagram has only %d flows.\nThe index is zero-based."
- % (connect[0], len(self.diagrams[prior].flows)))
- if connect[1] >= n:
- raise ValueError(
- "The connection index to this diagram is %d, but this diagram"
- "has only %d flows.\n The index is zero-based."
- % (connect[1], n))
- if self.diagrams[prior].angles[connect[0]] is None:
- raise ValueError(
- "The connection cannot be made. Check that the magnitude "
- "of flow %d of diagram %d is greater than or equal to the "
- "specified tolerance." % (connect[0], prior))
- flow_error = (self.diagrams[prior].flows[connect[0]] +
- flows[connect[1]])
- if abs(flow_error) >= self.tolerance:
- raise ValueError(
- "The scaled sum of the connected flows is %f, which is not "
- "within the tolerance (%f)." % (flow_error, self.tolerance))
-
- # Determine if the flows are inputs.
- are_inputs = [None] * n
- for i, flow in enumerate(flows):
- if flow >= self.tolerance:
- are_inputs[i] = True
- elif flow <= -self.tolerance:
- are_inputs[i] = False
- else:
- _log.info(
- "The magnitude of flow %d (%f) is below the "
- "tolerance (%f).\nIt will not be shown, and it "
- "cannot be used in a connection."
- % (i, flow, self.tolerance))
-
- # Determine the angles of the arrows (before rotation).
- angles = [None] * n
- for i, (orient, is_input) in enumerate(zip(orientations, are_inputs)):
- if orient == 1:
- if is_input:
- angles[i] = DOWN
- elif not is_input:
- # Be specific since is_input can be None.
- angles[i] = UP
- elif orient == 0:
- if is_input is not None:
- angles[i] = RIGHT
- else:
- if orient != -1:
- raise ValueError(
- "The value of orientations[%d] is %d, "
- "but it must be [ -1 | 0 | 1 ]." % (i, orient))
- if is_input:
- angles[i] = UP
- elif not is_input:
- angles[i] = DOWN
-
- # Justify the lengths of the paths.
- if iterable(pathlengths):
- if len(pathlengths) != n:
- raise ValueError(
- "If pathlengths is a list, then pathlengths and flows must "
- "have the same length.\npathlengths has length %d, but flows "
- "has length %d." % (len(pathlengths), n))
- else: # Make pathlengths into a list.
- urlength = pathlengths
- ullength = pathlengths
- lrlength = pathlengths
- lllength = pathlengths
- d = dict(RIGHT=pathlengths)
- pathlengths = [d.get(angle, 0) for angle in angles]
- # Determine the lengths of the top-side arrows
- # from the middle outwards.
- for i, (angle, is_input, flow) in enumerate(zip(angles, are_inputs,
- scaled_flows)):
- if angle == DOWN and is_input:
- pathlengths[i] = ullength
- ullength += flow
- elif angle == UP and not is_input:
- pathlengths[i] = urlength
- urlength -= flow # Flow is negative for outputs.
- # Determine the lengths of the bottom-side arrows
- # from the middle outwards.
- for i, (angle, is_input, flow) in enumerate(reversed(list(zip(
- angles, are_inputs, scaled_flows)))):
- if angle == UP and is_input:
- pathlengths[n - i - 1] = lllength
- lllength += flow
- elif angle == DOWN and not is_input:
- pathlengths[n - i - 1] = lrlength
- lrlength -= flow
- # Determine the lengths of the left-side arrows
- # from the bottom upwards.
- has_left_input = False
- for i, (angle, is_input, spec) in enumerate(reversed(list(zip(
- angles, are_inputs, zip(scaled_flows, pathlengths))))):
- if angle == RIGHT:
- if is_input:
- if has_left_input:
- pathlengths[n - i - 1] = 0
- else:
- has_left_input = True
- # Determine the lengths of the right-side arrows
- # from the top downwards.
- has_right_output = False
- for i, (angle, is_input, spec) in enumerate(zip(
- angles, are_inputs, list(zip(scaled_flows, pathlengths)))):
- if angle == RIGHT:
- if not is_input:
- if has_right_output:
- pathlengths[i] = 0
- else:
- has_right_output = True
-
- # Begin the subpaths, and smooth the transition if the sum of the flows
- # is nonzero.
- urpath = [(Path.MOVETO, [(self.gap - trunklength / 2.0), # Upper right
- gain / 2.0]),
- (Path.LINETO, [(self.gap - trunklength / 2.0) / 2.0,
- gain / 2.0]),
- (Path.CURVE4, [(self.gap - trunklength / 2.0) / 8.0,
- gain / 2.0]),
- (Path.CURVE4, [(trunklength / 2.0 - self.gap) / 8.0,
- -loss / 2.0]),
- (Path.LINETO, [(trunklength / 2.0 - self.gap) / 2.0,
- -loss / 2.0]),
- (Path.LINETO, [(trunklength / 2.0 - self.gap),
- -loss / 2.0])]
- llpath = [(Path.LINETO, [(trunklength / 2.0 - self.gap), # Lower left
- loss / 2.0]),
- (Path.LINETO, [(trunklength / 2.0 - self.gap) / 2.0,
- loss / 2.0]),
- (Path.CURVE4, [(trunklength / 2.0 - self.gap) / 8.0,
- loss / 2.0]),
- (Path.CURVE4, [(self.gap - trunklength / 2.0) / 8.0,
- -gain / 2.0]),
- (Path.LINETO, [(self.gap - trunklength / 2.0) / 2.0,
- -gain / 2.0]),
- (Path.LINETO, [(self.gap - trunklength / 2.0),
- -gain / 2.0])]
- lrpath = [(Path.LINETO, [(trunklength / 2.0 - self.gap), # Lower right
- loss / 2.0])]
- ulpath = [(Path.LINETO, [self.gap - trunklength / 2.0, # Upper left
- gain / 2.0])]
-
- # Add the subpaths and assign the locations of the tips and labels.
- tips = np.zeros((n, 2))
- label_locations = np.zeros((n, 2))
- # Add the top-side inputs and outputs from the middle outwards.
- for i, (angle, is_input, spec) in enumerate(zip(
- angles, are_inputs, list(zip(scaled_flows, pathlengths)))):
- if angle == DOWN and is_input:
- tips[i, :], label_locations[i, :] = self._add_input(
- ulpath, angle, *spec)
- elif angle == UP and not is_input:
- tips[i, :], label_locations[i, :] = self._add_output(
- urpath, angle, *spec)
- # Add the bottom-side inputs and outputs from the middle outwards.
- for i, (angle, is_input, spec) in enumerate(reversed(list(zip(
- angles, are_inputs, list(zip(scaled_flows, pathlengths)))))):
- if angle == UP and is_input:
- tip, label_location = self._add_input(llpath, angle, *spec)
- tips[n - i - 1, :] = tip
- label_locations[n - i - 1, :] = label_location
- elif angle == DOWN and not is_input:
- tip, label_location = self._add_output(lrpath, angle, *spec)
- tips[n - i - 1, :] = tip
- label_locations[n - i - 1, :] = label_location
- # Add the left-side inputs from the bottom upwards.
- has_left_input = False
- for i, (angle, is_input, spec) in enumerate(reversed(list(zip(
- angles, are_inputs, list(zip(scaled_flows, pathlengths)))))):
- if angle == RIGHT and is_input:
- if not has_left_input:
- # Make sure the lower path extends
- # at least as far as the upper one.
- if llpath[-1][1][0] > ulpath[-1][1][0]:
- llpath.append((Path.LINETO, [ulpath[-1][1][0],
- llpath[-1][1][1]]))
- has_left_input = True
- tip, label_location = self._add_input(llpath, angle, *spec)
- tips[n - i - 1, :] = tip
- label_locations[n - i - 1, :] = label_location
- # Add the right-side outputs from the top downwards.
- has_right_output = False
- for i, (angle, is_input, spec) in enumerate(zip(
- angles, are_inputs, list(zip(scaled_flows, pathlengths)))):
- if angle == RIGHT and not is_input:
- if not has_right_output:
- # Make sure the upper path extends
- # at least as far as the lower one.
- if urpath[-1][1][0] < lrpath[-1][1][0]:
- urpath.append((Path.LINETO, [lrpath[-1][1][0],
- urpath[-1][1][1]]))
- has_right_output = True
- tips[i, :], label_locations[i, :] = self._add_output(
- urpath, angle, *spec)
- # Trim any hanging vertices.
- if not has_left_input:
- ulpath.pop()
- llpath.pop()
- if not has_right_output:
- lrpath.pop()
- urpath.pop()
-
- # Concatenate the subpaths in the correct order (clockwise from top).
- path = (urpath + self._revert(lrpath) + llpath + self._revert(ulpath) +
- [(Path.CLOSEPOLY, urpath[0][1])])
-
- # Create a patch with the Sankey outline.
- codes, vertices = zip(*path)
- vertices = np.array(vertices)
-
- def _get_angle(a, r):
- if a is None:
- return None
- else:
- return a + r
-
- if prior is None:
- if rotation != 0: # By default, none of this is needed.
- angles = [_get_angle(angle, rotation) for angle in angles]
- rotate = Affine2D().rotate_deg(rotation * 90).transform_affine
- tips = rotate(tips)
- label_locations = rotate(label_locations)
- vertices = rotate(vertices)
- text = self.ax.text(0, 0, s=patchlabel, ha='center', va='center')
- else:
- rotation = (self.diagrams[prior].angles[connect[0]] -
- angles[connect[1]])
- angles = [_get_angle(angle, rotation) for angle in angles]
- rotate = Affine2D().rotate_deg(rotation * 90).transform_affine
- tips = rotate(tips)
- offset = self.diagrams[prior].tips[connect[0]] - tips[connect[1]]
- translate = Affine2D().translate(*offset).transform_affine
- tips = translate(tips)
- label_locations = translate(rotate(label_locations))
- vertices = translate(rotate(vertices))
- kwds = dict(s=patchlabel, ha='center', va='center')
- text = self.ax.text(*offset, **kwds)
- if rcParams['_internal.classic_mode']:
- fc = kwargs.pop('fc', kwargs.pop('facecolor', '#bfd1d4'))
- lw = kwargs.pop('lw', kwargs.pop('linewidth', 0.5))
- else:
- fc = kwargs.pop('fc', kwargs.pop('facecolor', None))
- lw = kwargs.pop('lw', kwargs.pop('linewidth', None))
- if fc is None:
- fc = next(self.ax._get_patches_for_fill.prop_cycler)['color']
- patch = PathPatch(Path(vertices, codes), fc=fc, lw=lw, **kwargs)
- self.ax.add_patch(patch)
-
- # Add the path labels.
- 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 label != '':
- label += "\n"
- label += quantity
- texts.append(self.ax.text(x=location[0], y=location[1],
- 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
- # user wants to provide labels later.
-
- # Expand the size of the diagram if necessary.
- self.extent = (min(np.min(vertices[:, 0]),
- np.min(label_locations[:, 0]),
- self.extent[0]),
- max(np.max(vertices[:, 0]),
- np.max(label_locations[:, 0]),
- self.extent[1]),
- min(np.min(vertices[:, 1]),
- np.min(label_locations[:, 1]),
- self.extent[2]),
- max(np.max(vertices[:, 1]),
- np.max(label_locations[:, 1]),
- self.extent[3]))
- # Include both vertices _and_ label locations in the extents; there are
- # where either could determine the margins (e.g., arrow shoulders).
-
- # Add this diagram as a subdiagram.
- self.diagrams.append(
- SimpleNamespace(patch=patch, flows=flows, angles=angles, tips=tips,
- text=text, texts=texts))
-
- # Allow a daisy-chained call structure (see docstring for the class).
- return self
-
- def finish(self):
- """
- Adjust the axes and return a list of information about the Sankey
- subdiagram(s).
-
- Return value is a list of subdiagrams represented with the following
- fields:
-
- =============== ===================================================
- Field Description
- =============== ===================================================
- *patch* Sankey outline (an instance of
- :class:`~maplotlib.patches.PathPatch`)
- *flows* values of the flows (positive for input, negative
- for output)
- *angles* list of angles of the arrows [deg/90]
- For example, if the diagram has not been rotated,
- an input to the top side will have an angle of 3
- (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*.
- *tips* array in which each row is an [x, y] pair
- indicating the positions of the tips (or "dips") of
- the flow paths
- If the magnitude of a flow is less the *tolerance*
- for the instance of :class:`Sankey`, the flow is
- skipped and its tip will be at the center of the
- diagram.
- *text* :class:`~matplotlib.text.Text` instance for the
- label of the diagram
- *texts* list of :class:`~matplotlib.text.Text` instances
- for the labels of flows
- =============== ===================================================
-
- .. seealso::
-
- :meth:`add`
- """
- self.ax.axis([self.extent[0] - self.margin,
- self.extent[1] + self.margin,
- self.extent[2] - self.margin,
- self.extent[3] + self.margin])
- self.ax.set_aspect('equal', adjustable='datalim')
- return self.diagrams
|