Creating Beautiful Sankey Diagrams with floWeaver | by Allison Stafford | Towards Data Science

A while back, I wrote “The What, Why, and How of Sankey Diagrams.” In the process of exploring Matplotlib’s Sankey library, I found a lovely alternative: floWeaver. If you’re looking to document flows with a Sankey-style diagram with a different style, this blog is for you.
Honestly, I was hoping that floWeaver would be infinitely “easier” to use than matplotlib sankey. I ended up finding that it is different, more than easier/harder. I do think that the diagrams appeal to me a little more with the default settings, where you have to have a vision and make the design choices happen a lot more intentionally in matplotlib sankey.
Here are some of the cases when you might choose floWeaver over matplotlib sankey:
You want to track flows across nodes (eg. marking the original input source, even as it travels through several steps)
You are not attached to arrowheads and directional endpoints (eg. flows that point up or down)
You are ok with nice-looking defaults and do not need a lot of control over things like font size
Find everything you need to know on the floWeaver GitHub. Or here’s the rundown for conda users:
Install floWeaver itself:
conda install -c conda-forge floweaver
Install ipysankeywidget to display the sankey diagrams in Jupyter Notebooks:
conda install -c conda-forge ipysankeywidget
Enable ipywidgets. This may be already enabled, but it doesn’t hurt to confirm:
jupyter nbextension enable --py widgetsnbextension --sys-prefix
It seems like floWeaver is expecting columns/features “source”, “target”, and “value”. From there, you can add additional features with whatever column names you want. These sources and targets are going to be one option for your end nodes. If you actually want to plot one of your other features as your end nodes, that works too, but one of your options for end nodes still should be in the “target” column (if my deduction is correct)
Here are the data tables for the examples I’ll use below:
First, find this and more in Rick Lubton and team’s tutorials, documentation, and examples. I highly recommend you use them (docs, GitHub repo) to investigate further.
Start with the imports and importing your dataset. Here we’re using the simple refugee data set shown above.
Start with the imports and importing your dataset. Here we’re using the simple refugee data set shown above.
import pandas as pdfrom floweaver import *refugees = pd.read_csv('refugees.csv')
2. Set your nodes, order, bundles, and partitions
Nodes are the gathering spots of your Sankey diagram. The stationary “locations” that your flows are going to flow into and out of. In this example, we have Syria on the left and the seven destinations on the right. You can name the node groups as you see fit. Each “ProcessGroup” is going to appear as a set of nodes aligned vertically. Note the ProcessGroup takes a list as its argument, but is flexible in whether the node names are repeated.
nodes = {    'start': ProcessGroup(['Syria']), # one (Syria) at the start     'end': ProcessGroup(list(refugees['target'])), # 7 at the end}
The order is what order (left to right) that you want the node ProcessGroups to appear in. Note that it is a list of lists of ProcessGroup/node names:
ordering = [[‘start’], [‘end’]]
The bundles are the set of connections that you want to show. This gives you the option to leave off any unnecessary connections. The bundles should be a list of floWeaver.Bundle objects that contain the connection information.
bundles = [Bundle('start', 'end')]
Even though you have defined the nodes, we need to set the partition in order to indicate how exactly we want them to be split. Using .unique() keeps us from having to list each destination individually.
nodes['start'].partition = Partition.Simple('source', ['Syria'])nodes['end'].partition = Partition.Simple('target',                                         refugees['target'].unique())
3. Create and display your Sankey Diagram
Create the Sankey using your nodes, ordering, and bundles from above:
sdd = SankeyDefinition(nodes, bundles, ordering)
To display the Sankey, use floWeaver’s weave(sankey_definition, dataset, measures=’value’, link_width=None, link_color=None, palette=None) in combination with .to_widget() to display in the Jupyter Notebook. If you want to save the image, add .auto_save_png(filename) after .to_widget().
weave(sdd, refugees).to_widget().auto_save_png('syria1.png')
As you see above, we have some pretty small flows. You can bundle them together using the same code as above except for modifying the end nodes partition. Instead of listing all of the nodes on the same level, create a tuple of (name, list of nodes) for any groupings you’d like to make.
nodes[‘end’].partition = Partition.Simple(‘target’,                             [‘Turkey’,                              ‘Lebanon’,                              ‘Jordan’,                              ‘Europe’,                             (‘Other’, [‘Iraq’, ‘Egypt’, ‘USA’])])
You’ll need to recreate the Sankey Definition and re-weave the Sankey:
sdd = SankeyDefinition(nodes, bundles, ordering)weave(sdd, refugees).to_widget().auto_save_png('syria_other.png')
To highlight a non-target feature, again, you adjust the partition. Here, we’re partitioning by region.
# set end partition to regionnodes[‘end’].partition = Partition.Simple(‘region’,                                        refugees[‘region’].unique())# create SankeyDefintion, display and save imagesdd = SankeyDefinition(nodes, bundles, ordering)weave(sdd, refugees).to_widget().auto_save_png('syria_region.png')
Another way to highlight different features or steps in the process is through waypoints (additional sets of nodes).
To add in this waypoint, we need to
Create a new partition for our waypoint (region)Add the waypoint to our node dictionary, feeding Waypoint the new region partitionAdd the waypoint to our ordering list (in the middle)Add the waypoint to our Bundle by adding waypoints=[node_name] as a parameter.
Create a new partition for our waypoint (region)
Add the waypoint to our node dictionary, feeding Waypoint the new region partition
Add the waypoint to our ordering list (in the middle)
Add the waypoint to our Bundle by adding waypoints=[node_name] as a parameter.
# create region partitionregion = Partition.Simple('region', refugees['region'].unique())# set the "nodes"nodes = {    'start': ProcessGroup(['Syria']),    'middle': Waypoint(region),    'end': ProcessGroup(list(refugees['target'].unique())), }# set the order of the nodes left to rightordering = [['start'], ['middle'], ['end']]# set the "bundle" of connectionsbundles = [Bundle('start', 'end', waypoints=['middle'])]# partition the groups for displaynodes['start'].partition = Partition.Simple('source', ['Syria'])nodes['end'].partition = Partition.Simple('target',                                         refugees['target'].unique())# set color palettepalette = {'North America': '#ED533B', 'Middle East': '#3CAEA3',            'Europe': '#F6D55C'}# create sankey with flow_partitionsdd = SankeyDefinition(nodes, bundles, ordering,                        flow_partition=region)# display sankey (with color palette) and save as pngweave(sdd, refugees, palette=palette).to_widget().auto_save_png('syria_waypt.png')
To adjust the color, create a color palette dictionary, add the partition you want to color on as the “flow_partition” in the SankeyDefinition, and add your palette to weave(). Note: this code uses the nodes, bundles, ordering, and dataset from above.
# set color palettepalette = {'North America': '#ED533B', 'Middle East': '#3CAEA3',            'Europe': '#F6D55C'}# create sankey with flow_partition for colorsdd = SankeyDefinition(nodes, bundles, ordering,                        flow_partition=region)# display sankey (with color palette) and save as pngweave(sdd, refugees, palette=palette) \     .to_widget().auto_save_png('syria_waypt.png')
To adjust the size of the Sankey output, create a size dictionary and add it as an argument for .to_widget()
# image sizesize = dict(width=570, height=300)# create sankey diagramweave(sdd, refugees).to_widget(**size).auto_save_png('syria_sm.png')
Here’s one last example, to show it all together.
# set the nodesnodes = {    'start': ProcessGroup(list(household['source'])),    'category': Waypoint(Partition.Simple('category',                          household['category'].unique())),    'end': ProcessGroup(list(household['target'])),}# set the order of the nodes left to rightordering = [['start'],             ['category'],            ['end']]# set the "bundle" of connections you want to showbundles = [Bundle('start', 'end', waypoints=['category'])]# add the partitionspartner = Partition.Simple('source', household['source'].unique())nodes['start'].partition = partnernodes['end'].partition = Partition.Simple('target',                          household['target'].unique())# add the color palettepalette = {'partner 1': '#027368', 'partner 2': '#58A4B0'}# create the sankey diagramsdd = SankeyDefinition(nodes, bundles, ordering,                        flow_partition=partner)# display the sankey diagramweave(sdd, household, palette=palette).to_widget().auto_save_png('./images/household.png')
You can see above that the color partitioning doesn’t always produce exactly the result I would choose, but at this point, I’m not sure how to adjust it. For me, coloring flows based on the waypoint looks the cutest, but doesn’t necessarily achieve the goal of the Sankey diagram — don’t forget to keep that at the forefront!
As always, check out the full code in the GitHub repo. Happy coding!