Skip to content

dbgnn

generate_bipartite_edge_index

Generate edge_index for bipartite graph connecting nodes of a second-order graph to first-order nodes.

Source code in src/pathpyG/utils/dbgnn.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def generate_bipartite_edge_index(
    g: Graph, g2: Graph, mapping: str = "last", device: Optional[torch.device] = None
) -> torch.Tensor:
    """Generate edge_index for bipartite graph connecting nodes of a second-order graph to first-order nodes."""

    if mapping == "last":
        bipartide_edge_index = torch.tensor([list(range(g2.n)), [v[1] for v in g2.data.node_sequence]], device=device)

    elif mapping == "first":
        bipartide_edge_index = torch.tensor([list(range(g2.n)), [v[0] for v in g2.data.node_sequence]], device=device)
    else:
        bipartide_edge_index = torch.tensor(
            [
                list(range(g2.n)) + list(range(g2.n)),
                [v[0] for v in g2.data.node_sequence] + [v[1] for v in g2.data.node_sequence],
            ],
            device=device,
        )

    return bipartide_edge_index