nx-cugraph: add weakly connected components

python/nx-cugraph/_nx_cugraph/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-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
@@ -70,6 +70,7 @@
         "in_degree_centrality",
         "is_connected",
         "is_isolate",
+        "is_weakly_connected",
         "isolates",
         "k_truss",
         "karate_club_graph",
@@ -85,6 +86,7 @@
         "number_connected_components",
         "number_of_isolates",
         "number_of_selfloops",
+        "number_weakly_connected_components",
         "octahedral_graph",
         "out_degree_centrality",
         "pagerank",
@@ -102,6 +104,7 @@
         "truncated_tetrahedron_graph",
         "turan_graph",
         "tutte_graph",
+        "weakly_connected_components",
         "wheel_graph",
         # END: functions
     },

python/nx-cugraph/nx_cugraph/algorithms/components/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-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
@@ -11,3 +11,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from .connected import *
+from .weakly_connected import *

python/nx-cugraph/nx_cugraph/algorithms/components/connected.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-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
@@ -32,7 +32,12 @@
 @not_implemented_for("directed")
 @networkx_algorithm
 def number_connected_components(G):
-    return sum(1 for _ in connected_components(G))
+    G = _to_undirected_graph(G)
+    return _number_connected_components(G)
+
+
+def _number_connected_components(G, symmetrize=None):
+    return sum(1 for _ in _connected_components(G, symmetrize=symmetrize))
     # PREFERRED IMPLEMENTATION, BUT PLC DOES NOT HANDLE ISOLATED VERTICES WELL
     # G = _to_undirected_graph(G)
     # unused_node_ids, labels = plc.weakly_connected_components(
@@ -60,12 +65,16 @@ def _(G):
 @networkx_algorithm
 def connected_components(G):
     G = _to_undirected_graph(G)
+    return _connected_components(G)
+
+
+def _connected_components(G, symmetrize=None):
     if G.src_indices.size == 0:
         # TODO: PLC doesn't handle empty graphs (or isolated nodes) gracefully!
         return [{key} for key in G._nodeiter_to_iter(range(len(G)))]
     node_ids, labels = plc.weakly_connected_components(
         resource_handle=plc.ResourceHandle(),
-        graph=G._get_plc_graph(),
+        graph=G._get_plc_graph(symmetrize=symmetrize),
         offsets=None,
         indices=None,
         weights=None,
@@ -75,7 +84,7 @@ def connected_components(G):
     groups = _groupby(labels, node_ids)
     it = (G._nodearray_to_set(connected_ids) for connected_ids in groups.values())
     # TODO: PLC doesn't handle isolated vertices yet, so this is a temporary fix
-    isolates = _isolates(G)
+    isolates = _isolates(G, symmetrize=symmetrize)
     if isolates.size > 0:
         isolates = isolates[isolates > node_ids.max()]
         if isolates.size > 0:
@@ -89,11 +98,15 @@ def connected_components(G):
 @networkx_algorithm
 def is_connected(G):
     G = _to_undirected_graph(G)
+    return _is_connected(G)
+
+
+def _is_connected(G, symmetrize=None):
     if len(G) == 0:
         raise nx.NetworkXPointlessConcept(
             "Connectivity is undefined for the null graph."
         )
-    for community in connected_components(G):
+    for community in _connected_components(G, symmetrize=symmetrize):
         return len(community) == len(G)
     raise RuntimeError  # pragma: no cover
     # PREFERRED IMPLEMENTATION, BUT PLC DOES NOT HANDLE ISOLATED VERTICES WELL

python/nx-cugraph/nx_cugraph/algorithms/components/weakly_connected.py

@@ -0,0 +1,47 @@
+# 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 nx_cugraph.convert import _to_directed_graph
+from nx_cugraph.utils import networkx_algorithm, not_implemented_for
+
+from .connected import (
+    _connected_components,
+    _is_connected,
+    _number_connected_components,
+)
+
+__all__ = [
+    "number_weakly_connected_components",
+    "weakly_connected_components",
+    "is_weakly_connected",
+]
+
+
+@not_implemented_for("undirected")
+@networkx_algorithm
+def weakly_connected_components(G):
+    G = _to_directed_graph(G)
+    return _connected_components(G, symmetrize="union")
+
+
+@not_implemented_for("undirected")
+@networkx_algorithm
+def number_weakly_connected_components(G):
+    G = _to_directed_graph(G)
+    return _number_connected_components(G, symmetrize="union")
+
+
+@not_implemented_for("undirected")
+@networkx_algorithm
+def is_weakly_connected(G):
+    G = _to_directed_graph(G)
+    return _is_connected(G, symmetrize="union")

python/nx-cugraph/nx_cugraph/algorithms/isolate.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-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
@@ -15,9 +15,10 @@
 from typing import TYPE_CHECKING
 
 import cupy as cp
+import numpy as np
 
 from nx_cugraph.convert import _to_graph
-from nx_cugraph.utils import networkx_algorithm
+from nx_cugraph.utils import index_dtype, networkx_algorithm
 
 if TYPE_CHECKING:  # pragma: no cover
     from nx_cugraph.typing import IndexValue
@@ -36,19 +37,28 @@ def is_isolate(G, n):
     )
 
 
