DAGCircuit: Add get_causal_cone method

qiskit/dagcircuit/dagcircuit.py

@@ -20,7 +20,7 @@
 composed, and modified. Some natural properties like depth can be computed
 directly from the graph.
 """
-from collections import OrderedDict, defaultdict
+from collections import OrderedDict, defaultdict, deque
 import copy
 import itertools
 import math
@@ -1898,6 +1898,64 @@ def count_ops_longest_path(self):
                 op_dict[name] += 1
         return op_dict
 
+    def _get_qubit_input_output_node(self, qubit, in_or_out=False):
+        """Returns qubit and input or output node from an index."""
+        nodes = self.output_map if in_or_out else self.input_map
+        # Check if index is passed
+        if isinstance(qubit, int):
+            if qubit >= self.num_qubits():
+                raise DAGCircuitError(f"Qubit index {qubit} is out of range")
+            qubit = list(nodes.keys())[qubit]
+        if qubit not in self.qubits:
+            raise DAGCircuitError("Qubit was not found in circuit")
+        return (qubit, nodes.get(qubit, None))
+
+    def get_qubit_input_node(self, qubit):
+        """Returns qubit and input node from a qubit index."""
+        return self._get_qubit_input_output_node(qubit)
+
+    def get_qubit_output_node(self, qubit):
+        """Returns qubit and output node from a qubit index."""
+        return self._get_qubit_input_output_node(qubit, True)
+
+    def get_causal_cone(self, qubit_index):
+        """Returns causal cone of a qubit."""
+        # Check if the qubit index is in range, else throw an error.
+        if qubit_index >= self.num_qubits():
+            raise DAGCircuitError(f"Qubit index {qubit_index} is out of range")
+
+        # Retrieve the output node and the qubit
+        qubit, output_node = self.get_qubit_output_node(qubit_index)
+        # Add the qubit to the causal cone.
+        qubits_to_check = set({qubit})
+        # Add predecessors of output node to the queue.
+        queue = deque(self.predecessors(output_node))
+
+        # While queue isn't empty
+        while queue:
+            # Pop first element.
+            node_to_check = queue.popleft()
+            # Check whether element is input or output node.
+            if not (isinstance(node_to_check, (DAGInNode, DAGOutNode))):
+                # 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 not node_to_check.op.name == "barrier"
+                ):
+                    # 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.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
+
     def properties(self):
         """Return a dictionary of circuit properties."""
         summary = {

releasenotes/notes/add-dag-causal-cone-5a19311e40fbb3af.yaml

@@ -0,0 +1,42 @@
+---
+features:
+  - |
+    Added :meth:`.DAGCircuit.get_causal_cone` to obtain the causal cone of a qubit
+    in a DAGCircuit.
+    The following example shows its correct usage::
+
+        from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
+        from qiskit.circuit.library import CXGate, CZGate
+        from qiskit.dagcircuit import DAGCircuit
+
+        # Build a DAGCircuit
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+        dag.apply_operation_back(CXGate(), qreg[[1, 2]], [])
+        dag.apply_operation_back(CXGate(), qreg[[0, 3]], [])
+        dag.apply_operation_back(CZGate(), qreg[[1, 4]], [])
+        dag.apply_operation_back(CZGate(), qreg[[2, 4]], [])
+        dag.apply_operation_back(CXGate(), qreg[[3, 4]], [])
+
+        # Get the causal node of qubit at index 0
+        result = dag.get_causal_cone(0)
+
+  - |
+    Added :meth:`.DAGCircuit.get_qubit_input_node` and 
+    :meth:`.DAGCircuit.get_qubit_output_node` to obtain the input/output node of a 
+    qubit in the `DAGCircuit`. The following example shows its correct usage::
+
+        from qiskit import QuantumRegister, ClassicalRegister
+        from qiskit.dagcircuit import DAGCircuit
+
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        # Get the output node of qubit at index 3
+        result = dag.get_qubit_output_node(3)

test/python/dagcircuit/test_dagcircuit.py

@@ -2322,5 +2322,183 @@ def test_2q_swap_partially_connected_prepost_spectators(self):
         self.assertEqual(dag, expected)
 
 
+class TestDagCausalCone(QiskitTestCase):
+    """Test `get_causal_node` function"""
+
+    def test_causal_cone_regular_circuit(self):
+        """Test causal cone with a regular circuit"""
+
+        # q_0: ───────■─────────────
+        #             │
+        # q_1: ──■────┼───■─────────
+        #      ┌─┴─┐  │   │
+        # q_2: ┤ X ├──┼───┼──■──────
+        #      └───┘┌─┴─┐ │  │
+        # q_3: ─────┤ X ├─┼──┼───■──
+        #           └───┘ │  │ ┌─┴─┐
+        # q_4: ───────────■──■─┤ X ├
+        #                      └───┘
+        # c: 5/═════════════════════
+
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+        dag.apply_operation_back(CXGate(), qreg[[1, 2]], [])
+        dag.apply_operation_back(CXGate(), qreg[[0, 3]], [])
+        dag.apply_operation_back(CZGate(), qreg[[1, 4]], [])
+        dag.apply_operation_back(CZGate(), qreg[[2, 4]], [])
+        dag.apply_operation_back(CXGate(), qreg[[3, 4]], [])
+
+        # Get causal cone of qubit at index 0
+        result = dag.get_causal_cone(0)
+
+        # Expected result
+        expected = set(qreg[[0, 3]])
+        self.assertEqual(result, expected)
+
+    def test_causal_cone_invalid_index(self):
+        """Test causal cone with invalid index"""
+
+        # q_0: ───────■─────────────
+        #             │
+        # q_1: ──■────┼───■─────────
+        #      ┌─┴─┐  │   │
+        # q_2: ┤ X ├──┼───┼──■──────
+        #      └───┘┌─┴─┐ │  │
+        # q_3: ─────┤ X ├─┼──┼───■──
+        #           └───┘ │  │ ┌─┴─┐
+        # q_4: ───────────■──■─┤ X ├
+        #                      └───┘
+        # c: 5/═════════════════════
+
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+        dag.apply_operation_back(CXGate(), qreg[[1, 2]], [])
+        dag.apply_operation_back(CXGate(), qreg[[0, 3]], [])
+        dag.apply_operation_back(CZGate(), qreg[[1, 4]], [])
+        dag.apply_operation_back(CZGate(), qreg[[2, 4]], [])
+        dag.apply_operation_back(CXGate(), qreg[[3, 4]], [])
+
+        # Raise error due to invalid index
+        self.assertRaises(DAGCircuitError, dag.get_causal_cone, 5)
+
+    def test_causal_cone_no_neighbor(self):
+        """Test causal cone with no neighbor"""
+
+        # q_0: ───────────
+
+        # q_1: ──■───■────
+        #      ┌─┴─┐ │
+        # q_2: ┤ X ├─┼──■─
+        #      ├───┤ │  │
+        # q_3: ┤ X ├─┼──┼─
+        #      └───┘ │  │
+        # q_4: ──────■──■─
+
+        # c: 5/═══════════
+
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+        dag.apply_operation_back(CXGate(), qreg[[1, 2]], [])
+        dag.apply_operation_back(CZGate(), qreg[[1, 4]], [])
+        dag.apply_operation_back(CZGate(), qreg[[2, 4]], [])
+        dag.apply_operation_back(XGate(), qreg[[3]], [])
+
+        # Get causal cone of Qubit at index 3.
+        result = dag.get_causal_cone(3)
+        # Expect only a set with Qubit at index 3
+        expected = set(qreg[[3]])
+        self.assertEqual(result, expected)
+
+    def test_causal_cone_empty_circuit(self):
+        """Test causal cone for circuit with no operations"""
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        # Get causal cone of qubit at index 4
+        result = dag.get_causal_cone(4)
+        # Expect only a set with Qubit at index 4
+        expected = set(qreg[[4]])
+
+        self.assertEqual(result, expected)
+
+
+class TestDagInputOutputNode(QiskitTestCase):
+    """Tests get_qubit_input_node and get_qubit_output_node"""
+
+    def test_input_node_on_dag_index(self):
+        """Test input node on empty dag"""
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        result = dag.get_qubit_input_node(2)
+        expected = qreg[2], dag.input_map[qreg[2]]
+
+        self.assertEqual(result, expected)
+
+    def test_output_node_on_dag_index(self):
+        """Test output node on empty dag"""
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        result = dag.get_qubit_output_node(3)
+        expected = qreg[3], dag.output_map[qreg[3]]
+
+        self.assertEqual(result, expected)
+
+    def test_input_node_on_dag_qubit(self):
+        """Test input node on empty dag"""
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        result = dag.get_qubit_input_node(qreg[1])
+        expected = qreg[1], dag.input_map[qreg[1]]
+
+        self.assertEqual(result, expected)
+
+    def test_output_node_on_dag_qubit(self):
+        """Test output node on empty dag"""
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        result = dag.get_qubit_output_node(qreg[0])
+        expected = qreg[0], dag.output_map[qreg[0]]
+
+        self.assertEqual(result, expected)
+
+    def test_input_node_on_invalid_index(self):
+        """Test output node on invalid node"""
+        dag = DAGCircuit()
+        qreg = QuantumRegister(5)
+        creg = ClassicalRegister(5)
+        dag.add_qreg(qreg)
+        dag.add_creg(creg)
+
+        self.assertRaises(DAGCircuitError, dag.get_qubit_output_node, 6)
+
+
 if __name__ == "__main__":
     unittest.main()