test_pylab.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791
  1. """Unit tests for matplotlib drawing functions."""
  2. import itertools
  3. import os
  4. import warnings
  5. import pytest
  6. mpl = pytest.importorskip("matplotlib")
  7. np = pytest.importorskip("numpy")
  8. mpl.use("PS")
  9. plt = pytest.importorskip("matplotlib.pyplot")
  10. plt.rcParams["text.usetex"] = False
  11. import networkx as nx
  12. barbell = nx.barbell_graph(4, 6)
  13. def test_draw():
  14. try:
  15. functions = [
  16. nx.draw_circular,
  17. nx.draw_kamada_kawai,
  18. nx.draw_planar,
  19. nx.draw_random,
  20. nx.draw_spectral,
  21. nx.draw_spring,
  22. nx.draw_shell,
  23. ]
  24. options = [{"node_color": "black", "node_size": 100, "width": 3}]
  25. for function, option in itertools.product(functions, options):
  26. function(barbell, **option)
  27. plt.savefig("test.ps")
  28. finally:
  29. try:
  30. os.unlink("test.ps")
  31. except OSError:
  32. pass
  33. def test_draw_shell_nlist():
  34. try:
  35. nlist = [list(range(4)), list(range(4, 10)), list(range(10, 14))]
  36. nx.draw_shell(barbell, nlist=nlist)
  37. plt.savefig("test.ps")
  38. finally:
  39. try:
  40. os.unlink("test.ps")
  41. except OSError:
  42. pass
  43. def test_edge_colormap():
  44. colors = range(barbell.number_of_edges())
  45. nx.draw_spring(
  46. barbell, edge_color=colors, width=4, edge_cmap=plt.cm.Blues, with_labels=True
  47. )
  48. # plt.show()
  49. def test_arrows():
  50. nx.draw_spring(barbell.to_directed())
  51. # plt.show()
  52. @pytest.mark.parametrize(
  53. ("edge_color", "expected"),
  54. (
  55. (None, "black"), # Default
  56. ("r", "red"), # Non-default color string
  57. (["r"], "red"), # Single non-default color in a list
  58. ((1.0, 1.0, 0.0), "yellow"), # single color as rgb tuple
  59. ([(1.0, 1.0, 0.0)], "yellow"), # single color as rgb tuple in list
  60. ((0, 1, 0, 1), "lime"), # single color as rgba tuple
  61. ([(0, 1, 0, 1)], "lime"), # single color as rgba tuple in list
  62. ("#0000ff", "blue"), # single color hex code
  63. (["#0000ff"], "blue"), # hex code in list
  64. ),
  65. )
  66. @pytest.mark.parametrize("edgelist", (None, [(0, 1)]))
  67. def test_single_edge_color_undirected(edge_color, expected, edgelist):
  68. """Tests ways of specifying all edges have a single color for edges
  69. drawn with a LineCollection"""
  70. G = nx.path_graph(3)
  71. drawn_edges = nx.draw_networkx_edges(
  72. G, pos=nx.random_layout(G), edgelist=edgelist, edge_color=edge_color
  73. )
  74. assert mpl.colors.same_color(drawn_edges.get_color(), expected)
  75. @pytest.mark.parametrize(
  76. ("edge_color", "expected"),
  77. (
  78. (None, "black"), # Default
  79. ("r", "red"), # Non-default color string
  80. (["r"], "red"), # Single non-default color in a list
  81. ((1.0, 1.0, 0.0), "yellow"), # single color as rgb tuple
  82. ([(1.0, 1.0, 0.0)], "yellow"), # single color as rgb tuple in list
  83. ((0, 1, 0, 1), "lime"), # single color as rgba tuple
  84. ([(0, 1, 0, 1)], "lime"), # single color as rgba tuple in list
  85. ("#0000ff", "blue"), # single color hex code
  86. (["#0000ff"], "blue"), # hex code in list
  87. ),
  88. )
  89. @pytest.mark.parametrize("edgelist", (None, [(0, 1)]))
  90. def test_single_edge_color_directed(edge_color, expected, edgelist):
  91. """Tests ways of specifying all edges have a single color for edges drawn
  92. with FancyArrowPatches"""
  93. G = nx.path_graph(3, create_using=nx.DiGraph)
  94. drawn_edges = nx.draw_networkx_edges(
  95. G, pos=nx.random_layout(G), edgelist=edgelist, edge_color=edge_color
  96. )
  97. for fap in drawn_edges:
  98. assert mpl.colors.same_color(fap.get_edgecolor(), expected)
  99. def test_edge_color_tuple_interpretation():
  100. """If edge_color is a sequence with the same length as edgelist, then each
  101. value in edge_color is mapped onto each edge via colormap."""
  102. G = nx.path_graph(6, create_using=nx.DiGraph)
  103. pos = {n: (n, n) for n in range(len(G))}
  104. # num edges != 3 or 4 --> edge_color interpreted as rgb(a)
  105. for ec in ((0, 0, 1), (0, 0, 1, 1)):
  106. # More than 4 edges
  107. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=ec)
  108. for fap in drawn_edges:
  109. assert mpl.colors.same_color(fap.get_edgecolor(), ec)
  110. # Fewer than 3 edges
  111. drawn_edges = nx.draw_networkx_edges(
  112. G, pos, edgelist=[(0, 1), (1, 2)], edge_color=ec
  113. )
  114. for fap in drawn_edges:
  115. assert mpl.colors.same_color(fap.get_edgecolor(), ec)
  116. # num edges == 3, len(edge_color) == 4: interpreted as rgba
  117. drawn_edges = nx.draw_networkx_edges(
  118. G, pos, edgelist=[(0, 1), (1, 2), (2, 3)], edge_color=(0, 0, 1, 1)
  119. )
  120. for fap in drawn_edges:
  121. assert mpl.colors.same_color(fap.get_edgecolor(), "blue")
  122. # num edges == 4, len(edge_color) == 3: interpreted as rgb
  123. drawn_edges = nx.draw_networkx_edges(
  124. G, pos, edgelist=[(0, 1), (1, 2), (2, 3), (3, 4)], edge_color=(0, 0, 1)
  125. )
  126. for fap in drawn_edges:
  127. assert mpl.colors.same_color(fap.get_edgecolor(), "blue")
  128. # num edges == len(edge_color) == 3: interpreted with cmap, *not* as rgb
  129. drawn_edges = nx.draw_networkx_edges(
  130. G, pos, edgelist=[(0, 1), (1, 2), (2, 3)], edge_color=(0, 0, 1)
  131. )
  132. assert mpl.colors.same_color(
  133. drawn_edges[0].get_edgecolor(), drawn_edges[1].get_edgecolor()
  134. )
  135. for fap in drawn_edges:
  136. assert not mpl.colors.same_color(fap.get_edgecolor(), "blue")
  137. # num edges == len(edge_color) == 4: interpreted with cmap, *not* as rgba
  138. drawn_edges = nx.draw_networkx_edges(
  139. G, pos, edgelist=[(0, 1), (1, 2), (2, 3), (3, 4)], edge_color=(0, 0, 1, 1)
  140. )
  141. assert mpl.colors.same_color(
  142. drawn_edges[0].get_edgecolor(), drawn_edges[1].get_edgecolor()
  143. )
  144. assert mpl.colors.same_color(
  145. drawn_edges[2].get_edgecolor(), drawn_edges[3].get_edgecolor()
  146. )
  147. for fap in drawn_edges:
  148. assert not mpl.colors.same_color(fap.get_edgecolor(), "blue")
  149. def test_fewer_edge_colors_than_num_edges_directed():
  150. """Test that the edge colors are cycled when there are fewer specified
  151. colors than edges."""
  152. G = barbell.to_directed()
  153. pos = nx.random_layout(barbell)
  154. edgecolors = ("r", "g", "b")
  155. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=edgecolors)
  156. for fap, expected in zip(drawn_edges, itertools.cycle(edgecolors)):
  157. assert mpl.colors.same_color(fap.get_edgecolor(), expected)
  158. def test_more_edge_colors_than_num_edges_directed():
  159. """Test that extra edge colors are ignored when there are more specified
  160. colors than edges."""
  161. G = nx.path_graph(4, create_using=nx.DiGraph) # 3 edges
  162. pos = nx.random_layout(barbell)
  163. edgecolors = ("r", "g", "b", "c") # 4 edge colors
  164. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=edgecolors)
  165. for fap, expected in zip(drawn_edges, edgecolors[:-1]):
  166. assert mpl.colors.same_color(fap.get_edgecolor(), expected)
  167. def test_edge_color_string_with_global_alpha_undirected():
  168. edge_collection = nx.draw_networkx_edges(
  169. barbell,
  170. pos=nx.random_layout(barbell),
  171. edgelist=[(0, 1), (1, 2)],
  172. edge_color="purple",
  173. alpha=0.2,
  174. )
  175. ec = edge_collection.get_color().squeeze() # as rgba tuple
  176. assert len(edge_collection.get_paths()) == 2
  177. assert mpl.colors.same_color(ec[:-1], "purple")
  178. assert ec[-1] == 0.2
  179. def test_edge_color_string_with_global_alpha_directed():
  180. drawn_edges = nx.draw_networkx_edges(
  181. barbell.to_directed(),
  182. pos=nx.random_layout(barbell),
  183. edgelist=[(0, 1), (1, 2)],
  184. edge_color="purple",
  185. alpha=0.2,
  186. )
  187. assert len(drawn_edges) == 2
  188. for fap in drawn_edges:
  189. ec = fap.get_edgecolor() # As rgba tuple
  190. assert mpl.colors.same_color(ec[:-1], "purple")
  191. assert ec[-1] == 0.2
  192. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  193. def test_edge_width_default_value(graph_type):
  194. """Test the default linewidth for edges drawn either via LineCollection or
  195. FancyArrowPatches."""
  196. G = nx.path_graph(2, create_using=graph_type)
  197. pos = {n: (n, n) for n in range(len(G))}
  198. drawn_edges = nx.draw_networkx_edges(G, pos)
  199. if isinstance(drawn_edges, list): # directed case: list of FancyArrowPatch
  200. drawn_edges = drawn_edges[0]
  201. assert drawn_edges.get_linewidth() == 1
  202. @pytest.mark.parametrize(
  203. ("edgewidth", "expected"),
  204. (
  205. (3, 3), # single-value, non-default
  206. ([3], 3), # Single value as a list
  207. ),
  208. )
  209. def test_edge_width_single_value_undirected(edgewidth, expected):
  210. G = nx.path_graph(4)
  211. pos = {n: (n, n) for n in range(len(G))}
  212. drawn_edges = nx.draw_networkx_edges(G, pos, width=edgewidth)
  213. assert len(drawn_edges.get_paths()) == 3
  214. assert drawn_edges.get_linewidth() == expected
  215. @pytest.mark.parametrize(
  216. ("edgewidth", "expected"),
  217. (
  218. (3, 3), # single-value, non-default
  219. ([3], 3), # Single value as a list
  220. ),
  221. )
  222. def test_edge_width_single_value_directed(edgewidth, expected):
  223. G = nx.path_graph(4, create_using=nx.DiGraph)
  224. pos = {n: (n, n) for n in range(len(G))}
  225. drawn_edges = nx.draw_networkx_edges(G, pos, width=edgewidth)
  226. assert len(drawn_edges) == 3
  227. for fap in drawn_edges:
  228. assert fap.get_linewidth() == expected
  229. @pytest.mark.parametrize(
  230. "edgelist",
  231. (
  232. [(0, 1), (1, 2), (2, 3)], # one width specification per edge
  233. None, # fewer widths than edges - widths cycle
  234. [(0, 1), (1, 2)], # More widths than edges - unused widths ignored
  235. ),
  236. )
  237. def test_edge_width_sequence(edgelist):
  238. G = barbell.to_directed()
  239. pos = nx.random_layout(G)
  240. widths = (0.5, 2.0, 12.0)
  241. drawn_edges = nx.draw_networkx_edges(G, pos, edgelist=edgelist, width=widths)
  242. for fap, expected_width in zip(drawn_edges, itertools.cycle(widths)):
  243. assert fap.get_linewidth() == expected_width
  244. def test_edge_color_with_edge_vmin_vmax():
  245. """Test that edge_vmin and edge_vmax properly set the dynamic range of the
  246. color map when num edges == len(edge_colors)."""
  247. G = nx.path_graph(3, create_using=nx.DiGraph)
  248. pos = nx.random_layout(G)
  249. # Extract colors from the original (unscaled) colormap
  250. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=[0, 1.0])
  251. orig_colors = [e.get_edgecolor() for e in drawn_edges]
  252. # Colors from scaled colormap
  253. drawn_edges = nx.draw_networkx_edges(
  254. G, pos, edge_color=[0.2, 0.8], edge_vmin=0.2, edge_vmax=0.8
  255. )
  256. scaled_colors = [e.get_edgecolor() for e in drawn_edges]
  257. assert mpl.colors.same_color(orig_colors, scaled_colors)
  258. def test_directed_edges_linestyle_default():
  259. """Test default linestyle for edges drawn with FancyArrowPatches."""
  260. G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
  261. pos = {n: (n, n) for n in range(len(G))}
  262. # edge with default style
  263. drawn_edges = nx.draw_networkx_edges(G, pos)
  264. assert len(drawn_edges) == 3
  265. for fap in drawn_edges:
  266. assert fap.get_linestyle() == "solid"
  267. @pytest.mark.parametrize(
  268. "style",
  269. (
  270. "dashed", # edge with string style
  271. "--", # edge with simplified string style
  272. (1, (1, 1)), # edge with (offset, onoffseq) style
  273. ),
  274. )
  275. def test_directed_edges_linestyle_single_value(style):
  276. """Tests support for specifying linestyles with a single value to be applied to
  277. all edges in ``draw_networkx_edges`` for FancyArrowPatch outputs
  278. (e.g. directed edges)."""
  279. G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
  280. pos = {n: (n, n) for n in range(len(G))}
  281. drawn_edges = nx.draw_networkx_edges(G, pos, style=style)
  282. assert len(drawn_edges) == 3
  283. for fap in drawn_edges:
  284. assert fap.get_linestyle() == style
  285. @pytest.mark.parametrize(
  286. "style_seq",
  287. (
  288. ["dashed"], # edge with string style in list
  289. ["--"], # edge with simplified string style in list
  290. [(1, (1, 1))], # edge with (offset, onoffseq) style in list
  291. ["--", "-", ":"], # edges with styles for each edge
  292. ["--", "-"], # edges with fewer styles than edges (styles cycle)
  293. ["--", "-", ":", "-."], # edges with more styles than edges (extra unused)
  294. ),
  295. )
  296. def test_directed_edges_linestyle_sequence(style_seq):
  297. """Tests support for specifying linestyles with sequences in
  298. ``draw_networkx_edges`` for FancyArrowPatch outputs (e.g. directed edges)."""
  299. G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
  300. pos = {n: (n, n) for n in range(len(G))}
  301. drawn_edges = nx.draw_networkx_edges(G, pos, style=style_seq)
  302. assert len(drawn_edges) == 3
  303. for fap, style in zip(drawn_edges, itertools.cycle(style_seq)):
  304. assert fap.get_linestyle() == style
  305. def test_labels_and_colors():
  306. G = nx.cubical_graph()
  307. pos = nx.spring_layout(G) # positions for all nodes
  308. # nodes
  309. nx.draw_networkx_nodes(
  310. G, pos, nodelist=[0, 1, 2, 3], node_color="r", node_size=500, alpha=0.75
  311. )
  312. nx.draw_networkx_nodes(
  313. G,
  314. pos,
  315. nodelist=[4, 5, 6, 7],
  316. node_color="b",
  317. node_size=500,
  318. alpha=[0.25, 0.5, 0.75, 1.0],
  319. )
  320. # edges
  321. nx.draw_networkx_edges(G, pos, width=1.0, alpha=0.5)
  322. nx.draw_networkx_edges(
  323. G,
  324. pos,
  325. edgelist=[(0, 1), (1, 2), (2, 3), (3, 0)],
  326. width=8,
  327. alpha=0.5,
  328. edge_color="r",
  329. )
  330. nx.draw_networkx_edges(
  331. G,
  332. pos,
  333. edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
  334. width=8,
  335. alpha=0.5,
  336. edge_color="b",
  337. )
  338. nx.draw_networkx_edges(
  339. G,
  340. pos,
  341. edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
  342. arrows=True,
  343. min_source_margin=0.5,
  344. min_target_margin=0.75,
  345. width=8,
  346. edge_color="b",
  347. )
  348. # some math labels
  349. labels = {}
  350. labels[0] = r"$a$"
  351. labels[1] = r"$b$"
  352. labels[2] = r"$c$"
  353. labels[3] = r"$d$"
  354. labels[4] = r"$\alpha$"
  355. labels[5] = r"$\beta$"
  356. labels[6] = r"$\gamma$"
  357. labels[7] = r"$\delta$"
  358. nx.draw_networkx_labels(G, pos, labels, font_size=16)
  359. nx.draw_networkx_edge_labels(G, pos, edge_labels=None, rotate=False)
  360. nx.draw_networkx_edge_labels(G, pos, edge_labels={(4, 5): "4-5"})
  361. # plt.show()
  362. @pytest.mark.mpl_image_compare
  363. def test_house_with_colors():
  364. G = nx.house_graph()
  365. # explicitly set positions
  366. fig, ax = plt.subplots()
  367. pos = {0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1), 4: (0.5, 2.0)}
  368. # Plot nodes with different properties for the "wall" and "roof" nodes
  369. nx.draw_networkx_nodes(
  370. G,
  371. pos,
  372. node_size=3000,
  373. nodelist=[0, 1, 2, 3],
  374. node_color="tab:blue",
  375. )
  376. nx.draw_networkx_nodes(
  377. G, pos, node_size=2000, nodelist=[4], node_color="tab:orange"
  378. )
  379. nx.draw_networkx_edges(G, pos, alpha=0.5, width=6)
  380. # Customize axes
  381. ax.margins(0.11)
  382. plt.tight_layout()
  383. plt.axis("off")
  384. return fig
  385. def test_axes():
  386. fig, ax = plt.subplots()
  387. nx.draw(barbell, ax=ax)
  388. nx.draw_networkx_edge_labels(barbell, nx.circular_layout(barbell), ax=ax)
  389. def test_empty_graph():
  390. G = nx.Graph()
  391. nx.draw(G)
  392. def test_draw_empty_nodes_return_values():
  393. # See Issue #3833
  394. import matplotlib.collections # call as mpl.collections
  395. G = nx.Graph([(1, 2), (2, 3)])
  396. DG = nx.DiGraph([(1, 2), (2, 3)])
  397. pos = nx.circular_layout(G)
  398. assert isinstance(
  399. nx.draw_networkx_nodes(G, pos, nodelist=[]), mpl.collections.PathCollection
  400. )
  401. assert isinstance(
  402. nx.draw_networkx_nodes(DG, pos, nodelist=[]), mpl.collections.PathCollection
  403. )
  404. # drawing empty edges used to return an empty LineCollection or empty list.
  405. # Now it is always an empty list (because edges are now lists of FancyArrows)
  406. assert nx.draw_networkx_edges(G, pos, edgelist=[], arrows=True) == []
  407. assert nx.draw_networkx_edges(G, pos, edgelist=[], arrows=False) == []
  408. assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=False) == []
  409. assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=True) == []
  410. def test_multigraph_edgelist_tuples():
  411. # See Issue #3295
  412. G = nx.path_graph(3, create_using=nx.MultiDiGraph)
  413. nx.draw_networkx(G, edgelist=[(0, 1, 0)])
  414. nx.draw_networkx(G, edgelist=[(0, 1, 0)], node_size=[10, 20, 0])
  415. def test_alpha_iter():
  416. pos = nx.random_layout(barbell)
  417. fig = plt.figure()
  418. # with fewer alpha elements than nodes
  419. fig.add_subplot(131) # Each test in a new axis object
  420. nx.draw_networkx_nodes(barbell, pos, alpha=[0.1, 0.2])
  421. # with equal alpha elements and nodes
  422. num_nodes = len(barbell.nodes)
  423. alpha = [x / num_nodes for x in range(num_nodes)]
  424. colors = range(num_nodes)
  425. fig.add_subplot(132)
  426. nx.draw_networkx_nodes(barbell, pos, node_color=colors, alpha=alpha)
  427. # with more alpha elements than nodes
  428. alpha.append(1)
  429. fig.add_subplot(133)
  430. nx.draw_networkx_nodes(barbell, pos, alpha=alpha)
  431. def test_error_invalid_kwds():
  432. with pytest.raises(ValueError, match="Received invalid argument"):
  433. nx.draw(barbell, foo="bar")
  434. def test_draw_networkx_arrowsize_incorrect_size():
  435. G = nx.DiGraph([(0, 1), (0, 2), (0, 3), (1, 3)])
  436. arrowsize = [1, 2, 3]
  437. with pytest.raises(
  438. ValueError, match="arrowsize should have the same length as edgelist"
  439. ):
  440. nx.draw(G, arrowsize=arrowsize)
  441. @pytest.mark.parametrize("arrowsize", (30, [10, 20, 30]))
  442. def test_draw_edges_arrowsize(arrowsize):
  443. G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
  444. pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
  445. edges = nx.draw_networkx_edges(G, pos=pos, arrowsize=arrowsize)
  446. arrowsize = itertools.repeat(arrowsize) if isinstance(arrowsize, int) else arrowsize
  447. for fap, expected in zip(edges, arrowsize):
  448. assert isinstance(fap, mpl.patches.FancyArrowPatch)
  449. assert fap.get_mutation_scale() == expected
  450. def test_np_edgelist():
  451. # see issue #4129
  452. nx.draw_networkx(barbell, edgelist=np.array([(0, 2), (0, 3)]))
  453. def test_draw_nodes_missing_node_from_position():
  454. G = nx.path_graph(3)
  455. pos = {0: (0, 0), 1: (1, 1)} # No position for node 2
  456. with pytest.raises(nx.NetworkXError, match="has no position"):
  457. nx.draw_networkx_nodes(G, pos)
  458. # NOTE: parametrizing on marker to test both branches of internal
  459. # nx.draw_networkx_edges.to_marker_edge function
  460. @pytest.mark.parametrize("node_shape", ("o", "s"))
  461. def test_draw_edges_min_source_target_margins(node_shape):
  462. """Test that there is a wider gap between the node and the start of an
  463. incident edge when min_source_margin is specified.
  464. This test checks that the use of min_{source/target}_margin kwargs result
  465. in shorter (more padding) between the edges and source and target nodes.
  466. As a crude visual example, let 's' and 't' represent source and target
  467. nodes, respectively:
  468. Default:
  469. s-----------------------------t
  470. With margins:
  471. s ----------------------- t
  472. """
  473. # Create a single axis object to get consistent pixel coords across
  474. # multiple draws
  475. fig, ax = plt.subplots()
  476. G = nx.DiGraph([(0, 1)])
  477. pos = {0: (0, 0), 1: (1, 0)} # horizontal layout
  478. # Get leftmost and rightmost points of the FancyArrowPatch object
  479. # representing the edge between nodes 0 and 1 (in pixel coordinates)
  480. default_patch = nx.draw_networkx_edges(G, pos, ax=ax, node_shape=node_shape)[0]
  481. default_extent = default_patch.get_extents().corners()[::2, 0]
  482. # Now, do the same but with "padding" for the source and target via the
  483. # min_{source/target}_margin kwargs
  484. padded_patch = nx.draw_networkx_edges(
  485. G,
  486. pos,
  487. ax=ax,
  488. node_shape=node_shape,
  489. min_source_margin=100,
  490. min_target_margin=100,
  491. )[0]
  492. padded_extent = padded_patch.get_extents().corners()[::2, 0]
  493. # With padding, the left-most extent of the edge should be further to the
  494. # right
  495. assert padded_extent[0] > default_extent[0]
  496. # And the rightmost extent of the edge, further to the left
  497. assert padded_extent[1] < default_extent[1]
  498. def test_nonzero_selfloop_with_single_node():
  499. """Ensure that selfloop extent is non-zero when there is only one node."""
  500. # Create explicit axis object for test
  501. fig, ax = plt.subplots()
  502. # Graph with single node + self loop
  503. G = nx.DiGraph()
  504. G.add_node(0)
  505. G.add_edge(0, 0)
  506. # Draw
  507. patch = nx.draw_networkx_edges(G, {0: (0, 0)})[0]
  508. # The resulting patch must have non-zero extent
  509. bbox = patch.get_extents()
  510. assert bbox.width > 0 and bbox.height > 0
  511. # Cleanup
  512. plt.delaxes(ax)
  513. def test_nonzero_selfloop_with_single_edge_in_edgelist():
  514. """Ensure that selfloop extent is non-zero when only a single edge is
  515. specified in the edgelist.
  516. """
  517. # Create explicit axis object for test
  518. fig, ax = plt.subplots()
  519. # Graph with selfloop
  520. G = nx.path_graph(2, create_using=nx.DiGraph)
  521. G.add_edge(1, 1)
  522. pos = {n: (n, n) for n in G.nodes}
  523. # Draw only the selfloop edge via the `edgelist` kwarg
  524. patch = nx.draw_networkx_edges(G, pos, edgelist=[(1, 1)])[0]
  525. # The resulting patch must have non-zero extent
  526. bbox = patch.get_extents()
  527. assert bbox.width > 0 and bbox.height > 0
  528. # Cleanup
  529. plt.delaxes(ax)
  530. def test_apply_alpha():
  531. """Test apply_alpha when there is a mismatch between the number of
  532. supplied colors and elements.
  533. """
  534. nodelist = [0, 1, 2]
  535. colorlist = ["r", "g", "b"]
  536. alpha = 0.5
  537. rgba_colors = nx.drawing.nx_pylab.apply_alpha(colorlist, alpha, nodelist)
  538. assert all(rgba_colors[:, -1] == alpha)
  539. def test_draw_edges_toggling_with_arrows_kwarg():
  540. """
  541. The `arrows` keyword argument is used as a 3-way switch to select which
  542. type of object to use for drawing edges:
  543. - ``arrows=None`` -> default (FancyArrowPatches for directed, else LineCollection)
  544. - ``arrows=True`` -> FancyArrowPatches
  545. - ``arrows=False`` -> LineCollection
  546. """
  547. import matplotlib.collections
  548. import matplotlib.patches
  549. UG = nx.path_graph(3)
  550. DG = nx.path_graph(3, create_using=nx.DiGraph)
  551. pos = {n: (n, n) for n in UG}
  552. # Use FancyArrowPatches when arrows=True, regardless of graph type
  553. for G in (UG, DG):
  554. edges = nx.draw_networkx_edges(G, pos, arrows=True)
  555. assert len(edges) == len(G.edges)
  556. assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
  557. # Use LineCollection when arrows=False, regardless of graph type
  558. for G in (UG, DG):
  559. edges = nx.draw_networkx_edges(G, pos, arrows=False)
  560. assert isinstance(edges, mpl.collections.LineCollection)
  561. # Default behavior when arrows=None: FAPs for directed, LC's for undirected
  562. edges = nx.draw_networkx_edges(UG, pos)
  563. assert isinstance(edges, mpl.collections.LineCollection)
  564. edges = nx.draw_networkx_edges(DG, pos)
  565. assert len(edges) == len(G.edges)
  566. assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
  567. @pytest.mark.parametrize("drawing_func", (nx.draw, nx.draw_networkx))
  568. def test_draw_networkx_arrows_default_undirected(drawing_func):
  569. import matplotlib.collections
  570. G = nx.path_graph(3)
  571. fig, ax = plt.subplots()
  572. drawing_func(G, ax=ax)
  573. assert any(isinstance(c, mpl.collections.LineCollection) for c in ax.collections)
  574. assert not ax.patches
  575. plt.delaxes(ax)
  576. @pytest.mark.parametrize("drawing_func", (nx.draw, nx.draw_networkx))
  577. def test_draw_networkx_arrows_default_directed(drawing_func):
  578. import matplotlib.collections
  579. G = nx.path_graph(3, create_using=nx.DiGraph)
  580. fig, ax = plt.subplots()
  581. drawing_func(G, ax=ax)
  582. assert not any(
  583. isinstance(c, mpl.collections.LineCollection) for c in ax.collections
  584. )
  585. assert ax.patches
  586. plt.delaxes(ax)
  587. def test_edgelist_kwarg_not_ignored():
  588. # See gh-4994
  589. G = nx.path_graph(3)
  590. G.add_edge(0, 0)
  591. fig, ax = plt.subplots()
  592. nx.draw(G, edgelist=[(0, 1), (1, 2)], ax=ax) # Exclude self-loop from edgelist
  593. assert not ax.patches
  594. plt.delaxes(ax)
  595. def test_draw_networkx_edge_label_multiedge_exception():
  596. """
  597. draw_networkx_edge_labels should raise an informative error message when
  598. the edge label includes keys
  599. """
  600. exception_msg = "draw_networkx_edge_labels does not support multiedges"
  601. G = nx.MultiGraph()
  602. G.add_edge(0, 1, weight=10)
  603. G.add_edge(0, 1, weight=20)
  604. edge_labels = nx.get_edge_attributes(G, "weight") # Includes edge keys
  605. pos = {n: (n, n) for n in G}
  606. with pytest.raises(nx.NetworkXError, match=exception_msg):
  607. nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
  608. def test_draw_networkx_edge_label_empty_dict():
  609. """Regression test for draw_networkx_edge_labels with empty dict. See
  610. gh-5372."""
  611. G = nx.path_graph(3)
  612. pos = {n: (n, n) for n in G.nodes}
  613. assert nx.draw_networkx_edge_labels(G, pos, edge_labels={}) == {}
  614. def test_draw_networkx_edges_undirected_selfloop_colors():
  615. """When an edgelist is supplied along with a sequence of colors, check that
  616. the self-loops have the correct colors."""
  617. fig, ax = plt.subplots()
  618. # Edge list and corresponding colors
  619. edgelist = [(1, 3), (1, 2), (2, 3), (1, 1), (3, 3), (2, 2)]
  620. edge_colors = ["pink", "cyan", "black", "red", "blue", "green"]
  621. G = nx.Graph(edgelist)
  622. pos = {n: (n, n) for n in G.nodes}
  623. nx.draw_networkx_edges(G, pos, ax=ax, edgelist=edgelist, edge_color=edge_colors)
  624. # Verify that there are three fancy arrow patches (1 per self loop)
  625. assert len(ax.patches) == 3
  626. # These are points that should be contained in the self loops. For example,
  627. # sl_points[0] will be (1, 1.1), which is inside the "path" of the first
  628. # self-loop but outside the others
  629. sl_points = np.array(edgelist[-3:]) + np.array([0, 0.1])
  630. # Check that the mapping between self-loop locations and their colors is
  631. # correct
  632. for fap, clr, slp in zip(ax.patches, edge_colors[-3:], sl_points):
  633. assert fap.get_path().contains_point(slp)
  634. assert mpl.colors.same_color(fap.get_edgecolor(), clr)
  635. plt.delaxes(ax)
  636. @pytest.mark.parametrize(
  637. "fap_only_kwarg", # Non-default values for kwargs that only apply to FAPs
  638. (
  639. {"arrowstyle": "-"},
  640. {"arrowsize": 20},
  641. {"connectionstyle": "arc3,rad=0.2"},
  642. {"min_source_margin": 10},
  643. {"min_target_margin": 10},
  644. ),
  645. )
  646. def test_user_warnings_for_unused_edge_drawing_kwargs(fap_only_kwarg):
  647. """Users should get a warning when they specify a non-default value for
  648. one of the kwargs that applies only to edges drawn with FancyArrowPatches,
  649. but FancyArrowPatches aren't being used under the hood."""
  650. G = nx.path_graph(3)
  651. pos = {n: (n, n) for n in G}
  652. fig, ax = plt.subplots()
  653. # By default, an undirected graph will use LineCollection to represent
  654. # the edges
  655. kwarg_name = list(fap_only_kwarg.keys())[0]
  656. with pytest.warns(
  657. UserWarning, match=f"\n\nThe {kwarg_name} keyword argument is not applicable"
  658. ):
  659. nx.draw_networkx_edges(G, pos, ax=ax, **fap_only_kwarg)
  660. # FancyArrowPatches are always used when `arrows=True` is specified.
  661. # Check that warnings are *not* raised in this case
  662. with warnings.catch_warnings():
  663. # Escalate warnings -> errors so tests fail if warnings are raised
  664. warnings.simplefilter("error")
  665. nx.draw_networkx_edges(G, pos, ax=ax, arrows=True, **fap_only_kwarg)
  666. plt.delaxes(ax)