Add edges to device (quantumlib#3993)

Closes quantumlib#3696

Author: daxfohl

cirq/__init__.py

@@ -90,6 +90,7 @@
     NO_NOISE,
     NOISE_MODEL_LIKE,
     NoiseModel,
+    SymmetricalQidPair,
     UNCONSTRAINED_DEVICE,
 )
 

cirq/contrib/graph_device/graph_device.py

@@ -159,6 +159,15 @@ def qubit_set(self) -> FrozenSet['cirq.Qid']:
     def edges(self):
         return tuple(sorted(self.device_graph.edges))
 
+    def qid_pairs(self) -> FrozenSet['cirq.SymmetricalQidPair']:
+        return frozenset(
+            [
+                devices.SymmetricalQidPair(*edge)  # type: ignore
+                for edge in self.device_graph.edges
+                if len(edge) == 2 and all(isinstance(q, ops.Qid) for q in edge)
+            ]
+        )
+
     @property
     def labelled_edges(self):
         return self.device_graph.labelled_edges

cirq/contrib/graph_device/graph_device_test.py

@@ -180,3 +180,10 @@ def test_qubit_set():
     device_graph = ccgd.UndirectedHypergraph(labelled_edges={(a, b): None, (c, d): None})
     device = ccgd.UndirectedGraphDevice(device_graph=device_graph)
     assert device.qubit_set() == {a, b, c, d}
+
+
+def test_qid_pairs():
+    a, b, c, d = cirq.LineQubit.range(4)
+    device_graph = ccgd.UndirectedHypergraph(labelled_edges={(a, b): None, (c, d): None})
+    device = ccgd.UndirectedGraphDevice(device_graph=device_graph)
+    assert len(device.qid_pairs()) == 2

cirq/devices/__init__.py

