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():
nodes = {'a': ProcessGroup(selection=['a1']), }
bundles = {}
order = [[], ['a'], []] # not at min/max rank
vd = SankeyDefinition(nodes, bundles, order)
new_waypoints, new_bundles = elsewhere_bundles(vd)
assert set(new_waypoints.keys()) == {'__a>', '__>a'}
assert set(new_bundles.values()) == {
Bundle('a', Elsewhere, waypoints=['__a>']),
Bundle(Elsewhere, 'a', waypoints=['__>a']),
}
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_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_elsewhere_bundles_not_added_at_minmax_rank_when_one_bundle_defined():
nodes = {'a': ProcessGroup(selection=['a1'])}
bundles = {0: Bundle('a', Elsewhere)}
order = [['a']]
vd = SankeyDefinition(nodes, bundles, order)
new_waypoints, new_bundles = elsewhere_bundles(vd)
assert len(new_waypoints) == 0
assert len(new_bundles) == 0
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_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_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_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_bundle_hashable():
assert hash(Bundle('a', 'b'))