Created routing utilities subdirectory in cirq-core/transformers and added MappingManager module

cirq-core/cirq/transformers/routing/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+
+"""Routing utilities in Cirq."""
+
+from cirq.transformers.routing.mapping_manager import MappingManager

cirq-core/cirq/transformers/routing/mapping_manager.py

@@ -0,0 +1,124 @@
+# 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, Sequence, TYPE_CHECKING
+from cirq._compat import cached_method
+import networkx as nx
+
+from cirq import protocols
+
+if TYPE_CHECKING:
+    import cirq
+
+
+class MappingManager:
+    """Class that keeps track of the mapping of logical to physical qubits.
+
+    Convenience methods over distance and mapping queries of the physical qubits are also provided.
+    All such public methods of this class expect logical qubits.
+    """
+
+    def __init__(
+        self, device_graph: nx.Graph, initial_mapping: Dict['cirq.Qid', 'cirq.Qid']
+    ) -> None:
+        """Initializes MappingManager.
+
+        Args:
+            device_graph: connectivity graph of qubits in the hardware device.
+            initial_mapping: the initial mapping of logical (keys) to physical qubits (values).
+        """
+        self.device_graph = device_graph
+        self._map = initial_mapping.copy()
+        self._inverse_map = {v: k for k, v in self._map.items()}
+        self._induced_subgraph = nx.induced_subgraph(self.device_graph, self._map.values())
+
+    @property
+    def map(self) -> Dict['cirq.Qid', 'cirq.Qid']:
+        """The mapping of logical qubits (keys) to physical qubits (values)."""
+        return self._map
+
+    @property
+    def inverse_map(self) -> Dict['cirq.Qid', 'cirq.Qid']:
+        """The mapping of physical qubits (keys) to logical qubits (values)."""
+        return self._inverse_map
+
+    @property
+    def induced_subgraph(self) -> nx.Graph:
+        """The induced subgraph on the set of physical qubits which are part of `self.map`."""
+        return self._induced_subgraph
+
+    def dist_on_device(self, lq1: 'cirq.Qid', lq2: 'cirq.Qid') -> int:
+        """Finds distance between logical qubits 'lq1' and 'lq2' on the device.
+
+        Args:
+            lq1: the first logical qubit.
+            lq2: the second logical qubit.
+
+        Returns:
+            The shortest path distance.
+        """
+        return len(self._physical_shortest_path(self._map[lq1], self._map[lq2])) - 1
+
+    def can_execute(self, op: 'cirq.Operation') -> bool:
+        """Finds whether the given operation acts on qubits that are adjacent on the device.
+
+        Args:
+            op: an operation on logical qubits.
+
+        Returns:
+            True, if physical qubits corresponding to logical qubits `op.qubits` are adjacent on
+            the device.
+        """
+        return protocols.num_qubits(op) < 2 or self.dist_on_device(*op.qubits) == 1
+
+    def apply_swap(self, lq1: 'cirq.Qid', lq2: 'cirq.Qid') -> None:
+        """Updates the mapping to simulate inserting a swap operation between `lq1` and `lq2`.
+
+        Args:
+            lq1: the first logical qubit.
+            lq2: the second logical qubit.
+
+        Raises:
+            ValueError: whenever lq1 and lq2 are no adjacent on the device.
+        """
+        if self.dist_on_device(lq1, lq2) > 1:
+            raise ValueError(
+                f"q1: {lq1} and q2: {lq2} are not adjacent on the device. Cannot swap them."
+            )
+
+        pq1, pq2 = self._map[lq1], self._map[lq2]
+        self._map[lq1], self._map[lq2] = self._map[lq2], self._map[lq1]
+
+        self._inverse_map[pq1], self._inverse_map[pq2] = (
+            self._inverse_map[pq2],
+            self._inverse_map[pq1],
+        )
+
+    def mapped_op(self, op: 'cirq.Operation') -> 'cirq.Operation':
+        """Transforms the given operation with the qubits in self._map.
+
+        Args:
+            op: an operation on logical qubits.
+
+        Returns:
+            The same operation on corresponding physical qubits."""
+        return op.transform_qubits(self._map)
+
+    def shortest_path(self, lq1: 'cirq.Qid', lq2: 'cirq.Qid') -> Sequence['cirq.Qid']:
+        """Find the shortest path between two logical qubits on the device given their mapping."""
+        return self._physical_shortest_path(self._map[lq1], self._map[lq2])
+
+    @cached_method
+    def _physical_shortest_path(self, pq1: 'cirq.Qid', pq2: 'cirq.Qid') -> Sequence['cirq.Qid']:
+        return nx.shortest_path(self._induced_subgraph, pq1, pq2)

cirq-core/cirq/transformers/routing/mapping_manager_test.py

