Added visualization for qubit permutations for routed circuits (quantumlib#5848)

ammareltigani · web-flow · commit 7d2dfadaf4b7 · 2022-09-12T15:01:28.000-07:00
* added circuit visualization code

* made print gate private and added more tests

* removed unused import and added gate to private list

* addressed comments

* added safeguard for wrong use of RoutingSwapTag

* addressed nits

* added return type for

* added _SwapPrintGate again to  and fixed type issue
diff --git a/cirq/__init__.py b/cirq/__init__.py
@@ -367,6 +367,7 @@
     prepare_two_qubit_state_using_cz,
     prepare_two_qubit_state_using_sqrt_iswap,
     RouteCQC,
+    routed_circuit_with_mapping,
     SqrtIswapTargetGateset,
     single_qubit_matrix_to_gates,
     single_qubit_matrix_to_pauli_rotations,
diff --git a/cirq/ops/gate_operation_test.py b/cirq/ops/gate_operation_test.py
@@ -494,6 +494,7 @@ def all_subclasses(cls):
         cirq.Pauli,
         # Private gates.
         cirq.transformers.analytical_decompositions.two_qubit_to_fsim._BGate,
+        cirq.transformers.routing.visualize_routed_circuit._SwapPrintGate,
         cirq.ops.raw_types._InverseCompositeGate,
         cirq.circuits.qasm_output.QasmTwoQubitGate,
         cirq.ops.MSGate,
diff --git a/cirq/transformers/__init__.py b/cirq/transformers/__init__.py
@@ -50,6 +50,7 @@
     LineInitialMapper,
     MappingManager,
     RouteCQC,
+    routed_circuit_with_mapping,
 )
 
 from cirq.transformers.target_gatesets import (
diff --git a/cirq/transformers/routing/__init__.py b/cirq/transformers/routing/__init__.py
@@ -18,3 +18,4 @@
 from cirq.transformers.routing.mapping_manager import MappingManager
 from cirq.transformers.routing.line_initial_mapper import LineInitialMapper
 from cirq.transformers.routing.route_circuit_cqc import RouteCQC
+from cirq.transformers.routing.visualize_routed_circuit import routed_circuit_with_mapping
diff --git a/cirq/transformers/routing/visualize_routed_circuit.py b/cirq/transformers/routing/visualize_routed_circuit.py
@@ -0,0 +1,78 @@
+# Copyright 2022 The Cirq Developers
+#
+# 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
+#
+#     https://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 typing import Dict, Optional, Tuple, TYPE_CHECKING
+from cirq import circuits, ops
+
+if TYPE_CHECKING:
+    import cirq
+
+
+class _SwapPrintGate(ops.Gate):
+    """A gate that displays the string representation of each qubits on the circuit."""
+
+    def __init__(self, qubits: Tuple[Tuple['cirq.Qid', 'cirq.Qid'], ...]) -> None:
+        self.qubits = qubits
+
+    def num_qubits(self):
+        return len(self.qubits)
+
+    def _circuit_diagram_info_(self, args: 'cirq.CircuitDiagramInfoArgs') -> Tuple[str, ...]:
+        return tuple(f'{str(q[1])}' for q in self.qubits)
+
+
+def routed_circuit_with_mapping(
+    routed_circuit: 'cirq.AbstractCircuit',
+    initial_map: Optional[Dict['cirq.Qid', 'cirq.Qid']] = None,
+) -> 'cirq.AbstractCircuit':
+    """Returns the same circuits with information about the permutation of qubits after each swap.
+
+    Args:
+        routed_circuit: a routed circuit that potentially has inserted swaps tagged with a
+            RoutingSwapTag.
+        initial_map: the initial mapping from logical to physical qubits. If this is not specified
+            then the identity mapping of the qubits in routed_circuit will be used as initial_map.
+
+    Raises:
+        ValueError: if a non-SWAP gate is tagged with a RoutingSwapTag.
+    """
+    all_qubits = sorted(routed_circuit.all_qubits())
+    qdict = {q: q for q in all_qubits}
+    if initial_map is None:
+        initial_map = qdict.copy()
+    inverse_map = {v: k for k, v in initial_map.items()}
+
+    def swap_print_moment() -> 'cirq.Operation':
+        return _SwapPrintGate(
+            tuple(zip(qdict.values(), [inverse_map[x] for x in qdict.values()]))
+        ).on(*all_qubits)
+
+    ret_circuit = circuits.Circuit(swap_print_moment())
+    for m in routed_circuit:
+        swap_in_moment = False
+        for op in m:
+            if ops.RoutingSwapTag() in op.tags:
+                if type(op.gate) != ops.swap_gates.SwapPowGate:
+                    raise ValueError(
+                        "Invalid circuit. A non-SWAP gate cannot be tagged a RoutingSwapTag."
+                    )
+                swap_in_moment = True
+                q1, q2 = op.qubits
+                qdict[q1], qdict[q2] = qdict[q2], qdict[q1]
+
+        ret_circuit.append(m)
+        if swap_in_moment:
+            ret_circuit.append(swap_print_moment())
+
+    return ret_circuit
diff --git a/cirq/transformers/routing/visualize_routed_circuit_test.py b/cirq/transformers/routing/visualize_routed_circuit_test.py
@@ -0,0 +1,89 @@
+# Copyright 2022 The Cirq Developers
+#
+# 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
+#
+#     https://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 pytest
+import cirq
+
+
+def test_routed_circuit_with_mapping_simple():
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit([cirq.Moment(cirq.SWAP(q[0], q[1]).with_tags(cirq.RoutingSwapTag()))])
+    expected_diagram = """
+0: ───q(0)───×[cirq.RoutingSwapTag()]───q(1)───
+      │      │                          │
+1: ───q(1)───×──────────────────────────q(0)───"""
+    cirq.testing.assert_has_diagram(cirq.routed_circuit_with_mapping(circuit), expected_diagram)
+
+    expected_diagram_with_initial_mapping = """
+0: ───a───×[cirq.RoutingSwapTag()]───b───
+      │   │                          │
+1: ───b───×──────────────────────────a───"""
+    cirq.testing.assert_has_diagram(
+        cirq.routed_circuit_with_mapping(
+            circuit, {cirq.NamedQubit("a"): q[0], cirq.NamedQubit("b"): q[1]}
+        ),
+        expected_diagram_with_initial_mapping,
+    )
+
+    # if swap is untagged should not affect the mapping
+    circuit = cirq.Circuit([cirq.Moment(cirq.SWAP(q[0], q[1]))])
+    expected_diagram = """
+0: ───q(0)───×───
+      │      │
+1: ───q(1)───×───"""
+    cirq.testing.assert_has_diagram(cirq.routed_circuit_with_mapping(circuit), expected_diagram)
+
+    circuit = cirq.Circuit(
+        [
+            cirq.Moment(cirq.X(q[0]).with_tags(cirq.RoutingSwapTag())),
+            cirq.Moment(cirq.SWAP(q[0], q[1])),
+        ]
+    )
+    with pytest.raises(
+        ValueError, match="Invalid circuit. A non-SWAP gate cannot be tagged a RoutingSwapTag."
+    ):
+        cirq.routed_circuit_with_mapping(circuit)
+
+
+def test_routed_circuit_with_mapping_multi_swaps():
+    q = cirq.LineQubit.range(6)
+    circuit = cirq.Circuit(
+        [
+            cirq.Moment(cirq.CNOT(q[3], q[4])),
+            cirq.Moment(cirq.CNOT(q[5], q[4]), cirq.CNOT(q[2], q[3])),
+            cirq.Moment(
+                cirq.CNOT(q[2], q[1]), cirq.SWAP(q[4], q[3]).with_tags(cirq.RoutingSwapTag())
+            ),
+            cirq.Moment(
+                cirq.SWAP(q[0], q[1]).with_tags(cirq.RoutingSwapTag()),
+                cirq.SWAP(q[3], q[2]).with_tags(cirq.RoutingSwapTag()),
+            ),
+            cirq.Moment(cirq.CNOT(q[2], q[1])),
+            cirq.Moment(cirq.CNOT(q[1], q[0])),
+        ]
+    )
+    expected_diagram = """
+0: ───q(0)──────────────────────────────────────q(0)───×[cirq.RoutingSwapTag()]───q(1)───────X───
+      │                                         │      │                          │          │
+1: ───q(1)───────────X──────────────────────────q(1)───×──────────────────────────q(0)───X───@───
+      │              │                          │                                 │      │
+2: ───q(2)───────@───@──────────────────────────q(2)───×──────────────────────────q(4)───@───────
+      │          │                              │      │                          │
+3: ───q(3)───@───X───×──────────────────────────q(4)───×[cirq.RoutingSwapTag()]───q(2)───────────
+      │      │       │                          │                                 │
+4: ───q(4)───X───X───×[cirq.RoutingSwapTag()]───q(3)──────────────────────────────q(3)───────────
+      │          │                              │                                 │
+5: ───q(5)───────@──────────────────────────────q(5)──────────────────────────────q(5)───────────
+"""
+    cirq.testing.assert_has_diagram(cirq.routed_circuit_with_mapping(circuit), expected_diagram)

Original file line number	Diff line number	Diff line change
`@@ -50,6 +50,7 @@`
`50`	`50`	`LineInitialMapper,`
`51`	`51`	`MappingManager,`
`52`	`52`	`RouteCQC,`
	`53`	`+ routed_circuit_with_mapping,`
`53`	`54`	`)`
`54`	`55`
`55`	`56`	`from cirq.transformers.target_gatesets import (`