improving quantum_causal_cone method in python

qiskit/dagcircuit/dagcircuit.py

@@ -2264,36 +2264,44 @@ def quantum_causal_cone(self, qubit):
         output_node = self.output_map.get(qubit, None)
         if not output_node:
             raise DAGCircuitError(f"Qubit {qubit} is not part of this circuit.")
-        # Add the qubit to the causal cone.
-        qubits_to_check = {qubit}
-        # Add predecessors of output node to the queue.
-        queue = deque(self.predecessors(output_node))
 
-        # While queue isn't empty
+        qubits_in_cone = {qubit}
+        queue = deque(self.quantum_predecessors(output_node))
+
+        # The processed_non_directive_nodes stores the set of processed non-directive nodes.
+        # This is an optimization to avoid considering the same non-directive node multiple
+        # times when reached from different paths.
+        # The directive nodes (such as barriers or measures) are trickier since when processing
+        # them we only add their predecessors that intersect qubits_in_cone. Hence, directive
+        # nodes have to be considered multiple times.
+        processed_non_directive_nodes = set()
+
         while queue:
-            # Pop first element.
             node_to_check = queue.popleft()
-            # Check whether element is input or output node.
+
             if isinstance(node_to_check, DAGOpNode):
-                # Keep all the qubits in the operation inside a set.
-                qubit_set = set(node_to_check.qargs)
-                # Check if there are any qubits in common and that the operation is not a barrier.
-                if (
-                    len(qubit_set.intersection(qubits_to_check)) > 0
-                    and node_to_check.op.name != "barrier"
-                    and not getattr(node_to_check.op, "_directive")
-                ):
-                    # If so, add all the qubits to the causal cone.
-                    qubits_to_check = qubits_to_check.union(qubit_set)
-            # For each predecessor of the current node, filter input/output nodes,
-            # also make sure it has at least one qubit in common. Then append.
-            for node in self.quantum_predecessors(node_to_check):
-                if (
-                    isinstance(node, DAGOpNode)
-                    and len(qubits_to_check.intersection(set(node.qargs))) > 0
-                ):
-                    queue.append(node)
-        return qubits_to_check
+                # If the operation is not a directive (in particular not a barrier nor a measure),
+                # we do not do anything if it was already processed. Otherwise, we add its qubits
+                # to qubits_in_cone, and append its predecessors to queue.
+                if not getattr(node_to_check.op, "_directive"):
+                    if node_to_check in processed_non_directive_nodes:
+                        continue
+                    qubits_in_cone = qubits_in_cone.union(set(node_to_check.qargs))
+                    processed_non_directive_nodes.add(node_to_check)
+                    for pred in self.quantum_predecessors(node_to_check):
+                        if isinstance(pred, DAGOpNode):
+                            queue.append(pred)
+                else:
+                    # Directives (such as barriers and measures) may be defined over all the qubits,
+                    # yet not all of these qubits should be considered in the causal cone. So we
+                    # only add those predecessors that have qubits in common with qubits_in_cone.
+                    for pred in self.quantum_predecessors(node_to_check):
+                        if isinstance(pred, DAGOpNode) and not qubits_in_cone.isdisjoint(
+                            set(pred.qargs)
+                        ):
+                            queue.append(pred)
+
+        return qubits_in_cone
 
     def properties(self):
         """Return a dictionary of circuit properties."""

releasenotes/notes/improve-quantum-causal-cone-f63eaaa9ab658811.yaml

@@ -0,0 +1,5 @@
+---
+features_circuits:
+  - |
+    Improved performance of the method :meth:`.DAGCircuit.quantum_causal_cone` by not examining
+    the same non-directive node multiple times when reached from different paths.

test/python/dagcircuit/test_dagcircuit.py

@@ -3475,6 +3475,103 @@ def test_causal_cone_barriers(self):
 
         self.assertEqual(result, expected)
 
