From 9a261ffd1f5e055e9c8b751c6009f99e9c39d0c1 Mon Sep 17 00:00:00 2001
From: Erik Welch <erik.n.welch@gmail.com>
Date: Thu, 25 Jan 2024 14:05:45 -0600
Subject: [PATCH] nx-cugraph: add `complement` and `reverse` (#4103)

We apparently already had `G.reverse()`, which made that function extra easy :)

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

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

URL: https://github.com/rapidsai/cugraph/pull/4103
---
 python/nx-cugraph/_nx_cugraph/__init__.py     |  2 +
 .../nx_cugraph/algorithms/__init__.py         |  2 +
 .../algorithms/operators/__init__.py          | 13 +++++
 .../nx_cugraph/algorithms/operators/unary.py  | 55 +++++++++++++++++++
 4 files changed, 72 insertions(+)
 create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py
 create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py

diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py
index 9bca031a2f0..2f283aa153c 100644
--- a/python/nx-cugraph/_nx_cugraph/__init__.py
+++ b/python/nx-cugraph/_nx_cugraph/__init__.py
@@ -43,6 +43,7 @@
         "chvatal_graph",
         "circular_ladder_graph",
         "clustering",
+        "complement",
         "complete_bipartite_graph",
         "complete_graph",
         "complete_multipartite_graph",
@@ -105,6 +106,7 @@
         "path_graph",
         "petersen_graph",
         "reciprocity",
+        "reverse",
         "sedgewick_maze_graph",
         "single_source_shortest_path_length",
         "single_target_shortest_path_length",
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
index 08658ad94cb..7aafa85f5b7 100644
--- a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
+++ b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py
@@ -17,6 +17,7 @@
     community,
     components,
     link_analysis,
+    operators,
     shortest_paths,
     traversal,
     tree,
@@ -29,6 +30,7 @@
 from .dag import *
 from .isolate import *
 from .link_analysis import *
+from .operators import *
 from .reciprocity import *
 from .shortest_paths import *
 from .traversal import *
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/operators/__init__.py
new file mode 100644
index 00000000000..32fd45f5726
--- /dev/null
+++ b/python/nx-cugraph/nx_cugraph/algorithms/operators/__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 .unary import *
diff --git a/python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py b/python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py
new file mode 100644
index 00000000000..08abc9f2872
--- /dev/null
+++ b/python/nx-cugraph/nx_cugraph/algorithms/operators/unary.py
@@ -0,0 +1,55 @@
+# 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 numpy as np
+
+import nx_cugraph as nxcg
+from nx_cugraph.convert import _to_graph
+from nx_cugraph.utils import index_dtype, networkx_algorithm
+
+__all__ = ["complement", "reverse"]
+
+
+@networkx_algorithm(version_added="24.02")
+def complement(G):
+    G = _to_graph(G)
+    N = G._N
+    # Upcast to int64 so indices don't overflow.
+    edges_a_b = N * G.src_indices.astype(np.int64) + G.dst_indices
+    # Now compute flattened indices for all edges except self-loops
+    # Alt (slower):
+    # edges_full = np.arange(N * N)
+    # edges_full = edges_full[(edges_full % (N + 1)).astype(bool)]
+    edges_full = cp.arange(1, N * (N - 1) + 1) + cp.repeat(cp.arange(N - 1), N)
+    edges_comp = cp.setdiff1d(
+        edges_full,
+        edges_a_b,
+        assume_unique=not G.is_multigraph(),
+    )
+    src_indices, dst_indices = cp.divmod(edges_comp, N)
+    return G.__class__.from_coo(
+        N,
+        src_indices.astype(index_dtype),
+        dst_indices.astype(index_dtype),
+        key_to_id=G.key_to_id,
+    )
+
+
+@networkx_algorithm(version_added="24.02")
+def reverse(G, copy=True):
+    if not G.is_directed():
+        raise nx.NetworkXError("Cannot reverse an undirected graph.")
+    if isinstance(G, nx.Graph):
+        G = nxcg.from_networkx(G, preserve_all_attrs=True)
+    return G.reverse(copy=copy)