Repository: networkx/networkx

Issue:
multipartite_layout() fails when some node labels are not float
<!--- Provide a general summary of the issue in the Title above -->
Dear contributors,
I am facing currently the following issue which I believe is a bug.
Calling `multipartite_layout()` on a graph which has some nodes with labels which are not float (or integers) will result in an error.

### Current Behavior
Calling `multipartite_layout()` on a graph which has some nodes with hashable labels which are not float (or integers) will result in the following exception:
```
numpy.core._exceptions.UFuncTypeError: ufunc 'subtract' did not contain a loop with signature matching types (dtype('<U32'), dtype('<U32')) -> dtype('<U32')
```
The traceback points at line 1101 of drawing/layout.py, in multipartite_layout:  `layer_pos = np.column_stack([xs, ys]) - offset`

### Expected Behavior
Calling `multipartite_layout()` on a graph which has some nodes with hashable labels which are not float (or integers) should return a `pos` allowing to draw the graph as a multipartite graph.

### Steps to Reproduce
Here is a minimum non-working example:
```python
G = nx.Graph()
G.add_node(0, subset='s0')
G.add_node("b", subset='s0')
G.add_node(2, subset='s1')
G.add_node("d", subset='s1')
G.add_edges_from([(0, 2), (0, "d"), ("b", 2)])
pos = nx.multipartite_layout(G)  # <-- raise the exception mentioned above
```

### Environment
<!--- Please provide details about your local environment -->
Python version: 3.8.11
NetworkX version: 2.6.3
NumPy version: 1.20.3


### Additional context
I fixed this problem with the following workaround (from `multipartite_layout` code, in drawing/layout.py):
```python
import numpy as np

if align not in ("vertical", "horizontal"):
    msg = "align must be either vertical or horizontal."
    raise ValueError(msg)

G, center = _process_params(G, center=center, dim=2)
if len(G) == 0:
    return {}

layers = {}
for v, data in G.nodes(data=True):
    try:
        layer = data[subset_key]
    except KeyError:
        msg = "all nodes must have subset_key (default='subset') as data"
        raise ValueError(msg)
    layers[layer] = [v] + layers.get(layer, [])

pos = None
nodes = []

width = len(layers)
# for i, layer in layers.items():  # <-- original code
for i, layer in enumerate(layers.values()):  # <-- modification
    height = len(layer)
    xs = np.repeat(i, height)
    ys = np.arange(0, height, dtype=float)
    offset = ((width - 1) / 2, (height - 1) / 2)
    layer_pos = np.column_stack([xs, ys]) - offset
    if pos is None:
        pos = layer_pos
    else:
        pos = np.concatenate([pos, layer_pos])
    nodes.extend(layer)
pos = rescale_layout(pos, scale=scale) + center
if align == "horizontal":
    pos = np.flip(pos, 1)
pos = dict(zip(nodes, pos))
return pos
```
However, being new to using NetworkX, I cannot guarantee that it is a reliable fix.

Discussion:

Patch:
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
index e08b785f1..f43b11a86 100644
--- a/CONTRIBUTING.rst
+++ b/CONTRIBUTING.rst
@@ -390,3 +390,14 @@ Bugs
 ----
 
 Please `report bugs on GitHub <https://github.com/networkx/networkx/issues>`_.