+    def test_causal_cone_more_barriers(self):
+        """Test causal cone for circuit with barriers. This example shows
+        why barriers may need to be examined multiple times."""
+
+        # q0_0: ──■────────░────────────────────────
+        #       ┌─┴─┐      ░
+        # q0_1: ┤ X ├──────░───■────────────────────
+        #       ├───┤      ░ ┌─┴─┐┌───┐┌───┐┌───┐
+        # q0_2: ┤ H ├──────░─┤ X ├┤ H ├┤ H ├┤ H ├─X─
+        #       ├───┤┌───┐ ░ └───┘└───┘└───┘└───┘ │
+        # q0_3: ┤ H ├┤ X ├─░──────────────────────X─
+        #       ├───┤└─┬─┘ ░
+        # q0_4: ┤ X ├──■───░────────────────────────
+        #       └─┬─┘      ░
+        # q0_5: ──■────────░────────────────────────
+
+        qreg = QuantumRegister(6)
+        qc = QuantumCircuit(qreg)
+        qc.cx(0, 1)
+        qc.h(2)
+        qc.cx(5, 4)
+        qc.h(3)
+        qc.cx(4, 3)
+        qc.barrier()
+        qc.cx(1, 2)
+
+        qc.h(2)
+        qc.h(2)
+        qc.h(2)
+        qc.swap(2, 3)
+
+        dag = circuit_to_dag(qc)
+
+        result = dag.quantum_causal_cone(qreg[2])
+        expected = {qreg[0], qreg[1], qreg[2], qreg[3], qreg[4], qreg[5]}
+
+        self.assertEqual(result, expected)
+
+    def test_causal_cone_measure(self):
+        """Test causal cone with measures."""
+
+        #         ┌───┐ ░ ┌─┐
+        #    q_0: ┤ H ├─░─┤M├────────────
+        #         ├───┤ ░ └╥┘┌─┐
+        #    q_1: ┤ H ├─░──╫─┤M├─────────
+        #         ├───┤ ░  ║ └╥┘┌─┐
+        #    q_2: ┤ H ├─░──╫──╫─┤M├──────
+        #         ├───┤ ░  ║  ║ └╥┘┌─┐
+        #    q_3: ┤ H ├─░──╫──╫──╫─┤M├───
+        #         ├───┤ ░  ║  ║  ║ └╥┘┌─┐
+        #    q_4: ┤ H ├─░──╫──╫──╫──╫─┤M├
+        #         └───┘ ░  ║  ║  ║  ║ └╥┘
+        #    c: 5/═════════╬══╬══╬══╬══╬═
+        #                  ║  ║  ║  ║  ║
+        # meas: 5/═════════╩══╩══╩══╩══╩═
+        #                  0  1  2  3  4
+
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        circuit = QuantumCircuit(qreg, creg)
+        for i in range(5):
+            circuit.h(i)
+        circuit.measure_all()
+
+        dag = circuit_to_dag(circuit)
+
+        result = dag.quantum_causal_cone(dag.qubits[1])
+        expected = {qreg[1]}
+        self.assertEqual(result, expected)
+
+    def test_reconvergent_paths(self):
+        """Test circuit with reconvergent paths."""
+
+        # q0_0: ──■─────────■─────────■─────────■─────────■─────────■───────
+        #       ┌─┴─┐     ┌─┴─┐     ┌─┴─┐     ┌─┴─┐     ┌─┴─┐     ┌─┴─┐
+        # q0_1: ┤ X ├──■──┤ X ├──■──┤ X ├──■──┤ X ├──■──┤ X ├──■──┤ X ├──■──
+        #       └───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐
+        # q0_2: ──■──┤ X ├──■──┤ X ├──■──┤ X ├──■──┤ X ├──■──┤ X ├──■──┤ X ├
+        #       ┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘┌─┴─┐└───┘
+        # q0_3: ┤ X ├─────┤ X ├─────┤ X ├─────┤ X ├─────┤ X ├─────┤ X ├─────
+        #       └───┘     └───┘     └───┘     └───┘     └───┘     └───┘
+        # q0_4: ────────────────────────────────────────────────────────────
+
+        qreg = QuantumRegister(5)
+        circuit = QuantumCircuit(qreg)
+
+        for _ in range(6):
+            circuit.cx(0, 1)
+            circuit.cx(2, 3)
+            circuit.cx(1, 2)
+
+        dag = circuit_to_dag(circuit)
+
+        result = dag.quantum_causal_cone(dag.qubits[1])
+        expected = {qreg[0], qreg[1], qreg[2], qreg[3]}
+        self.assertEqual(result, expected)
+
 
 if __name__ == "__main__":
     unittest.main()