def test_view_graph_Elsewhere_bundles(): nodes = { 'a': ProcessGroup(selection=('a', )), 'b': ProcessGroup(selection=('b', )), } order = [[['a']], [['b']]] bundles = [ Bundle('a', 'b'), Bundle(Elsewhere, 'b'), ] G = view_graph(SankeyDefinition(nodes, bundles, order)) assert sorted(G.nodes(data=True)) == [ ('a', { 'node': ProcessGroup(selection=('a', )) }), ('b', { 'node': ProcessGroup(selection=('b', )), 'from_elsewhere_bundles': [1] }), ] assert sorted(G.edges(data=True)) == [ ('a', 'b', { 'bundles': [0] }), ] # Now with a Waypoint on the Elsewhere bundle nodes['w'] = Waypoint() bundles[1] = Bundle(Elsewhere, 'b', waypoints=['w']) order = [[['a', 'w']], [['b']]] G2 = view_graph(SankeyDefinition(nodes, bundles, order)) assert sorted(G2.nodes(data=True)) == [ ('a', { 'node': ProcessGroup(selection=('a', )) }), ('b', { 'node': ProcessGroup(selection=('b', )) }), ('w', { 'node': Waypoint() }), ] assert sorted(G2.edges(data=True)) == [ ('a', 'b', { 'bundles': [0] }), ('w', 'b', { 'bundles': [1] }), ]
def test_view_graph_adds_waypoints(): nodes = { 'n1': ProcessGroup(selection=['n1']), 'n2': ProcessGroup(selection=['n2']), 'w1': Waypoint(), } bundles = [ Bundle('n1', 'n2', waypoints=['w1']), ] order0 = [['n1'], [], ['w1'], [], [], ['n2']] G = view_graph(SankeyDefinition(nodes, bundles, order0)) assert sorted(nodes_ignoring_elsewhere(G, data=True)) == [ ('__n1_w1_1', { 'node': Waypoint(title='') }), ('__w1_n2_3', { 'node': Waypoint(title='') }), ('__w1_n2_4', { 'node': Waypoint(title='') }), ('n1', { 'node': ProcessGroup(selection=['n1']) }), ('n2', { 'node': ProcessGroup(selection=['n2']) }), ('w1', { 'node': Waypoint() }), ] assert sorted(edges_ignoring_elsewhere(G, data=True)) == [ ('__n1_w1_1', 'w1', { 'bundles': [0] }), ('__w1_n2_3', '__w1_n2_4', { 'bundles': [0] }), ('__w1_n2_4', 'n2', { 'bundles': [0] }), ('n1', '__n1_w1_1', { 'bundles': [0] }), ('w1', '__w1_n2_3', { 'bundles': [0] }), ] assert G.ordering == Ordering([[['n1']], [['__n1_w1_1']], [['w1']], [['__w1_n2_3']], [['__w1_n2_4']], [['n2']]])
def test_elsewhere_bundles_does_not_duplicate(): nodes = { 'a': ProcessGroup(selection=('a1')), 'in': Waypoint(), 'out': Waypoint() } bundles = { 0: Bundle(Elsewhere, 'a', waypoints=['in']), 1: Bundle('a', Elsewhere, waypoints=['out']), } order = [['in'], ['a'], ['out']] # not at min/max rank vd = SankeyDefinition(nodes, bundles, order) new_waypoints, new_bundles = elsewhere_bundles(vd) assert new_bundles == {}
def test_view_graph_bundle_flow_partitions_must_be_equal(): material_partition_mn = Partition.Simple('material', ['m', 'n']) material_partition_XY = Partition.Simple('material', ['X', 'Y']) nodes = { 'a': ProcessGroup(selection=['a1']), 'b': ProcessGroup(selection=['b1']), 'c': ProcessGroup(selection=['c1']), 'via': Waypoint(), } order = [['a', 'b'], ['via'], ['c']] bundles = [ Bundle('a', 'c', waypoints=['via'], flow_partition=material_partition_mn), Bundle('b', 'c', waypoints=['via'], flow_partition=material_partition_XY), ] # Do partition based on flows stored in bundles with pytest.raises(ValueError): G = view_graph(SankeyDefinition(nodes, bundles, order)) bundles[1] = Bundle('b', 'c', waypoints=['via'], flow_partition=material_partition_mn) assert view_graph(SankeyDefinition(nodes, bundles, order))
def test_view_graph_merges_bundles_between_same_nodes(): nodes = { 'n1': ProcessGroup(selection=['n1']), 'n2': ProcessGroup(selection=['n2']), 'n3': ProcessGroup(selection=['n3']), 'via': Waypoint(), } order0 = [['n1', 'n2'], ['via'], ['n3']] bundles = [ Bundle('n1', 'n3', waypoints=['via']), Bundle('n2', 'n3', waypoints=['via']), ] G = view_graph(SankeyDefinition(nodes, bundles, order0)) assert G.nodes['n3'] == {'node': nodes['n3']} assert sorted(edges_ignoring_elsewhere(G, data=True)) == [ ('n1', 'via', { 'bundles': [0] }), ('n2', 'via', { 'bundles': [1] }), ('via', 'n3', { 'bundles': [0, 1] }), ]
def test_sankey_definition_as_script_with_partitions(): nodes = { 'a': ProcessGroup(selection=['a1', 'a2']), 'b': ProcessGroup(selection=['b1']), 'c': ProcessGroup(selection=['c1', 'c2'], partition=Partition.Simple('process', ['c1', 'c2'])), 'via': Waypoint(partition=Partition.Simple('material', ['m', 'n'])), } bundles = [ Bundle('a', 'c', waypoints=['via']), Bundle('b', 'c', waypoints=['via']), ] ordering = [[['a', 'b']], [['via']], [['c']]] sdd = SankeyDefinition(nodes, bundles, ordering, flow_partition=Partition.Simple( 'material', ['m', 'n'])) code = sdd.to_code() # Check roundtrip ctx = {} exec(code, ctx) assert ctx["sdd"] == sdd
def test_elsewhere_bundles_not_added_to_waypoints(): nodes = {'waypoint': Waypoint(), } bundles = {} order = [[], ['waypoint'], []] vd = SankeyDefinition(nodes, bundles, order) new_waypoints, new_bundles = elsewhere_bundles(vd) assert new_waypoints == {} assert new_bundles == {}
def test_elsewhere_bundles_are_added_when_no_bundles_defined(): # make it easier to get started nodes = {'a': ProcessGroup(selection=['a1'])} bundles = {} order = [['a']] vd = SankeyDefinition(nodes, bundles, order) new_waypoints, new_bundles = elsewhere_bundles(vd) assert len(new_bundles) == 2 assert new_waypoints == { '__>a': Waypoint(title='→'), '__a>': Waypoint(title='→'), } # when direction is to left nodes['a'] = ProcessGroup(selection=['a1'], direction='L') vd = SankeyDefinition(nodes, bundles, order) new_waypoints, new_bundles = elsewhere_bundles(vd) assert new_waypoints == { '__>a': Waypoint(direction='L', title='←'), '__a>': Waypoint(direction='L', title='←'), }
def test_augment_waypoint_alignment(): # j -- a -- x # b # k -- c -- y # # should insert "from b" betwen x and y # and "to b" between j and k G = LayeredGraph() G.add_nodes_from([ ('a', {'node': ProcessGroup()}), ('b', {'node': ProcessGroup(selection=['b1'])}), ('c', {'node': ProcessGroup()}), ('x', {'node': ProcessGroup()}), ('y', {'node': ProcessGroup()}), ('j', {'node': ProcessGroup()}), ('k', {'node': ProcessGroup()}), ]) G.add_edges_from([ ('a', 'x', {'bundles': [2]}), ('k', 'c', {'bundles': [1]}), ('j', 'a', {'bundles': [0]}), ('c', 'y', {'bundles': [3]}), ]) G.ordering = Ordering([[['j', 'k']], [['a', 'b', 'c']], [['x', 'y']]]) new_waypoints = { 'from b': Waypoint(), 'to b': Waypoint(), } new_bundles = { 'b>': Bundle('b', Elsewhere, waypoints=['from b']), '>b': Bundle(Elsewhere, 'b', waypoints=['to b']), } G2 = augment(G, new_waypoints, new_bundles) assert set(G2.nodes()).difference(G.nodes()) == {'from b', 'to b'} assert G2.ordering == Ordering([ [['j', 'to b', 'k']], [['a', 'b', 'c']], [['x', 'from b', 'y']] ])
def test_sankey_definition_as_script(): nodes = { 'a': ProcessGroup(selection=['a1']), 'b': ProcessGroup(selection=['b1']), 'waypoint': Waypoint(), } ordering = [['a'], ['waypoint'], ['b']] bundles = [Bundle('a', 'b')] sdd = SankeyDefinition(nodes, bundles, ordering) code = sdd.to_code() assert code == dedent(""" from floweaver import ( ProcessGroup, Waypoint, Partition, Group, Elsewhere, Bundle, SankeyDefinition, ) nodes = { 'a': ProcessGroup(selection=['a1']), 'b': ProcessGroup(selection=['b1']), 'waypoint': Waypoint(), } ordering = [ [['a']], [['waypoint']], [['b']], ] bundles = [ Bundle(source='a', target='b'), ] sdd = SankeyDefinition(nodes, bundles, ordering) """) # Check roundtrip ctx = {} exec(code, ctx) assert ctx["sdd"] == sdd
def test_view_graph_adds_waypoints_partition(): nodes = { 'n1': ProcessGroup(selection=['n1']), 'n2': ProcessGroup(selection=['n2']), } g = Partition.Simple('test', ['x']) bundles = [ Bundle('n1', 'n2', default_partition=g), ] order0 = [['n1'], [], ['n2']] G = view_graph(SankeyDefinition(nodes, bundles, order0)) assert sorted(nodes_ignoring_elsewhere(G, data=True)) == [ ('__n1_n2_1', { 'node': Waypoint(title='', partition=g) }), ('n1', { 'node': ProcessGroup(selection=['n1']) }), ('n2', { 'node': ProcessGroup(selection=['n2']) }), ]
def test_sankey_definition_checks_bundles(): nodes = { 'a': ProcessGroup(selection=('a1')), 'b': ProcessGroup(selection=('b1')), 'waypoint': Waypoint(), } ordering = Ordering([]) with pytest.raises(ValueError): bundles = {0: Bundle('waypoint', 'b')} SankeyDefinition(nodes, bundles, ordering) with pytest.raises(ValueError): bundles = {0: Bundle('b', 'waypoint')} SankeyDefinition(nodes, bundles, ordering) # should work bundles = {0: Bundle('a', 'b')} assert SankeyDefinition(nodes, bundles, ordering) # also accepts a list bundles = [Bundle('a', 'b')] assert SankeyDefinition(nodes, bundles, ordering).bundles \ == {0: Bundle('a', 'b')}
def test_weave_results(): nodes = { 'a': ProcessGroup(selection=['a1', 'a2']), 'b': ProcessGroup(selection=['b1']), 'c': ProcessGroup(selection=['c1', 'c2'], partition=Partition.Simple('process', ['c1', 'c2'])), 'via': Waypoint(partition=Partition.Simple('material', ['m', 'n'])), } bundles = [ Bundle('a', 'c', waypoints=['via']), Bundle('b', 'c', waypoints=['via']), ] ordering = [[['a', 'b']], [['via']], [['c']]] sdd = SankeyDefinition(nodes, bundles, ordering) # Dataset flows = pd.DataFrame.from_records([ ('a1', 'c1', 'm', 3), ('a2', 'c1', 'n', 1), ('b1', 'c1', 'm', 1), ('b1', 'c2', 'm', 2), ('b1', 'c2', 'n', 1), ], columns=('source', 'target', 'material', 'value')) dim_process = pd.DataFrame({ 'id': list(flows.source.unique()) + list(flows.target.unique()) }).set_index('id') dataset = Dataset(flows, dim_process) result = weave(sdd, dataset) def link(src, tgt, original_flows, value, link_type='*', color='#FBB4AE'): return SankeyLink(source=src, target=tgt, type=link_type, time='*', data={'value': value}, title=link_type, color=color, original_flows=original_flows) assert set(n.id for n in result.nodes) == { 'a^*', 'b^*', 'via^m', 'via^n', 'c^c1', 'c^c2' } assert sorted(result.links) == [ link('a^*', 'via^m', [0], 3), link('a^*', 'via^n', [1], 1), link('b^*', 'via^m', [2, 3], 3), link('b^*', 'via^n', [4], 1), link('via^m', 'c^c1', [0, 2], 4), link('via^m', 'c^c2', [3], 2), link('via^n', 'c^c1', [1], 1), link('via^n', 'c^c2', [4], 1), ] assert result.ordering == Ordering([ [['a^*', 'b^*']], [['via^m', 'via^n']], [['c^c1', 'c^c2']], ]) assert result.groups == [ { 'id': 'via', 'title': '', 'type': 'group', 'nodes': ['via^m', 'via^n'] }, { 'id': 'c', 'title': '', 'type': 'process', 'nodes': ['c^c1', 'c^c2'] }, ] # Can also set flow_partition for all bundles at once sdd2 = SankeyDefinition(nodes, bundles, ordering, flow_partition=Partition.Simple( 'material', ['m', 'n'])) scale = CategoricalScale('type', palette=['red', 'blue']) scale.set_domain(['m', 'n']) result = weave(sdd2, dataset, link_color=scale) assert sorted(result.links) == [ link('a^*', 'via^m', [0], 3, 'm', 'red'), link('a^*', 'via^n', [1], 1, 'n', 'blue'), link('b^*', 'via^m', [2, 3], 3, 'm', 'red'), link('b^*', 'via^n', [4], 1, 'n', 'blue'), link('via^m', 'c^c1', [0, 2], 4, 'm', 'red'), link('via^m', 'c^c2', [3], 2, 'm', 'red'), link('via^n', 'c^c1', [1], 1, 'n', 'blue'), link('via^n', 'c^c2', [4], 1, 'n', 'blue'), ]
def test_results_graph_overall(): material_partition = Partition.Simple('material', ['m', 'n']) c_partition = Partition.Simple('process', ['c1', 'c2']) view_graph = LayeredGraph() view_graph.add_node('a', node=ProcessGroup(title='Node a')) view_graph.add_node('b', node=ProcessGroup()) view_graph.add_node('c', node=ProcessGroup(partition=c_partition)) view_graph.add_node('via', node=Waypoint(partition=material_partition)) view_graph.add_edges_from([ ('a', 'via', { 'bundles': [0], 'flow_partition': material_partition }), ('b', 'via', { 'bundles': [1], 'flow_partition': material_partition }), ('via', 'c', { 'bundles': [0, 1], 'flow_partition': material_partition }), ]) view_graph.ordering = Ordering([[['a', 'b']], [['via']], [['c']]]) # Mock flow data bundle_flows = { 0: pd.DataFrame.from_records([ ('a1', 'c1', 'm', 3), ('a2', 'c1', 'n', 1), ], index=(0, 1), columns=('source', 'target', 'material', 'value')), 1: pd.DataFrame.from_records([ ('b1', 'c1', 'm', 1), ('b1', 'c2', 'm', 2), ('b1', 'c2', 'n', 1), ], index=(2, 3, 4), columns=('source', 'target', 'material', 'value')) } # Do partition based on flows stored in bundles Gr, groups = results_graph(view_graph, bundle_flows) assert sorted(Gr.nodes(data=True)) == [ ('a^*', { 'direction': 'R', 'type': 'process', 'title': 'Node a' }), ('b^*', { 'direction': 'R', 'type': 'process', 'title': 'b' }), ('c^c1', { 'direction': 'R', 'type': 'process', 'title': 'c1' }), ('c^c2', { 'direction': 'R', 'type': 'process', 'title': 'c2' }), ('via^m', { 'direction': 'R', 'type': 'group', 'title': 'm' }), ('via^n', { 'direction': 'R', 'type': 'group', 'title': 'n' }), ] assert sorted(Gr.edges(keys=True, data=True)) == [ ('a^*', 'via^m', ('m', '*'), { 'measures': { 'value': 3 }, 'original_flows': [0], 'bundles': [0] }), ('a^*', 'via^n', ('n', '*'), { 'measures': { 'value': 1 }, 'original_flows': [1], 'bundles': [0] }), ('b^*', 'via^m', ('m', '*'), { 'measures': { 'value': 3 }, 'original_flows': [2, 3], 'bundles': [1] }), ('b^*', 'via^n', ('n', '*'), { 'measures': { 'value': 1 }, 'original_flows': [4], 'bundles': [1] }), ('via^m', 'c^c1', ('m', '*'), { 'measures': { 'value': 4 }, 'original_flows': [0, 2], 'bundles': [0, 1] }), ('via^m', 'c^c2', ('m', '*'), { 'measures': { 'value': 2 }, 'original_flows': [3], 'bundles': [0, 1] }), ('via^n', 'c^c1', ('n', '*'), { 'measures': { 'value': 1 }, 'original_flows': [1], 'bundles': [0, 1] }), ('via^n', 'c^c2', ('n', '*'), { 'measures': { 'value': 1 }, 'original_flows': [4], 'bundles': [0, 1] }), ] assert Gr.ordering == Ordering([ [['a^*', 'b^*']], [['via^m', 'via^n']], [['c^c1', 'c^c2']], ]) # Only includes groups where the title is not the same as the single node # title assert groups == [ { 'id': 'via', 'title': '', 'type': 'group', 'nodes': ['via^m', 'via^n'] }, { 'id': 'c', 'title': '', 'type': 'process', 'nodes': ['c^c1', 'c^c2'] }, ]