.. _api-sankey_demo_old:

api example code: sankey_demo_old.py
====================================



.. plot:: /builddir/build/BUILD/matplotlib-1.2.0/doc/mpl_examples/api/sankey_demo_old.py

::

    #!/usr/bin/env python
    
    from __future__ import print_function
    
    __author__ = "Yannick Copin <ycopin@ipnl.in2p3.fr>"
    __version__ = "Time-stamp: <10/02/2010 16:49 ycopin@lyopc548.in2p3.fr>"
    
    import numpy as np
    
    
    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 = np.absolute(outputs)
        outsigns = np.sign(outputs)
        outsigns[-1] = 0  # Last output
    
        ins = np.absolute(inputs)
        insigns = np.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)*np.tan(outangle/180. * np.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)*np.tan(inangle/180. * np.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 reversed(list(zip(ins, insigns))):
            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 = np.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 plt
    
        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 = plt.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')
    
        plt.show()
    

Keywords: python, matplotlib, pylab, example, codex (see :ref:`how-to-search-examples`)