def test_results_graph_with_extra_or_not_enough_groups():
# Mock flow data
bundle_flows = {
0: pd.DataFrame.from_records(
[
('a1', 'b1', 'm', 3),
('a2', 'b1', 'm', 1),
],
columns=('source', 'target', 'material', 'value'))
}
# Group 'a3' not used. ProcessGroup 'a2' isn't in any group.
node_a = ProcessGroup(partition=Partition.Simple('process', ['a1', 'a3']))
node_b = ProcessGroup(partition=Partition.Simple('process', ['b1']))
view_graph = LayeredGraph()
view_graph.add_node('a', node=node_a)
view_graph.add_node('b', node=node_b)
view_graph.add_edges_from([('a', 'b', {'bundles': [0]}), ])
view_graph.ordering = Ordering([[['a']], [['b']]])
# Do partition based on flows stored in bundles
Gr, groups = results_graph(view_graph, bundle_flows)
assert set(Gr.nodes()) == {'a^a1', 'a^_', 'b^b1'}
assert sorted(Gr.edges(keys=True, data=True)) == [
('a^_', 'b^b1', ('*', '*'), {'value': 1, 'measures': {},
'bundles': [0]}),
('a^a1', 'b^b1', ('*', '*'), {'value': 3, 'measures': {},
'bundles': [0]}),
]
assert Gr.ordering == Ordering([
[['a^a1', 'a^_']],
[['b^b1']],
])
def test_results_graph_bands():
bundles = [Bundle('a', 'b'), ]
# Mock flow data
bundle_flows = {
bundles[0]: pd.DataFrame.from_records(
[
('a1', 'b1', 'm', 3),
],
columns=('source', 'target', 'material', 'value'))
}
view_graph = LayeredGraph()
view_graph.add_node('a', node=ProcessGroup())
view_graph.add_node('b', node=ProcessGroup())
view_graph.add_edges_from([('a', 'b', {'bundles': bundles}), ])
view_graph.ordering = Ordering([
[['a'], []],
[[], ['b']],
])
# Do partition based on flows stored in bundles
Gr, groups = results_graph(view_graph, bundle_flows)
assert Gr.ordering == Ordering([
# rank 1
[['a^*'], []],
# rank 2
[[], ['b^*']],
])
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.node['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_view_graph_does_short_bundles_last():
"""Return loops are inserted immediately below the source node, so work from
the outside in."""
#
# ,a -- b -- c-,
# | `----`|
# `-----------'
#
nodes = {
'a': ProcessGroup(selection=('a', )),
'b': ProcessGroup(selection=('b', )),
'c': ProcessGroup(selection=('c', )),
}
order = [[['a']], [['b']], [['c']]]
bundles = [
Bundle('a', 'b'),
Bundle('b', 'c'),
Bundle('c', 'b'),
Bundle('c', 'a'),
]
G = view_graph(SankeyDefinition(nodes, bundles, order))
assert G.ordering == Ordering([
[['a', '__c_a_0']],
[['b', '__c_b_1', '__c_a_1']],
[['c', '__c_b_2', '__c_a_2']],
])
# order of bundles doesn't affect it
G2 = view_graph(SankeyDefinition(nodes, bundles[::-1], order))
assert G.ordering == G2.ordering
def test_view_graph_does_non_dummy_bundles_first():
"""It's important to do bundles that don't require adding dummy nodes first, so
when it comes to return loops, they are better placed."""
nodes = {
'a': ProcessGroup(selection=('a', )),
'b': ProcessGroup(selection=('b', )),
'c': ProcessGroup(selection=('c', )),
'd': ProcessGroup(selection=('d', )),
}
order = [[['a', 'c']], [['b', 'd']]]
bundles = [
Bundle('a', 'b'),
Bundle('c', 'd'),
Bundle('b', 'a'),
]
G = view_graph(SankeyDefinition(nodes, bundles, order))
assert G.ordering == Ordering([
[['a', '__b_a_0', 'c']],
[['b', '__b_a_1', 'd']],
])
# order of bundles doesn't affect it
G2 = view_graph(SankeyDefinition(nodes, bundles[::-1], order))
assert G2.ordering == G.ordering
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 _twonode_viewgraph():
view_graph = LayeredGraph()
view_graph.add_node('a', node=ProcessGroup())
view_graph.add_node('b', node=ProcessGroup())
view_graph.add_edge('a', 'b', {'bundles': [0]})
view_graph.ordering = Ordering([
[['a']],
[['b']],
])
return view_graph
def _twonodes(xrank, xdir, yrank, ydir, **kwargs):
G = LayeredGraph()
G.add_node('x', node=ProcessGroup(direction=xdir))
G.add_node('y', node=ProcessGroup(direction=ydir))
layers = [[[]] for i in range(max(xrank, yrank) + 1)]
layers[xrank][0].append('x')
layers[yrank][0].append('y')
G.ordering = Ordering(layers)
kwargs.setdefault('bundle_key', None)
return add_dummy_nodes(G, 'x', 'y', **kwargs)
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_dummy_nodes_merge_bundles():
G = LayeredGraph()
G.add_node('a', node=ProcessGroup())
G.add_node('b', node=ProcessGroup())
G.ordering = Ordering([[['a']], [['b']]])
G = add_dummy_nodes(G, 'a', 'b', bundle_key=1)
assert G['a']['b']['bundles'] == [1]
G = add_dummy_nodes(G, 'a', 'b', bundle_key=2)
assert G['a']['b']['bundles'] == [1, 2]
assert set(G.nodes()) == {'a', 'b'}
assert set(G.edges()) == {('a', 'b')}
assert G.ordering == Ordering([[['a']], [['b']]])
def test_sankey_definition_checks_nodes_exist():
nodes = {
'a': ProcessGroup(selection=('a1')),
'b': ProcessGroup(selection=('b1')),
'waypoint': ProcessGroup(),
}
ordering = Ordering([])
with pytest.raises(ValueError):
bundles = [Bundle('does not exist', 'b')]
SankeyDefinition(nodes, bundles, ordering)
with pytest.raises(ValueError):
bundles = [Bundle('a', 'b', waypoints=['does not exist'])]
SankeyDefinition(nodes, bundles, ordering)
def test_sankey_view_accepts_dataframe_as_dataset():
nodes = {
'a': ProcessGroup(selection=['a']),
'b': ProcessGroup(selection=['b']),
}
bundles = [
Bundle('a', 'b'),
]
ordering = [['a'], ['b']]
vd = SankeyDefinition(nodes, bundles, ordering)
flows = pd.DataFrame.from_records([('a', 'b', 'm', 3)],
columns=('source', 'target', 'material',
'value'))
GR, groups = sankey_view(vd, flows)
def test_results_graph_time_partition():
time_partition = Partition.Simple('time', [1, 2])
view_graph = LayeredGraph()
view_graph.add_node('a', node=ProcessGroup())
view_graph.add_node('b', node=ProcessGroup())
view_graph.add_edges_from([('a', 'b', {'bundles': [0]}), ])
view_graph.ordering = Ordering([[['a']], [['b']]])
# Mock flow data
bundle_flows = {
0: pd.DataFrame.from_records(
[
('a1', 'b1', 'm', 1, 3),
('a2', 'b1', 'n', 1, 1),
('a2', 'b2', 'n', 1, 2),
('a1', 'b1', 'm', 2, 1),
('a1', 'b1', 'n', 2, 3),
],
columns=('source', 'target', 'material', 'time', 'value')),
}
# Do partition based on flows stored in bundles
Gr, groups = results_graph(view_graph,
bundle_flows,
time_partition=time_partition)
assert sorted(Gr.edges(keys=True, data=True)) == [
('a^*', 'b^*', ('*', '1'), {'value': 6, 'measures': {},
'bundles': [0]}),
('a^*', 'b^*', ('*', '2'), {'value': 4, 'measures': {},
'bundles': [0]}),
]
# Now add a material partition too
material_partition = Partition.Simple('material', ['m', 'n'])
Gr, groups = results_graph(view_graph, bundle_flows, material_partition,
time_partition)
assert sorted(Gr.edges(keys=True, data=True)) == [
('a^*', 'b^*', ('m', '1'), {'value': 3, 'measures': {},
'bundles': [0]}),
('a^*', 'b^*', ('m', '2'), {'value': 1, 'measures': {},
'bundles': [0]}),
('a^*', 'b^*', ('n', '1'), {'value': 3, 'measures': {},
'bundles': [0]}),
('a^*', 'b^*', ('n', '2'), {'value': 3, 'measures': {},
'bundles': [0]}),
]
def test_view_graph_does_not_mutate_definition():
nodes = {
'n1': ProcessGroup(selection=['n1']),
'n2': ProcessGroup(selection=['n2']),
}
bundles = [
Bundle('n1', 'n2'),
]
order0 = [['n1'], [], ['n2']]
vd = SankeyDefinition(nodes, bundles, order0)
G = view_graph(vd)
assert vd.nodes == {
'n1': ProcessGroup(selection=['n1']),
'n2': ProcessGroup(selection=['n2']),
}
assert vd.bundles == {
0: Bundle('n1', 'n2'),
}
assert vd.ordering == Ordering([[['n1']], [[]], [['n2']]])
def test_sankey_view_results_time_partition():
nodes = {
'a': ProcessGroup(selection=['a1']),
'b': ProcessGroup(selection=['b1']),
}
bundles = [Bundle('a', 'b')]
ordering = [[['a']], [['b']]]
time_partition = Partition.Simple('time', [1, 2])
vd = SankeyDefinition(nodes,
bundles,
ordering,
time_partition=time_partition)
# Dataset
flows = pd.DataFrame.from_records([
('a1', 'b1', 'm', 1, 3),
('a1', 'b1', 'm', 2, 2),
],
columns=('source', 'target', 'material',
'time', 'value'))
dim_process = pd.DataFrame({'id': ['a1', 'b1']}).set_index('id')
dataset = Dataset(flows, dim_process)
GR, groups = sankey_view(vd, dataset)
assert set(GR.nodes()) == {'a^*', 'b^*'}
assert sorted(GR.edges(keys=True, data=True)) == [
('a^*', 'b^*', ('*', '1'), {
'value': 3,
'measures': {},
'bundles': [0]
}),
('a^*', 'b^*', ('*', '2'), {
'value': 2,
'measures': {},
'bundles': [0]
}),
]
assert GR.ordering == Ordering([[['a^*']], [['b^*']]])
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_results_graph_material_key():
# Mock flow data
flows = pd.DataFrame.from_records(
[
('a1', 'c1', 'm', 'long', 3),
('a1', 'c1', 'n', 'long', 1),
],
columns=('source', 'target', 'material_type', 'shape', 'value'))
view_graph = LayeredGraph()
view_graph.add_node('a', node=ProcessGroup())
view_graph.add_node('c', node=ProcessGroup())
view_graph.add_edge('a', 'c', bundles=[0])
view_graph.ordering = Ordering([[['a']], [['c']]])
bundle_flows = {0: flows}
material_partition = Partition.Simple('material_type', ['m', 'n'])
shape_partition = Partition.Simple('shape', ['long', 'thin'])
# Partition based on material_type
view_graph.edge['a']['c']['flow_partition'] = material_partition
Gr, groups = results_graph(view_graph, bundle_flows)
assert sorted(Gr.edges(keys=True, data=True)) == [
('a^*', 'c^*', ('m', '*'), {'value': 3, 'measures': {},
'bundles': [0]}),
('a^*', 'c^*', ('n', '*'), {'value': 1, 'measures': {},
'bundles': [0]}),
]
# Partition based on shape
view_graph.edge['a']['c']['flow_partition'] = shape_partition
Gr, groups = results_graph(view_graph, bundle_flows)
assert sorted(Gr.edges(keys=True, data=True)) == [
('a^*', 'c^*', ('long', '*'), {'value': 4, 'measures': {},
'bundles': [0]}),
]
def test_dummy_nodes_order_dependence():
#
# a b
# c d
#
# bundles a-b, c-d, b-a
G = LayeredGraph()
G.add_nodes_from('abcd', node=ProcessGroup())
G.ordering = Ordering([[['a', 'c']], [['b', 'd']]])
# Correct G.order: a-b, c-d, b-a
G1 = _apply_bundles(G, ('ab', 'cd', 'ba'))
assert G1.ordering == Ordering([[['a', '__b_a_0', 'c']],
[['b', '__b_a_1', 'd']]])
# Incorrect G.order: b-a first
G2 = _apply_bundles(G, ('ba', 'ab', 'cd'))
assert G2.ordering == Ordering([[['a', 'c', '__b_a_0']],
[['b', '__b_a_1', 'd']]])
def test_sankey_view_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']]]
vd = 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)
GR, groups = sankey_view(vd, dataset)
assert set(GR.nodes()) == {'a^*', 'b^*', 'via^m', 'via^n', 'c^c1', 'c^c2'}
assert sorted(GR.edges(keys=True, data=True)) == [
('a^*', 'via^m', ('*', '*'), {
'value': 3,
'measures': {},
'bundles': [0]
}),
('a^*', 'via^n', ('*', '*'), {
'value': 1,
'measures': {},
'bundles': [0]
}),
('b^*', 'via^m', ('*', '*'), {
'value': 3,
'measures': {},
'bundles': [1]
}),
('b^*', 'via^n', ('*', '*'), {
'value': 1,
'measures': {},
'bundles': [1]
}),
('via^m', 'c^c1', ('*', '*'), {
'value': 4,
'measures': {},
'bundles': [0, 1]
}),
('via^m', 'c^c2', ('*', '*'), {
'value': 2,
'measures': {},
'bundles': [0, 1]
}),
('via^n', 'c^c1', ('*', '*'), {
'value': 1,
'measures': {},
'bundles': [0, 1]
}),
('via^n', 'c^c2', ('*', '*'), {
'value': 1,
'measures': {},
'bundles': [0, 1]
}),
]
assert GR.ordering == Ordering([
[['a^*', 'b^*']],
[['via^m', 'via^n']],
[['c^c1', 'c^c2']],
])
assert groups == [
{
'id': 'a',
'title': '',
'type': 'process',
'nodes': ['a^*']
},
{
'id': 'b',
'title': '',
'type': 'process',
'nodes': ['b^*']
},
{
'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
vd2 = SankeyDefinition(nodes,
bundles,
ordering,
flow_partition=Partition.Simple(
'material', ['m', 'n']))
GR, groups = sankey_view(vd2, dataset)
assert sorted(GR.edges(keys=True, data=True)) == [
('a^*', 'via^m', ('m', '*'), {
'value': 3,
'measures': {},
'bundles': [0]
}),
('a^*', 'via^n', ('n', '*'), {
'value': 1,
'measures': {},
'bundles': [0]
}),
('b^*', 'via^m', ('m', '*'), {
'value': 3,
'measures': {},
'bundles': [1]
}),
('b^*', 'via^n', ('n', '*'), {
'value': 1,
'measures': {},
'bundles': [1]
}),
('via^m', 'c^c1', ('m', '*'), {
'value': 4,
'measures': {},
'bundles': [0, 1]
}),
('via^m', 'c^c2', ('m', '*'), {
'value': 2,
'measures': {},
'bundles': [0, 1]
}),
('via^n', 'c^c1', ('n', '*'), {
'value': 1,
'measures': {},
'bundles': [0, 1]
}),
('via^n', 'c^c2', ('n', '*'), {
'value': 1,
'measures': {},
'bundles': [0, 1]
}),
]
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),
],
columns=('source', 'target', 'material', 'value')),
1: pd.DataFrame.from_records(
[
('b1', 'c1', 'm', 1),
('b1', 'c2', 'm', 2),
('b1', 'c2', 'n', 1),
],
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', '*'), {'value': 3, 'measures': {},
'bundles': [0]}),
('a^*', 'via^n', ('n', '*'), {'value': 1, 'measures': {},
'bundles': [0]}),
('b^*', 'via^m', ('m', '*'), {'value': 3, 'measures': {},
'bundles': [1]}),
('b^*', 'via^n', ('n', '*'), {'value': 1, 'measures': {},
'bundles': [1]}),
('via^m', 'c^c1', ('m', '*'), {'value': 4, 'measures': {},
'bundles': [0, 1]}),
('via^m', 'c^c2', ('m', '*'), {'value': 2, 'measures': {},
'bundles': [0, 1]}),
('via^n', 'c^c1', ('n', '*'), {'value': 1, 'measures': {},
'bundles': [0, 1]}),
('via^n', 'c^c2', ('n', '*'), {'value': 1, 'measures': {},
'bundles': [0, 1]}),
]
assert Gr.ordering == Ordering([
[['a^*', 'b^*']],
[['via^m', 'via^n']],
[['c^c1', 'c^c2']],
])
assert groups == [
{'id': 'a',
'title': 'Node a',
'type': 'process',
'nodes': ['a^*']},
{'id': 'b',
'title': '',
'type': 'process',
'nodes': ['b^*']},
{'id': 'via',
'title': '',
'type': 'group',
'nodes': ['via^m', 'via^n']},
{'id': 'c',
'title': '',
'type': 'process',
'nodes': ['c^c1', 'c^c2']},
]