floweaver/debug_c_catchall.py at master · ricklupton/floweaver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
"""Debug c^_ node creation."""
import pandas as pd
from floweaver import (
    SankeyDefinition,
    ProcessGroup,
    Waypoint,
    Bundle,
    Partition,
    Dataset,
)
from floweaver.weave import weave_compiled

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']))

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_compiled(sdd, dataset)

print("\n=== NODES ===")
for node in result.nodes:
    print(f"{node.id}: hidden={node.hidden}")

print("\n=== LINKS involving c^_ ===")
c_catchall_links = [link for link in result.links if 'c^_' in (link.source, link.target)]
for link in c_catchall_links:
    print(f"{link.source} -> {link.target}: value={link.link_width}, flows={link.original_flows}")

print(f"\nTotal links involving c^_: {len(c_catchall_links)}")