"vscode:/vscode.git/clone" did not exist on "a9f8f2589c84707ea8714bb7e3756f2b729f591d"
graph_cases.py 1010 Bytes
Newer Older
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
from collections import defaultdict
import dgl
import networkx as nx

case_registry = defaultdict(list)

def register_case(labels):
    def wrapper(fn):
        for lbl in labels:
            case_registry[lbl].append(fn)
        return fn
    return wrapper

def get_cases(labels=None, exclude=None):
    cases = set()
    if labels is None:
        # get all the cases
        labels = case_registry.keys()
    for lbl in labels:
        if exclude is not None and lbl in exclude:
            continue
        cases.update(case_registry[lbl])
    return [fn() for fn in cases]

@register_case(['dglgraph', 'path', 'small'])
def dglgraph_path():
    return dgl.DGLGraph(nx.path_graph(5))

@register_case(['bipartite', 'small', 'hetero', 'zero-degree'])
def bipartite1():
    return dgl.bipartite([(0, 0), (0, 1), (0, 4), (2, 1), (2, 4), (3, 3)])

@register_case(['bipartite', 'small', 'hetero'])
def bipartite_full():
    return dgl.bipartite([(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3)])