Replace NetworkX by RetworkX (#1791)

* Add retworkx as pl-requirements

* Add RX support to qaoa/cycle.py

* Update RX in qaoa/cycle.py

* Add RX support to qaoa/cost.py

* Update RX in qaoa/cost.py

* Add RX support to qaoa/mixers.py

* Update RX support in qaoa

* Add RX support to CircuitGraph and update QAOA tests

* Update CircuitGraph

* Update RX in circuit_graph

* Update CircuitGraph tests

* Update RX @ qaoa

* Update formatting

* Update circuit_graph

* Add RX support to CircuitGraph

* Update formatting

* Apply codecov suggestions

* Add lollipop_graph_rx

* Apply code review suggestions

* Update doc/releases/changelog-dev.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

* Add more comments

* Update has_path

Co-authored-by: Josh Izaac <josh146@gmail.com>

doc/releases/changelog-dev.md

@@ -122,6 +122,14 @@
 * Interferometer is now a class with `shape` method.
   [(#1946)](https://github.com/PennyLaneAI/pennylane/pull/1946)
 
+* The `CircuitGraph`, used to represent circuits via directed acyclic graphs, now
+  uses RetworkX for its internal representation. This results in significant speedup
+  for algorithms that rely on a directed acyclic graph representation.
+  [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791)
+
+* The QAOA module now accepts both NetworkX and RetworkX graphs as function inputs.
+  [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791)
+
 <h3>Breaking changes</h3>
 
 <h3>Bug fixes</h3>
@@ -151,4 +159,4 @@
 
 This release contains contributions from (in alphabetical order):
 
-Juan Miguel Arrazola, Esther Cruz, Olivia Di Matteo, Diego Guala, Ankit Khandelwal, Antal Száva, David Wierichs, Shaoming Zhang
+Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Olivia Di Matteo, Diego Guala, Ankit Khandelwal, Antal Száva, David Wierichs, Shaoming Zhang

pennylane/circuit_graph.py

@@ -18,7 +18,7 @@
 # pylint: disable=too-many-branches,too-many-arguments,too-many-instance-attributes
 from collections import Counter, OrderedDict, namedtuple
 
-import networkx as nx
+import retworkx as rx
 
 import pennylane as qml
 import numpy as np
@@ -167,22 +167,35 @@ def __init__(self, ops, obs, wires, par_info=None, trainable_params=None):
 
         # TODO: State preparations demolish the incoming state entirely, and therefore should have no incoming edges.
 
-        self._graph = nx.DiGraph()  #: nx.DiGraph: DAG representation of the quantum circuit
+        self._graph = rx.PyDiGraph(
+            multigraph=False
+        )  #: rx.PyDiGraph: DAG representation of the quantum circuit
+
         # Iterate over each (populated) wire in the grid
         for wire in self._grid.values():
             # Add the first operator on the wire to the graph
             # This operator does not depend on any others
-            self._graph.add_node(wire[0])
+
+            # Check if wire[0] in self._grid.values()
+            # is already added to the graph; this
+            # condition avoids adding new nodes with
+            # the same value but different indexes
+            if wire[0] not in self._graph.nodes():
+                self._graph.add_node(wire[0])
 
             for i in range(1, len(wire)):
                 # For subsequent operators on the wire:
-                if wire[i] not in self._graph:
+                if wire[i] not in self._graph.nodes():
                     # Add them to the graph if they are not already
                     # in the graph (multi-qubit operators might already have been placed)
                     self._graph.add_node(wire[i])
 
                 # Create an edge between this and the previous operator
-                self._graph.add_edge(wire[i - 1], wire[i])
+                # There isn't any default value for the edge-data in
+                # rx.PyDiGraph.add_edge(); this is set to an empty string
+                self._graph.add_edge(
+                    self._graph.nodes().index(wire[i - 1]), self._graph.nodes().index(wire[i]), ""
+                )
 
         # For computing depth; want only a graph with the operations, not
         # including the observables
@@ -264,7 +277,7 @@ def observables_in_order(self):
         Returns:
             list[Observable]: observables
         """
-        nodes = [node for node in self._graph.nodes if _is_observable(node)]
+        nodes = [node for node in self._graph.nodes() if _is_observable(node)]
         return sorted(nodes, key=_by_idx)
 
     @property
@@ -284,7 +297,7 @@ def operations_in_order(self):
         Returns:
             list[Operation]: operations
         """
-        nodes = [node for node in self._graph.nodes if not _is_observable(node)]
+        nodes = [node for node in self._graph.nodes() if not _is_observable(node)]
         return sorted(nodes, key=_by_idx)
 
     @property
@@ -300,7 +313,7 @@ def graph(self):
         and directed edges pointing from nodes to their immediate dependents/successors.
 
         Returns:
-            networkx.DiGraph: the directed acyclic graph representing the quantum circuit
+            retworkx.PyDiGraph: the directed acyclic graph representing the quantum circuit
         """
         return self._graph
 
@@ -324,7 +337,14 @@ def ancestors(self, ops):
         Returns:
             set[Operator]: ancestors of the given operators
         """
-        return set().union(*(nx.dag.ancestors(self._graph, o) for o in ops)) - set(ops)
+        anc = set(
+            self._graph.get_node_data(n)
+            for n in set().union(
+                # rx.ancestors() returns node indexes instead of node-values
+                *(rx.ancestors(self._graph, self._graph.nodes().index(o)) for o in ops)
+            )
+        )
+        return anc - set(ops)
 
     def descendants(self, ops):
         """Descendants of a given set of operators.
@@ -335,7 +355,14 @@ def descendants(self, ops):
         Returns:
             set[Operator]: descendants of the given operators
         """
-        return set().union(*(nx.dag.descendants(self._graph, o) for o in ops)) - set(ops)
+        des = set(
+            self._graph.get_node_data(n)
+            for n in set().union(
+                # rx.descendants() returns node indexes instead of node-values
+                *(rx.descendants(self._graph, self._graph.nodes().index(o)) for o in ops)
+            )
+        )
+        return des - set(ops)
 
     def _in_topological_order(self, ops):
         """Sorts a set of operators in the circuit in a topological order.
@@ -346,8 +373,9 @@ def _in_topological_order(self, ops):
         Returns:
             Iterable[Operator]: same set of operators, topologically ordered
         """
-        G = nx.DiGraph(self._graph.subgraph(ops))
-        return nx.dag.topological_sort(G)
+        G = self._graph.subgraph(list(self._graph.nodes().index(o) for o in ops))
+        indexes = rx.topological_sort(G)
+        return list(G[x] for x in indexes)
 
     def ancestors_in_order(self, ops):
         """Operator ancestors in a topological order.
@@ -360,8 +388,7 @@ def ancestors_in_order(self, ops):
         Returns:
             list[Operator]: ancestors of the given operators, topologically ordered
         """
-        # return self._in_topological_order(self.ancestors(ops))  # an abitrary topological order
-        return sorted(self.ancestors(ops), key=_by_idx)
+        return sorted(self.ancestors(ops), key=_by_idx)  # an abitrary topological order
 
     def descendants_in_order(self, ops):
         """Operator descendants in a topological order.
@@ -585,8 +612,10 @@ def update_node(self, old, new):
         # NOTE Does not alter the graph edges in any way. variable_deps is not changed, _grid is not changed. Dangerous!
         if new.wires != old.wires:
             raise ValueError("The new Operator must act on the same wires as the old one.")
+
         new.queue_idx = old.queue_idx
-        nx.relabel_nodes(self._graph, {old: new}, copy=False)  # change the graph in place
+        self._graph[self._graph.nodes().index(old)] = new
+
         self._operations = self.operations_in_order
         self._observables = self.observables_in_order
 
@@ -632,9 +661,10 @@ def get_depth(self):
         # expressed in terms of edges, and we want it in terms of nodes).
         if self._depth is None and self.operations:
             if self._operation_graph is None:
-                self._operation_graph = self.graph.subgraph(self.operations)
-                self._depth = nx.dag_longest_path_length(self._operation_graph) + 1
-
+                self._operation_graph = self._graph.subgraph(
+                    list(self._graph.nodes().index(node) for node in self.operations)
+                )
+                self._depth = rx.dag_longest_path_length(self._operation_graph) + 1
         return self._depth
 
     def has_path(self, a, b):
@@ -647,7 +677,22 @@ def has_path(self, a, b):
         Returns:
             bool: returns ``True`` if a path exists
         """
-        return nx.has_path(self._graph, a, b)
+        if a == b:
+            return True
+
+        return (
+            len(
+                rx.digraph_dijkstra_shortest_paths(
+                    self._graph,
+                    self._graph.nodes().index(a),
+                    self._graph.nodes().index(b),
+                    weight_fn=None,
+                    default_weight=1.0,
+                    as_undirected=False,
+                )
+            )
+            != 0
+        )
 
     @property
     def max_simultaneous_measurements(self):