test_attrmatrix.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. import pytest
  2. np = pytest.importorskip("numpy")
  3. import networkx as nx
  4. def test_attr_matrix():
  5. G = nx.Graph()
  6. G.add_edge(0, 1, thickness=1, weight=3)
  7. G.add_edge(0, 1, thickness=1, weight=3)
  8. G.add_edge(0, 2, thickness=2)
  9. G.add_edge(1, 2, thickness=3)
  10. def node_attr(u):
  11. return G.nodes[u].get("size", 0.5) * 3
  12. def edge_attr(u, v):
  13. return G[u][v].get("thickness", 0.5)
  14. M = nx.attr_matrix(G, edge_attr=edge_attr, node_attr=node_attr)
  15. np.testing.assert_equal(M[0], np.array([[6.0]]))
  16. assert M[1] == [1.5]
  17. def test_attr_matrix_directed():
  18. G = nx.DiGraph()
  19. G.add_edge(0, 1, thickness=1, weight=3)
  20. G.add_edge(0, 1, thickness=1, weight=3)
  21. G.add_edge(0, 2, thickness=2)
  22. G.add_edge(1, 2, thickness=3)
  23. M = nx.attr_matrix(G, rc_order=[0, 1, 2])
  24. # fmt: off
  25. data = np.array(
  26. [[0., 1., 1.],
  27. [0., 0., 1.],
  28. [0., 0., 0.]]
  29. )
  30. # fmt: on
  31. np.testing.assert_equal(M, np.array(data))
  32. def test_attr_matrix_multigraph():
  33. G = nx.MultiGraph()
  34. G.add_edge(0, 1, thickness=1, weight=3)
  35. G.add_edge(0, 1, thickness=1, weight=3)
  36. G.add_edge(0, 1, thickness=1, weight=3)
  37. G.add_edge(0, 2, thickness=2)
  38. G.add_edge(1, 2, thickness=3)
  39. M = nx.attr_matrix(G, rc_order=[0, 1, 2])
  40. # fmt: off
  41. data = np.array(
  42. [[0., 3., 1.],
  43. [3., 0., 1.],
  44. [1., 1., 0.]]
  45. )
  46. # fmt: on
  47. np.testing.assert_equal(M, np.array(data))
  48. M = nx.attr_matrix(G, edge_attr="weight", rc_order=[0, 1, 2])
  49. # fmt: off
  50. data = np.array(
  51. [[0., 9., 1.],
  52. [9., 0., 1.],
  53. [1., 1., 0.]]
  54. )
  55. # fmt: on
  56. np.testing.assert_equal(M, np.array(data))
  57. M = nx.attr_matrix(G, edge_attr="thickness", rc_order=[0, 1, 2])
  58. # fmt: off
  59. data = np.array(
  60. [[0., 3., 2.],
  61. [3., 0., 3.],
  62. [2., 3., 0.]]
  63. )
  64. # fmt: on
  65. np.testing.assert_equal(M, np.array(data))
  66. def test_attr_sparse_matrix():
  67. pytest.importorskip("scipy")
  68. G = nx.Graph()
  69. G.add_edge(0, 1, thickness=1, weight=3)
  70. G.add_edge(0, 2, thickness=2)
  71. G.add_edge(1, 2, thickness=3)
  72. M = nx.attr_sparse_matrix(G)
  73. mtx = M[0]
  74. data = np.ones((3, 3), float)
  75. np.fill_diagonal(data, 0)
  76. np.testing.assert_equal(mtx.todense(), np.array(data))
  77. assert M[1] == [0, 1, 2]
  78. def test_attr_sparse_matrix_directed():
  79. pytest.importorskip("scipy")
  80. G = nx.DiGraph()
  81. G.add_edge(0, 1, thickness=1, weight=3)
  82. G.add_edge(0, 1, thickness=1, weight=3)
  83. G.add_edge(0, 2, thickness=2)
  84. G.add_edge(1, 2, thickness=3)
  85. M = nx.attr_sparse_matrix(G, rc_order=[0, 1, 2])
  86. # fmt: off
  87. data = np.array(
  88. [[0., 1., 1.],
  89. [0., 0., 1.],
  90. [0., 0., 0.]]
  91. )
  92. # fmt: on
  93. np.testing.assert_equal(M.todense(), np.array(data))