Add validation for couplers to GridDevice

cirq-google/cirq_google/devices/grid_device.py

@@ -37,6 +37,7 @@
 from cirq_google import transformers
 from cirq_google.api import v2
 from cirq_google.devices import known_devices
+from cirq_google.devices.coupler import Coupler
 from cirq_google.experimental import ops as experimental_ops
 
 
@@ -622,7 +623,12 @@ def validate_operation(self, operation: cirq.Operation) -> None:
             raise ValueError(f'Operation {operation} contains a gate which is not supported.')
 
         for q in operation.qubits:
-            if q not in self._metadata.qubit_set:
+            if isinstance(q, Coupler):
+                if any(qc not in self._metadata.qubit_set for qc in q.qubits):
+                    raise ValueError(f'Qubits on coupler not on device: {q.qubits}.')
+                if frozenset(q.qubits) not in self._metadata.qubit_pairs:
+                    raise ValueError(f'Coupler pair is not valid on device: {q.qubits}.')
+            elif q not in self._metadata.qubit_set:
                 raise ValueError(f'Qubit not on device: {q!r}.')
 
         if (

cirq-google/cirq_google/devices/grid_device_test.py

@@ -21,7 +21,7 @@
 import cirq
 import cirq_google
 from cirq_google.api import v2
-from cirq_google.devices import grid_device
+from cirq_google.devices import Coupler, grid_device
 
 
 GRID_HEIGHT = 5
@@ -465,6 +465,47 @@ def test_grid_device_validate_operations_negative():
         )
 
 
+def test_grid_device_validate_operation_coupler_for_horizontal_couplings():
+    """Tests coupler device on a device spec that only
+    has horizontal couplings."""
+    _, spec = _create_device_spec_with_horizontal_couplings()
+    device = cirq_google.GridDevice.from_proto(spec)
+
+    g = cirq_google.InternalGate(
+        gate_name="DetuneCoupler", gate_module='internal_module', num_qubits=1, freq=5.5
+    )
+    for y in range(GRID_HEIGHT):
+        # Valid couplers
+        coupler = Coupler(cirq.GridQubit(y, 0), cirq.GridQubit(y, 1))
+        device.validate_operation(g(coupler))
+        coupler = Coupler(cirq.GridQubit(y, 1), cirq.GridQubit(y, 0))
+        device.validate_operation(g(coupler))
+        # One coupler off grid
+        coupler = Coupler(cirq.GridQubit(y, 1), cirq.GridQubit(y, 2))
+        with pytest.raises(ValueError, match="Qubits on coupler not on device"):
+            device.validate_operation(g(coupler))
+        # Both couplers off grid
+        coupler = Coupler(cirq.GridQubit(y, 2), cirq.GridQubit(y, 3))
+        with pytest.raises(ValueError, match="Qubits on coupler not on device"):
+            device.validate_operation(g(coupler))
+        # Vertical Coupler (not on device)
+        coupler = Coupler(cirq.GridQubit(y, 0), cirq.GridQubit((y + 1) % GRID_HEIGHT, 0))
+        with pytest.raises(ValueError, match="Coupler pair is not valid on device"):
+            device.validate_operation(g(coupler))
+
+
+def test_grid_device_validate_operation_coupler_for_vertical_couplings():
+    gateset = cirq.Gateset(cirq.GateFamily(cirq_google.InternalGate))
+    device = grid_device.GridDevice._from_device_information(
+        qubit_pairs=[(cirq.GridQubit(1, 0), cirq.GridQubit(0, 0))], gateset=gateset
+    )
+    g = cirq_google.InternalGate(
+        gate_name="DetuneCoupler", gate_module='internal_module', num_qubits=1, freq=5.5
+    )
+    coupler = Coupler(cirq.GridQubit(1, 0), cirq.GridQubit(0, 0))
+    device.validate_operation(g(coupler))
+
+
 @pytest.mark.parametrize(
     'spec, error_match',
     [