258 lines
8.6 KiB
Python
258 lines
8.6 KiB
Python
|
"""Unit tests for matplotlib drawing functions."""
|
||
|
|
import os
|
||
|
|
import itertools
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
mpl = pytest.importorskip("matplotlib")
|
||
|
|
mpl.use("PS")
|
||
|
|
plt = pytest.importorskip("matplotlib.pyplot")
|
||
|
|
plt.rcParams["text.usetex"] = False
|
||
|
|
|
||
|
|
import networkx as nx
|
||
|
|
|
||
|
|
|
||
|
|
class TestPylab:
|
||
|
|
@classmethod
|
||
|
|
def setup_class(cls):
|
||
|
|
cls.G = nx.barbell_graph(4, 6)
|
||
|
|
|
||
|
|
def test_draw(self):
|
||
|
|
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(self.G, **option)
|
||
|
|
plt.savefig("test.ps")
|
||
|
|
|
||
|
|
finally:
|
||
|
|
try:
|
||
|
|
os.unlink("test.ps")
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def test_draw_shell_nlist(self):
|
||
|
|
try:
|
||
|
|
nlist = [list(range(4)), list(range(4, 10)), list(range(10, 14))]
|
||
|
|
nx.draw_shell(self.G, nlist=nlist)
|
||
|
|
plt.savefig("test.ps")
|
||
|
|
finally:
|
||
|
|
try:
|
||
|
|
os.unlink("test.ps")
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def test_edge_colormap(self):
|
||
|
|
colors = range(self.G.number_of_edges())
|
||
|
|
nx.draw_spring(
|
||
|
|
self.G, edge_color=colors, width=4, edge_cmap=plt.cm.Blues, with_labels=True
|
||
|
|
)
|
||
|
|
# plt.show()
|
||
|
|
|
||
|
|
def test_arrows(self):
|
||
|
|
nx.draw_spring(self.G.to_directed())
|
||
|
|
# plt.show()
|
||
|
|
|
||
|
|
def test_edge_colors_and_widths(self):
|
||
|
|
pos = nx.circular_layout(self.G)
|
||
|
|
for G in (self.G, self.G.to_directed()):
|
||
|
|
nx.draw_networkx_nodes(G, pos, node_color=[(1.0, 1.0, 0.2, 0.5)])
|
||
|
|
nx.draw_networkx_labels(G, pos)
|
||
|
|
# edge with default color and width
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(0, 1)], width=None, edge_color=None
|
||
|
|
)
|
||
|
|
# edges with global color strings and widths in lists
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(0, 2), (0, 3)], width=[3], edge_color=["r"]
|
||
|
|
)
|
||
|
|
# edges with color strings and widths for each edge
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(0, 2), (0, 3)], width=[1, 3], edge_color=["r", "b"]
|
||
|
|
)
|
||
|
|
# edges with fewer color strings and widths than edges
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G,
|
||
|
|
pos,
|
||
|
|
edgelist=[(1, 2), (1, 3), (2, 3), (3, 4)],
|
||
|
|
width=[1, 3],
|
||
|
|
edge_color=["g", "m", "c"],
|
||
|
|
)
|
||
|
|
# edges with more color strings and widths than edges
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G,
|
||
|
|
pos,
|
||
|
|
edgelist=[(3, 4)],
|
||
|
|
width=[1, 2, 3, 4],
|
||
|
|
edge_color=["r", "b", "g", "k"],
|
||
|
|
)
|
||
|
|
# with rgb tuple and 3 edges - is interpreted with cmap
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(4, 5), (5, 6), (6, 7)], edge_color=(1.0, 0.4, 0.3)
|
||
|
|
)
|
||
|
|
# with rgb tuple in list
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(7, 8), (8, 9)], edge_color=[(0.4, 1.0, 0.0)]
|
||
|
|
)
|
||
|
|
# with rgba tuple and 4 edges - is interpretted with cmap
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G,
|
||
|
|
pos,
|
||
|
|
edgelist=[(9, 10), (10, 11), (10, 12), (10, 13)],
|
||
|
|
edge_color=(0.0, 1.0, 1.0, 0.5),
|
||
|
|
)
|
||
|
|
# with rgba tuple in list
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G,
|
||
|
|
pos,
|
||
|
|
edgelist=[(9, 10), (10, 11), (10, 12), (10, 13)],
|
||
|
|
edge_color=[(0.0, 1.0, 1.0, 0.5)],
|
||
|
|
)
|
||
|
|
# with color string and global alpha
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(11, 12), (11, 13)], edge_color="purple", alpha=0.2
|
||
|
|
)
|
||
|
|
# with color string in a list
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G, pos, edgelist=[(11, 12), (11, 13)], edge_color=["purple"]
|
||
|
|
)
|
||
|
|
# with single edge and hex color string
|
||
|
|
nx.draw_networkx_edges(G, pos, edgelist=[(12, 13)], edge_color="#1f78b4f0")
|
||
|
|
|
||
|
|
# edge_color as numeric using vmin, vmax
|
||
|
|
nx.draw_networkx_edges(
|
||
|
|
G,
|
||
|
|
pos,
|
||
|
|
edgelist=[(7, 8), (8, 9)],
|
||
|
|
edge_color=[0.2, 0.5],
|
||
|
|
edge_vmin=0.1,
|
||
|
|
edge_vmax=0.6,
|
||
|
|
)
|
||
|
|
|
||
|
|
# plt.show()
|
||
|
|
|
||
|
|
def test_labels_and_colors(self):
|
||
|
|
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)],
|
||
|
|
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$"
|
||
|
|
nx.draw_networkx_labels(G, pos, labels, font_size=16)
|
||
|
|
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()
|
||
|
|
|
||
|
|
def test_axes(self):
|
||
|
|
fig, ax = plt.subplots()
|
||
|
|
nx.draw(self.G, ax=ax)
|
||
|
|
|
||
|
|
def test_empty_graph(self):
|
||
|
|
G = nx.Graph()
|
||
|
|
nx.draw(G)
|
||
|
|
|
||
|
|
def test_draw_empty_nodes_return_values(self):
|
||
|
|
# See Issue #3833
|
||
|
|
from matplotlib.collections import PathCollection, LineCollection
|
||
|
|
|
||
|
|
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=[]), PathCollection)
|
||
|
|
assert isinstance(nx.draw_networkx_nodes(DG, pos, nodelist=[]), PathCollection)
|
||
|
|
|
||
|
|
# drawing empty edges either return an empty LineCollection or empty list.
|
||
|
|
assert isinstance(
|
||
|
|
nx.draw_networkx_edges(G, pos, edgelist=[], arrows=True), LineCollection
|
||
|
|
)
|
||
|
|
assert isinstance(
|
||
|
|
nx.draw_networkx_edges(G, pos, edgelist=[], arrows=False), LineCollection
|
||
|
|
)
|
||
|
|
assert isinstance(
|
||
|
|
nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=False), LineCollection
|
||
|
|
)
|
||
|
|
assert nx.draw_networkx_edges(DG, pos, edgelist=[], arrows=True) == []
|
||
|
|
|
||
|
|
def test_multigraph_edgelist_tuples(self):
|
||
|
|
# 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(self):
|
||
|
|
pos = nx.random_layout(self.G)
|
||
|
|
# with fewer alpha elements than nodes
|
||
|
|
plt.subplot(131)
|
||
|
|
nx.draw_networkx_nodes(self.G, pos, alpha=[0.1, 0.2])
|
||
|
|
# with equal alpha elements and nodes
|
||
|
|
num_nodes = len(self.G.nodes)
|
||
|
|
alpha = [x / num_nodes for x in range(num_nodes)]
|
||
|
|
colors = range(num_nodes)
|
||
|
|
plt.subplot(132)
|
||
|
|
nx.draw_networkx_nodes(self.G, pos, node_color=colors, alpha=alpha)
|
||
|
|
# with more alpha elements than nodes
|
||
|
|
alpha.append(1)
|
||
|
|
plt.subplot(133)
|
||
|
|
nx.draw_networkx_nodes(self.G, pos, alpha=alpha)
|
||
|
|
|
||
|
|
def test_error_invalid_kwds(self):
|
||
|
|
with pytest.raises(ValueError, match="Received invalid argument"):
|
||
|
|
nx.draw(self.G, foo="bar")
|
||
|
|
|
||
|
|
def test_np_edgelist(self):
|
||
|
|
# see issue #4129
|
||
|
|
np = pytest.importorskip("numpy")
|
||
|
|
nx.draw_networkx(self.G, edgelist=np.array([(0, 2), (0, 3)]))
|