109 lines
2.7 KiB
Python
109 lines
2.7 KiB
Python
|
import pytest
|
||
|
|
||
|
np = pytest.importorskip("numpy")
|
||
|
import numpy.testing as npt
|
||
|
|
||
|
import networkx as nx
|
||
|
|
||
|
|
||
|
def test_attr_matrix():
|
||
|
G = nx.Graph()
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 2, thickness=2)
|
||
|
G.add_edge(1, 2, thickness=3)
|
||
|
|
||
|
def node_attr(u):
|
||
|
return G.nodes[u].get("size", 0.5) * 3
|
||
|
|
||
|
def edge_attr(u, v):
|
||
|
return G[u][v].get("thickness", 0.5)
|
||
|
|
||
|
M = nx.attr_matrix(G, edge_attr=edge_attr, node_attr=node_attr)
|
||
|
npt.assert_equal(M[0], np.array([[6.0]]))
|
||
|
assert M[1] == [1.5]
|
||
|
|
||
|
|
||
|
def test_attr_matrix_directed():
|
||
|
G = nx.DiGraph()
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 2, thickness=2)
|
||
|
G.add_edge(1, 2, thickness=3)
|
||
|
M = nx.attr_matrix(G, rc_order=[0, 1, 2])
|
||
|
# fmt: off
|
||
|
data = np.array(
|
||
|
[[0., 1., 1.],
|
||
|
[0., 0., 1.],
|
||
|
[0., 0., 0.]]
|
||
|
)
|
||
|
# fmt: on
|
||
|
npt.assert_equal(M, np.array(data))
|
||
|
|
||
|
|
||
|
def test_attr_matrix_multigraph():
|
||
|
G = nx.MultiGraph()
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 2, thickness=2)
|
||
|
G.add_edge(1, 2, thickness=3)
|
||
|
M = nx.attr_matrix(G, rc_order=[0, 1, 2])
|
||
|
# fmt: off
|
||
|
data = np.array(
|
||
|
[[0., 3., 1.],
|
||
|
[3., 0., 1.],
|
||
|
[1., 1., 0.]]
|
||
|
)
|
||
|
# fmt: on
|
||
|
npt.assert_equal(M, np.array(data))
|
||
|
M = nx.attr_matrix(G, edge_attr="weight", rc_order=[0, 1, 2])
|
||
|
# fmt: off
|
||
|
data = np.array(
|
||
|
[[0., 9., 1.],
|
||
|
[9., 0., 1.],
|
||
|
[1., 1., 0.]]
|
||
|
)
|
||
|
# fmt: on
|
||
|
npt.assert_equal(M, np.array(data))
|
||
|
M = nx.attr_matrix(G, edge_attr="thickness", rc_order=[0, 1, 2])
|
||
|
# fmt: off
|
||
|
data = np.array(
|
||
|
[[0., 3., 2.],
|
||
|
[3., 0., 3.],
|
||
|
[2., 3., 0.]]
|
||
|
)
|
||
|
# fmt: on
|
||
|
npt.assert_equal(M, np.array(data))
|
||
|
|
||
|
|
||
|
def test_attr_sparse_matrix():
|
||
|
pytest.importorskip("scipy")
|
||
|
G = nx.Graph()
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 2, thickness=2)
|
||
|
G.add_edge(1, 2, thickness=3)
|
||
|
M = nx.attr_sparse_matrix(G)
|
||
|
mtx = M[0]
|
||
|
data = np.ones((3, 3), float)
|
||
|
np.fill_diagonal(data, 0)
|
||
|
npt.assert_equal(mtx.todense(), np.array(data))
|
||
|
assert M[1] == [0, 1, 2]
|
||
|
|
||
|
|
||
|
def test_attr_sparse_matrix_directed():
|
||
|
G = nx.DiGraph()
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 1, thickness=1, weight=3)
|
||
|
G.add_edge(0, 2, thickness=2)
|
||
|
G.add_edge(1, 2, thickness=3)
|
||
|
M = nx.attr_sparse_matrix(G, rc_order=[0, 1, 2])
|
||
|
# fmt: off
|
||
|
data = np.array(
|
||
|
[[0., 1., 1.],
|
||
|
[0., 0., 1.],
|
||
|
[0., 0., 0.]]
|
||
|
)
|
||
|
# fmt: on
|
||
|
npt.assert_equal(M.todense(), np.array(data))
|