Fix edge drawing performance regression (#4825)

* Add back line coll. style drawing and encapsulate.

Adds back drawing edges via LineCollection that was removed in
#4360.
Encapsulates the two different drawing methods in internal
functions; one for FancyArrowPatches and one for LineCollection.

* WIP: add new kwarg, mv up fn defns.

* apply axis scaling to both edge drawing fns.

* use fancy arrow patch for digraphs + arrow=True (default).

* draw selfloops even for undirected graphs.

* TST: Add test for desired behavior of arrows kwarg.

Tests 3-way switching between default behavior, forcing
fancy arrow patches, and forcing line collections.

* Implement 3-way behavior switch w/ arrows kwarg.

Co-Authored-By: Dan Schult <dschult@colgate.edu>

* Update docstring and add example.

* Fix docstring

Co-authored-by: Dan Schult <dschult@colgate.edu>

Co-authored-by: Dan Schult <dschult@colgate.edu>

2 files changed, 219 insertions, 127 deletions

diff --git a/networkx/drawing/nx_pylab.py b/networkx/drawing/nx_pylab.py
index 3d35b6b2..823457da 100644
--- a/networkx/drawing/nx_pylab.py
+++ b/networkx/drawing/nx_pylab.py
@@ -495,13 +495,13 @@ def draw_networkx_edges(
     edge_color="k",
     style="solid",
     alpha=None,
-    arrowstyle=None,
+    arrowstyle="-|>",
     arrowsize=10,
     edge_cmap=None,
     edge_vmin=None,
     edge_vmax=None,
     ax=None,
-    arrows=True,
+    arrows=None,
     label=None,
     node_size=300,
     nodelist=None,
@@ -552,15 +552,17 @@ def draw_networkx_edges(
     ax : Matplotlib Axes object, optional
         Draw the graph in the specified Matplotlib axes.
 
-    arrows : bool, optional (default=True)
-        For directed graphs, if True set default to drawing arrowheads.
-        Otherwise set default to no arrowheads. Ignored if `arrowstyle` is set.
+    arrows : bool or None, optional (default=None)
+        If `None`, directed graphs draw arrowheads with
+        `~matplotlib.patches.FancyArrowPatch`, while undirected graphs draw edges
+        via `~matplotlib.collections.LineCollection` for speed.
+        If `True`, draw arrowheads with FancyArrowPatches (bendable and stylish).
+        If `False`, draw edges using LineCollection (linear and fast).
 
-        Note: Arrows will be the same color as edges.
+        Note: Arrowheads will be the same color as edges.
 
-    arrowstyle : str (default='-\|>' if directed else '-')
+    arrowstyle : str (default='-\|>')
         For directed graphs and `arrows==True` defaults to '-\|>',
-        otherwise defaults to '-'.
 
         See `matplotlib.patches.ArrowStyle` for more options.
 
@@ -597,8 +599,11 @@ def draw_networkx_edges(
 
     Returns
     -------
-    list of matplotlib.patches.FancyArrowPatch
-        `FancyArrowPatch` instances of the directed edges
+     matplotlib.colections.LineCollection or a list of matplotlib.patches.FancyArrowPatch
+        If ``arrows=True``, a list of FancyArrowPatches is returned.
+        If ``arrows=False``, a LineCollection is returned.
+        If ``arrows=None`` (the default), then a LineCollection is returned if
+        `G` is undirected, otherwise returns a list of FancyArrowPatches.
 
     Notes
     -----
@@ -609,6 +614,13 @@ def draw_networkx_edges(
     Be sure to include `node_size` as a keyword argument; arrows are
     drawn considering the size of nodes.
 
+    Self-loops are always drawn with `~matplotlib.patches.FancyArrowPatch`
+    regardless of the value of `arrows` or whether `G` is directed.
+    When ``arrows=False`` or ``arrows=None`` and `G` is undirected, the
+    FancyArrowPatches corresponding to the self-loops are not explicitly
+    returned. They should instead be accessed via the ``Axes.patches``
+    attribute (see examples).
+
     Examples
     --------
     >>> G = nx.dodecahedral_graph()
@@ -621,6 +633,16 @@ def draw_networkx_edges(
     >>> for i, arc in enumerate(arcs):  # change alpha values of arcs
     ...     arc.set_alpha(alphas[i])
 
+    The FancyArrowPatches corresponding to self-loops are not always
+    returned, but can always be accessed via the ``patches`` attribute of the
+    `matplotlib.Axes` object.
+
+    >>> import matplotlib.pyplot as plt
+    >>> fig, ax = plt.subplots()
+    >>> G = nx.Graph([(0, 1), (0, 0)])  # Self-loop at node 0
+    >>> edge_collection = nx.draw_networkx_edges(G, pos=nx.circular_layout(G), ax=ax)
+    >>> self_loop_fap = ax.patches[0]
+
     Also see the NetworkX drawing examples at
     https://networkx.org/documentation/latest/auto_examples/index.html
 
@@ -637,14 +659,17 @@ def draw_networkx_edges(
     import matplotlib as mpl
     import matplotlib.colors  # call as mpl.colors
     import matplotlib.patches  # call as mpl.patches
+    import matplotlib.collections  # call as mpl.collections
     import matplotlib.path  # call as mpl.path
     import matplotlib.pyplot as plt
 
-    if arrowstyle is None:
-        if G.is_directed() and arrows:
-            arrowstyle = "-|>"
-        else:
-            arrowstyle = "-"
+    # The default behavior is to use LineCollection to draw edges for
+    # undirected graphs (for performance reasons) and use FancyArrowPatches
+    # for directed graphs.
+    # The `arrows` keyword can be used to override the default behavior
+    use_linecollection = not G.is_directed()
+    if arrows in (True, False):
+        use_linecollection = not arrows
 
     if ax is None:
         ax = plt.gca()
@@ -683,21 +708,135 @@ def draw_networkx_edges(
         color_normal = mpl.colors.Normalize(vmin=edge_vmin, vmax=edge_vmax)
         edge_color = [edge_cmap(color_normal(e)) for e in edge_color]
 
-    # Note: Waiting for someone to implement arrow to intersection with
-    # marker.  Meanwhile, this works well for polygons with more than 4
-    # sides and circle.
-
-    def to_marker_edge(marker_size, marker):
-        if marker in "s^>v<d":  # `large` markers need extra space
-            return np.sqrt(2 * marker_size) / 2
-        else:
-            return np.sqrt(marker_size) / 2
-
-    # Draw arrows with `matplotlib.patches.FancyarrowPatch`
-    arrow_collection = []
-    mutation_scale = arrowsize  # scale factor of arrow head
-
-    # compute view
+    def _draw_networkx_edges_line_collection():
+        edge_collection = mpl.collections.LineCollection(
+            edge_pos,
+            colors=edge_color,
+            linewidths=width,
+            antialiaseds=(1,),
+            linestyle=style,
+            transOffset=ax.transData,
+            alpha=alpha,
+        )
+        edge_collection.set_cmap(edge_cmap)
+        edge_collection.set_clim(edge_vmin, edge_vmax)
+        edge_collection.set_zorder(1)  # edges go behind nodes
+        edge_collection.set_label(label)
+        ax.add_collection(edge_collection)
+
+        return edge_collection
+
+    def _draw_networkx_edges_fancy_arrow_patch():
+        # Note: Waiting for someone to implement arrow to intersection with
+        # marker.  Meanwhile, this works well for polygons with more than 4
+        # sides and circle.
+
+        def to_marker_edge(marker_size, marker):
+            if marker in "s^>v<d":  # `large` markers need extra space
+                return np.sqrt(2 * marker_size) / 2
+            else:
+                return np.sqrt(marker_size) / 2
+
+        # Draw arrows with `matplotlib.patches.FancyarrowPatch`
+        arrow_collection = []
+        mutation_scale = arrowsize  # scale factor of arrow head
+
+        base_connection_style = mpl.patches.ConnectionStyle(connectionstyle)
+
+        # Fallback for self-loop scale. Left outside of _connectionstyle so it is
+        # only computed once
+        max_nodesize = np.array(node_size).max()
+
+        def _connectionstyle(posA, posB, *args, **kwargs):
+            # check if we need to do a self-loop
+            if np.all(posA == posB):
+                # Self-loops are scaled by view extent, except in cases the extent
+                # is 0, e.g. for a single node. In this case, fall back to scaling
+                # by the maximum node size
+                selfloop_ht = 0.005 * max_nodesize if h == 0 else h
+                # this is called with _screen space_ values so covert back
+                # to data space
+                data_loc = ax.transData.inverted().transform(posA)
+                v_shift = 0.1 * selfloop_ht
+                h_shift = v_shift * 0.5
+                # put the top of the loop first so arrow is not hidden by node
+                path = [
+                    # 1
+                    data_loc + np.asarray([0, v_shift]),
+                    # 4 4 4
+                    data_loc + np.asarray([h_shift, v_shift]),
+                    data_loc + np.asarray([h_shift, 0]),
+                    data_loc,
+                    # 4 4 4
+                    data_loc + np.asarray([-h_shift, 0]),
+                    data_loc + np.asarray([-h_shift, v_shift]),
+                    data_loc + np.asarray([0, v_shift]),
+                ]
+
+                ret = mpl.path.Path(ax.transData.transform(path), [1, 4, 4, 4, 4, 4, 4])
+            # if not, fall back to the user specified behavior
+            else:
+                ret = base_connection_style(posA, posB, *args, **kwargs)
+
+            return ret
+
+        # FancyArrowPatch doesn't handle color strings
+        arrow_colors = mpl.colors.colorConverter.to_rgba_array(edge_color, alpha)
+        for i, (src, dst) in enumerate(edge_pos):
+            x1, y1 = src
+            x2, y2 = dst
+            shrink_source = 0  # space from source to tail
+            shrink_target = 0  # space from  head to target
+            if np.iterable(node_size):  # many node sizes
+                source, target = edgelist[i][:2]
+                source_node_size = node_size[nodelist.index(source)]
+                target_node_size = node_size[nodelist.index(target)]
+                shrink_source = to_marker_edge(source_node_size, node_shape)
+                shrink_target = to_marker_edge(target_node_size, node_shape)
+            else:
+                shrink_source = shrink_target = to_marker_edge(node_size, node_shape)
+
+            if shrink_source < min_source_margin:
+                shrink_source = min_source_margin
+
+            if shrink_target < min_target_margin:
+                shrink_target = min_target_margin
+
+            if len(arrow_colors) == len(edge_pos):
+                arrow_color = arrow_colors[i]
+            elif len(arrow_colors) == 1:
+                arrow_color = arrow_colors[0]
+            else:  # Cycle through colors
+                arrow_color = arrow_colors[i % len(arrow_colors)]
+
+            if np.iterable(width):
+                if len(width) == len(edge_pos):
+                    line_width = width[i]
+                else:
+                    line_width = width[i % len(width)]
+            else:
+                line_width = width
+
+            arrow = mpl.patches.FancyArrowPatch(
+                (x1, y1),
+                (x2, y2),
+                arrowstyle=arrowstyle,
+                shrinkA=shrink_source,
+                shrinkB=shrink_target,
+                mutation_scale=mutation_scale,
+                color=arrow_color,
+                linewidth=line_width,
+                connectionstyle=_connectionstyle,
+                linestyle=style,
+                zorder=1,
+            )  # arrows go behind nodes
+
+            arrow_collection.append(arrow)
+            ax.add_patch(arrow)
+
+        return arrow_collection
+
+    # compute initial view
     minx = np.amin(np.ravel(edge_pos[:, :, 0]))
     maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
     miny = np.amin(np.ravel(edge_pos[:, :, 1]))
@@ -705,100 +844,19 @@ def draw_networkx_edges(
     w = maxx - minx
     h = maxy - miny
 
-    base_connection_style = mpl.patches.ConnectionStyle(connectionstyle)
-
-    # Fallback for self-loop scale. Left outside of _connectionstyle so it is
-    # only computed once
-    max_nodesize = np.array(node_size).max()
-
-    def _connectionstyle(posA, posB, *args, **kwargs):
-        # check if we need to do a self-loop
-        if np.all(posA == posB):
-            # Self-loops are scaled by view extent, except in cases the extent
-            # is 0, e.g. for a single node. In this case, fall back to scaling
-            # by the maximum node size
-            selfloop_ht = 0.005 * max_nodesize if h == 0 else h
-            # this is called with _screen space_ values so covert back
-            # to data space
-            data_loc = ax.transData.inverted().transform(posA)
-            v_shift = 0.1 * selfloop_ht
-            h_shift = v_shift * 0.5
-            # put the top of the loop first so arrow is not hidden by node
-            path = [
-                # 1
-                data_loc + np.asarray([0, v_shift]),
-                # 4 4 4
-                data_loc + np.asarray([h_shift, v_shift]),
-                data_loc + np.asarray([h_shift, 0]),
-                data_loc,
-                # 4 4 4
-                data_loc + np.asarray([-h_shift, 0]),
-                data_loc + np.asarray([-h_shift, v_shift]),
-                data_loc + np.asarray([0, v_shift]),
-            ]
-
-            ret = mpl.path.Path(ax.transData.transform(path), [1, 4, 4, 4, 4, 4, 4])
-        # if not, fall back to the user specified behavior
-        else:
-            ret = base_connection_style(posA, posB, *args, **kwargs)
-
-        return ret
-
-    # FancyArrowPatch doesn't handle color strings
-    arrow_colors = mpl.colors.colorConverter.to_rgba_array(edge_color, alpha)
-    for i, (src, dst) in enumerate(edge_pos):
-        x1, y1 = src
-        x2, y2 = dst
-        shrink_source = 0  # space from source to tail
-        shrink_target = 0  # space from  head to target
-        if np.iterable(node_size):  # many node sizes
-            source, target = edgelist[i][:2]
-            source_node_size = node_size[nodelist.index(source)]
-            target_node_size = node_size[nodelist.index(target)]
-            shrink_source = to_marker_edge(source_node_size, node_shape)
-            shrink_target = to_marker_edge(target_node_size, node_shape)
-        else:
-            shrink_source = shrink_target = to_marker_edge(node_size, node_shape)
-
-        if shrink_source < min_source_margin:
-            shrink_source = min_source_margin
-
-        if shrink_target < min_target_margin:
-            shrink_target = min_target_margin
-
-        if len(arrow_colors) == len(edge_pos):
-            arrow_color = arrow_colors[i]
-        elif len(arrow_colors) == 1:
-            arrow_color = arrow_colors[0]
-        else:  # Cycle through colors
-            arrow_color = arrow_colors[i % len(arrow_colors)]
-
-        if np.iterable(width):
-            if len(width) == len(edge_pos):
-                line_width = width[i]
-            else:
-                line_width = width[i % len(width)]
-        else:
-            line_width = width
-
-        arrow = mpl.patches.FancyArrowPatch(
-            (x1, y1),
-            (x2, y2),
-            arrowstyle=arrowstyle,
-            shrinkA=shrink_source,
-            shrinkB=shrink_target,
-            mutation_scale=mutation_scale,
-            color=arrow_color,
-            linewidth=line_width,
-            connectionstyle=_connectionstyle,
-            linestyle=style,
-            zorder=1,
-        )  # arrows go behind nodes
-
-        arrow_collection.append(arrow)
-        ax.add_patch(arrow)
+    # Draw the edges
+    if use_linecollection:
+        edge_viz_obj = _draw_networkx_edges_line_collection()
+        # Make sure selfloop edges are also drawn.
+        edgelist = list(nx.selfloop_edges(G))
+        if edgelist:
+            edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
+            arrowstyle = "-"
+            _draw_networkx_edges_fancy_arrow_patch()
+    else:
+        edge_viz_obj = _draw_networkx_edges_fancy_arrow_patch()
 
-    # update view
+    # update view after drawing
     padx, pady = 0.05 * w, 0.05 * h
     corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
     ax.update_datalim(corners)
@@ -813,7 +871,7 @@ def draw_networkx_edges(
         labelleft=False,
     )
 
-    return arrow_collection
+    return edge_viz_obj
 
 
 def draw_networkx_labels(
diff --git a/networkx/drawing/tests/test_pylab.py b/networkx/drawing/tests/test_pylab.py
index 1a3f5083..6d94ddd9 100644
--- a/networkx/drawing/tests/test_pylab.py
+++ b/networkx/drawing/tests/test_pylab.py
@@ -297,7 +297,7 @@ def test_draw_edges_min_source_target_margins(node_shape):
     # Create a single axis object to get consistent pixel coords across
     # multiple draws
     fig, ax = plt.subplots()
-    G = nx.Graph([(0, 1)])
+    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)
@@ -327,7 +327,7 @@ def test_nonzero_selfloop_with_single_node():
     # Create explicit axis object for test
     fig, ax = plt.subplots()
     # Graph with single node + self loop
-    G = nx.Graph()
+    G = nx.DiGraph()
     G.add_node(0)
     G.add_edge(0, 0)
     # Draw
@@ -346,7 +346,7 @@ def test_nonzero_selfloop_with_single_edge_in_edgelist():
     # Create explicit axis object for test
     fig, ax = plt.subplots()
     # Graph with selfloop
-    G = nx.path_graph(2)
+    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
@@ -367,3 +367,37 @@ def test_apply_alpha():
     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.patches
+    import matplotlib.collections
+
+    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)