test_pylab.py 54 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582
  1. """Unit tests for matplotlib drawing functions."""
  2. import itertools
  3. import os
  4. import warnings
  5. import pytest
  6. import networkx as nx
  7. mpl = pytest.importorskip("matplotlib")
  8. np = pytest.importorskip("numpy")
  9. mpl.use("PS")
  10. plt = pytest.importorskip("matplotlib.pyplot")
  11. plt.rcParams["text.usetex"] = False
  12. barbell = nx.barbell_graph(4, 6)
  13. defaults = {
  14. "node_pos": None,
  15. "node_visible": True,
  16. "node_color": "#1f78b4",
  17. "node_size": 300,
  18. "node_label": {
  19. "size": 12,
  20. "color": "#000000",
  21. "family": "sans-serif",
  22. "weight": "normal",
  23. "alpha": 1.0,
  24. "background_color": None,
  25. "background_alpha": None,
  26. "h_align": "center",
  27. "v_align": "center",
  28. "bbox": None,
  29. },
  30. "node_shape": "o",
  31. "node_alpha": 1.0,
  32. "node_border_width": 1.0,
  33. "node_border_color": "face",
  34. "edge_visible": True,
  35. "edge_width": 1.0,
  36. "edge_color": "#000000",
  37. "edge_label": {
  38. "size": 12,
  39. "color": "#000000",
  40. "family": "sans-serif",
  41. "weight": "normal",
  42. "alpha": 1.0,
  43. "bbox": {"boxstyle": "round", "ec": (1.0, 1.0, 1.0), "fc": (1.0, 1.0, 1.0)},
  44. "h_align": "center",
  45. "v_align": "center",
  46. "pos": 0.5,
  47. "rotate": True,
  48. },
  49. "edge_style": "-",
  50. "edge_alpha": 1.0,
  51. # These are for undirected-graphs. Directed graphs shouls use "-|>" and 10, respectively
  52. "edge_arrowstyle": "-",
  53. "edge_arrowsize": 0,
  54. "edge_curvature": "arc3",
  55. "edge_source_margin": 0,
  56. "edge_target_margin": 0,
  57. }
  58. @pytest.mark.parametrize(
  59. ("param_name", "param_value", "expected"),
  60. (
  61. ("node_color", None, defaults["node_color"]),
  62. ("node_color", "#FF0000", "red"),
  63. ("node_color", "color", "lime"),
  64. ),
  65. )
  66. def test_display_arg_handling_node_color(param_name, param_value, expected):
  67. G = nx.path_graph(4)
  68. nx.set_node_attributes(G, "#00FF00", "color")
  69. canvas = plt.figure().add_subplot(111)
  70. nx.display(G, canvas=canvas, **{param_name: param_value})
  71. assert mpl.colors.same_color(canvas.get_children()[0].get_edgecolors()[0], expected)
  72. plt.close()
  73. @pytest.mark.parametrize(
  74. ("param_value", "expected"),
  75. (
  76. (None, (1, 1, 1, 1)), # default value
  77. (0.5, (0.5, 0.5, 0.5, 0.5)),
  78. ("n_alpha", (1.0, 1 / 2, 1 / 3, 0.25)),
  79. ),
  80. )
  81. def test_display_arg_handling_node_alpha(param_value, expected):
  82. G = nx.path_graph(4)
  83. nx.set_node_attributes(G, {n: 1 / (n + 1) for n in G.nodes()}, "n_alpha")
  84. canvas = plt.figure().add_subplot(111)
  85. nx.display(G, canvas=canvas, node_alpha=param_value)
  86. assert all(
  87. canvas.get_children()[0].get_fc()[:, 3] == expected
  88. ) # Extract just the alpha from the node colors
  89. plt.close()
  90. def test_display_node_position():
  91. G = nx.path_graph(4)
  92. nx.set_node_attributes(G, {n: (n, n) for n in G.nodes()}, "pos")
  93. canvas = plt.figure().add_subplot(111)
  94. nx.display(G, canvas=canvas, node_pos="pos")
  95. assert np.all(
  96. canvas.get_children()[0].get_offsets().data == [[0, 0], [1, 1], [2, 2], [3, 3]]
  97. )
  98. plt.close()
  99. def test_display_line_collection():
  100. G = nx.karate_club_graph()
  101. nx.set_edge_attributes(
  102. G, {(u, v): "-|>" if (u + v) % 2 else "-" for u, v in G.edges()}, "arrowstyle"
  103. )
  104. canvas = plt.figure().add_subplot(111)
  105. nx.display(G, canvas=canvas, edge_arrowsize=10)
  106. # There should only be one line collection in any given visualization
  107. lc = [
  108. l
  109. for l in canvas.get_children()
  110. if isinstance(l, mpl.collections.LineCollection)
  111. ][0]
  112. assert len(lc.get_paths()) == sum([1 for u, v in G.edges() if (u + v) % 2])
  113. plt.close()
  114. @pytest.mark.parametrize(
  115. ("edge_color", "expected"),
  116. (
  117. (None, "black"),
  118. ("r", "red"),
  119. ((1.0, 1.0, 0.0), "yellow"),
  120. ((0, 1, 0, 1), "lime"),
  121. ("color", "blue"),
  122. ("#0000FF", "blue"),
  123. ),
  124. )
  125. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  126. def test_display_edge_single_color(edge_color, expected, graph_type):
  127. G = nx.path_graph(3, create_using=graph_type)
  128. nx.set_edge_attributes(G, "#0000FF", "color")
  129. canvas = plt.figure().add_subplot(111)
  130. nx.display(G, edge_color=edge_color, canvas=canvas)
  131. if G.is_directed():
  132. colors = [
  133. f.get_fc()
  134. for f in canvas.get_children()
  135. if isinstance(f, mpl.patches.FancyArrowPatch)
  136. ]
  137. else:
  138. colors = [
  139. l
  140. for l in canvas.collections
  141. if isinstance(l, mpl.collections.LineCollection)
  142. ][0].get_colors()
  143. assert all(mpl.colors.same_color(c, expected) for c in colors)
  144. plt.close()
  145. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  146. def test_display_edge_multiple_colors(graph_type):
  147. G = nx.path_graph(3, create_using=graph_type)
  148. nx.set_edge_attributes(G, {(0, 1): "#FF0000", (1, 2): (0, 0, 1)}, "color")
  149. ax = plt.figure().add_subplot(111)
  150. nx.display(G, canvas=ax)
  151. expected = ["red", "blue"]
  152. if G.is_directed():
  153. colors = [
  154. f.get_fc()
  155. for f in ax.get_children()
  156. if isinstance(f, mpl.patches.FancyArrowPatch)
  157. ]
  158. else:
  159. colors = [
  160. l for l in ax.collections if isinstance(l, mpl.collections.LineCollection)
  161. ][0].get_colors()
  162. assert mpl.colors.same_color(colors, expected)
  163. plt.close()
  164. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  165. def test_display_edge_position(graph_type):
  166. G = nx.path_graph(3, create_using=graph_type)
  167. nx.set_node_attributes(G, {n: (n, n) for n in G.nodes()}, "pos")
  168. ax = plt.figure().add_subplot(111)
  169. nx.display(G, canvas=ax)
  170. if G.is_directed():
  171. end_points = [
  172. (f.get_path().vertices[0, :], f.get_path().vertices[-2, :])
  173. for f in ax.get_children()
  174. if isinstance(f, mpl.patches.FancyArrowPatch)
  175. ]
  176. else:
  177. line_collection = [
  178. l for l in ax.collections if isinstance(l, mpl.collections.LineCollection)
  179. ][0]
  180. end_points = [
  181. (p.vertices[0, :], p.vertices[-1, :]) for p in line_collection.get_paths()
  182. ]
  183. expected = [((0, 0), (1, 1)), ((1, 1), (2, 2))]
  184. # Use the threshold to account for slight shifts in FancyArrowPatch margins to
  185. # avoid covering the arrow head with the node.
  186. threshold = 0.05
  187. for a, e in zip(end_points, expected):
  188. act_start, act_end = a
  189. exp_start, exp_end = e
  190. assert all(abs(act_start - exp_start) < (threshold, threshold)) and all(
  191. abs(act_end - exp_end) < (threshold, threshold)
  192. )
  193. plt.close()
  194. def test_display_position_function():
  195. G = nx.karate_club_graph()
  196. def fixed_layout(G):
  197. return nx.spring_layout(G, seed=314159)
  198. pos = fixed_layout(G)
  199. ax = plt.figure().add_subplot(111)
  200. nx.display(G, node_pos=fixed_layout, canvas=ax)
  201. # rebuild the position dictionary from the canvas
  202. act_pos = {
  203. n: tuple(p) for n, p in zip(G.nodes(), ax.get_children()[0].get_offsets().data)
  204. }
  205. for n in G.nodes():
  206. assert all(pos[n] == act_pos[n])
  207. plt.close()
  208. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  209. def test_display_edge_colormaps(graph_type):
  210. G = nx.path_graph(3, create_using=graph_type)
  211. nx.set_edge_attributes(G, {(0, 1): 0, (1, 2): 1}, "weight")
  212. cmap = mpl.colormaps["plasma"]
  213. nx.apply_matplotlib_colors(G, "weight", "color", cmap, nodes=False)
  214. canvas = plt.figure().add_subplot(111)
  215. nx.display(G, canvas=canvas)
  216. mapper = mpl.cm.ScalarMappable(cmap=cmap)
  217. mapper.set_clim(0, 1)
  218. expected = [mapper.to_rgba(0), mapper.to_rgba(1)]
  219. if G.is_directed():
  220. colors = [
  221. f.get_facecolor()
  222. for f in canvas.get_children()
  223. if isinstance(f, mpl.patches.FancyArrowPatch)
  224. ]
  225. else:
  226. colors = [
  227. l
  228. for l in canvas.collections
  229. if isinstance(l, mpl.collections.LineCollection)
  230. ][0].get_colors()
  231. assert mpl.colors.same_color(expected[0], G.edges[0, 1]["color"])
  232. assert mpl.colors.same_color(expected[1], G.edges[1, 2]["color"])
  233. assert mpl.colors.same_color(expected, colors)
  234. plt.close()
  235. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  236. def test_display_node_colormaps(graph_type):
  237. G = nx.path_graph(3, create_using=graph_type)
  238. nx.set_node_attributes(G, {0: 0, 1: 0.5, 2: 1}, "weight")
  239. cmap = mpl.colormaps["plasma"]
  240. nx.apply_matplotlib_colors(G, "weight", "color", cmap)
  241. canvas = plt.figure().add_subplot(111)
  242. nx.display(G, canvas=canvas)
  243. mapper = mpl.cm.ScalarMappable(cmap=cmap)
  244. mapper.set_clim(0, 1)
  245. expected = [mapper.to_rgba(0), mapper.to_rgba(0.5), mapper.to_rgba(1)]
  246. colors = [
  247. s for s in canvas.collections if isinstance(s, mpl.collections.PathCollection)
  248. ][0].get_edgecolors()
  249. assert mpl.colors.same_color(expected[0], G.nodes[0]["color"])
  250. assert mpl.colors.same_color(expected[1], G.nodes[1]["color"])
  251. assert mpl.colors.same_color(expected[2], G.nodes[2]["color"])
  252. assert mpl.colors.same_color(expected, colors)
  253. plt.close()
  254. @pytest.mark.parametrize(
  255. ("param_value", "expected"),
  256. (
  257. (None, [defaults["edge_width"], defaults["edge_width"]]),
  258. (5, [5, 5]),
  259. ("width", [5, 10]),
  260. ),
  261. )
  262. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  263. def test_display_edge_width(param_value, expected, graph_type):
  264. G = nx.path_graph(3, create_using=graph_type)
  265. nx.set_edge_attributes(G, {(0, 1): 5, (1, 2): 10}, "width")
  266. canvas = plt.figure().add_subplot(111)
  267. nx.display(G, edge_width=param_value, canvas=canvas)
  268. if G.is_directed():
  269. widths = [
  270. f.get_linewidth()
  271. for f in canvas.get_children()
  272. if isinstance(f, mpl.patches.FancyArrowPatch)
  273. ]
  274. else:
  275. widths = list(
  276. [
  277. l
  278. for l in canvas.collections
  279. if isinstance(l, mpl.collections.LineCollection)
  280. ][0].get_linewidths()
  281. )
  282. assert widths == expected
  283. @pytest.mark.parametrize(
  284. ("param_value", "expected"),
  285. (
  286. (None, [defaults["edge_style"], defaults["edge_style"]]),
  287. (":", [":", ":"]),
  288. ("style", ["-", ":"]),
  289. ),
  290. )
  291. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  292. def test_display_edge_style(param_value, expected, graph_type):
  293. G = nx.path_graph(3, create_using=graph_type)
  294. nx.set_edge_attributes(G, {(0, 1): "-", (1, 2): ":"}, "style")
  295. canvas = plt.figure().add_subplot(111)
  296. nx.display(G, edge_style=param_value, canvas=canvas)
  297. if G.is_directed():
  298. styles = [
  299. f.get_linestyle()
  300. for f in canvas.get_children()
  301. if isinstance(f, mpl.patches.FancyArrowPatch)
  302. ]
  303. else:
  304. # Convert back from tuple description to character form
  305. linestyles = {(0, None): "-", (0, (1, 1.65)): ":"}
  306. styles = [
  307. linestyles[(s[0], tuple(s[1]) if s[1] is not None else None)]
  308. for s in [
  309. l
  310. for l in canvas.collections
  311. if isinstance(l, mpl.collections.LineCollection)
  312. ][0].get_linestyles()
  313. ]
  314. assert styles == expected
  315. plt.close()
  316. def test_display_node_labels():
  317. G = nx.path_graph(4)
  318. canvas = plt.figure().add_subplot(111)
  319. nx.display(G, canvas=canvas, node_label={"size": 20})
  320. labels = [t for t in canvas.get_children() if isinstance(t, mpl.text.Text)]
  321. for n, l in zip(G.nodes(), labels):
  322. assert l.get_text() == str(n)
  323. assert l.get_size() == 20.0
  324. plt.close()
  325. def test_display_edge_labels():
  326. G = nx.path_graph(4)
  327. canvas = plt.figure().add_subplot(111)
  328. # While we can pass in dicts for edge label defaults without errors,
  329. # this isn't helpful unless we want one label for all edges.
  330. nx.set_edge_attributes(G, {(u, v): {"label": u + v} for u, v in G.edges()})
  331. nx.display(G, canvas=canvas, edge_label={"color": "r"}, node_label=None)
  332. labels = [t for t in canvas.get_children() if isinstance(t, mpl.text.Text)]
  333. print(labels)
  334. for e, l in zip(G.edges(), labels):
  335. assert l.get_text() == str(e[0] + e[1])
  336. assert l.get_color() == "r"
  337. plt.close()
  338. def test_display_multigraph_non_integer_keys():
  339. G = nx.MultiGraph()
  340. G.add_nodes_from(["A", "B", "C", "D"])
  341. G.add_edges_from(
  342. [
  343. ("A", "B", "0"),
  344. ("A", "B", "1"),
  345. ("B", "C", "-1"),
  346. ("B", "C", "1"),
  347. ("C", "D", "-1"),
  348. ("C", "D", "0"),
  349. ]
  350. )
  351. nx.set_edge_attributes(
  352. G, {e: f"arc3,rad={0.2 * int(e[2])}" for e in G.edges(keys=True)}, "curvature"
  353. )
  354. canvas = plt.figure().add_subplot(111)
  355. nx.display(G, canvas=canvas)
  356. rads = [
  357. f.get_connectionstyle().rad
  358. for f in canvas.get_children()
  359. if isinstance(f, mpl.patches.FancyArrowPatch)
  360. ]
  361. assert rads == [0.0, 0.2, -0.2, 0.2, -0.2, 0.0]
  362. plt.close()
  363. def test_display_raises_for_bad_arg():
  364. G = nx.karate_club_graph()
  365. with pytest.raises(nx.NetworkXError):
  366. nx.display(G, bad_arg="bad_arg")
  367. plt.close()
  368. def test_display_arrow_size():
  369. G = nx.path_graph(4, create_using=nx.DiGraph)
  370. nx.set_edge_attributes(
  371. G, {(u, v): (u + v + 2) ** 2 for u, v in G.edges()}, "arrowsize"
  372. )
  373. ax = plt.axes()
  374. nx.display(G, canvas=ax)
  375. assert [9, 25, 49] == [
  376. f.get_mutation_scale()
  377. for f in ax.get_children()
  378. if isinstance(f, mpl.patches.FancyArrowPatch)
  379. ]
  380. plt.close()
  381. def test_display_mismatched_edge_position():
  382. """
  383. This test ensures that a error is raised for incomplete position data.
  384. """
  385. G = nx.path_graph(5)
  386. # Notice that there is no position for node 3
  387. nx.set_node_attributes(G, {0: (0, 0), 1: (1, 1), 2: (2, 2), 4: (4, 4)}, "pos")
  388. # But that's not a problem since we don't want to show node 4, right?
  389. nx.set_node_attributes(G, {n: n < 4 for n in G.nodes()}, "visible")
  390. # However, if we try to visualize every edge (including 3 -> 4)...
  391. # That's a problem since node 4 doesn't have a position
  392. with pytest.raises(nx.NetworkXError):
  393. nx.display(G)
  394. # NOTE: parametrizing on marker to test both branches of internal
  395. # to_marker_edge function
  396. @pytest.mark.parametrize("node_shape", ("o", "s"))
  397. def test_display_edge_margins(node_shape):
  398. """
  399. Test that there is a wider gap between the node and the start of an
  400. incident edge when min_source_margin is specified.
  401. This test checks that the use os min_{source/target}_margin edge
  402. attributes result is shorter (more padding) between the edges and
  403. source and target nodes.
  404. As a crude visual example, let 's' and 't' represent source and target
  405. nodes, respectively:
  406. Default:
  407. s-----------------------------t
  408. With margins:
  409. s ----------------------- t
  410. """
  411. ax = plt.figure().add_subplot(111)
  412. G = nx.DiGraph([(0, 1)])
  413. nx.set_node_attributes(G, {0: (0, 0), 1: (1, 1)}, "pos")
  414. # Get the default patches from the regular visualization
  415. nx.display(G, canvas=ax, node_shape=node_shape)
  416. default_arrow = [
  417. f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
  418. ][0]
  419. default_extent = default_arrow.get_extents().corners()[::2, 0]
  420. # Now plot again with margins
  421. ax = plt.figure().add_subplot(111)
  422. nx.display(
  423. G,
  424. canvas=ax,
  425. edge_source_margin=100,
  426. edge_target_margin=100,
  427. node_shape=node_shape,
  428. )
  429. padded_arrow = [
  430. f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
  431. ][0]
  432. padded_extent = padded_arrow.get_extents().corners()[::2, 0]
  433. # With padding, the left-most extent of the edge should be further to the right
  434. assert padded_extent[0] > default_extent[0]
  435. # And the rightmost extent of the edge, further to the left
  436. assert padded_extent[1] < default_extent[1]
  437. plt.close()
  438. @pytest.mark.parametrize("ticks", [False, True])
  439. def test_display_hide_ticks(ticks):
  440. G = nx.path_graph(3)
  441. nx.set_node_attributes(G, {n: (n, n) for n in G.nodes()}, "pos")
  442. ax = plt.axes()
  443. nx.display(G, hide_ticks=ticks)
  444. for axis in [ax.xaxis, ax.yaxis]:
  445. assert bool(axis.get_ticklabels()) != ticks
  446. plt.close()
  447. def test_display_self_loop():
  448. ax = plt.axes()
  449. G = nx.DiGraph()
  450. G.add_node(0)
  451. G.add_edge(0, 0)
  452. nx.set_node_attributes(G, {0: (0, 0)}, "pos")
  453. nx.display(G, canvas=ax)
  454. arrow = [
  455. f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
  456. ][0]
  457. bbox = arrow.get_extents()
  458. print(bbox.width)
  459. print(bbox.height)
  460. assert bbox.width > 0 and bbox.height > 0
  461. plt.delaxes(ax)
  462. plt.close()
  463. def test_display_remove_pos_attr():
  464. """
  465. If the pos attribute isn't provided or is a function, display computes the layout
  466. and adds it to the graph. We need to ensure that this new attribute is removed from
  467. the returned graph.
  468. """
  469. G = nx.karate_club_graph()
  470. nx.display(G)
  471. assert nx.get_node_attributes(G, "display's position attribute name") == {}
  472. @pytest.fixture
  473. def subplots():
  474. fig, ax = plt.subplots()
  475. yield fig, ax
  476. plt.delaxes(ax)
  477. plt.close()
  478. @pytest.mark.parametrize(
  479. "function",
  480. [
  481. nx.draw_circular,
  482. nx.draw_kamada_kawai,
  483. nx.draw_planar,
  484. nx.draw_random,
  485. nx.draw_spectral,
  486. nx.draw_spring,
  487. nx.draw_shell,
  488. nx.draw_forceatlas2,
  489. ],
  490. )
  491. def test_draw(function, subplots, tmp_path):
  492. if function == nx.draw_kamada_kawai:
  493. pytest.importorskip("scipy", reason="draw_kamada_kawai requires scipy")
  494. fig, _ = subplots
  495. options = {"node_color": "black", "node_size": 100, "width": 3}
  496. function(barbell, **options)
  497. fig.savefig(tmp_path / "test.ps")
  498. def test_draw_shell_nlist(subplots, tmp_path):
  499. fig, _ = subplots
  500. nlist = [list(range(4)), list(range(4, 10)), list(range(10, 14))]
  501. nx.draw_shell(barbell, nlist=nlist)
  502. fig.savefig(tmp_path / "test.ps")
  503. def test_draw_bipartite(subplots, tmp_path):
  504. fig, _ = subplots
  505. G = nx.complete_bipartite_graph(2, 5)
  506. nx.draw_bipartite(G)
  507. fig.savefig(tmp_path / "test.ps")
  508. def test_edge_colormap():
  509. colors = range(barbell.number_of_edges())
  510. nx.draw_spring(
  511. barbell, edge_color=colors, width=4, edge_cmap=plt.cm.Blues, with_labels=True
  512. )
  513. # plt.show()
  514. def test_draw_networkx_edge_labels(subplots, tmp_path):
  515. fig, _ = subplots
  516. edge = (0, 1)
  517. G = nx.DiGraph([edge])
  518. pos = {n: (n, n) for n in G}
  519. nx.draw(G, pos=pos)
  520. nx.draw_networkx_edge_labels(G, pos, edge_labels={edge: "edge"})
  521. fig.savefig(tmp_path / "test.ps")
  522. def test_arrows():
  523. nx.draw_spring(barbell.to_directed())
  524. # plt.show()
  525. @pytest.mark.parametrize(
  526. ("edge_color", "expected"),
  527. (
  528. (None, "black"), # Default
  529. ("r", "red"), # Non-default color string
  530. (["r"], "red"), # Single non-default color in a list
  531. ((1.0, 1.0, 0.0), "yellow"), # single color as rgb tuple
  532. ([(1.0, 1.0, 0.0)], "yellow"), # single color as rgb tuple in list
  533. ((0, 1, 0, 1), "lime"), # single color as rgba tuple
  534. ([(0, 1, 0, 1)], "lime"), # single color as rgba tuple in list
  535. ("#0000ff", "blue"), # single color hex code
  536. (["#0000ff"], "blue"), # hex code in list
  537. ),
  538. )
  539. @pytest.mark.parametrize("edgelist", (None, [(0, 1)]))
  540. def test_single_edge_color_undirected(edge_color, expected, edgelist):
  541. """Tests ways of specifying all edges have a single color for edges
  542. drawn with a LineCollection"""
  543. G = nx.path_graph(3)
  544. drawn_edges = nx.draw_networkx_edges(
  545. G, pos=nx.random_layout(G), edgelist=edgelist, edge_color=edge_color
  546. )
  547. assert mpl.colors.same_color(drawn_edges.get_color(), expected)
  548. @pytest.mark.parametrize(
  549. ("edge_color", "expected"),
  550. (
  551. (None, "black"), # Default
  552. ("r", "red"), # Non-default color string
  553. (["r"], "red"), # Single non-default color in a list
  554. ((1.0, 1.0, 0.0), "yellow"), # single color as rgb tuple
  555. ([(1.0, 1.0, 0.0)], "yellow"), # single color as rgb tuple in list
  556. ((0, 1, 0, 1), "lime"), # single color as rgba tuple
  557. ([(0, 1, 0, 1)], "lime"), # single color as rgba tuple in list
  558. ("#0000ff", "blue"), # single color hex code
  559. (["#0000ff"], "blue"), # hex code in list
  560. ),
  561. )
  562. @pytest.mark.parametrize("edgelist", (None, [(0, 1)]))
  563. def test_single_edge_color_directed(edge_color, expected, edgelist):
  564. """Tests ways of specifying all edges have a single color for edges drawn
  565. with FancyArrowPatches"""
  566. G = nx.path_graph(3, create_using=nx.DiGraph)
  567. drawn_edges = nx.draw_networkx_edges(
  568. G, pos=nx.random_layout(G), edgelist=edgelist, edge_color=edge_color
  569. )
  570. for fap in drawn_edges:
  571. assert mpl.colors.same_color(fap.get_edgecolor(), expected)
  572. def test_edge_color_tuple_interpretation():
  573. """If edge_color is a sequence with the same length as edgelist, then each
  574. value in edge_color is mapped onto each edge via colormap."""
  575. G = nx.path_graph(6, create_using=nx.DiGraph)
  576. pos = {n: (n, n) for n in range(len(G))}
  577. # num edges != 3 or 4 --> edge_color interpreted as rgb(a)
  578. for ec in ((0, 0, 1), (0, 0, 1, 1)):
  579. # More than 4 edges
  580. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=ec)
  581. for fap in drawn_edges:
  582. assert mpl.colors.same_color(fap.get_edgecolor(), ec)
  583. # Fewer than 3 edges
  584. drawn_edges = nx.draw_networkx_edges(
  585. G, pos, edgelist=[(0, 1), (1, 2)], edge_color=ec
  586. )
  587. for fap in drawn_edges:
  588. assert mpl.colors.same_color(fap.get_edgecolor(), ec)
  589. # num edges == 3, len(edge_color) == 4: interpreted as rgba
  590. drawn_edges = nx.draw_networkx_edges(
  591. G, pos, edgelist=[(0, 1), (1, 2), (2, 3)], edge_color=(0, 0, 1, 1)
  592. )
  593. for fap in drawn_edges:
  594. assert mpl.colors.same_color(fap.get_edgecolor(), "blue")
  595. # num edges == 4, len(edge_color) == 3: interpreted as rgb
  596. drawn_edges = nx.draw_networkx_edges(
  597. G, pos, edgelist=[(0, 1), (1, 2), (2, 3), (3, 4)], edge_color=(0, 0, 1)
  598. )
  599. for fap in drawn_edges:
  600. assert mpl.colors.same_color(fap.get_edgecolor(), "blue")
  601. # num edges == len(edge_color) == 3: interpreted with cmap, *not* as rgb
  602. drawn_edges = nx.draw_networkx_edges(
  603. G, pos, edgelist=[(0, 1), (1, 2), (2, 3)], edge_color=(0, 0, 1)
  604. )
  605. assert mpl.colors.same_color(
  606. drawn_edges[0].get_edgecolor(), drawn_edges[1].get_edgecolor()
  607. )
  608. for fap in drawn_edges:
  609. assert not mpl.colors.same_color(fap.get_edgecolor(), "blue")
  610. # num edges == len(edge_color) == 4: interpreted with cmap, *not* as rgba
  611. drawn_edges = nx.draw_networkx_edges(
  612. G, pos, edgelist=[(0, 1), (1, 2), (2, 3), (3, 4)], edge_color=(0, 0, 1, 1)
  613. )
  614. assert mpl.colors.same_color(
  615. drawn_edges[0].get_edgecolor(), drawn_edges[1].get_edgecolor()
  616. )
  617. assert mpl.colors.same_color(
  618. drawn_edges[2].get_edgecolor(), drawn_edges[3].get_edgecolor()
  619. )
  620. for fap in drawn_edges:
  621. assert not mpl.colors.same_color(fap.get_edgecolor(), "blue")
  622. def test_fewer_edge_colors_than_num_edges_directed():
  623. """Test that the edge colors are cycled when there are fewer specified
  624. colors than edges."""
  625. G = barbell.to_directed()
  626. pos = nx.random_layout(barbell)
  627. edgecolors = ("r", "g", "b")
  628. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=edgecolors)
  629. for fap, expected in zip(drawn_edges, itertools.cycle(edgecolors)):
  630. assert mpl.colors.same_color(fap.get_edgecolor(), expected)
  631. def test_more_edge_colors_than_num_edges_directed():
  632. """Test that extra edge colors are ignored when there are more specified
  633. colors than edges."""
  634. G = nx.path_graph(4, create_using=nx.DiGraph) # 3 edges
  635. pos = nx.random_layout(barbell)
  636. edgecolors = ("r", "g", "b", "c") # 4 edge colors
  637. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=edgecolors)
  638. for fap, expected in zip(drawn_edges, edgecolors[:-1]):
  639. assert mpl.colors.same_color(fap.get_edgecolor(), expected)
  640. def test_edge_color_string_with_global_alpha_undirected():
  641. edge_collection = nx.draw_networkx_edges(
  642. barbell,
  643. pos=nx.random_layout(barbell),
  644. edgelist=[(0, 1), (1, 2)],
  645. edge_color="purple",
  646. alpha=0.2,
  647. )
  648. ec = edge_collection.get_color().squeeze() # as rgba tuple
  649. assert len(edge_collection.get_paths()) == 2
  650. assert mpl.colors.same_color(ec[:-1], "purple")
  651. assert ec[-1] == 0.2
  652. def test_edge_color_string_with_global_alpha_directed():
  653. drawn_edges = nx.draw_networkx_edges(
  654. barbell.to_directed(),
  655. pos=nx.random_layout(barbell),
  656. edgelist=[(0, 1), (1, 2)],
  657. edge_color="purple",
  658. alpha=0.2,
  659. )
  660. assert len(drawn_edges) == 2
  661. for fap in drawn_edges:
  662. ec = fap.get_edgecolor() # As rgba tuple
  663. assert mpl.colors.same_color(ec[:-1], "purple")
  664. assert ec[-1] == 0.2
  665. @pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
  666. def test_edge_width_default_value(graph_type):
  667. """Test the default linewidth for edges drawn either via LineCollection or
  668. FancyArrowPatches."""
  669. G = nx.path_graph(2, create_using=graph_type)
  670. pos = {n: (n, n) for n in range(len(G))}
  671. drawn_edges = nx.draw_networkx_edges(G, pos)
  672. if isinstance(drawn_edges, list): # directed case: list of FancyArrowPatch
  673. drawn_edges = drawn_edges[0]
  674. assert drawn_edges.get_linewidth() == 1
  675. @pytest.mark.parametrize(
  676. ("edgewidth", "expected"),
  677. (
  678. (3, 3), # single-value, non-default
  679. ([3], 3), # Single value as a list
  680. ),
  681. )
  682. def test_edge_width_single_value_undirected(edgewidth, expected):
  683. G = nx.path_graph(4)
  684. pos = {n: (n, n) for n in range(len(G))}
  685. drawn_edges = nx.draw_networkx_edges(G, pos, width=edgewidth)
  686. assert len(drawn_edges.get_paths()) == 3
  687. assert drawn_edges.get_linewidth() == expected
  688. @pytest.mark.parametrize(
  689. ("edgewidth", "expected"),
  690. (
  691. (3, 3), # single-value, non-default
  692. ([3], 3), # Single value as a list
  693. ),
  694. )
  695. def test_edge_width_single_value_directed(edgewidth, expected):
  696. G = nx.path_graph(4, create_using=nx.DiGraph)
  697. pos = {n: (n, n) for n in range(len(G))}
  698. drawn_edges = nx.draw_networkx_edges(G, pos, width=edgewidth)
  699. assert len(drawn_edges) == 3
  700. for fap in drawn_edges:
  701. assert fap.get_linewidth() == expected
  702. @pytest.mark.parametrize(
  703. "edgelist",
  704. (
  705. [(0, 1), (1, 2), (2, 3)], # one width specification per edge
  706. None, # fewer widths than edges - widths cycle
  707. [(0, 1), (1, 2)], # More widths than edges - unused widths ignored
  708. ),
  709. )
  710. def test_edge_width_sequence(edgelist):
  711. G = barbell.to_directed()
  712. pos = nx.random_layout(G)
  713. widths = (0.5, 2.0, 12.0)
  714. drawn_edges = nx.draw_networkx_edges(G, pos, edgelist=edgelist, width=widths)
  715. for fap, expected_width in zip(drawn_edges, itertools.cycle(widths)):
  716. assert fap.get_linewidth() == expected_width
  717. def test_edge_color_with_edge_vmin_vmax():
  718. """Test that edge_vmin and edge_vmax properly set the dynamic range of the
  719. color map when num edges == len(edge_colors)."""
  720. G = nx.path_graph(3, create_using=nx.DiGraph)
  721. pos = nx.random_layout(G)
  722. # Extract colors from the original (unscaled) colormap
  723. drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=[0, 1.0])
  724. orig_colors = [e.get_edgecolor() for e in drawn_edges]
  725. # Colors from scaled colormap
  726. drawn_edges = nx.draw_networkx_edges(
  727. G, pos, edge_color=[0.2, 0.8], edge_vmin=0.2, edge_vmax=0.8
  728. )
  729. scaled_colors = [e.get_edgecolor() for e in drawn_edges]
  730. assert mpl.colors.same_color(orig_colors, scaled_colors)
  731. def test_directed_edges_linestyle_default():
  732. """Test default linestyle for edges drawn with FancyArrowPatches."""
  733. G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
  734. pos = {n: (n, n) for n in range(len(G))}
  735. # edge with default style
  736. drawn_edges = nx.draw_networkx_edges(G, pos)
  737. assert len(drawn_edges) == 3
  738. for fap in drawn_edges:
  739. assert fap.get_linestyle() == "solid"
  740. @pytest.mark.parametrize(
  741. "style",
  742. (
  743. "dashed", # edge with string style
  744. "--", # edge with simplified string style
  745. (1, (1, 1)), # edge with (offset, onoffseq) style
  746. ),
  747. )
  748. def test_directed_edges_linestyle_single_value(style):
  749. """Tests support for specifying linestyles with a single value to be applied to
  750. all edges in ``draw_networkx_edges`` for FancyArrowPatch outputs
  751. (e.g. directed edges)."""
  752. G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
  753. pos = {n: (n, n) for n in range(len(G))}
  754. drawn_edges = nx.draw_networkx_edges(G, pos, style=style)
  755. assert len(drawn_edges) == 3
  756. for fap in drawn_edges:
  757. assert fap.get_linestyle() == style
  758. @pytest.mark.parametrize(
  759. "style_seq",
  760. (
  761. ["dashed"], # edge with string style in list
  762. ["--"], # edge with simplified string style in list
  763. [(1, (1, 1))], # edge with (offset, onoffseq) style in list
  764. ["--", "-", ":"], # edges with styles for each edge
  765. ["--", "-"], # edges with fewer styles than edges (styles cycle)
  766. ["--", "-", ":", "-."], # edges with more styles than edges (extra unused)
  767. ),
  768. )
  769. def test_directed_edges_linestyle_sequence(style_seq):
  770. """Tests support for specifying linestyles with sequences in
  771. ``draw_networkx_edges`` for FancyArrowPatch outputs (e.g. directed edges)."""
  772. G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
  773. pos = {n: (n, n) for n in range(len(G))}
  774. drawn_edges = nx.draw_networkx_edges(G, pos, style=style_seq)
  775. assert len(drawn_edges) == 3
  776. for fap, style in zip(drawn_edges, itertools.cycle(style_seq)):
  777. assert fap.get_linestyle() == style
  778. def test_return_types():
  779. from matplotlib.collections import LineCollection, PathCollection
  780. from matplotlib.patches import FancyArrowPatch
  781. G = nx.frucht_graph(create_using=nx.Graph)
  782. dG = nx.frucht_graph(create_using=nx.DiGraph)
  783. pos = nx.spring_layout(G, seed=42)
  784. dpos = nx.spring_layout(dG, seed=42)
  785. # nodes
  786. nodes = nx.draw_networkx_nodes(G, pos)
  787. assert isinstance(nodes, PathCollection)
  788. # edges
  789. edges = nx.draw_networkx_edges(dG, dpos, arrows=True)
  790. assert isinstance(edges, list)
  791. if len(edges) > 0:
  792. assert isinstance(edges[0], FancyArrowPatch)
  793. edges = nx.draw_networkx_edges(dG, dpos, arrows=False)
  794. assert isinstance(edges, LineCollection)
  795. edges = nx.draw_networkx_edges(G, dpos, arrows=None)
  796. assert isinstance(edges, LineCollection)
  797. edges = nx.draw_networkx_edges(dG, pos, arrows=None)
  798. assert isinstance(edges, list)
  799. if len(edges) > 0:
  800. assert isinstance(edges[0], FancyArrowPatch)
  801. def test_labels_and_colors():
  802. G = nx.cubical_graph()
  803. pos = nx.spring_layout(G, seed=42) # positions for all nodes
  804. # nodes
  805. nx.draw_networkx_nodes(
  806. G, pos, nodelist=[0, 1, 2, 3], node_color="r", node_size=500, alpha=0.75
  807. )
  808. nx.draw_networkx_nodes(
  809. G,
  810. pos,
  811. nodelist=[4, 5, 6, 7],
  812. node_color="b",
  813. node_size=500,
  814. alpha=[0.25, 0.5, 0.75, 1.0],
  815. )
  816. # edges
  817. nx.draw_networkx_edges(G, pos, width=1.0, alpha=0.5)
  818. nx.draw_networkx_edges(
  819. G,
  820. pos,
  821. edgelist=[(0, 1), (1, 2), (2, 3), (3, 0)],
  822. width=8,
  823. alpha=0.5,
  824. edge_color="r",
  825. )
  826. nx.draw_networkx_edges(
  827. G,
  828. pos,
  829. edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
  830. width=8,
  831. alpha=0.5,
  832. edge_color="b",
  833. )
  834. nx.draw_networkx_edges(
  835. G,
  836. pos,
  837. edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
  838. arrows=True,
  839. min_source_margin=0.5,
  840. min_target_margin=0.75,
  841. width=8,
  842. edge_color="b",
  843. )
  844. # some math labels
  845. labels = {}
  846. labels[0] = r"$a$"
  847. labels[1] = r"$b$"
  848. labels[2] = r"$c$"
  849. labels[3] = r"$d$"
  850. labels[4] = r"$\alpha$"
  851. labels[5] = r"$\beta$"
  852. labels[6] = r"$\gamma$"
  853. labels[7] = r"$\delta$"
  854. colors = {n: "k" if n % 2 == 0 else "r" for n in range(8)}
  855. nx.draw_networkx_labels(G, pos, labels, font_size=16)
  856. nx.draw_networkx_labels(G, pos, labels, font_size=16, font_color=colors)
  857. nx.draw_networkx_edge_labels(G, pos, edge_labels=None, rotate=False)
  858. nx.draw_networkx_edge_labels(G, pos, edge_labels={(4, 5): "4-5"})
  859. # plt.show()
  860. def test_axes(subplots):
  861. fig, ax = subplots
  862. nx.draw(barbell, ax=ax)
  863. nx.draw_networkx_edge_labels(barbell, nx.circular_layout(barbell), ax=ax)
  864. def test_empty_graph():
  865. G = nx.Graph()
  866. nx.draw(G)
  867. def test_draw_empty_nodes_return_values():
  868. # See Issue #3833
  869. import matplotlib.collections # call as mpl.collections
  870. G = nx.Graph([(1, 2), (2, 3)])
  871. DG = nx.DiGraph([(1, 2), (2, 3)])
  872. pos = nx.circular_layout(G)
  873. assert isinstance(
  874. nx.draw_networkx_nodes(G, pos, nodelist=[]), mpl.collections.PathCollection
  875. )
  876. assert isinstance(
  877. nx.draw_networkx_nodes(DG, pos, nodelist=[]), mpl.collections.PathCollection
  878. )
  879. # drawing empty edges used to return an empty LineCollection or empty list.
  880. # Now it is always an empty list (because edges are now lists of FancyArrows)
  881. assert nx.draw_networkx_edges(G, pos, edgelist=[], arrows=True) == []
  882. assert nx.draw_networkx_edges(G, pos, edgelist=[], arrows=False) == []
  883. assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=False) == []
  884. assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=True) == []
  885. def test_multigraph_edgelist_tuples():
  886. # See Issue #3295
  887. G = nx.path_graph(3, create_using=nx.MultiDiGraph)
  888. nx.draw_networkx(G, edgelist=[(0, 1, 0)])
  889. nx.draw_networkx(G, edgelist=[(0, 1, 0)], node_size=[10, 20, 0])
  890. def test_alpha_iter():
  891. pos = nx.random_layout(barbell)
  892. fig = plt.figure()
  893. # with fewer alpha elements than nodes
  894. fig.add_subplot(131) # Each test in a new axis object
  895. nx.draw_networkx_nodes(barbell, pos, alpha=[0.1, 0.2])
  896. # with equal alpha elements and nodes
  897. num_nodes = len(barbell.nodes)
  898. alpha = [x / num_nodes for x in range(num_nodes)]
  899. colors = range(num_nodes)
  900. fig.add_subplot(132)
  901. nx.draw_networkx_nodes(barbell, pos, node_color=colors, alpha=alpha)
  902. # with more alpha elements than nodes
  903. alpha.append(1)
  904. fig.add_subplot(133)
  905. nx.draw_networkx_nodes(barbell, pos, alpha=alpha)
  906. def test_multiple_node_shapes(subplots):
  907. fig, ax = subplots
  908. G = nx.path_graph(4)
  909. nx.draw(G, node_shape=["o", "h", "s", "^"], ax=ax)
  910. scatters = [
  911. s for s in ax.get_children() if isinstance(s, mpl.collections.PathCollection)
  912. ]
  913. assert len(scatters) == 4
  914. def test_individualized_font_attributes(subplots):
  915. G = nx.karate_club_graph()
  916. fig, ax = subplots
  917. nx.draw(
  918. G,
  919. ax=ax,
  920. font_color={n: "k" if n % 2 else "r" for n in G.nodes()},
  921. font_size={n: int(n / (34 / 15) + 5) for n in G.nodes()},
  922. )
  923. for n, t in zip(
  924. G.nodes(),
  925. [
  926. t
  927. for t in ax.get_children()
  928. if isinstance(t, mpl.text.Text) and len(t.get_text()) > 0
  929. ],
  930. ):
  931. expected = "black" if n % 2 else "red"
  932. assert mpl.colors.same_color(t.get_color(), expected)
  933. assert int(n / (34 / 15) + 5) == t.get_size()
  934. def test_individualized_edge_attributes(subplots):
  935. G = nx.karate_club_graph()
  936. fig, ax = subplots
  937. arrowstyles = ["-|>" if (u + v) % 2 == 0 else "-[" for u, v in G.edges()]
  938. arrowsizes = [10 * (u % 2 + v % 2) + 10 for u, v in G.edges()]
  939. nx.draw(G, ax=ax, arrows=True, arrowstyle=arrowstyles, arrowsize=arrowsizes)
  940. arrows = [
  941. f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
  942. ]
  943. for e, a in zip(G.edges(), arrows):
  944. assert a.get_mutation_scale() == 10 * (e[0] % 2 + e[1] % 2) + 10
  945. expected = (
  946. mpl.patches.ArrowStyle.BracketB
  947. if sum(e) % 2
  948. else mpl.patches.ArrowStyle.CurveFilledB
  949. )
  950. assert isinstance(a.get_arrowstyle(), expected)
  951. def test_error_invalid_kwds():
  952. with pytest.raises(ValueError, match="Received invalid argument"):
  953. nx.draw(barbell, foo="bar")
  954. def test_draw_networkx_arrowsize_incorrect_size():
  955. G = nx.DiGraph([(0, 1), (0, 2), (0, 3), (1, 3)])
  956. arrowsize = [1, 2, 3]
  957. with pytest.raises(
  958. ValueError, match="arrowsize should have the same length as edgelist"
  959. ):
  960. nx.draw(G, arrowsize=arrowsize)
  961. @pytest.mark.parametrize("arrowsize", (30, [10, 20, 30]))
  962. def test_draw_edges_arrowsize(arrowsize):
  963. G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
  964. pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
  965. edges = nx.draw_networkx_edges(G, pos=pos, arrowsize=arrowsize)
  966. arrowsize = itertools.repeat(arrowsize) if isinstance(arrowsize, int) else arrowsize
  967. for fap, expected in zip(edges, arrowsize):
  968. assert isinstance(fap, mpl.patches.FancyArrowPatch)
  969. assert fap.get_mutation_scale() == expected
  970. @pytest.mark.parametrize("arrowstyle", ("-|>", ["-|>", "-[", "<|-|>"]))
  971. def test_draw_edges_arrowstyle(arrowstyle):
  972. G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
  973. pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
  974. edges = nx.draw_networkx_edges(G, pos=pos, arrowstyle=arrowstyle)
  975. arrowstyle = (
  976. itertools.repeat(arrowstyle) if isinstance(arrowstyle, str) else arrowstyle
  977. )
  978. arrow_objects = {
  979. "-|>": mpl.patches.ArrowStyle.CurveFilledB,
  980. "-[": mpl.patches.ArrowStyle.BracketB,
  981. "<|-|>": mpl.patches.ArrowStyle.CurveFilledAB,
  982. }
  983. for fap, expected in zip(edges, arrowstyle):
  984. assert isinstance(fap, mpl.patches.FancyArrowPatch)
  985. assert isinstance(fap.get_arrowstyle(), arrow_objects[expected])
  986. def test_np_edgelist():
  987. # see issue #4129
  988. nx.draw_networkx(barbell, edgelist=np.array([(0, 2), (0, 3)]))
  989. def test_draw_nodes_missing_node_from_position():
  990. G = nx.path_graph(3)
  991. pos = {0: (0, 0), 1: (1, 1)} # No position for node 2
  992. with pytest.raises(nx.NetworkXError, match="has no position"):
  993. nx.draw_networkx_nodes(G, pos)
  994. def test_draw_networkx_nodes_node_shape_list_with_scalar_color(subplots):
  995. """Ensure draw_networkx_nodes works when node_shape is a Python list.
  996. This covers the case where node_shape is a sequence (list) and node_color
  997. is a single scalar color, which should be supported.
  998. """
  999. fig, ax = subplots
  1000. G = nx.empty_graph(5)
  1001. pos = {i: (i, i) for i in G}
  1002. shapes = ["o", "^", "o", "^", "o"]
  1003. nodes = nx.draw_networkx_nodes(
  1004. G,
  1005. pos,
  1006. node_color="red", # scalar color (supported)
  1007. node_shape=shapes, # list of shapes – this used to be buggy
  1008. ax=ax,
  1009. )
  1010. # Should get a PathCollection with an element in it (same as with numpy arrays)
  1011. assert len(nodes.get_offsets()) > 0
  1012. # NOTE: When node_shape is a sequence, draw_networkx_nodes internally calls
  1013. # ax.scatter multiple times and returns only the last PathCollection.
  1014. # Therefore, we do NOT assert a value for len(nodes.get_offsets()) here.
  1015. # NOTE: parametrizing on marker to test both branches of internal
  1016. # nx.draw_networkx_edges.to_marker_edge function
  1017. @pytest.mark.parametrize("node_shape", ("o", "s"))
  1018. def test_draw_edges_min_source_target_margins(node_shape, subplots):
  1019. """Test that there is a wider gap between the node and the start of an
  1020. incident edge when min_source_margin is specified.
  1021. This test checks that the use of min_{source/target}_margin kwargs result
  1022. in shorter (more padding) between the edges and source and target nodes.
  1023. As a crude visual example, let 's' and 't' represent source and target
  1024. nodes, respectively:
  1025. Default:
  1026. s-----------------------------t
  1027. With margins:
  1028. s ----------------------- t
  1029. """
  1030. # Create a single axis object to get consistent pixel coords across
  1031. # multiple draws
  1032. fig, ax = subplots
  1033. G = nx.DiGraph([(0, 1)])
  1034. pos = {0: (0, 0), 1: (1, 0)} # horizontal layout
  1035. # Get leftmost and rightmost points of the FancyArrowPatch object
  1036. # representing the edge between nodes 0 and 1 (in pixel coordinates)
  1037. default_patch = nx.draw_networkx_edges(G, pos, ax=ax, node_shape=node_shape)[0]
  1038. default_extent = default_patch.get_extents().corners()[::2, 0]
  1039. # Now, do the same but with "padding" for the source and target via the
  1040. # min_{source/target}_margin kwargs
  1041. padded_patch = nx.draw_networkx_edges(
  1042. G,
  1043. pos,
  1044. ax=ax,
  1045. node_shape=node_shape,
  1046. min_source_margin=100,
  1047. min_target_margin=100,
  1048. )[0]
  1049. padded_extent = padded_patch.get_extents().corners()[::2, 0]
  1050. # With padding, the left-most extent of the edge should be further to the
  1051. # right
  1052. assert padded_extent[0] > default_extent[0]
  1053. # And the rightmost extent of the edge, further to the left
  1054. assert padded_extent[1] < default_extent[1]
  1055. # NOTE: parametrizing on marker to test both branches of internal
  1056. # nx.draw_networkx_edges.to_marker_edge function
  1057. @pytest.mark.parametrize("node_shape", ("o", "s"))
  1058. def test_draw_edges_min_source_target_margins_individual(node_shape, subplots):
  1059. """Test that there is a wider gap between the node and the start of an
  1060. incident edge when min_source_margin is specified.
  1061. This test checks that the use of min_{source/target}_margin kwargs result
  1062. in shorter (more padding) between the edges and source and target nodes.
  1063. As a crude visual example, let 's' and 't' represent source and target
  1064. nodes, respectively:
  1065. Default:
  1066. s-----------------------------t
  1067. With margins:
  1068. s ----------------------- t
  1069. """
  1070. # Create a single axis object to get consistent pixel coords across
  1071. # multiple draws
  1072. fig, ax = subplots
  1073. G = nx.DiGraph([(0, 1), (1, 2)])
  1074. pos = {0: (0, 0), 1: (1, 0), 2: (2, 0)} # horizontal layout
  1075. # Get leftmost and rightmost points of the FancyArrowPatch object
  1076. # representing the edge between nodes 0 and 1 (in pixel coordinates)
  1077. default_patch = nx.draw_networkx_edges(G, pos, ax=ax, node_shape=node_shape)
  1078. default_extent = [d.get_extents().corners()[::2, 0] for d in default_patch]
  1079. # Now, do the same but with "padding" for the source and target via the
  1080. # min_{source/target}_margin kwargs
  1081. padded_patch = nx.draw_networkx_edges(
  1082. G,
  1083. pos,
  1084. ax=ax,
  1085. node_shape=node_shape,
  1086. min_source_margin=[98, 102],
  1087. min_target_margin=[98, 102],
  1088. )
  1089. padded_extent = [p.get_extents().corners()[::2, 0] for p in padded_patch]
  1090. for d, p in zip(default_extent, padded_extent):
  1091. # With padding, the left-most extent of the edge should be further to the
  1092. # right
  1093. assert p[0] > d[0]
  1094. # And the rightmost extent of the edge, further to the left
  1095. assert p[1] < d[1]
  1096. def test_nonzero_selfloop_with_single_node(subplots):
  1097. """Ensure that selfloop extent is non-zero when there is only one node."""
  1098. # Create explicit axis object for test
  1099. fig, ax = subplots
  1100. # Graph with single node + self loop
  1101. G = nx.DiGraph()
  1102. G.add_node(0)
  1103. G.add_edge(0, 0)
  1104. # Draw
  1105. patch = nx.draw_networkx_edges(G, {0: (0, 0)})[0]
  1106. # The resulting patch must have non-zero extent
  1107. bbox = patch.get_extents()
  1108. assert bbox.width > 0 and bbox.height > 0
  1109. def test_nonzero_selfloop_with_single_edge_in_edgelist(subplots):
  1110. """Ensure that selfloop extent is non-zero when only a single edge is
  1111. specified in the edgelist.
  1112. """
  1113. # Create explicit axis object for test
  1114. fig, ax = subplots
  1115. # Graph with selfloop
  1116. G = nx.path_graph(2, create_using=nx.DiGraph)
  1117. G.add_edge(1, 1)
  1118. pos = {n: (n, n) for n in G.nodes}
  1119. # Draw only the selfloop edge via the `edgelist` kwarg
  1120. patch = nx.draw_networkx_edges(G, pos, edgelist=[(1, 1)])[0]
  1121. # The resulting patch must have non-zero extent
  1122. bbox = patch.get_extents()
  1123. assert bbox.width > 0 and bbox.height > 0
  1124. def test_apply_alpha():
  1125. """Test apply_alpha when there is a mismatch between the number of
  1126. supplied colors and elements.
  1127. """
  1128. nodelist = [0, 1, 2]
  1129. colorlist = ["r", "g", "b"]
  1130. alpha = 0.5
  1131. rgba_colors = nx.drawing.nx_pylab.apply_alpha(colorlist, alpha, nodelist)
  1132. assert all(rgba_colors[:, -1] == alpha)
  1133. def test_draw_edges_toggling_with_arrows_kwarg():
  1134. """
  1135. The `arrows` keyword argument is used as a 3-way switch to select which
  1136. type of object to use for drawing edges:
  1137. - ``arrows=None`` -> default (FancyArrowPatches for directed, else LineCollection)
  1138. - ``arrows=True`` -> FancyArrowPatches
  1139. - ``arrows=False`` -> LineCollection
  1140. """
  1141. import matplotlib.collections
  1142. import matplotlib.patches
  1143. UG = nx.path_graph(3)
  1144. DG = nx.path_graph(3, create_using=nx.DiGraph)
  1145. pos = {n: (n, n) for n in UG}
  1146. # Use FancyArrowPatches when arrows=True, regardless of graph type
  1147. for G in (UG, DG):
  1148. edges = nx.draw_networkx_edges(G, pos, arrows=True)
  1149. assert len(edges) == len(G.edges)
  1150. assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
  1151. # Use LineCollection when arrows=False, regardless of graph type
  1152. for G in (UG, DG):
  1153. edges = nx.draw_networkx_edges(G, pos, arrows=False)
  1154. assert isinstance(edges, mpl.collections.LineCollection)
  1155. # Default behavior when arrows=None: FAPs for directed, LC's for undirected
  1156. edges = nx.draw_networkx_edges(UG, pos)
  1157. assert isinstance(edges, mpl.collections.LineCollection)
  1158. edges = nx.draw_networkx_edges(DG, pos)
  1159. assert len(edges) == len(G.edges)
  1160. assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
  1161. @pytest.mark.parametrize("drawing_func", (nx.draw, nx.draw_networkx))
  1162. def test_draw_networkx_arrows_default_undirected(drawing_func, subplots):
  1163. import matplotlib.collections
  1164. G = nx.path_graph(3)
  1165. fig, ax = subplots
  1166. drawing_func(G, ax=ax)
  1167. assert any(isinstance(c, mpl.collections.LineCollection) for c in ax.collections)
  1168. assert not ax.patches
  1169. @pytest.mark.parametrize("drawing_func", (nx.draw, nx.draw_networkx))
  1170. def test_draw_networkx_arrows_default_directed(drawing_func, subplots):
  1171. import matplotlib.collections
  1172. G = nx.path_graph(3, create_using=nx.DiGraph)
  1173. fig, ax = subplots
  1174. drawing_func(G, ax=ax)
  1175. assert not any(
  1176. isinstance(c, mpl.collections.LineCollection) for c in ax.collections
  1177. )
  1178. assert ax.patches
  1179. def test_edgelist_kwarg_not_ignored(subplots):
  1180. # See gh-4994
  1181. G = nx.path_graph(3)
  1182. G.add_edge(0, 0)
  1183. fig, ax = subplots
  1184. nx.draw(G, edgelist=[(0, 1), (1, 2)], ax=ax) # Exclude self-loop from edgelist
  1185. assert not ax.patches
  1186. @pytest.mark.parametrize(
  1187. ("G", "expected_n_edges"),
  1188. ([nx.DiGraph(), 2], [nx.MultiGraph(), 4], [nx.MultiDiGraph(), 4]),
  1189. )
  1190. def test_draw_networkx_edges_multiedge_connectionstyle(G, expected_n_edges):
  1191. """Draws edges correctly for 3 types of graphs and checks for valid length"""
  1192. for i, (u, v) in enumerate([(0, 1), (0, 1), (0, 1), (0, 2)]):
  1193. G.add_edge(u, v, weight=round(i / 3, 2))
  1194. pos = {n: (n, n) for n in G}
  1195. # Raises on insufficient connectionstyle length
  1196. for conn_style in [
  1197. "arc3,rad=0.1",
  1198. ["arc3,rad=0.1", "arc3,rad=0.1"],
  1199. ["arc3,rad=0.1", "arc3,rad=0.1", "arc3,rad=0.2"],
  1200. ]:
  1201. nx.draw_networkx_edges(G, pos, connectionstyle=conn_style)
  1202. arrows = nx.draw_networkx_edges(G, pos, connectionstyle=conn_style)
  1203. assert len(arrows) == expected_n_edges
  1204. @pytest.mark.parametrize(
  1205. ("G", "expected_n_edges"),
  1206. ([nx.DiGraph(), 2], [nx.MultiGraph(), 4], [nx.MultiDiGraph(), 4]),
  1207. )
  1208. def test_draw_networkx_edge_labels_multiedge_connectionstyle(G, expected_n_edges):
  1209. """Draws labels correctly for 3 types of graphs and checks for valid length and class names"""
  1210. for i, (u, v) in enumerate([(0, 1), (0, 1), (0, 1), (0, 2)]):
  1211. G.add_edge(u, v, weight=round(i / 3, 2))
  1212. pos = {n: (n, n) for n in G}
  1213. # Raises on insufficient connectionstyle length
  1214. arrows = nx.draw_networkx_edges(
  1215. G, pos, connectionstyle=["arc3,rad=0.1", "arc3,rad=0.1", "arc3,rad=0.1"]
  1216. )
  1217. for conn_style in [
  1218. "arc3,rad=0.1",
  1219. ["arc3,rad=0.1", "arc3,rad=0.2"],
  1220. ["arc3,rad=0.1", "arc3,rad=0.1", "arc3,rad=0.1"],
  1221. ]:
  1222. text_items = nx.draw_networkx_edge_labels(G, pos, connectionstyle=conn_style)
  1223. assert len(text_items) == expected_n_edges
  1224. for ti in text_items.values():
  1225. assert ti.__class__.__name__ == "CurvedArrowText"
  1226. def test_draw_networkx_edge_label_multiedge():
  1227. G = nx.MultiGraph()
  1228. G.add_edge(0, 1, weight=10)
  1229. G.add_edge(0, 1, weight=20)
  1230. edge_labels = nx.get_edge_attributes(G, "weight") # Includes edge keys
  1231. pos = {n: (n, n) for n in G}
  1232. text_items = nx.draw_networkx_edge_labels(
  1233. G,
  1234. pos,
  1235. edge_labels=edge_labels,
  1236. connectionstyle=["arc3,rad=0.1", "arc3,rad=0.2"],
  1237. )
  1238. assert len(text_items) == 2
  1239. def test_draw_networkx_edge_label_empty_dict():
  1240. """Regression test for draw_networkx_edge_labels with empty dict. See
  1241. gh-5372."""
  1242. G = nx.path_graph(3)
  1243. pos = {n: (n, n) for n in G.nodes}
  1244. assert nx.draw_networkx_edge_labels(G, pos, edge_labels={}) == {}
  1245. def test_draw_networkx_edges_undirected_selfloop_colors(subplots):
  1246. """When an edgelist is supplied along with a sequence of colors, check that
  1247. the self-loops have the correct colors."""
  1248. fig, ax = subplots
  1249. # Edge list and corresponding colors
  1250. edgelist = [(1, 3), (1, 2), (2, 3), (1, 1), (3, 3), (2, 2)]
  1251. edge_colors = ["pink", "cyan", "black", "red", "blue", "green"]
  1252. G = nx.Graph(edgelist)
  1253. pos = {n: (n, n) for n in G.nodes}
  1254. nx.draw_networkx_edges(G, pos, ax=ax, edgelist=edgelist, edge_color=edge_colors)
  1255. # Verify that there are three fancy arrow patches (1 per self loop)
  1256. assert len(ax.patches) == 3
  1257. # These are points that should be contained in the self loops. For example,
  1258. # sl_points[0] will be (1, 1.1), which is inside the "path" of the first
  1259. # self-loop but outside the others
  1260. sl_points = np.array(edgelist[-3:]) + np.array([0, 0.1])
  1261. # Check that the mapping between self-loop locations and their colors is
  1262. # correct
  1263. for fap, clr, slp in zip(ax.patches, edge_colors[-3:], sl_points):
  1264. assert fap.get_path().contains_point(slp)
  1265. assert mpl.colors.same_color(fap.get_edgecolor(), clr)
  1266. @pytest.mark.parametrize(
  1267. "fap_only_kwarg", # Non-default values for kwargs that only apply to FAPs
  1268. (
  1269. {"arrowstyle": "-"},
  1270. {"arrowsize": 20},
  1271. {"connectionstyle": "arc3,rad=0.2"},
  1272. {"min_source_margin": 10},
  1273. {"min_target_margin": 10},
  1274. ),
  1275. )
  1276. def test_user_warnings_for_unused_edge_drawing_kwargs(fap_only_kwarg, subplots):
  1277. """Users should get a warning when they specify a non-default value for
  1278. one of the kwargs that applies only to edges drawn with FancyArrowPatches,
  1279. but FancyArrowPatches aren't being used under the hood."""
  1280. G = nx.path_graph(3)
  1281. pos = {n: (n, n) for n in G}
  1282. fig, ax = subplots
  1283. # By default, an undirected graph will use LineCollection to represent
  1284. # the edges
  1285. kwarg_name = list(fap_only_kwarg.keys())[0]
  1286. with pytest.warns(
  1287. UserWarning, match=f"\n\nThe {kwarg_name} keyword argument is not applicable"
  1288. ):
  1289. nx.draw_networkx_edges(G, pos, ax=ax, **fap_only_kwarg)
  1290. # FancyArrowPatches are always used when `arrows=True` is specified.
  1291. # Check that warnings are *not* raised in this case
  1292. with warnings.catch_warnings():
  1293. # Escalate warnings -> errors so tests fail if warnings are raised
  1294. warnings.simplefilter("error")
  1295. warnings.filterwarnings("ignore", category=DeprecationWarning)
  1296. nx.draw_networkx_edges(G, pos, ax=ax, arrows=True, **fap_only_kwarg)
  1297. @pytest.mark.parametrize("draw_fn", (nx.draw, nx.draw_circular))
  1298. def test_no_warning_on_default_draw_arrowstyle(draw_fn, subplots):
  1299. # See gh-7284
  1300. fig, ax = subplots
  1301. G = nx.cycle_graph(5)
  1302. with warnings.catch_warnings(record=True) as w:
  1303. draw_fn(G, ax=ax)
  1304. assert len(w) == 0
  1305. @pytest.mark.parametrize("hide_ticks", [False, True])
  1306. @pytest.mark.parametrize(
  1307. "method",
  1308. [
  1309. nx.draw_networkx,
  1310. nx.draw_networkx_edge_labels,
  1311. nx.draw_networkx_edges,
  1312. nx.draw_networkx_labels,
  1313. nx.draw_networkx_nodes,
  1314. ],
  1315. )
  1316. def test_hide_ticks(method, hide_ticks, subplots):
  1317. G = nx.path_graph(3)
  1318. pos = {n: (n, n) for n in G.nodes}
  1319. _, ax = subplots
  1320. method(G, pos=pos, ax=ax, hide_ticks=hide_ticks)
  1321. for axis in [ax.xaxis, ax.yaxis]:
  1322. assert bool(axis.get_ticklabels()) != hide_ticks
  1323. @pytest.mark.parametrize(
  1324. "style", ["angle", "angle3", "arc", "arc3,rad=0.0", "bar,fraction=0.1"]
  1325. )
  1326. def test_edge_label_all_connectionstyles(subplots, style):
  1327. """
  1328. Check that FancyArrowPatches with all `connectionstyle`s are supported
  1329. in edge label rendering. See gh-7735 and gh-8106.
  1330. """
  1331. fig, ax = subplots
  1332. edge = (0, 1)
  1333. G = nx.DiGraph([edge])
  1334. pos = {n: (n, 0) for n in G}
  1335. name = style.split(",")[0]
  1336. labels = nx.draw_networkx_edge_labels(
  1337. G, pos, edge_labels={edge: "edge"}, connectionstyle=style
  1338. )
  1339. hmid = (pos[0][0] + pos[1][0]) / 2
  1340. vmid = (pos[0][1] + pos[1][1]) / 2
  1341. if name in {"arc", "arc3"}: # The label should be at roughly the midpoint.
  1342. assert labels[edge].x, labels[edge].y == pytest.approx((hmid, vmid))
  1343. elif name == "bar": # The label should be below the vertical midpoint.
  1344. assert labels[edge].y < vmid
  1345. @pytest.mark.parametrize("label_pos", [-0.1, 1.1])
  1346. def test_edge_label_label_pos(subplots, label_pos):
  1347. """
  1348. Check that label positions can be extrapolated outside [0, 1].
  1349. """
  1350. fig, ax = subplots
  1351. edge = (0, 1)
  1352. G = nx.DiGraph([edge])
  1353. pos = {n: (n, n) for n in G}
  1354. lbl = nx.draw_networkx_edge_labels(
  1355. G, pos, edge_labels={edge: "edge"}, label_pos=label_pos, connectionstyle="angle"
  1356. )
  1357. assert lbl[edge].x, lbl[edge].y == pytest.approx((label_pos, label_pos))