Bugfix: degrees and clustering for a subset of nodes

nngt/analysis/clustering.py

@@ -24,6 +24,7 @@
 import numpy as np
 
 import nngt
+from nngt.lib import nonstring_container
 
 
 __all__ = [
@@ -172,7 +173,10 @@ def local_clustering_binary_undirected(g, nodes=None):
     triangles = triangle_count(g, weights=None, nodes=nodes, directed=False)
     triplets  = triplet_count(g, weights=None, nodes=nodes, directed=False)
 
-    triplets[triangles == 0] = 1
+    if nonstring_container(triangles):
+        triplets[triangles == 0] = 1
+    elif triangles == 0:
+        return 0
 
     return triangles / triplets
 
@@ -295,22 +299,23 @@ def local_clustering(g, nodes=None, directed=True, weights=None,
     directed *= g.is_directed()
     weighted  = weights not in (None, False)
 
+    triplets, triangles = None, None
+
     if not directed and not weighted:
         # undirected binary clustering uses the library method
         return local_clustering_binary_undirected(g, nodes=nodes)
     elif not weighted:
         # directed clustering
         triangles = triangle_count(g, nodes=nodes, mode=mode)
         triplets  = triplet_count(g, nodes, mode=mode).astype(float)
+    else:
+        triangles, triplets = _triangles_and_triplets(
+            g, directed, weights, method, mode, combine_weights, nodes)
 
+    if nonstring_container(triplets):
         triplets[triangles == 0] = 1
-
-        return triangles / triplets
-
-    triangles, triplets = _triangles_and_triplets(
-        g, directed, weights, method, mode, combine_weights, nodes)
-
-    triplets[triangles == 0] = 1
+    elif triangles == 0:
+        return 0
 
     return triangles / triplets
 
@@ -468,11 +473,17 @@ def triplet_count(g, nodes=None, directed=True, weights=None,
                 _, adjsym = _get_matrices(g, directed, None, False,
                                           combine_weights)
 
-                deg = adjsym.sum(axis=0).A1
+                if nodes is None:
+                    deg = adjsym.sum(axis=0).A1
+                else:
+                    deg = adjsym.sum(axis=0).A1[nodes]
             else:
                 deg = g.get_degrees(nodes=nodes)
 
-            return (0.5*deg*(deg - 1)).astype(int)
+            if nodes is None or nonstring_container(nodes):
+                return (0.5*deg*(deg - 1)).astype(int)
+
+            return 0.5*deg*(deg - 1)
 
         # directed
         if mode in ("total", "cycle", "middleman"):

nngt/analysis/nx_functions.py

@@ -72,12 +72,23 @@ def local_clustering_binary_undirected(g, nodes=None):
     ----------
     .. [nx-local-clustering] :nxdoc:`algorithms.cluster.clustering`
     '''
+    num_nodes = g.node_nb()
+
+    if nonstring_container(nodes):
+        num_nodes = len(nodes)
+    elif nodes is not None:
+        num_nodes = 1
+
     lc = nx.clustering(g.graph.to_undirected(as_view=True), nodes=nodes,
                        weight=None)
-
-    lc = np.array([lc[i] for i in range(g.node_nb())], dtype=float)
 
-    return lc
+    if num_nodes == 1:
+        return lc
+
+    if nodes is None:
+        nodes = list(range(num_nodes))
+
+    return np.array([lc[n] for n in nodes], dtype=float)
 
 
 def assortativity(g, degree, weights=None):

nngt/core/nngt_graph.py

@@ -717,9 +717,12 @@ def get_degrees(self, mode="total", nodes=None, weights=None):
         if nodes is None:
             num_nodes = self.node_nb()
             nodes = slice(num_nodes)
-        else:
+        elif nonstring_container(nodes):
             nodes = list(nodes)
             num_nodes = len(nodes)
+        else:
+            nodes = [nodes]
+            num_nodes = 1
 
         # weighted
         if nonstring_container(weights) or weights in self._eattr:
@@ -750,6 +753,9 @@ def get_degrees(self, mode="total", nodes=None, weights=None):
             else:
                 degrees += [g._out_deg[i] for i in nodes]
 
+        if num_nodes == 1:
+            return degrees[0]
+
         return degrees
 
     def neighbours(self, node, mode="all"):

nngt/core/nx_graph.py

@@ -644,7 +644,10 @@ def get_degrees(self, mode="total", nodes=None, weights=None):
         else:
             raise ValueError("Unknown `mode` '{}'".format(mode))
 
-        return np.array([di_deg[i] for i in nodes], dtype=dtype)
+        if nonstring_container(nodes):
+            return np.array([di_deg[i] for i in nodes], dtype=dtype)
+
+        return di_deg
 
     def is_connected(self, mode="strong"):
         '''

testing/test_analysis.py

@@ -325,12 +325,25 @@ def test_clustering_parameters():
     cc_bu = na.local_clustering_binary_undirected(g)
 
     for m in methods:
+        # combine
         for combine in ("sum", "max", "min", "mean"):
             cc_wu = na.local_clustering(g, directed=False, weights=True,
                                         combine_weights=combine, method=m)
 
             assert np.all(np.isclose(cc_wu, cc_bu))
 
+        # subset of nodes
+        nodes = [1, [1, 2]]
+
+        for n in nodes:
+            cc = na.local_clustering_binary_undirected(g, nodes=n)
+
+            assert np.all(np.isclose(cc, cc_bu[n]))
+
+            cc = na.local_clustering(g, nodes=n, method=m)
+
+            assert np.all(np.isclose(cc, cc_bu[n]))
+
     # check different but equivalent matrices
     for m in methods:
         W = np.array([
@@ -358,6 +371,16 @@ def test_clustering_parameters():
 
         assert np.all(np.isclose(cc_sum, cc_sym))
 
+        # subset of nodes
+        nodes = [1, [1, 2]]
+
+        cc_dir = na.local_clustering(g, weights="weight", method=m)
+
+        for n in nodes:
+            cc = na.local_clustering(g, nodes=n, weights="weight", method=m)
+
+            assert np.all(np.isclose(cc, cc_dir[n]))
+
 
 @pytest.mark.mpi_skip
 def test_partial_directed_clustering():
@@ -617,4 +640,5 @@ def test_swp():
         test_iedges()
         test_swp()
         test_partial_directed_clustering()
+        nngt.use_backend("nngt")
         test_clustering_parameters()