From f20219ed62aaade9ff21e5867ea828ee20baef9b Mon Sep 17 00:00:00 2001
From: Erik Welch <erik.n.welch@gmail.com>
Date: Mon, 22 Jan 2024 10:31:48 -0600
Subject: [PATCH] nx-cugraph: add `is_tree`, etc. (#4097)

These are now possible b/c we have `connected_components` and `weakly_connected_components` (and their `is_*` equivalents).

Authors:
  - Erik Welch (https://github.com/eriknw)
  - Rick Ratzel (https://github.com/rlratzel)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)

URL: https://github.com/rapidsai/cugraph/pull/4097
---
 python/nx-cugraph/_nx_cugraph/__init__.py     |  4 +
 .../nx_cugraph/algorithms/__init__.py         |  2 +
 .../nx_cugraph/algorithms/isolate.py          |  2 +-
 .../nx_cugraph/algorithms/tree/__init__.py    | 13 ++++
 .../nx_cugraph/algorithms/tree/recognition.py | 74 +++++++++++++++++++
 .../nx-cugraph/nx_cugraph/classes/digraph.py  | 12 +--
 .../nx-cugraph/nx_cugraph/classes/function.py |  4 +-
 python/nx-cugraph/nx_cugraph/classes/graph.py |  4 +-
 8 files changed, 107 insertions(+), 8 deletions(-)
 create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/tree/__init__.py
 create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py

diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py
index 8deac55f4ad..9bca031a2f0 100644
--- a/python/nx-cugraph/_nx_cugraph/__init__.py
+++ b/python/nx-cugraph/_nx_cugraph/__init__.py
@@ -71,10 +71,14 @@
         "house_x_graph",
         "icosahedral_graph",
         "in_degree_centrality",
+        "is_arborescence",
         "is_bipartite",
+        "is_branching",
         "is_connected",
+        "is_forest",
         "is_isolate",
         "is_strongly_connected",
+        "is_tree",
         "is_weakly_connected",
         "isolates",
         "k_truss",
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
index de4e9466ba0..08658ad94cb 100644
--- a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
+++ b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
@@ -19,6 +19,7 @@
     link_analysis,
     shortest_paths,
     traversal,
+    tree,
 )
 from .bipartite import complete_bipartite_graph, is_bipartite
 from .centrality import *
@@ -31,3 +32,4 @@
 from .reciprocity import *
 from .shortest_paths import *
 from .traversal import *
+from .tree.recognition import *
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/isolate.py b/python/nx-cugraph/nx_cugraph/algorithms/isolate.py
index 62b47a9b354..9621fbeaa9d 100644
--- a/python/nx-cugraph/nx_cugraph/algorithms/isolate.py
+++ b/python/nx-cugraph/nx_cugraph/algorithms/isolate.py
@@ -70,4 +70,4 @@ def isolates(G):
 @networkx_algorithm(version_added="23.10")
 def number_of_isolates(G):
     G = _to_graph(G)
-    return _mark_isolates(G).sum().tolist()
+    return int(cp.count_nonzero(_mark_isolates(G)))
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/tree/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/tree/__init__.py
new file mode 100644
index 00000000000..91bf72417be
--- /dev/null
+++ b/python/nx-cugraph/nx_cugraph/algorithms/tree/__init__.py
@@ -0,0 +1,13 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .recognition import *
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py b/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py
new file mode 100644
index 00000000000..0b82f079d43
--- /dev/null
+++ b/python/nx-cugraph/nx_cugraph/algorithms/tree/recognition.py
@@ -0,0 +1,74 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import cupy as cp
+import networkx as nx
+
+import nx_cugraph as nxcg
+from nx_cugraph.convert import _to_directed_graph, _to_graph
+from nx_cugraph.utils import networkx_algorithm, not_implemented_for
+
+__all__ = ["is_arborescence", "is_branching", "is_forest", "is_tree"]
+
+
+@not_implemented_for("undirected")
+@networkx_algorithm(plc="weakly_connected_components", version_added="24.02")
+def is_arborescence(G):
+    G = _to_directed_graph(G)
+    return is_tree(G) and int(G._in_degrees_array().max()) <= 1
+
+
+@not_implemented_for("undirected")
+@networkx_algorithm(plc="weakly_connected_components", version_added="24.02")
+def is_branching(G):
+    G = _to_directed_graph(G)
+    return is_forest(G) and int(G._in_degrees_array().max()) <= 1
+
+
+@networkx_algorithm(plc="weakly_connected_components", version_added="24.02")
+def is_forest(G):
+    G = _to_graph(G)
+    if len(G) == 0:
+        raise nx.NetworkXPointlessConcept("G has no nodes.")
+    if is_directed := G.is_directed():
+        connected_components = nxcg.weakly_connected_components
+    else:
+        connected_components = nxcg.connected_components
+    for components in connected_components(G):
+        node_ids = G._list_to_nodearray(list(components))
+        # TODO: create utilities for creating subgraphs
+        mask = cp.isin(G.src_indices, node_ids) & cp.isin(G.dst_indices, node_ids)
+        # A tree must have an edge count equal to the number of nodes minus the
+        # tree's root node.
+        if is_directed:
+            if int(cp.count_nonzero(mask)) != len(components) - 1:
+                return False
+        else:
+            src_indices = G.src_indices[mask]
+            dst_indices = G.dst_indices[mask]
+            if int(cp.count_nonzero(src_indices <= dst_indices)) != len(components) - 1:
+                return False
+    return True
+
+
+@networkx_algorithm(plc="weakly_connected_components", version_added="24.02")
+def is_tree(G):
+    G = _to_graph(G)
+    if len(G) == 0:
+        raise nx.NetworkXPointlessConcept("G has no nodes.")
+    if G.is_directed():
+        is_connected = nxcg.is_weakly_connected
+    else:
+        is_connected = nxcg.is_connected
+    # A tree must have an edge count equal to the number of nodes minus the
+    # tree's root node.
+    return len(G) - 1 == G.number_of_edges() and is_connected(G)
diff --git a/python/nx-cugraph/nx_cugraph/classes/digraph.py b/python/nx-cugraph/nx_cugraph/classes/digraph.py
index f8217a2c79f..169815eb067 100644
--- a/python/nx-cugraph/nx_cugraph/classes/digraph.py
+++ b/python/nx-cugraph/nx_cugraph/classes/digraph.py
@@ -25,7 +25,7 @@
 from .graph import Graph
 
 if TYPE_CHECKING:  # pragma: no cover
-    from nx_cugraph.typing import NodeKey
+    from nx_cugraph.typing import AttrKey
 
 __all__ = ["DiGraph"]
 
@@ -47,10 +47,8 @@ def to_networkx_class(cls) -> type[nx.DiGraph]:
         return nx.DiGraph
 
     @networkx_api
-    def number_of_edges(
-        self, u: NodeKey | None = None, v: NodeKey | None = None
-    ) -> int:
-        if u is not None or v is not None:
+    def size(self, weight: AttrKey | None = None) -> int:
+        if weight is not None:
             raise NotImplementedError
         return self.src_indices.size
 
@@ -182,6 +180,8 @@ def _in_degrees_array(self, *, ignore_selfloops=False):
         if ignore_selfloops:
             not_selfloops = self.src_indices != dst_indices
             dst_indices = dst_indices[not_selfloops]
+        if dst_indices.size == 0:
+            return cp.zeros(self._N, dtype=np.int64)
         return cp.bincount(dst_indices, minlength=self._N)
 
     def _out_degrees_array(self, *, ignore_selfloops=False):
@@ -189,4 +189,6 @@ def _out_degrees_array(self, *, ignore_selfloops=False):
         if ignore_selfloops:
             not_selfloops = src_indices != self.dst_indices
             src_indices = src_indices[not_selfloops]
+        if src_indices.size == 0:
+            return cp.zeros(self._N, dtype=np.int64)
         return cp.bincount(src_indices, minlength=self._N)
diff --git a/python/nx-cugraph/nx_cugraph/classes/function.py b/python/nx-cugraph/nx_cugraph/classes/function.py
index 435dfe37239..7212a4d2da9 100644
--- a/python/nx-cugraph/nx_cugraph/classes/function.py
+++ b/python/nx-cugraph/nx_cugraph/classes/function.py
@@ -10,6 +10,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import cupy as cp
+
 from nx_cugraph.convert import _to_graph
 from nx_cugraph.utils import networkx_algorithm
 
@@ -20,4 +22,4 @@
 def number_of_selfloops(G):
     G = _to_graph(G)
     is_selfloop = G.src_indices == G.dst_indices
-    return is_selfloop.sum().tolist()
+    return int(cp.count_nonzero(is_selfloop))
diff --git a/python/nx-cugraph/nx_cugraph/classes/graph.py b/python/nx-cugraph/nx_cugraph/classes/graph.py
index 4aa2de1538e..f697668750d 100644
--- a/python/nx-cugraph/nx_cugraph/classes/graph.py
+++ b/python/nx-cugraph/nx_cugraph/classes/graph.py
@@ -522,7 +522,7 @@ def size(self, weight: AttrKey | None = None) -> int:
         if weight is not None:
             raise NotImplementedError
         # If no self-edges, then `self.src_indices.size // 2`
-        return int((self.src_indices <= self.dst_indices).sum())
+        return int(cp.count_nonzero(self.src_indices <= self.dst_indices))
 
     @networkx_api
     def to_directed(self, as_view: bool = False) -> nxcg.DiGraph:
@@ -740,6 +740,8 @@ def _degrees_array(self, *, ignore_selfloops=False):
             src_indices = src_indices[not_selfloops]
             if self.is_directed():
                 dst_indices = dst_indices[not_selfloops]
+        if src_indices.size == 0:
+            return cp.zeros(self._N, dtype=np.int64)
         degrees = cp.bincount(src_indices, minlength=self._N)
         if self.is_directed():
             degrees += cp.bincount(dst_indices, minlength=self._N)