@@ -0,0 +1,154 @@
+# 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 networkx.utils.misc import graphs_equal
+import pytest
+import networkx as nx
+
+import cirq
+
+
+def construct_device_graph_and_mapping():
+    device_graph = nx.Graph(
+        [
+            (cirq.NamedQubit("a"), cirq.NamedQubit("b")),
+            (cirq.NamedQubit("b"), cirq.NamedQubit("c")),
+            (cirq.NamedQubit("c"), cirq.NamedQubit("d")),
+            (cirq.NamedQubit("a"), cirq.NamedQubit("e")),
+            (cirq.NamedQubit("e"), cirq.NamedQubit("d")),
+        ]
+    )
+    q = cirq.LineQubit.range(5)
+    initial_mapping = {
+        q[1]: cirq.NamedQubit("a"),
+        q[3]: cirq.NamedQubit("b"),
+        q[2]: cirq.NamedQubit("c"),
+        q[4]: cirq.NamedQubit("d"),
+    }
+    return device_graph, initial_mapping, q
+
+
+def test_induced_subgraph():
+    device_graph, initial_mapping, _ = construct_device_graph_and_mapping()
+    mm = cirq.transformers.routing.MappingManager(device_graph, initial_mapping)
+
+    expected_induced_subgraph = nx.Graph(
+        [
+            (cirq.NamedQubit("a"), cirq.NamedQubit("b")),
+            (cirq.NamedQubit("b"), cirq.NamedQubit("c")),
+            (cirq.NamedQubit("c"), cirq.NamedQubit("d")),
+        ]
+    )
+    assert graphs_equal(mm.induced_subgraph, expected_induced_subgraph)
+
+
+def test_mapped_op():
+    device_graph, initial_mapping, q = construct_device_graph_and_mapping()
+    mm = cirq.transformers.routing.MappingManager(device_graph, initial_mapping)
+
+    assert mm.mapped_op(cirq.CNOT(q[1], q[3])).qubits == (
+        cirq.NamedQubit("a"),
+        cirq.NamedQubit("b"),
+    )
+    # does not fail if qubits non-adjacent
+    assert mm.mapped_op(cirq.CNOT(q[3], q[4])).qubits == (
+        cirq.NamedQubit("b"),
+        cirq.NamedQubit("d"),
+    )
+
+    # correctly changes mapped qubits when swapped
+    mm.apply_swap(q[2], q[3])
+    assert mm.mapped_op(cirq.CNOT(q[1], q[2])).qubits == (
+        cirq.NamedQubit("a"),
+        cirq.NamedQubit("b"),
+    )
+    # does not fial if qubits non-adjacent
+    assert mm.mapped_op(cirq.CNOT(q[1], q[3])).qubits == (
+        cirq.NamedQubit("a"),
+        cirq.NamedQubit("c"),
+    )
+
+
+def test_distance_on_device_and_can_execute():
+    device_graph, initial_mapping, q = construct_device_graph_and_mapping()
+    mm = cirq.transformers.routing.MappingManager(device_graph, initial_mapping)
+
+    # adjacent qubits have distance 1 and are thus executable
+    assert mm.dist_on_device(q[1], q[3]) == 1
+    assert mm.can_execute(cirq.CNOT(q[1], q[3]))
+
+    # non-adjacent qubits with distance > 1 are not executable
+    assert mm.dist_on_device(q[1], q[2]) == 2
+    assert mm.can_execute(cirq.CNOT(q[1], q[2])) is False
+
+    # 'dist_on_device' does not use cirq.NamedQubit("e") to find shorter shortest path
+    assert mm.dist_on_device(q[1], q[4]) == 3
+
+    # distance changes after applying swap
+    mm.apply_swap(q[2], q[3])
+    assert mm.dist_on_device(q[1], q[3]) == 2
+    assert mm.can_execute(cirq.CNOT(q[1], q[3])) is False
+    assert mm.dist_on_device(q[1], q[2]) == 1
+    assert mm.can_execute(cirq.CNOT(q[1], q[2]))
+
+    # distance between other qubits doesn't change
+    assert mm.dist_on_device(q[1], q[4]) == 3
+
+
+def test_apply_swap():
+    device_graph, initial_mapping, q = construct_device_graph_and_mapping()
+    mm = cirq.transformers.routing.MappingManager(device_graph, initial_mapping)
+
+    # swapping non-adjacent qubits raises error
+    with pytest.raises(ValueError):
+        mm.apply_swap(q[1], q[2])
+
+    # applying swap on same qubit does nothing
+    map_before_swap = mm.map.copy()
+    mm.apply_swap(q[1], q[1])
+    assert map_before_swap == mm.map
+
+    # applying same swap twice does nothing
+    mm.apply_swap(q[1], q[3])
+    mm.apply_swap(q[1], q[3])
+    assert map_before_swap == mm.map
+
+    # qubits in inverse map get swapped correctly
+    assert mm.inverse_map == {v: k for k, v in mm.map.items()}
+
+
+def test_shortest_path():
+    device_graph, initial_mapping, q = construct_device_graph_and_mapping()
+    mm = cirq.transformers.routing.MappingManager(device_graph, initial_mapping)
+
+    assert mm.shortest_path(q[1], q[2]) == [
+        cirq.NamedQubit("a"),
+        cirq.NamedQubit("b"),
+        cirq.NamedQubit("c"),
+    ]
+    shortest_one_to_four = [
+        cirq.NamedQubit("a"),
+        cirq.NamedQubit("b"),
+        cirq.NamedQubit("c"),
+        cirq.NamedQubit("d"),
+    ]
+    assert mm.shortest_path(q[1], q[4]) == shortest_one_to_four
+    # shortest path on symmetric qubit reverses the list
+    assert mm.shortest_path(q[4], q[1]) == shortest_one_to_four[::-1]
+
+    # swapping doesn't change shortest paths involving other qubits
+    mm.apply_swap(q[3], q[2])
+    assert mm.shortest_path(q[1], q[4]) == shortest_one_to_four
+    # swapping changes shortest paths involving the swapped qubits
+    assert mm.shortest_path(q[1], q[2]) == [cirq.NamedQubit("a"), cirq.NamedQubit("b")]