File: test_custom_motif.py

package info (click to toggle)
pytorch-geometric 2.6.1-7
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 12,904 kB
  • sloc: python: 127,155; sh: 338; cpp: 27; makefile: 18; javascript: 16
file content (38 lines) | stat: -rw-r--r-- 978 bytes parent folder | download
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
import pytest
import torch

from torch_geometric.data import Data
from torch_geometric.datasets.motif_generator import CustomMotif
from torch_geometric.testing import withPackage


def test_custom_motif_pyg_data():
    structure = Data(
        num_nodes=3,
        edge_index=torch.tensor([[0, 1, 2, 1, 2, 0], [1, 2, 0, 0, 1, 2]]),
    )

    motif_generator = CustomMotif(structure)
    assert str(motif_generator) == 'CustomMotif()'

    assert structure == motif_generator()


@withPackage('networkx')
def test_custom_motif_networkx():
    import networkx as nx

    structure = nx.gnm_random_graph(5, 10, seed=2000)

    motif_generator = CustomMotif(structure)
    assert str(motif_generator) == 'CustomMotif()'

    out = motif_generator()
    assert len(out) == 2
    assert out.num_nodes == 5
    assert out.num_edges == 20


def test_custom_motif_unknown():
    with pytest.raises(ValueError, match="motif structure of type"):
        CustomMotif(structure='unknown')