-def _mark_isolates(G) -> cp.ndarray[bool]:
+def _mark_isolates(G, symmetrize=None) -> cp.ndarray[bool]:
     """Return a boolean mask array indicating indices of isolated nodes."""
     mark_isolates = cp.ones(len(G), bool)
-    mark_isolates[G.src_indices] = False
-    if G.is_directed():
-        mark_isolates[G.dst_indices] = False
+    if G.is_directed() and symmetrize == "intersection":
+        N = G._N
+        # Upcast to int64 so indices don't overflow
+        src_dst = N * G.src_indices.astype(np.int64) + G.dst_indices
+        src_dst_T = G.src_indices + N * G.dst_indices.astype(np.int64)
+        src_dst_new = cp.intersect1d(src_dst, src_dst_T)
+        new_indices = cp.floor_divide(src_dst_new, N, dtype=index_dtype)
+        mark_isolates[new_indices] = False
+    else:
+        mark_isolates[G.src_indices] = False
+        if G.is_directed():
+            mark_isolates[G.dst_indices] = False
     return mark_isolates
 
 
-def _isolates(G) -> cp.ndarray[IndexValue]:
+def _isolates(G, symmetrize=None) -> cp.ndarray[IndexValue]:
     """Like isolates, but return an array of indices instead of an iterator of nodes."""
     G = _to_graph(G)
-    return cp.nonzero(_mark_isolates(G))[0]
+    return cp.nonzero(_mark_isolates(G, symmetrize=symmetrize))[0]
 
 
 @networkx_algorithm

python/nx-cugraph/nx_cugraph/classes/graph.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-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
@@ -579,6 +579,7 @@ def _get_plc_graph(
         store_transposed: bool = False,
         switch_indices: bool = False,
         edge_array: cp.ndarray[EdgeValue] | None = None,
+        symmetrize: str | None = None,
     ):
         if edge_array is not None or edge_attr is None:
             pass
@@ -637,11 +638,29 @@ def _get_plc_graph(
         dst_indices = self.dst_indices
         if switch_indices:
             src_indices, dst_indices = dst_indices, src_indices
+        if symmetrize is not None:
+            if edge_array is not None:
+                raise NotImplementedError(
+                    "edge_array must be None when symmetrizing the graph"
+                )
+            N = self._N
+            # Upcast to int64 so indices don't overflow
+            src_dst = N * src_indices.astype(np.int64) + dst_indices
+            src_dst_T = src_indices + N * dst_indices.astype(np.int64)
+            if symmetrize == "union":
+                src_dst_new = cp.union1d(src_dst, src_dst_T)
+            elif symmetrize == "intersection":
+                src_dst_new = cp.intersect1d(src_dst, src_dst_T)
+            else:
+                raise ValueError(
+                    f'symmetrize must be "union" or "intersection"; got "{symmetrize}"'
+                )
+            src_indices, dst_indices = cp.divmod(src_dst_new, N, dtype=index_dtype)
         return plc.SGGraph(
             resource_handle=plc.ResourceHandle(),
             graph_properties=plc.GraphProperties(
-                is_multigraph=self.is_multigraph(),
-                is_symmetric=not self.is_directed(),
+                is_multigraph=self.is_multigraph() and symmetrize is None,
+                is_symmetric=not self.is_directed() or symmetrize is not None,
             ),
             src_or_offset_array=src_indices,
             dst_or_index_array=dst_indices,