team-10/env/Lib/site-packages/networkx/drawing/tests/test_pylab.py
2025-08-02 07:34:44 +02:00

1756 lines
58 KiB
Python

"""Unit tests for matplotlib drawing functions."""
import itertools
import os
import warnings
import pytest
import networkx as nx
mpl = pytest.importorskip("matplotlib")
np = pytest.importorskip("numpy")
mpl.use("PS")
plt = pytest.importorskip("matplotlib.pyplot")
plt.rcParams["text.usetex"] = False
barbell = nx.barbell_graph(4, 6)
defaults = {
"node_pos": None,
"node_visible": True,
"node_color": "#1f78b4",
"node_size": 300,
"node_label": {
"size": 12,
"color": "#000000",
"family": "sans-serif",
"weight": "normal",
"alpha": 1.0,
"background_color": None,
"background_alpha": None,
"h_align": "center",
"v_align": "center",
"bbox": None,
},
"node_shape": "o",
"node_alpha": 1.0,
"node_border_width": 1.0,
"node_border_color": "face",
"edge_visible": True,
"edge_width": 1.0,
"edge_color": "#000000",
"edge_label": {
"size": 12,
"color": "#000000",
"family": "sans-serif",
"weight": "normal",
"alpha": 1.0,
"bbox": {"boxstyle": "round", "ec": (1.0, 1.0, 1.0), "fc": (1.0, 1.0, 1.0)},
"h_align": "center",
"v_align": "center",
"pos": 0.5,
"rotate": True,
},
"edge_style": "-",
"edge_alpha": 1.0,
# These are for undirected-graphs. Directed graphs shouls use "-|>" and 10, respectively
"edge_arrowstyle": "-",
"edge_arrowsize": 0,
"edge_curvature": "arc3",
"edge_source_margin": 0,
"edge_target_margin": 0,
}
@pytest.mark.parametrize(
("param_name", "param_value", "expected"),
(
("node_color", None, defaults["node_color"]),
("node_color", "#FF0000", "red"),
("node_color", "color", "lime"),
),
)
def test_display_arg_handling_node_color(param_name, param_value, expected):
G = nx.path_graph(4)
nx.set_node_attributes(G, "#00FF00", "color")
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas, **{param_name: param_value})
assert mpl.colors.same_color(canvas.get_children()[0].get_edgecolors()[0], expected)
plt.close()
@pytest.mark.parametrize(
("param_value", "expected"),
(
(None, (1, 1, 1, 1)), # default value
(0.5, (0.5, 0.5, 0.5, 0.5)),
("n_alpha", (1.0, 1 / 2, 1 / 3, 0.25)),
),
)
def test_display_arg_handling_node_alpha(param_value, expected):
G = nx.path_graph(4)
nx.set_node_attributes(G, {n: 1 / (n + 1) for n in G.nodes()}, "n_alpha")
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas, node_alpha=param_value)
assert all(
canvas.get_children()[0].get_fc()[:, 3] == expected
) # Extract just the alpha from the node colors
plt.close()
def test_display_node_position():
G = nx.path_graph(4)
nx.set_node_attributes(G, {n: (n, n) for n in G.nodes()}, "pos")
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas, node_pos="pos")
assert np.all(
canvas.get_children()[0].get_offsets().data == [[0, 0], [1, 1], [2, 2], [3, 3]]
)
plt.close()
@pytest.mark.mpl_image_compare
def test_display_house_with_colors():
"""
Originally, I wanted to use the exact samge image as test_house_with_colors.
But I can't seem to find the correct value for the margins to get the figures
to line up perfectly. To the human eye, these visualizations are basically the
same.
"""
G = nx.house_graph()
fig, ax = plt.subplots()
nx.set_node_attributes(
G, {0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1), 4: (0.5, 2.0)}, "pos"
)
nx.set_node_attributes(
G,
{
n: {
"size": 3000 if n != 4 else 2000,
"color": "tab:blue" if n != 4 else "tab:orange",
}
for n in G.nodes()
},
)
nx.display(
G,
node_pos="pos",
edge_alpha=0.5,
edge_width=6,
node_label=None,
node_border_color="k",
)
ax.margins(0.17)
plt.tight_layout()
plt.axis("off")
return fig
def test_display_line_collection():
G = nx.karate_club_graph()
nx.set_edge_attributes(
G, {(u, v): "-|>" if (u + v) % 2 else "-" for u, v in G.edges()}, "arrowstyle"
)
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas, edge_arrowsize=10)
# There should only be one line collection in any given visualization
lc = [
l
for l in canvas.get_children()
if isinstance(l, mpl.collections.LineCollection)
][0]
assert len(lc.get_paths()) == sum([1 for u, v in G.edges() if (u + v) % 2])
plt.close()
@pytest.mark.mpl_image_compare
def test_display_labels_and_colors():
"""See 'Labels and Colors' gallery example"""
fig, ax = plt.subplots()
G = nx.cubical_graph()
pos = nx.spring_layout(G, seed=3113794652) # positions for all nodes
nx.set_node_attributes(G, pos, "pos") # Will not be needed after PR 7571
labels = iter(
[
r"$a$",
r"$b$",
r"$c$",
r"$d$",
r"$\alpha$",
r"$\beta$",
r"$\gamma$",
r"$\delta$",
]
)
nx.set_node_attributes(
G,
{
n: {
"size": 800,
"alpha": 0.9,
"color": "tab:red" if n < 4 else "tab:blue",
"label": {"label": next(labels), "size": 22, "color": "whitesmoke"},
}
for n in G.nodes()
},
)
nx.display(G, node_pos="pos", edge_color="tab:grey")
# The tricky bit is the highlighted colors for the edges
edgelist = [(0, 1), (1, 2), (2, 3), (0, 3)]
nx.set_edge_attributes(
G,
{
(u, v): {
"width": 8,
"alpha": 0.5,
"color": "tab:red",
"visible": (u, v) in edgelist,
}
for u, v in G.edges()
},
)
nx.display(G, node_pos="pos", node_visible=False)
edgelist = [(4, 5), (5, 6), (6, 7), (4, 7)]
nx.set_edge_attributes(
G,
{
(u, v): {
"color": "tab:blue",
"visible": (u, v) in edgelist,
}
for u, v in G.edges()
},
)
nx.display(G, node_pos="pos", node_visible=False)
plt.tight_layout()
plt.axis("off")
return fig
@pytest.mark.mpl_image_compare
def test_display_complex():
import itertools as it
fig, ax = plt.subplots()
G = nx.MultiDiGraph()
nodes = "ABC"
prod = list(it.product(nodes, repeat=2)) * 4
G = nx.MultiDiGraph()
for i, (u, v) in enumerate(prod):
G.add_edge(u, v, w=round(i / 3, 2))
nx.set_node_attributes(G, nx.spring_layout(G, seed=3113794652), "pos")
csi = it.cycle([f"arc3,rad={r}" for r in it.accumulate([0.15] * 4)])
nx.set_edge_attributes(G, {e: next(csi) for e in G.edges(keys=True)}, "curvature")
nx.set_edge_attributes(
G,
{
tuple(e): {"label": w, "bbox": {"alpha": 0}}
for *e, w in G.edges(keys=True, data="w")
},
"label",
)
nx.apply_matplotlib_colors(G, "w", "color", mpl.colormaps["inferno"], nodes=False)
nx.display(G, canvas=ax, node_pos="pos")
plt.tight_layout()
plt.axis("off")
return fig
@pytest.mark.mpl_image_compare
def test_display_shortest_path():
fig, ax = plt.subplots()
G = nx.Graph()
G.add_nodes_from(["A", "B", "C", "D", "E", "F", "G", "H"])
G.add_edge("A", "B", weight=4)
G.add_edge("A", "H", weight=8)
G.add_edge("B", "C", weight=8)
G.add_edge("B", "H", weight=11)
G.add_edge("C", "D", weight=7)
G.add_edge("C", "F", weight=4)
G.add_edge("C", "I", weight=2)
G.add_edge("D", "E", weight=9)
G.add_edge("D", "F", weight=14)
G.add_edge("E", "F", weight=10)
G.add_edge("F", "G", weight=2)
G.add_edge("G", "H", weight=1)
G.add_edge("G", "I", weight=6)
G.add_edge("H", "I", weight=7)
# Find the shortest path from node A to node E
path = nx.shortest_path(G, "A", "E", weight="weight")
# Create a list of edges in the shortest path
path_edges = list(zip(path, path[1:]))
nx.set_node_attributes(G, nx.spring_layout(G, seed=37), "pos")
nx.set_edge_attributes(
G,
{
(u, v): {
"color": (
"red"
if (u, v) in path_edges or tuple(reversed((u, v))) in path_edges
else "black"
),
"label": {"label": d["weight"], "rotate": False},
}
for u, v, d in G.edges(data=True)
},
)
nx.display(G, canvas=ax)
plt.tight_layout()
plt.axis("off")
return fig
@pytest.mark.parametrize(
("edge_color", "expected"),
(
(None, "black"),
("r", "red"),
((1.0, 1.0, 0.0), "yellow"),
((0, 1, 0, 1), "lime"),
("color", "blue"),
("#0000FF", "blue"),
),
)
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_edge_single_color(edge_color, expected, graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_edge_attributes(G, "#0000FF", "color")
canvas = plt.figure().add_subplot(111)
nx.display(G, edge_color=edge_color, canvas=canvas)
if G.is_directed():
colors = [
f.get_fc()
for f in canvas.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
else:
colors = [
l
for l in canvas.collections
if isinstance(l, mpl.collections.LineCollection)
][0].get_colors()
assert all(mpl.colors.same_color(c, expected) for c in colors)
plt.close()
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_edge_multiple_colors(graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_edge_attributes(G, {(0, 1): "#FF0000", (1, 2): (0, 0, 1)}, "color")
ax = plt.figure().add_subplot(111)
nx.display(G, canvas=ax)
expected = ["red", "blue"]
if G.is_directed():
colors = [
f.get_fc()
for f in ax.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
else:
colors = [
l for l in ax.collections if isinstance(l, mpl.collections.LineCollection)
][0].get_colors()
assert mpl.colors.same_color(colors, expected)
plt.close()
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_edge_position(graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_node_attributes(G, {n: (n, n) for n in G.nodes()}, "pos")
ax = plt.figure().add_subplot(111)
nx.display(G, canvas=ax)
if G.is_directed():
end_points = [
(f.get_path().vertices[0, :], f.get_path().vertices[-2, :])
for f in ax.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
else:
line_collection = [
l for l in ax.collections if isinstance(l, mpl.collections.LineCollection)
][0]
end_points = [
(p.vertices[0, :], p.vertices[-1, :]) for p in line_collection.get_paths()
]
expected = [((0, 0), (1, 1)), ((1, 1), (2, 2))]
# Use the threshold to account for slight shifts in FancyArrowPatch margins to
# avoid covering the arrow head with the node.
threshold = 0.05
for a, e in zip(end_points, expected):
act_start, act_end = a
exp_start, exp_end = e
assert all(abs(act_start - exp_start) < (threshold, threshold)) and all(
abs(act_end - exp_end) < (threshold, threshold)
)
plt.close()
def test_display_position_function():
G = nx.karate_club_graph()
def fixed_layout(G):
return nx.spring_layout(G, seed=314159)
pos = fixed_layout(G)
ax = plt.figure().add_subplot(111)
nx.display(G, node_pos=fixed_layout, canvas=ax)
# rebuild the position dictionary from the canvas
act_pos = {
n: tuple(p) for n, p in zip(G.nodes(), ax.get_children()[0].get_offsets().data)
}
for n in G.nodes():
assert all(pos[n] == act_pos[n])
plt.close()
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_edge_colormaps(graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_edge_attributes(G, {(0, 1): 0, (1, 2): 1}, "weight")
cmap = mpl.colormaps["plasma"]
nx.apply_matplotlib_colors(G, "weight", "color", cmap, nodes=False)
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas)
mapper = mpl.cm.ScalarMappable(cmap=cmap)
mapper.set_clim(0, 1)
expected = [mapper.to_rgba(0), mapper.to_rgba(1)]
if G.is_directed():
colors = [
f.get_facecolor()
for f in canvas.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
else:
colors = [
l
for l in canvas.collections
if isinstance(l, mpl.collections.LineCollection)
][0].get_colors()
assert mpl.colors.same_color(expected[0], G.edges[0, 1]["color"])
assert mpl.colors.same_color(expected[1], G.edges[1, 2]["color"])
assert mpl.colors.same_color(expected, colors)
plt.close()
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_node_colormaps(graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_node_attributes(G, {0: 0, 1: 0.5, 2: 1}, "weight")
cmap = mpl.colormaps["plasma"]
nx.apply_matplotlib_colors(G, "weight", "color", cmap)
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas)
mapper = mpl.cm.ScalarMappable(cmap=cmap)
mapper.set_clim(0, 1)
expected = [mapper.to_rgba(0), mapper.to_rgba(0.5), mapper.to_rgba(1)]
colors = [
s for s in canvas.collections if isinstance(s, mpl.collections.PathCollection)
][0].get_edgecolors()
assert mpl.colors.same_color(expected[0], G.nodes[0]["color"])
assert mpl.colors.same_color(expected[1], G.nodes[1]["color"])
assert mpl.colors.same_color(expected[2], G.nodes[2]["color"])
assert mpl.colors.same_color(expected, colors)
plt.close()
@pytest.mark.parametrize(
("param_value", "expected"),
(
(None, [defaults["edge_width"], defaults["edge_width"]]),
(5, [5, 5]),
("width", [5, 10]),
),
)
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_edge_width(param_value, expected, graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_edge_attributes(G, {(0, 1): 5, (1, 2): 10}, "width")
canvas = plt.figure().add_subplot(111)
nx.display(G, edge_width=param_value, canvas=canvas)
if G.is_directed():
widths = [
f.get_linewidth()
for f in canvas.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
else:
widths = list(
[
l
for l in canvas.collections
if isinstance(l, mpl.collections.LineCollection)
][0].get_linewidths()
)
assert widths == expected
@pytest.mark.parametrize(
("param_value", "expected"),
(
(None, [defaults["edge_style"], defaults["edge_style"]]),
(":", [":", ":"]),
("style", ["-", ":"]),
),
)
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_display_edge_style(param_value, expected, graph_type):
G = nx.path_graph(3, create_using=graph_type)
nx.set_edge_attributes(G, {(0, 1): "-", (1, 2): ":"}, "style")
canvas = plt.figure().add_subplot(111)
nx.display(G, edge_style=param_value, canvas=canvas)
if G.is_directed():
styles = [
f.get_linestyle()
for f in canvas.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
else:
# Convert back from tuple description to character form
linestyles = {(0, None): "-", (0, (1, 1.65)): ":"}
styles = [
linestyles[(s[0], tuple(s[1]) if s[1] is not None else None)]
for s in [
l
for l in canvas.collections
if isinstance(l, mpl.collections.LineCollection)
][0].get_linestyles()
]
assert styles == expected
plt.close()
def test_display_node_labels():
G = nx.path_graph(4)
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas, node_label={"size": 20})
labels = [t for t in canvas.get_children() if isinstance(t, mpl.text.Text)]
for n, l in zip(G.nodes(), labels):
assert l.get_text() == str(n)
assert l.get_size() == 20.0
plt.close()
def test_display_edge_labels():
G = nx.path_graph(4)
canvas = plt.figure().add_subplot(111)
# While we can pass in dicts for edge label defaults without errors,
# this isn't helpful unless we want one label for all edges.
nx.set_edge_attributes(G, {(u, v): {"label": u + v} for u, v in G.edges()})
nx.display(G, canvas=canvas, edge_label={"color": "r"}, node_label=None)
labels = [t for t in canvas.get_children() if isinstance(t, mpl.text.Text)]
print(labels)
for e, l in zip(G.edges(), labels):
assert l.get_text() == str(e[0] + e[1])
assert l.get_color() == "r"
plt.close()
@pytest.mark.mpl_image_compare
def test_display_empty_graph():
G = nx.empty_graph()
fig, ax = plt.subplots()
nx.display(G, canvas=ax)
plt.tight_layout()
plt.axis("off")
return fig
def test_display_multigraph_non_integer_keys():
G = nx.MultiGraph()
G.add_nodes_from(["A", "B", "C", "D"])
G.add_edges_from(
[
("A", "B", "0"),
("A", "B", "1"),
("B", "C", "-1"),
("B", "C", "1"),
("C", "D", "-1"),
("C", "D", "0"),
]
)
nx.set_edge_attributes(
G, {e: f"arc3,rad={0.2 * int(e[2])}" for e in G.edges(keys=True)}, "curvature"
)
canvas = plt.figure().add_subplot(111)
nx.display(G, canvas=canvas)
rads = [
f.get_connectionstyle().rad
for f in canvas.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
assert rads == [0.0, 0.2, -0.2, 0.2, -0.2, 0.0]
plt.close()
def test_display_raises_for_bad_arg():
G = nx.karate_club_graph()
with pytest.raises(nx.NetworkXError):
nx.display(G, bad_arg="bad_arg")
plt.close()
def test_display_arrow_size():
G = nx.path_graph(4, create_using=nx.DiGraph)
nx.set_edge_attributes(
G, {(u, v): (u + v + 2) ** 2 for u, v in G.edges()}, "arrowsize"
)
ax = plt.axes()
nx.display(G, canvas=ax)
assert [9, 25, 49] == [
f.get_mutation_scale()
for f in ax.get_children()
if isinstance(f, mpl.patches.FancyArrowPatch)
]
plt.close()
def test_display_mismatched_edge_position():
"""
This test ensures that a error is raised for incomplete position data.
"""
G = nx.path_graph(5)
# Notice that there is no position for node 3
nx.set_node_attributes(G, {0: (0, 0), 1: (1, 1), 2: (2, 2), 4: (4, 4)}, "pos")
# But that's not a problem since we don't want to show node 4, right?
nx.set_node_attributes(G, {n: n < 4 for n in G.nodes()}, "visible")
# However, if we try to visualize every edge (including 3 -> 4)...
# That's a problem since node 4 doesn't have a position
with pytest.raises(nx.NetworkXError):
nx.display(G)
# NOTE: parametrizing on marker to test both branches of internal
# to_marker_edge function
@pytest.mark.parametrize("node_shape", ("o", "s"))
def test_display_edge_margins(node_shape):
"""
Test that there is a wider gap between the node and the start of an
incident edge when min_source_margin is specified.
This test checks that the use os min_{source/target}_margin edge
attributes result is shorter (more padding) between the edges and
source and target nodes.
As a crude visual example, let 's' and 't' represent source and target
nodes, respectively:
Default:
s-----------------------------t
With margins:
s ----------------------- t
"""
ax = plt.figure().add_subplot(111)
G = nx.DiGraph([(0, 1)])
nx.set_node_attributes(G, {0: (0, 0), 1: (1, 1)}, "pos")
# Get the default patches from the regular visualization
nx.display(G, canvas=ax, node_shape=node_shape)
default_arrow = [
f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
][0]
default_extent = default_arrow.get_extents().corners()[::2, 0]
# Now plot again with margins
ax = plt.figure().add_subplot(111)
nx.display(
G,
canvas=ax,
edge_source_margin=100,
edge_target_margin=100,
node_shape=node_shape,
)
padded_arrow = [
f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
][0]
padded_extent = padded_arrow.get_extents().corners()[::2, 0]
# With padding, the left-most extent of the edge should be further to the right
assert padded_extent[0] > default_extent[0]
# And the rightmost extent of the edge, further to the left
assert padded_extent[1] < default_extent[1]
plt.close()
@pytest.mark.parametrize("ticks", [False, True])
def test_display_hide_ticks(ticks):
G = nx.path_graph(3)
nx.set_node_attributes(G, {n: (n, n) for n in G.nodes()}, "pos")
ax = plt.axes()
nx.display(G, hide_ticks=ticks)
for axis in [ax.xaxis, ax.yaxis]:
assert bool(axis.get_ticklabels()) != ticks
plt.close()
def test_display_self_loop():
ax = plt.axes()
G = nx.DiGraph()
G.add_node(0)
G.add_edge(0, 0)
nx.set_node_attributes(G, {0: (0, 0)}, "pos")
nx.display(G, canvas=ax)
arrow = [
f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
][0]
bbox = arrow.get_extents()
print(bbox.width)
print(bbox.height)
assert bbox.width > 0 and bbox.height > 0
plt.delaxes(ax)
plt.close()
def test_display_remove_pos_attr():
"""
If the pos attribute isn't provided or is a function, display computes the layout
and adds it to the graph. We need to ensure that this new attribute is removed from
the returned graph.
"""
G = nx.karate_club_graph()
nx.display(G)
assert nx.get_node_attributes(G, "display's position attribute name") == {}
@pytest.fixture
def subplots():
fig, ax = plt.subplots()
yield fig, ax
plt.delaxes(ax)
plt.close()
def test_draw():
try:
functions = [
nx.draw_circular,
nx.draw_kamada_kawai,
nx.draw_planar,
nx.draw_random,
nx.draw_spectral,
nx.draw_spring,
nx.draw_shell,
]
options = [{"node_color": "black", "node_size": 100, "width": 3}]
for function, option in itertools.product(functions, options):
function(barbell, **option)
plt.savefig("test.ps")
except ModuleNotFoundError: # draw_kamada_kawai requires scipy
pass
finally:
try:
os.unlink("test.ps")
except OSError:
pass
def test_draw_shell_nlist():
try:
nlist = [list(range(4)), list(range(4, 10)), list(range(10, 14))]
nx.draw_shell(barbell, nlist=nlist)
plt.savefig("test.ps")
finally:
try:
os.unlink("test.ps")
except OSError:
pass
def test_draw_bipartite():
try:
G = nx.complete_bipartite_graph(2, 5)
nx.draw_bipartite(G)
plt.savefig("test.ps")
finally:
try:
os.unlink("test.ps")
except OSError:
pass
def test_edge_colormap():
colors = range(barbell.number_of_edges())
nx.draw_spring(
barbell, edge_color=colors, width=4, edge_cmap=plt.cm.Blues, with_labels=True
)
# plt.show()
def test_arrows():
nx.draw_spring(barbell.to_directed())
# plt.show()
@pytest.mark.parametrize(
("edge_color", "expected"),
(
(None, "black"), # Default
("r", "red"), # Non-default color string
(["r"], "red"), # Single non-default color in a list
((1.0, 1.0, 0.0), "yellow"), # single color as rgb tuple
([(1.0, 1.0, 0.0)], "yellow"), # single color as rgb tuple in list
((0, 1, 0, 1), "lime"), # single color as rgba tuple
([(0, 1, 0, 1)], "lime"), # single color as rgba tuple in list
("#0000ff", "blue"), # single color hex code
(["#0000ff"], "blue"), # hex code in list
),
)
@pytest.mark.parametrize("edgelist", (None, [(0, 1)]))
def test_single_edge_color_undirected(edge_color, expected, edgelist):
"""Tests ways of specifying all edges have a single color for edges
drawn with a LineCollection"""
G = nx.path_graph(3)
drawn_edges = nx.draw_networkx_edges(
G, pos=nx.random_layout(G), edgelist=edgelist, edge_color=edge_color
)
assert mpl.colors.same_color(drawn_edges.get_color(), expected)
@pytest.mark.parametrize(
("edge_color", "expected"),
(
(None, "black"), # Default
("r", "red"), # Non-default color string
(["r"], "red"), # Single non-default color in a list
((1.0, 1.0, 0.0), "yellow"), # single color as rgb tuple
([(1.0, 1.0, 0.0)], "yellow"), # single color as rgb tuple in list
((0, 1, 0, 1), "lime"), # single color as rgba tuple
([(0, 1, 0, 1)], "lime"), # single color as rgba tuple in list
("#0000ff", "blue"), # single color hex code
(["#0000ff"], "blue"), # hex code in list
),
)
@pytest.mark.parametrize("edgelist", (None, [(0, 1)]))
def test_single_edge_color_directed(edge_color, expected, edgelist):
"""Tests ways of specifying all edges have a single color for edges drawn
with FancyArrowPatches"""
G = nx.path_graph(3, create_using=nx.DiGraph)
drawn_edges = nx.draw_networkx_edges(
G, pos=nx.random_layout(G), edgelist=edgelist, edge_color=edge_color
)
for fap in drawn_edges:
assert mpl.colors.same_color(fap.get_edgecolor(), expected)
def test_edge_color_tuple_interpretation():
"""If edge_color is a sequence with the same length as edgelist, then each
value in edge_color is mapped onto each edge via colormap."""
G = nx.path_graph(6, create_using=nx.DiGraph)
pos = {n: (n, n) for n in range(len(G))}
# num edges != 3 or 4 --> edge_color interpreted as rgb(a)
for ec in ((0, 0, 1), (0, 0, 1, 1)):
# More than 4 edges
drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=ec)
for fap in drawn_edges:
assert mpl.colors.same_color(fap.get_edgecolor(), ec)
# Fewer than 3 edges
drawn_edges = nx.draw_networkx_edges(
G, pos, edgelist=[(0, 1), (1, 2)], edge_color=ec
)
for fap in drawn_edges:
assert mpl.colors.same_color(fap.get_edgecolor(), ec)
# num edges == 3, len(edge_color) == 4: interpreted as rgba
drawn_edges = nx.draw_networkx_edges(
G, pos, edgelist=[(0, 1), (1, 2), (2, 3)], edge_color=(0, 0, 1, 1)
)
for fap in drawn_edges:
assert mpl.colors.same_color(fap.get_edgecolor(), "blue")
# num edges == 4, len(edge_color) == 3: interpreted as rgb
drawn_edges = nx.draw_networkx_edges(
G, pos, edgelist=[(0, 1), (1, 2), (2, 3), (3, 4)], edge_color=(0, 0, 1)
)
for fap in drawn_edges:
assert mpl.colors.same_color(fap.get_edgecolor(), "blue")
# num edges == len(edge_color) == 3: interpreted with cmap, *not* as rgb
drawn_edges = nx.draw_networkx_edges(
G, pos, edgelist=[(0, 1), (1, 2), (2, 3)], edge_color=(0, 0, 1)
)
assert mpl.colors.same_color(
drawn_edges[0].get_edgecolor(), drawn_edges[1].get_edgecolor()
)
for fap in drawn_edges:
assert not mpl.colors.same_color(fap.get_edgecolor(), "blue")
# num edges == len(edge_color) == 4: interpreted with cmap, *not* as rgba
drawn_edges = nx.draw_networkx_edges(
G, pos, edgelist=[(0, 1), (1, 2), (2, 3), (3, 4)], edge_color=(0, 0, 1, 1)
)
assert mpl.colors.same_color(
drawn_edges[0].get_edgecolor(), drawn_edges[1].get_edgecolor()
)
assert mpl.colors.same_color(
drawn_edges[2].get_edgecolor(), drawn_edges[3].get_edgecolor()
)
for fap in drawn_edges:
assert not mpl.colors.same_color(fap.get_edgecolor(), "blue")
def test_fewer_edge_colors_than_num_edges_directed():
"""Test that the edge colors are cycled when there are fewer specified
colors than edges."""
G = barbell.to_directed()
pos = nx.random_layout(barbell)
edgecolors = ("r", "g", "b")
drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=edgecolors)
for fap, expected in zip(drawn_edges, itertools.cycle(edgecolors)):
assert mpl.colors.same_color(fap.get_edgecolor(), expected)
def test_more_edge_colors_than_num_edges_directed():
"""Test that extra edge colors are ignored when there are more specified
colors than edges."""
G = nx.path_graph(4, create_using=nx.DiGraph) # 3 edges
pos = nx.random_layout(barbell)
edgecolors = ("r", "g", "b", "c") # 4 edge colors
drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=edgecolors)
for fap, expected in zip(drawn_edges, edgecolors[:-1]):
assert mpl.colors.same_color(fap.get_edgecolor(), expected)
def test_edge_color_string_with_global_alpha_undirected():
edge_collection = nx.draw_networkx_edges(
barbell,
pos=nx.random_layout(barbell),
edgelist=[(0, 1), (1, 2)],
edge_color="purple",
alpha=0.2,
)
ec = edge_collection.get_color().squeeze() # as rgba tuple
assert len(edge_collection.get_paths()) == 2
assert mpl.colors.same_color(ec[:-1], "purple")
assert ec[-1] == 0.2
def test_edge_color_string_with_global_alpha_directed():
drawn_edges = nx.draw_networkx_edges(
barbell.to_directed(),
pos=nx.random_layout(barbell),
edgelist=[(0, 1), (1, 2)],
edge_color="purple",
alpha=0.2,
)
assert len(drawn_edges) == 2
for fap in drawn_edges:
ec = fap.get_edgecolor() # As rgba tuple
assert mpl.colors.same_color(ec[:-1], "purple")
assert ec[-1] == 0.2
@pytest.mark.parametrize("graph_type", (nx.Graph, nx.DiGraph))
def test_edge_width_default_value(graph_type):
"""Test the default linewidth for edges drawn either via LineCollection or
FancyArrowPatches."""
G = nx.path_graph(2, create_using=graph_type)
pos = {n: (n, n) for n in range(len(G))}
drawn_edges = nx.draw_networkx_edges(G, pos)
if isinstance(drawn_edges, list): # directed case: list of FancyArrowPatch
drawn_edges = drawn_edges[0]
assert drawn_edges.get_linewidth() == 1
@pytest.mark.parametrize(
("edgewidth", "expected"),
(
(3, 3), # single-value, non-default
([3], 3), # Single value as a list
),
)
def test_edge_width_single_value_undirected(edgewidth, expected):
G = nx.path_graph(4)
pos = {n: (n, n) for n in range(len(G))}
drawn_edges = nx.draw_networkx_edges(G, pos, width=edgewidth)
assert len(drawn_edges.get_paths()) == 3
assert drawn_edges.get_linewidth() == expected
@pytest.mark.parametrize(
("edgewidth", "expected"),
(
(3, 3), # single-value, non-default
([3], 3), # Single value as a list
),
)
def test_edge_width_single_value_directed(edgewidth, expected):
G = nx.path_graph(4, create_using=nx.DiGraph)
pos = {n: (n, n) for n in range(len(G))}
drawn_edges = nx.draw_networkx_edges(G, pos, width=edgewidth)
assert len(drawn_edges) == 3
for fap in drawn_edges:
assert fap.get_linewidth() == expected
@pytest.mark.parametrize(
"edgelist",
(
[(0, 1), (1, 2), (2, 3)], # one width specification per edge
None, # fewer widths than edges - widths cycle
[(0, 1), (1, 2)], # More widths than edges - unused widths ignored
),
)
def test_edge_width_sequence(edgelist):
G = barbell.to_directed()
pos = nx.random_layout(G)
widths = (0.5, 2.0, 12.0)
drawn_edges = nx.draw_networkx_edges(G, pos, edgelist=edgelist, width=widths)
for fap, expected_width in zip(drawn_edges, itertools.cycle(widths)):
assert fap.get_linewidth() == expected_width
def test_edge_color_with_edge_vmin_vmax():
"""Test that edge_vmin and edge_vmax properly set the dynamic range of the
color map when num edges == len(edge_colors)."""
G = nx.path_graph(3, create_using=nx.DiGraph)
pos = nx.random_layout(G)
# Extract colors from the original (unscaled) colormap
drawn_edges = nx.draw_networkx_edges(G, pos, edge_color=[0, 1.0])
orig_colors = [e.get_edgecolor() for e in drawn_edges]
# Colors from scaled colormap
drawn_edges = nx.draw_networkx_edges(
G, pos, edge_color=[0.2, 0.8], edge_vmin=0.2, edge_vmax=0.8
)
scaled_colors = [e.get_edgecolor() for e in drawn_edges]
assert mpl.colors.same_color(orig_colors, scaled_colors)
def test_directed_edges_linestyle_default():
"""Test default linestyle for edges drawn with FancyArrowPatches."""
G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
pos = {n: (n, n) for n in range(len(G))}
# edge with default style
drawn_edges = nx.draw_networkx_edges(G, pos)
assert len(drawn_edges) == 3
for fap in drawn_edges:
assert fap.get_linestyle() == "solid"
@pytest.mark.parametrize(
"style",
(
"dashed", # edge with string style
"--", # edge with simplified string style
(1, (1, 1)), # edge with (offset, onoffseq) style
),
)
def test_directed_edges_linestyle_single_value(style):
"""Tests support for specifying linestyles with a single value to be applied to
all edges in ``draw_networkx_edges`` for FancyArrowPatch outputs
(e.g. directed edges)."""
G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
pos = {n: (n, n) for n in range(len(G))}
drawn_edges = nx.draw_networkx_edges(G, pos, style=style)
assert len(drawn_edges) == 3
for fap in drawn_edges:
assert fap.get_linestyle() == style
@pytest.mark.parametrize(
"style_seq",
(
["dashed"], # edge with string style in list
["--"], # edge with simplified string style in list
[(1, (1, 1))], # edge with (offset, onoffseq) style in list
["--", "-", ":"], # edges with styles for each edge
["--", "-"], # edges with fewer styles than edges (styles cycle)
["--", "-", ":", "-."], # edges with more styles than edges (extra unused)
),
)
def test_directed_edges_linestyle_sequence(style_seq):
"""Tests support for specifying linestyles with sequences in
``draw_networkx_edges`` for FancyArrowPatch outputs (e.g. directed edges)."""
G = nx.path_graph(4, create_using=nx.DiGraph) # Graph with 3 edges
pos = {n: (n, n) for n in range(len(G))}
drawn_edges = nx.draw_networkx_edges(G, pos, style=style_seq)
assert len(drawn_edges) == 3
for fap, style in zip(drawn_edges, itertools.cycle(style_seq)):
assert fap.get_linestyle() == style
def test_return_types():
from matplotlib.collections import LineCollection, PathCollection
from matplotlib.patches import FancyArrowPatch
G = nx.frucht_graph(create_using=nx.Graph)
dG = nx.frucht_graph(create_using=nx.DiGraph)
pos = nx.spring_layout(G)
dpos = nx.spring_layout(dG)
# nodes
nodes = nx.draw_networkx_nodes(G, pos)
assert isinstance(nodes, PathCollection)
# edges
edges = nx.draw_networkx_edges(dG, dpos, arrows=True)
assert isinstance(edges, list)
if len(edges) > 0:
assert isinstance(edges[0], FancyArrowPatch)
edges = nx.draw_networkx_edges(dG, dpos, arrows=False)
assert isinstance(edges, LineCollection)
edges = nx.draw_networkx_edges(G, dpos, arrows=None)
assert isinstance(edges, LineCollection)
edges = nx.draw_networkx_edges(dG, pos, arrows=None)
assert isinstance(edges, list)
if len(edges) > 0:
assert isinstance(edges[0], FancyArrowPatch)
def test_labels_and_colors():
G = nx.cubical_graph()
pos = nx.spring_layout(G) # positions for all nodes
# nodes
nx.draw_networkx_nodes(
G, pos, nodelist=[0, 1, 2, 3], node_color="r", node_size=500, alpha=0.75
)
nx.draw_networkx_nodes(
G,
pos,
nodelist=[4, 5, 6, 7],
node_color="b",
node_size=500,
alpha=[0.25, 0.5, 0.75, 1.0],
)
# edges
nx.draw_networkx_edges(G, pos, width=1.0, alpha=0.5)
nx.draw_networkx_edges(
G,
pos,
edgelist=[(0, 1), (1, 2), (2, 3), (3, 0)],
width=8,
alpha=0.5,
edge_color="r",
)
nx.draw_networkx_edges(
G,
pos,
edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
width=8,
alpha=0.5,
edge_color="b",
)
nx.draw_networkx_edges(
G,
pos,
edgelist=[(4, 5), (5, 6), (6, 7), (7, 4)],
arrows=True,
min_source_margin=0.5,
min_target_margin=0.75,
width=8,
edge_color="b",
)
# some math labels
labels = {}
labels[0] = r"$a$"
labels[1] = r"$b$"
labels[2] = r"$c$"
labels[3] = r"$d$"
labels[4] = r"$\alpha$"
labels[5] = r"$\beta$"
labels[6] = r"$\gamma$"
labels[7] = r"$\delta$"
colors = {n: "k" if n % 2 == 0 else "r" for n in range(8)}
nx.draw_networkx_labels(G, pos, labels, font_size=16)
nx.draw_networkx_labels(G, pos, labels, font_size=16, font_color=colors)
nx.draw_networkx_edge_labels(G, pos, edge_labels=None, rotate=False)
nx.draw_networkx_edge_labels(G, pos, edge_labels={(4, 5): "4-5"})
# plt.show()
@pytest.mark.mpl_image_compare
def test_house_with_colors():
G = nx.house_graph()
# explicitly set positions
fig, ax = plt.subplots()
pos = {0: (0, 0), 1: (1, 0), 2: (0, 1), 3: (1, 1), 4: (0.5, 2.0)}
# Plot nodes with different properties for the "wall" and "roof" nodes
nx.draw_networkx_nodes(
G,
pos,
node_size=3000,
nodelist=[0, 1, 2, 3],
node_color="tab:blue",
)
nx.draw_networkx_nodes(
G, pos, node_size=2000, nodelist=[4], node_color="tab:orange"
)
nx.draw_networkx_edges(G, pos, alpha=0.5, width=6)
# Customize axes
ax.margins(0.11)
plt.tight_layout()
plt.axis("off")
return fig
def test_axes(subplots):
fig, ax = subplots
nx.draw(barbell, ax=ax)
nx.draw_networkx_edge_labels(barbell, nx.circular_layout(barbell), ax=ax)
def test_empty_graph():
G = nx.Graph()
nx.draw(G)
def test_draw_empty_nodes_return_values():
# See Issue #3833
import matplotlib.collections # call as mpl.collections
G = nx.Graph([(1, 2), (2, 3)])
DG = nx.DiGraph([(1, 2), (2, 3)])
pos = nx.circular_layout(G)
assert isinstance(
nx.draw_networkx_nodes(G, pos, nodelist=[]), mpl.collections.PathCollection
)
assert isinstance(
nx.draw_networkx_nodes(DG, pos, nodelist=[]), mpl.collections.PathCollection
)
# drawing empty edges used to return an empty LineCollection or empty list.
# Now it is always an empty list (because edges are now lists of FancyArrows)
assert nx.draw_networkx_edges(G, pos, edgelist=[], arrows=True) == []
assert nx.draw_networkx_edges(G, pos, edgelist=[], arrows=False) == []
assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=False) == []
assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=True) == []
def test_multigraph_edgelist_tuples():
# See Issue #3295
G = nx.path_graph(3, create_using=nx.MultiDiGraph)
nx.draw_networkx(G, edgelist=[(0, 1, 0)])
nx.draw_networkx(G, edgelist=[(0, 1, 0)], node_size=[10, 20, 0])
def test_alpha_iter():
pos = nx.random_layout(barbell)
fig = plt.figure()
# with fewer alpha elements than nodes
fig.add_subplot(131) # Each test in a new axis object
nx.draw_networkx_nodes(barbell, pos, alpha=[0.1, 0.2])
# with equal alpha elements and nodes
num_nodes = len(barbell.nodes)
alpha = [x / num_nodes for x in range(num_nodes)]
colors = range(num_nodes)
fig.add_subplot(132)
nx.draw_networkx_nodes(barbell, pos, node_color=colors, alpha=alpha)
# with more alpha elements than nodes
alpha.append(1)
fig.add_subplot(133)
nx.draw_networkx_nodes(barbell, pos, alpha=alpha)
def test_multiple_node_shapes(subplots):
fig, ax = subplots
G = nx.path_graph(4)
nx.draw(G, node_shape=["o", "h", "s", "^"], ax=ax)
scatters = [
s for s in ax.get_children() if isinstance(s, mpl.collections.PathCollection)
]
assert len(scatters) == 4
def test_individualized_font_attributes(subplots):
G = nx.karate_club_graph()
fig, ax = subplots
nx.draw(
G,
ax=ax,
font_color={n: "k" if n % 2 else "r" for n in G.nodes()},
font_size={n: int(n / (34 / 15) + 5) for n in G.nodes()},
)
for n, t in zip(
G.nodes(),
[
t
for t in ax.get_children()
if isinstance(t, mpl.text.Text) and len(t.get_text()) > 0
],
):
expected = "black" if n % 2 else "red"
assert mpl.colors.same_color(t.get_color(), expected)
assert int(n / (34 / 15) + 5) == t.get_size()
def test_individualized_edge_attributes(subplots):
G = nx.karate_club_graph()
fig, ax = subplots
arrowstyles = ["-|>" if (u + v) % 2 == 0 else "-[" for u, v in G.edges()]
arrowsizes = [10 * (u % 2 + v % 2) + 10 for u, v in G.edges()]
nx.draw(G, ax=ax, arrows=True, arrowstyle=arrowstyles, arrowsize=arrowsizes)
arrows = [
f for f in ax.get_children() if isinstance(f, mpl.patches.FancyArrowPatch)
]
for e, a in zip(G.edges(), arrows):
assert a.get_mutation_scale() == 10 * (e[0] % 2 + e[1] % 2) + 10
expected = (
mpl.patches.ArrowStyle.BracketB
if sum(e) % 2
else mpl.patches.ArrowStyle.CurveFilledB
)
assert isinstance(a.get_arrowstyle(), expected)
def test_error_invalid_kwds():
with pytest.raises(ValueError, match="Received invalid argument"):
nx.draw(barbell, foo="bar")
def test_draw_networkx_arrowsize_incorrect_size():
G = nx.DiGraph([(0, 1), (0, 2), (0, 3), (1, 3)])
arrowsize = [1, 2, 3]
with pytest.raises(
ValueError, match="arrowsize should have the same length as edgelist"
):
nx.draw(G, arrowsize=arrowsize)
@pytest.mark.parametrize("arrowsize", (30, [10, 20, 30]))
def test_draw_edges_arrowsize(arrowsize):
G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
edges = nx.draw_networkx_edges(G, pos=pos, arrowsize=arrowsize)
arrowsize = itertools.repeat(arrowsize) if isinstance(arrowsize, int) else arrowsize
for fap, expected in zip(edges, arrowsize):
assert isinstance(fap, mpl.patches.FancyArrowPatch)
assert fap.get_mutation_scale() == expected
@pytest.mark.parametrize("arrowstyle", ("-|>", ["-|>", "-[", "<|-|>"]))
def test_draw_edges_arrowstyle(arrowstyle):
G = nx.DiGraph([(0, 1), (0, 2), (1, 2)])
pos = {0: (0, 0), 1: (0, 1), 2: (1, 0)}
edges = nx.draw_networkx_edges(G, pos=pos, arrowstyle=arrowstyle)
arrowstyle = (
itertools.repeat(arrowstyle) if isinstance(arrowstyle, str) else arrowstyle
)
arrow_objects = {
"-|>": mpl.patches.ArrowStyle.CurveFilledB,
"-[": mpl.patches.ArrowStyle.BracketB,
"<|-|>": mpl.patches.ArrowStyle.CurveFilledAB,
}
for fap, expected in zip(edges, arrowstyle):
assert isinstance(fap, mpl.patches.FancyArrowPatch)
assert isinstance(fap.get_arrowstyle(), arrow_objects[expected])
def test_np_edgelist():
# see issue #4129
nx.draw_networkx(barbell, edgelist=np.array([(0, 2), (0, 3)]))
def test_draw_nodes_missing_node_from_position():
G = nx.path_graph(3)
pos = {0: (0, 0), 1: (1, 1)} # No position for node 2
with pytest.raises(nx.NetworkXError, match="has no position"):
nx.draw_networkx_nodes(G, pos)
# NOTE: parametrizing on marker to test both branches of internal
# nx.draw_networkx_edges.to_marker_edge function
@pytest.mark.parametrize("node_shape", ("o", "s"))
def test_draw_edges_min_source_target_margins(node_shape, subplots):
"""Test that there is a wider gap between the node and the start of an
incident edge when min_source_margin is specified.
This test checks that the use of min_{source/target}_margin kwargs result
in shorter (more padding) between the edges and source and target nodes.
As a crude visual example, let 's' and 't' represent source and target
nodes, respectively:
Default:
s-----------------------------t
With margins:
s ----------------------- t
"""
# Create a single axis object to get consistent pixel coords across
# multiple draws
fig, ax = subplots
G = nx.DiGraph([(0, 1)])
pos = {0: (0, 0), 1: (1, 0)} # horizontal layout
# Get leftmost and rightmost points of the FancyArrowPatch object
# representing the edge between nodes 0 and 1 (in pixel coordinates)
default_patch = nx.draw_networkx_edges(G, pos, ax=ax, node_shape=node_shape)[0]
default_extent = default_patch.get_extents().corners()[::2, 0]
# Now, do the same but with "padding" for the source and target via the
# min_{source/target}_margin kwargs
padded_patch = nx.draw_networkx_edges(
G,
pos,
ax=ax,
node_shape=node_shape,
min_source_margin=100,
min_target_margin=100,
)[0]
padded_extent = padded_patch.get_extents().corners()[::2, 0]
# With padding, the left-most extent of the edge should be further to the
# right
assert padded_extent[0] > default_extent[0]
# And the rightmost extent of the edge, further to the left
assert padded_extent[1] < default_extent[1]
# NOTE: parametrizing on marker to test both branches of internal
# nx.draw_networkx_edges.to_marker_edge function
@pytest.mark.parametrize("node_shape", ("o", "s"))
def test_draw_edges_min_source_target_margins_individual(node_shape, subplots):
"""Test that there is a wider gap between the node and the start of an
incident edge when min_source_margin is specified.
This test checks that the use of min_{source/target}_margin kwargs result
in shorter (more padding) between the edges and source and target nodes.
As a crude visual example, let 's' and 't' represent source and target
nodes, respectively:
Default:
s-----------------------------t
With margins:
s ----------------------- t
"""
# Create a single axis object to get consistent pixel coords across
# multiple draws
fig, ax = subplots
G = nx.DiGraph([(0, 1), (1, 2)])
pos = {0: (0, 0), 1: (1, 0), 2: (2, 0)} # horizontal layout
# Get leftmost and rightmost points of the FancyArrowPatch object
# representing the edge between nodes 0 and 1 (in pixel coordinates)
default_patch = nx.draw_networkx_edges(G, pos, ax=ax, node_shape=node_shape)
default_extent = [d.get_extents().corners()[::2, 0] for d in default_patch]
# Now, do the same but with "padding" for the source and target via the
# min_{source/target}_margin kwargs
padded_patch = nx.draw_networkx_edges(
G,
pos,
ax=ax,
node_shape=node_shape,
min_source_margin=[98, 102],
min_target_margin=[98, 102],
)
padded_extent = [p.get_extents().corners()[::2, 0] for p in padded_patch]
for d, p in zip(default_extent, padded_extent):
# With padding, the left-most extent of the edge should be further to the
# right
assert p[0] > d[0]
# And the rightmost extent of the edge, further to the left
assert p[1] < d[1]
def test_nonzero_selfloop_with_single_node(subplots):
"""Ensure that selfloop extent is non-zero when there is only one node."""
# Create explicit axis object for test
fig, ax = subplots
# Graph with single node + self loop
G = nx.DiGraph()
G.add_node(0)
G.add_edge(0, 0)
# Draw
patch = nx.draw_networkx_edges(G, {0: (0, 0)})[0]
# The resulting patch must have non-zero extent
bbox = patch.get_extents()
assert bbox.width > 0 and bbox.height > 0
def test_nonzero_selfloop_with_single_edge_in_edgelist(subplots):
"""Ensure that selfloop extent is non-zero when only a single edge is
specified in the edgelist.
"""
# Create explicit axis object for test
fig, ax = subplots
# Graph with selfloop
G = nx.path_graph(2, create_using=nx.DiGraph)
G.add_edge(1, 1)
pos = {n: (n, n) for n in G.nodes}
# Draw only the selfloop edge via the `edgelist` kwarg
patch = nx.draw_networkx_edges(G, pos, edgelist=[(1, 1)])[0]
# The resulting patch must have non-zero extent
bbox = patch.get_extents()
assert bbox.width > 0 and bbox.height > 0
def test_apply_alpha():
"""Test apply_alpha when there is a mismatch between the number of
supplied colors and elements.
"""
nodelist = [0, 1, 2]
colorlist = ["r", "g", "b"]
alpha = 0.5
rgba_colors = nx.drawing.nx_pylab.apply_alpha(colorlist, alpha, nodelist)
assert all(rgba_colors[:, -1] == alpha)
def test_draw_edges_toggling_with_arrows_kwarg():
"""
The `arrows` keyword argument is used as a 3-way switch to select which
type of object to use for drawing edges:
- ``arrows=None`` -> default (FancyArrowPatches for directed, else LineCollection)
- ``arrows=True`` -> FancyArrowPatches
- ``arrows=False`` -> LineCollection
"""
import matplotlib.collections
import matplotlib.patches
UG = nx.path_graph(3)
DG = nx.path_graph(3, create_using=nx.DiGraph)
pos = {n: (n, n) for n in UG}
# Use FancyArrowPatches when arrows=True, regardless of graph type
for G in (UG, DG):
edges = nx.draw_networkx_edges(G, pos, arrows=True)
assert len(edges) == len(G.edges)
assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
# Use LineCollection when arrows=False, regardless of graph type
for G in (UG, DG):
edges = nx.draw_networkx_edges(G, pos, arrows=False)
assert isinstance(edges, mpl.collections.LineCollection)
# Default behavior when arrows=None: FAPs for directed, LC's for undirected
edges = nx.draw_networkx_edges(UG, pos)
assert isinstance(edges, mpl.collections.LineCollection)
edges = nx.draw_networkx_edges(DG, pos)
assert len(edges) == len(G.edges)
assert isinstance(edges[0], mpl.patches.FancyArrowPatch)
@pytest.mark.parametrize("drawing_func", (nx.draw, nx.draw_networkx))
def test_draw_networkx_arrows_default_undirected(drawing_func, subplots):
import matplotlib.collections
G = nx.path_graph(3)
fig, ax = subplots
drawing_func(G, ax=ax)
assert any(isinstance(c, mpl.collections.LineCollection) for c in ax.collections)
assert not ax.patches
@pytest.mark.parametrize("drawing_func", (nx.draw, nx.draw_networkx))
def test_draw_networkx_arrows_default_directed(drawing_func, subplots):
import matplotlib.collections
G = nx.path_graph(3, create_using=nx.DiGraph)
fig, ax = subplots
drawing_func(G, ax=ax)
assert not any(
isinstance(c, mpl.collections.LineCollection) for c in ax.collections
)
assert ax.patches
def test_edgelist_kwarg_not_ignored(subplots):
# See gh-4994
G = nx.path_graph(3)
G.add_edge(0, 0)
fig, ax = subplots
nx.draw(G, edgelist=[(0, 1), (1, 2)], ax=ax) # Exclude self-loop from edgelist
assert not ax.patches
@pytest.mark.parametrize(
("G", "expected_n_edges"),
([nx.DiGraph(), 2], [nx.MultiGraph(), 4], [nx.MultiDiGraph(), 4]),
)
def test_draw_networkx_edges_multiedge_connectionstyle(G, expected_n_edges):
"""Draws edges correctly for 3 types of graphs and checks for valid length"""
for i, (u, v) in enumerate([(0, 1), (0, 1), (0, 1), (0, 2)]):
G.add_edge(u, v, weight=round(i / 3, 2))
pos = {n: (n, n) for n in G}
# Raises on insufficient connectionstyle length
for conn_style in [
"arc3,rad=0.1",
["arc3,rad=0.1", "arc3,rad=0.1"],
["arc3,rad=0.1", "arc3,rad=0.1", "arc3,rad=0.2"],
]:
nx.draw_networkx_edges(G, pos, connectionstyle=conn_style)
arrows = nx.draw_networkx_edges(G, pos, connectionstyle=conn_style)
assert len(arrows) == expected_n_edges
@pytest.mark.parametrize(
("G", "expected_n_edges"),
([nx.DiGraph(), 2], [nx.MultiGraph(), 4], [nx.MultiDiGraph(), 4]),
)
def test_draw_networkx_edge_labels_multiedge_connectionstyle(G, expected_n_edges):
"""Draws labels correctly for 3 types of graphs and checks for valid length and class names"""
for i, (u, v) in enumerate([(0, 1), (0, 1), (0, 1), (0, 2)]):
G.add_edge(u, v, weight=round(i / 3, 2))
pos = {n: (n, n) for n in G}
# Raises on insufficient connectionstyle length
arrows = nx.draw_networkx_edges(
G, pos, connectionstyle=["arc3,rad=0.1", "arc3,rad=0.1", "arc3,rad=0.1"]
)
for conn_style in [
"arc3,rad=0.1",
["arc3,rad=0.1", "arc3,rad=0.2"],
["arc3,rad=0.1", "arc3,rad=0.1", "arc3,rad=0.1"],
]:
text_items = nx.draw_networkx_edge_labels(G, pos, connectionstyle=conn_style)
assert len(text_items) == expected_n_edges
for ti in text_items.values():
assert ti.__class__.__name__ == "CurvedArrowText"
def test_draw_networkx_edge_label_multiedge():
G = nx.MultiGraph()
G.add_edge(0, 1, weight=10)
G.add_edge(0, 1, weight=20)
edge_labels = nx.get_edge_attributes(G, "weight") # Includes edge keys
pos = {n: (n, n) for n in G}
text_items = nx.draw_networkx_edge_labels(
G,
pos,
edge_labels=edge_labels,
connectionstyle=["arc3,rad=0.1", "arc3,rad=0.2"],
)
assert len(text_items) == 2
def test_draw_networkx_edge_label_empty_dict():
"""Regression test for draw_networkx_edge_labels with empty dict. See
gh-5372."""
G = nx.path_graph(3)
pos = {n: (n, n) for n in G.nodes}
assert nx.draw_networkx_edge_labels(G, pos, edge_labels={}) == {}
def test_draw_networkx_edges_undirected_selfloop_colors(subplots):
"""When an edgelist is supplied along with a sequence of colors, check that
the self-loops have the correct colors."""
fig, ax = subplots
# Edge list and corresponding colors
edgelist = [(1, 3), (1, 2), (2, 3), (1, 1), (3, 3), (2, 2)]
edge_colors = ["pink", "cyan", "black", "red", "blue", "green"]
G = nx.Graph(edgelist)
pos = {n: (n, n) for n in G.nodes}
nx.draw_networkx_edges(G, pos, ax=ax, edgelist=edgelist, edge_color=edge_colors)
# Verify that there are three fancy arrow patches (1 per self loop)
assert len(ax.patches) == 3
# These are points that should be contained in the self loops. For example,
# sl_points[0] will be (1, 1.1), which is inside the "path" of the first
# self-loop but outside the others
sl_points = np.array(edgelist[-3:]) + np.array([0, 0.1])
# Check that the mapping between self-loop locations and their colors is
# correct
for fap, clr, slp in zip(ax.patches, edge_colors[-3:], sl_points):
assert fap.get_path().contains_point(slp)
assert mpl.colors.same_color(fap.get_edgecolor(), clr)
@pytest.mark.parametrize(
"fap_only_kwarg", # Non-default values for kwargs that only apply to FAPs
(
{"arrowstyle": "-"},
{"arrowsize": 20},
{"connectionstyle": "arc3,rad=0.2"},
{"min_source_margin": 10},
{"min_target_margin": 10},
),
)
def test_user_warnings_for_unused_edge_drawing_kwargs(fap_only_kwarg, subplots):
"""Users should get a warning when they specify a non-default value for
one of the kwargs that applies only to edges drawn with FancyArrowPatches,
but FancyArrowPatches aren't being used under the hood."""
G = nx.path_graph(3)
pos = {n: (n, n) for n in G}
fig, ax = subplots
# By default, an undirected graph will use LineCollection to represent
# the edges
kwarg_name = list(fap_only_kwarg.keys())[0]
with pytest.warns(
UserWarning, match=f"\n\nThe {kwarg_name} keyword argument is not applicable"
):
nx.draw_networkx_edges(G, pos, ax=ax, **fap_only_kwarg)
# FancyArrowPatches are always used when `arrows=True` is specified.
# Check that warnings are *not* raised in this case
with warnings.catch_warnings():
# Escalate warnings -> errors so tests fail if warnings are raised
warnings.simplefilter("error")
warnings.filterwarnings("ignore", category=DeprecationWarning)
nx.draw_networkx_edges(G, pos, ax=ax, arrows=True, **fap_only_kwarg)
@pytest.mark.parametrize("draw_fn", (nx.draw, nx.draw_circular))
def test_no_warning_on_default_draw_arrowstyle(draw_fn, subplots):
# See gh-7284
fig, ax = subplots
G = nx.cycle_graph(5)
with warnings.catch_warnings(record=True) as w:
draw_fn(G, ax=ax)
assert len(w) == 0
@pytest.mark.parametrize("hide_ticks", [False, True])
@pytest.mark.parametrize(
"method",
[
nx.draw_networkx,
nx.draw_networkx_edge_labels,
nx.draw_networkx_edges,
nx.draw_networkx_labels,
nx.draw_networkx_nodes,
],
)
def test_hide_ticks(method, hide_ticks, subplots):
G = nx.path_graph(3)
pos = {n: (n, n) for n in G.nodes}
_, ax = subplots
method(G, pos=pos, ax=ax, hide_ticks=hide_ticks)
for axis in [ax.xaxis, ax.yaxis]:
assert bool(axis.get_ticklabels()) != hide_ticks
def test_edge_label_bar_connectionstyle(subplots):
"""Check that FancyArrowPatches with `bar` connectionstyle are also supported
in edge label rendering. See gh-7735."""
fig, ax = subplots
edge = (0, 1)
G = nx.DiGraph([edge])
pos = {n: (n, 0) for n in G} # Edge is horizontal line between (0, 0) and (1, 0)
style_arc = "arc3,rad=0.0"
style_bar = "bar,fraction=0.1"
arc_lbl = nx.draw_networkx_edge_labels(
G, pos, edge_labels={edge: "edge"}, connectionstyle=style_arc
)
# This would fail prior to gh-7739
bar_lbl = nx.draw_networkx_edge_labels(
G, pos, edge_labels={edge: "edge"}, connectionstyle=style_bar
)
# For the "arc" style, the label should be at roughly the midpoint
assert arc_lbl[edge].x, arc_lbl[edge].y == pytest.approx((0.5, 0))
# The label should be below the x-axis for the "bar" style
assert bar_lbl[edge].y < arc_lbl[edge].y