+
+Policies
+--------
+
+All interactions with the project are subject to the
+:doc:`NetworkX code of conduct <code_of_conduct>`.
+
+We also follow these policies:
+
+* :doc:`NetworkX deprecation policy <deprecations>`
+* :doc:`Python version support <nep-0029-deprecation_policy>`
diff --git a/doc/conf.py b/doc/conf.py
index 75b2fa5e0..aee59e44d 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -220,6 +220,7 @@ latex_appendices = ["tutorial"]
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
+    "neps": ("https://numpy.org/neps", None),
     "matplotlib": ("https://matplotlib.org/stable", None),
     "scipy": ("https://docs.scipy.org/doc/scipy/reference", None),
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable", None),
diff --git a/networkx/drawing/layout.py b/networkx/drawing/layout.py
index ca0f04373..c16446e2c 100644
--- a/networkx/drawing/layout.py
+++ b/networkx/drawing/layout.py
@@ -1092,7 +1092,7 @@ def multipartite_layout(G, subset_key="subset", align="vertical", scale=1, cente
     nodes = []
 
     width = len(layers)
-    for i, layer in layers.items():
+    for i, layer in enumerate(layers.values()):
         height = len(layer)
         xs = np.repeat(i, height)
         ys = np.arange(0, height, dtype=float)
diff --git a/networkx/drawing/nx_pylab.py b/networkx/drawing/nx_pylab.py
index 836ff85ea..16fc59337 100644
--- a/networkx/drawing/nx_pylab.py
+++ b/networkx/drawing/nx_pylab.py
@@ -155,10 +155,11 @@ def draw_networkx(G, pos=None, arrows=None, with_labels=True, **kwds):
         For directed graphs, choose the style of the arrowsheads.
         See `matplotlib.patches.ArrowStyle` for more options.
 
-    arrowsize : int (default=10)
+    arrowsize : int or list (default=10)
         For directed graphs, choose the size of the arrow head's length and
-        width. See `matplotlib.patches.FancyArrowPatch` for attribute
-        `mutation_scale` for more info.
+        width. A list of values can be passed in to assign a different size for arrow head's length and width.
+        See `matplotlib.patches.FancyArrowPatch` for attribute `mutation_scale`
+        for more info.
 
     with_labels :  bool (default=True)
         Set to True to draw labels on the nodes.
@@ -750,7 +751,12 @@ def draw_networkx_edges(
 
         # Draw arrows with `matplotlib.patches.FancyarrowPatch`
         arrow_collection = []
-        mutation_scale = arrowsize  # scale factor of arrow head
+
+        if isinstance(arrowsize, list):
+            if len(arrowsize) != len(edge_pos):
+                raise ValueError("arrowsize should have the same length as edgelist")
+        else:
+            mutation_scale = arrowsize  # scale factor of arrow head
 
         base_connection_style = mpl.patches.ConnectionStyle(connectionstyle)
 
@@ -798,6 +804,11 @@ def draw_networkx_edges(
             x2, y2 = dst
             shrink_source = 0  # space from source to tail
             shrink_target = 0  # space from  head to target
+
+            if isinstance(arrowsize, list):
+                # Scale each factor of each arrow based on arrowsize list
+                mutation_scale = arrowsize[i]
+
             if np.iterable(node_size):  # many node sizes
                 source, target = edgelist[i][:2]
                 source_node_size = node_size[nodelist.index(source)]


Test patch:
diff --git a/networkx/drawing/tests/test_layout.py b/networkx/drawing/tests/test_layout.py
index 183c9fbce..ab9730884 100644
--- a/networkx/drawing/tests/test_layout.py
+++ b/networkx/drawing/tests/test_layout.py
@@ -412,3 +412,15 @@ class TestLayout:
         }
         for k, v in expectation.items():
             assert (s_vpos[k] == v).all()
+
+
+def test_multipartite_layout_nonnumeric_partition_labels():
+    """See gh-5123."""
+    G = nx.Graph()
+    G.add_node(0, subset="s0")
+    G.add_node(1, subset="s0")
+    G.add_node(2, subset="s1")
+    G.add_node(3, subset="s1")
+    G.add_edges_from([(0, 2), (0, 3), (1, 2)])
+    pos = nx.multipartite_layout(G)
+    assert len(pos) == len(G)
diff --git a/networkx/drawing/tests/test_pylab.py b/networkx/drawing/tests/test_pylab.py
index f4a1a73e6..4795613dd 100644
--- a/networkx/drawing/tests/test_pylab.py
+++ b/networkx/drawing/tests/test_pylab.py
@@ -478,6 +478,28 @@ def test_error_invalid_kwds():
         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
+
+
 def test_np_edgelist():
     # see issue #4129
     np = pytest.importorskip("numpy")