@@ -15,6 +15,7 @@
 """Types for devices, device-specific qubits, and noise models."""
 from cirq.devices.device import (
     Device,
+    SymmetricalQidPair,
 )
 
 from cirq.devices.grid_qubit import (

cirq/devices/device.py

@@ -13,7 +13,11 @@
 # limitations under the License.
 
 import abc
-from typing import TYPE_CHECKING, Optional, AbstractSet
+from typing import TYPE_CHECKING, Optional, AbstractSet, cast, FrozenSet, Iterator
+
+from cirq import value
+from cirq.devices.grid_qubit import _BaseGridQid
+from cirq.devices.line_qubit import _BaseLineQid
 
 if TYPE_CHECKING:
     import cirq
@@ -46,6 +50,43 @@ def qubit_set(self) -> Optional[AbstractSet['cirq.Qid']]:
         # Default to the qubits being unknown.
         return None
 
+    def qid_pairs(self) -> Optional[FrozenSet['cirq.SymmetricalQidPair']]:
+        """Returns a set of qubit edges on the device, if possible.
+
+        This property can be overridden in child classes for special handling.
+        The default handling is: GridQids and LineQids will have neighbors as
+        edges, and others will be fully connected.
+
+        Returns:
+            If the device has a finite set of qubits, then a set of all edges
+            on the device is returned.
+
+            If the device has no well defined finite set of qubits (e.g.
+            `cirq.UnconstrainedDevice` has this property), then `None` is
+            returned.
+        """
+        qs = self.qubit_set()
+        if qs is None:
+            return None
+        if all(isinstance(q, _BaseGridQid) for q in qs):
+            return frozenset(
+                [
+                    SymmetricalQidPair(q, q2)
+                    for q in [cast(_BaseGridQid, q) for q in qs]
+                    for q2 in [q + (0, 1), q + (1, 0)]
+                    if q2 in qs
+                ]
+            )
+        if all(isinstance(q, _BaseLineQid) for q in qs):
+            return frozenset(
+                [
+                    SymmetricalQidPair(q, q + 1)
+                    for q in [cast(_BaseLineQid, q) for q in qs]
+                    if q + 1 in qs
+                ]
+            )
+        return frozenset([SymmetricalQidPair(q, q2) for q in qs for q2 in qs if q < q2])
+
     def decompose_operation(self, operation: 'cirq.Operation') -> 'cirq.OP_TREE':
         """Returns a device-valid decomposition for the given operation.
 
@@ -105,3 +146,36 @@ def can_add_operation_into_moment(
             Whether or not the moment will validate after adding the operation.
         """
         return not moment.operates_on(operation.qubits)
+
+
+@value.value_equality
+class SymmetricalQidPair:
+    def __init__(self, qid1: 'cirq.Qid', qid2: 'cirq.Qid'):
+        if qid1 == qid2:
+            raise ValueError('A QidPair cannot have identical qids.')
+        self.qids = frozenset([qid1, qid2])
+
+    def _value_equality_values_(self):
+        return self.qids
+
+    def __repr__(self):
+        return f'cirq.QidPair({repr(sorted(self.qids))[1:-1]})'
+
+    def _json_dict_(self):
+        return {
+            'qids': sorted(self.qids),
+            'cirq_type': self.__class__.__name__,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, qids, **kwargs):
+        return cls(qids[0], qids[1])
+
+    def __len__(self) -> int:
+        return 2
+
+    def __iter__(self) -> Iterator['cirq.Qid']:
+        yield from sorted(self.qids)
+
+    def __contains__(self, item: 'cirq.Qid') -> bool:
+        return item in self.qids

cirq/devices/device_test.py

@@ -1,3 +1,4 @@
+import pytest
 import cirq
 
 
@@ -30,3 +31,46 @@ def _qubits(self):
             return cirq.LineQubit.range(6)
 
     assert PrivateQubitMethodDevice().qubit_set() == frozenset(cirq.LineQubit.range(6))
+
+
+def test_qid_pairs():
+    class RawDevice(cirq.Device):
+        pass
+
+    assert RawDevice().qid_pairs() is None
+
+    class QubitFieldDevice(cirq.Device):
+        def __init__(self, qubits):
+            self.qubits = qubits
+
+    assert len(QubitFieldDevice(cirq.LineQubit.range(10)).qid_pairs()) == 9
+    assert len(QubitFieldDevice(cirq.GridQubit.rect(10, 10)).qid_pairs()) == 180
+    assert len(QubitFieldDevice([cirq.NamedQubit(str(s)) for s in range(10)]).qid_pairs()) == 45
+
+
+def test_qid_pair():
+    q0, q1, q2, q3 = cirq.LineQubit.range(4)
+    e1 = cirq.SymmetricalQidPair(q0, q1)
+    e2 = cirq.SymmetricalQidPair(q1, q0)
+    e3 = cirq.SymmetricalQidPair(q2, q3)
+    assert e1 == e2
+    assert e2 != e3
+    assert repr(e1) == "cirq.QidPair(cirq.LineQubit(0), cirq.LineQubit(1))"
+
+    assert len(e1) == 2
+    a, b = e1
+    assert (a, b) == (q0, q1)
+    a, b = e2
+    assert (a, b) == (q0, q1)
+
+    assert q0 in e1
+    assert q1 in e1
+    assert q2 not in e1
+
+    set1 = frozenset([e1, e2])
+    set2 = frozenset([e2, e3])
+    assert len(set1) == 1
+    assert len(set2) == 2
+
+    with pytest.raises(ValueError, match='A QidPair cannot have identical qids.'):
+        cirq.SymmetricalQidPair(q0, q0)

cirq/ion/ion_device.py

@@ -52,6 +52,15 @@ def __init__(
     def qubit_set(self) -> FrozenSet['cirq.LineQubit']:
         return self.qubits
 
+    def qid_pairs(self) -> FrozenSet['cirq.SymmetricalQidPair']:
+        """Qubits have all-to-all connectivity, so returns all pairs.
+
+        Returns:
+            All qubit pairs on the device.
+        """
+        qs = self.qubits
+        return frozenset([devices.SymmetricalQidPair(q, q2) for q in qs for q2 in qs if q < q2])
+
     def decompose_operation(self, operation: ops.Operation) -> ops.OP_TREE:
         return convert_to_ion_gates.ConvertToIonGates().convert_one(operation)
 

cirq/ion/ion_device_test.py

@@ -198,3 +198,7 @@ def test_at():
 
 def test_qubit_set():
     assert ion_device(3).qubit_set() == frozenset(cirq.LineQubit.range(3))
+
+
+def test_qid_pairs():
+    assert len(ion_device(10).qid_pairs()) == 45

cirq/json_resolver_cache.py

@@ -119,15 +119,16 @@ def two_qubit_matrix_gate(matrix):
         'PhasedISwapPowGate': cirq.PhasedISwapPowGate,
         'PhasedXPowGate': cirq.PhasedXPowGate,
         'PhasedXZGate': cirq.PhasedXZGate,
-        'QuantumFourierTransformGate': cirq.QuantumFourierTransformGate,
         'RandomGateChannel': cirq.RandomGateChannel,
+        'QuantumFourierTransformGate': cirq.QuantumFourierTransformGate,
         'RepetitionsStoppingCriteria': cirq.work.RepetitionsStoppingCriteria,
         'ResetChannel': cirq.ResetChannel,
         'SingleQubitMatrixGate': single_qubit_matrix_gate,
         'SingleQubitPauliStringGateOperation': cirq.SingleQubitPauliStringGateOperation,
         'SingleQubitReadoutCalibrationResult': cirq.experiments.SingleQubitReadoutCalibrationResult,
         'StabilizerStateChForm': cirq.StabilizerStateChForm,
         'SwapPowGate': cirq.SwapPowGate,
+        'SymmetricalQidPair': cirq.SymmetricalQidPair,
         'TaggedOperation': cirq.TaggedOperation,
         'ThreeDQubit': cirq.pasqal.ThreeDQubit,
         'Result': cirq.Result,

cirq/pasqal/pasqal_device.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import FrozenSet, Callable, List, Sequence, Any, Union, Dict
+
 import numpy as np
 
 import cirq
@@ -361,6 +362,22 @@ def _value_equality_values_(self) -> Any:
     def _json_dict_(self) -> Dict[str, Any]:
         return cirq.protocols.obj_to_dict_helper(self, ['control_radius', 'qubits'])
 
+    def qid_pairs(self) -> FrozenSet['cirq.SymmetricalQidPair']:
+        """Returns a list of qubit edges on the device.
+
+        Returns:
+            All qubit pairs that are less or equal to the control radius apart.
+        """
+        qs = self.qubits
+        return frozenset(
+            [
+                cirq.SymmetricalQidPair(q, q2)
+                for q in qs
+                for q2 in qs
+                if q < q2 and self.distance(q, q2) <= self.control_radius
+            ]
+        )
+
 
 class PasqalConverter(cirq.neutral_atoms.ConvertToNeutralAtomGates):
     """A gate converter for compatibility with Pasqal processors.

cirq/pasqal/pasqal_device_test.py

@@ -278,3 +278,27 @@ def test_to_json():
         "control_radius": 2,
         "qubits": [cirq.pasqal.TwoDQubit(0, 0)],
     }
+
+
+def test_qid_pairs():
+    dev = PasqalVirtualDevice(
+        1,
+        qubits=[
+            ThreeDQubit(0, 0, 0),
+            ThreeDQubit(1, 0, 0),
+            ThreeDQubit(0, 1, 0),
+            ThreeDQubit(1, 1, 0),
+            ThreeDQubit(1, 1, 1),
+        ],
+    )
+    assert len(dev.qid_pairs()) == 5
+    dev1 = PasqalVirtualDevice(
+        5,
+        qubits=[
+            TwoDQubit(0, 0),
+            TwoDQubit(3, 2),
+            TwoDQubit(3, 4),
+            TwoDQubit(3, 6),
+        ],
+    )
+    assert len(dev1.qid_pairs()) == 5

cirq/protocols/json_test_data/SymmetricalQidPair.json

@@ -0,0 +1,13 @@
+{
+  "cirq_type": "SymmetricalQidPair",
+  "qids": [
+    {
+      "cirq_type": "LineQubit",
+      "x": 0
+    },
+    {
+      "cirq_type": "LineQubit",
+      "x": 1
+    }
+  ]
+}

cirq/protocols/json_test_data/SymmetricalQidPair.repr

@@ -0,0 +1 @@
+cirq.SymmetricalQidPair(cirq.LineQubit(0), cirq.LineQubit(1))