105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
import sys
|
|
|
|
import easygraph as eg
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from easygraph.functions.basic.localassort import localAssort
|
|
|
|
|
|
class TestLocalAssort:
|
|
@classmethod
|
|
def setup_class(self):
|
|
self.G = eg.get_graph_karateclub()
|
|
edgelist = []
|
|
node_num = len(self.G.nodes)
|
|
for e in self.G.edges:
|
|
edgelist.append([e[0] - 1, e[1] - 1])
|
|
self.edgelist = np.int32(edgelist)
|
|
self.valuelist = np.arange(node_num, dtype=np.int32) % 6
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info.major <= 3 and sys.version_info.minor <= 7,
|
|
reason="python version should higher than 3.7",
|
|
)
|
|
def test_karateclub(self):
|
|
assortM, assortT, Z = eg.localAssort(
|
|
self.edgelist, self.valuelist, pr=np.arange(0, 1, 0.1)
|
|
)
|
|
|
|
_, assortT, Z = eg.functions.basic.localassort.localAssort(
|
|
self.edgelist, self.valuelist, pr=np.array([0.9])
|
|
)
|
|
|
|
|
|
def test_localassort_small_complete_graph():
|
|
G = eg.complete_graph(4)
|
|
edgelist = np.array(list(G.edges))
|
|
node_attr = np.array([0, 0, 1, 1])
|
|
assortM, assortT, Z = localAssort(edgelist, node_attr)
|
|
assert assortM.shape == (4, 10)
|
|
assert assortT.shape == (4,)
|
|
assert Z.shape == (4,)
|
|
assert np.all(Z >= 0) and np.all(Z <= 1)
|
|
|
|
|
|
def test_localassort_with_missing_attributes():
|
|
G = eg.path_graph(5)
|
|
edgelist = np.array(list(G.edges))
|
|
node_attr = np.array([0, -1, 1, -1, 1])
|
|
assortM, assortT, Z = localAssort(edgelist, node_attr, pr=np.array([0.5]))
|
|
assert assortT.shape == (5,)
|
|
assert Z.shape == (5,)
|
|
assert np.any(np.isnan(assortT))
|
|
|
|
|
|
def test_localassort_directed_graph():
|
|
G = eg.DiGraph()
|
|
G.add_edges_from([(0, 1), (1, 2), (2, 3)])
|
|
edgelist = np.array(list(G.edges))
|
|
node_attr = np.array([0, 1, 0, 1])
|
|
assortM, assortT, Z = localAssort(edgelist, node_attr, undir=False)
|
|
assert assortM.shape == (4, 10)
|
|
assert assortT.shape == (4,)
|
|
assert Z.shape == (4,)
|
|
|
|
|
|
def test_localassort_single_node_graph():
|
|
edgelist = np.empty((0, 2), dtype=int)
|
|
node_attr = np.array([0])
|
|
assortM, assortT, Z = localAssort(edgelist, node_attr)
|
|
assert assortM.shape == (1, 10)
|
|
assert np.all(np.isnan(assortM)) or np.allclose(assortM, 0, atol=1e-5)
|
|
assert np.all(np.isnan(assortT)) or np.allclose(assortT, 0, atol=1e-5)
|
|
assert np.all(np.isnan(Z)) or np.allclose(Z, 0, atol=1e-5)
|
|
|
|
|
|
def test_localassort_disconnected_graph():
|
|
G = eg.Graph()
|
|
G.add_nodes_from(range(5))
|
|
edgelist = np.empty((0, 2), dtype=int)
|
|
node_attr = np.array([0, 1, 0, 1, 1])
|
|
assortM, assortT, Z = localAssort(edgelist, node_attr)
|
|
assert assortM.shape == (5, 10)
|
|
assert np.all(np.isnan(assortM)) or np.allclose(assortM, 0, atol=1e-5)
|
|
assert np.all(np.isnan(assortT)) or np.allclose(assortT, 0, atol=1e-5)
|
|
assert np.all(np.isnan(Z)) or np.allclose(Z, 0, atol=1e-5)
|
|
|
|
|
|
def test_localassort_high_restart_probabilities():
|
|
G = eg.path_graph(5)
|
|
edgelist = np.array(list(G.edges))
|
|
node_attr = np.array([1, 0, 1, 0, 1])
|
|
pr = np.array([0.95, 0.99])
|
|
assortM, assortT, Z = localAssort(edgelist, node_attr, pr=pr)
|
|
assert assortM.shape == (5, 2)
|
|
assert assortT.shape == (5,)
|
|
assert Z.shape == (5,)
|
|
|
|
|
|
def test_localassort_invalid_attribute_length():
|
|
edgelist = np.array([[0, 1], [1, 2]])
|
|
node_attr = np.array([0, 1]) # too short
|
|
with pytest.raises(ValueError):
|
|
localAssort(edgelist, node_attr)
|