Fix numeric and degree assortativity coefficient calculation (#4928)

3 files changed, 38 insertions, 10 deletions

diff --git a/networkx/algorithms/assortativity/correlation.py b/networkx/algorithms/assortativity/correlation.py
index b885a22e..8a244fc7 100644
--- a/networkx/algorithms/assortativity/correlation.py
+++ b/networkx/algorithms/assortativity/correlation.py
@@ -73,8 +73,12 @@ def degree_assortativity_coefficient(G, x="out", y="in", weight=None, nodes=None
     .. [2] Foster, J.G., Foster, D.V., Grassberger, P. & Paczuski, M.
        Edge direction and the structure of networks, PNAS 107, 10815-20 (2010).
     """
-    M = degree_mixing_matrix(G, x=x, y=y, nodes=nodes, weight=weight)
-    return numeric_ac(M)
+    if nodes is None:
+        nodes = G.nodes
+    degrees = set([d for n, d in G.degree(nodes, weight=weight)])
+    mapping = {d: i for i, d, in enumerate(degrees)}
+    M = degree_mixing_matrix(G, x=x, y=y, nodes=nodes, weight=weight, mapping=mapping)
+    return numeric_ac(M, mapping=mapping)
 
 
 def degree_pearson_correlation_coefficient(G, x="out", y="in", weight=None, nodes=None):
@@ -223,8 +227,12 @@ def numeric_assortativity_coefficient(G, attribute, nodes=None):
     .. [1] M. E. J. Newman, Mixing patterns in networks
            Physical Review E, 67 026126, 2003
     """
-    a = numeric_mixing_matrix(G, attribute, nodes)
-    return numeric_ac(a)
+    if nodes is None:
+        nodes = G.nodes
+    vals = set(G.nodes[n][attribute] for n in nodes)
+    mapping = {d: i for i, d, in enumerate(vals)}
+    M = attribute_mixing_matrix(G, attribute, nodes, mapping)
+    return numeric_ac(M, mapping)
 
 
 def attribute_ac(M):
@@ -254,7 +262,7 @@ def attribute_ac(M):
     return r
 
 
-def numeric_ac(M):
+def numeric_ac(M, mapping):
     # M is a numpy matrix or array
     # numeric assortativity coefficient, pearsonr
     import numpy as np
@@ -262,12 +270,13 @@ def numeric_ac(M):
     if M.sum() != 1.0:
         M = M / float(M.sum())
     nx, ny = M.shape  # nx=ny
-    x = np.arange(nx)
-    y = np.arange(ny)
+    x = np.array(list(mapping.keys()))
+    y = x  # x and y have the same support
+    idx = list(mapping.values())
     a = M.sum(axis=0)
     b = M.sum(axis=1)
-    vara = (a * x ** 2).sum() - ((a * x).sum()) ** 2
-    varb = (b * x ** 2).sum() - ((b * x).sum()) ** 2
+    vara = (a[idx] * x ** 2).sum() - ((a[idx] * x).sum()) ** 2
+    varb = (b[idx] * y ** 2).sum() - ((b[idx] * y).sum()) ** 2
     xy = np.outer(x, y)
-    ab = np.outer(a, b)
+    ab = np.outer(a[idx], b[idx])
     return (xy * (M - ab)).sum() / np.sqrt(vara * varb)
diff --git a/networkx/algorithms/assortativity/tests/base_test.py b/networkx/algorithms/assortativity/tests/base_test.py
index 5b371231..ee5b126c 100644
--- a/networkx/algorithms/assortativity/tests/base_test.py
+++ b/networkx/algorithms/assortativity/tests/base_test.py
@@ -52,6 +52,10 @@ class BaseTestDegreeMixing:
         cls.W = nx.Graph()
         cls.W.add_edges_from([(0, 3), (1, 3), (2, 3)], weight=0.5)
         cls.W.add_edge(0, 2, weight=1)
+        S1 = nx.star_graph(4)
+        S2 = nx.star_graph(4)
+        cls.DS = nx.disjoint_union(S1, S2)
+        cls.DS.add_edge(4, 5)
 
 
 class BaseTestNumericMixing:
@@ -70,3 +74,10 @@ class BaseTestNumericMixing:
         F.add_edge(0, 2, weight=1)
         nx.set_node_attributes(F, dict(F.degree(weight="weight")), "margin")
         cls.F = F
+
+        M = nx.Graph()
+        M.add_nodes_from([1, 2], margin=-1)
+        M.add_nodes_from([3], margin=1)
+        M.add_nodes_from([4], margin=2)
+        M.add_edges_from([(3, 4), (1, 2), (1, 3)])
+        cls.M = M
diff --git a/networkx/algorithms/assortativity/tests/test_correlation.py b/networkx/algorithms/assortativity/tests/test_correlation.py
index 763f0a41..8da72652 100644
--- a/networkx/algorithms/assortativity/tests/test_correlation.py
+++ b/networkx/algorithms/assortativity/tests/test_correlation.py
@@ -42,6 +42,10 @@ class TestDegreeMixingCorrelation(BaseTestDegreeMixing):
         r = nx.degree_assortativity_coefficient(self.W, weight="weight")
         np.testing.assert_almost_equal(r, -0.1429, decimal=4)
 
+    def test_degree_assortativity_double_star(self):
+        r = nx.degree_assortativity_coefficient(self.DS)
+        np.testing.assert_almost_equal(r, -0.9339, decimal=4)
+
 
 class TestAttributeMixingCorrelation(BaseTestAttributeMixing):
     def test_attribute_assortativity_undirected(self):
@@ -91,3 +95,7 @@ class TestNumericMixingCorrelation(BaseTestNumericMixing):
     def test_numeric_assortativity_float(self):
         r = nx.numeric_assortativity_coefficient(self.F, "margin")
         np.testing.assert_almost_equal(r, -0.1429, decimal=4)
+
+    def test_numeric_assortativity_mixed(self):
+        r = nx.numeric_assortativity_coefficient(self.M, "margin")
+        np.testing.assert_almost_equal(r, 0.4340, decimal=4)