Repository: quantumlib/Cirq

Issue:
Move `durations_of` and `neigbors_of` methods in AQTDevice to metadata
**Description of the issue**

Follow-up to the changes requested in #5696.



Discussion:
The AQTDevice.neighbors_of() method can be removed and replaced by LineQubit.neighbor().

I am however not sure about the durations_of method.
The DeviceMetadata does not have an attribute for gate durations.
There is only one derived class, GridDeviceMetadata, which has a gate_durations dictionary,
but AQTDevice is not GridDevice so this class is not appropriate.

@MichaelBroughton , @dstrain115 - should I be adding a new AQTDeviceMetadata class with an extra gate_durations attribute? 

You will have to create a new derived class `AQTDeviceMetadata` which exposes the relevant properties / methods like `durations_of` and `neighbors_of` 

Resolving patch:
diff --git a/cirq-aqt/cirq_aqt/aqt_device.py b/cirq-aqt/cirq_aqt/aqt_device.py
--- a/cirq-aqt/cirq_aqt/aqt_device.py
+++ b/cirq-aqt/cirq_aqt/aqt_device.py
@@ -11,6 +11,7 @@
 # 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.
+
 """Current device parameters for the AQT/UIBK ion trap device
 
 The device is based on a linear calcium ion string with
@@ -30,7 +31,7 @@
 import numpy as np
 
 import cirq
-from cirq_aqt import aqt_target_gateset
+from cirq_aqt import aqt_device_metadata
 
 gate_dict = {'X': cirq.X, 'Y': cirq.Y, 'Z': cirq.Z, 'MS': cirq.XX, 'R': cirq.PhasedXPowGate}
 
@@ -245,41 +246,28 @@ def __init__(
         Raises:
             TypeError: If not all the qubits supplied are `cirq.LineQubit`s.
         """
-        self._measurement_duration = cirq.Duration(measurement_duration)
-        self._twoq_gates_duration = cirq.Duration(twoq_gates_duration)
-        self._oneq_gates_duration = cirq.Duration(oneq_gates_duration)
         if not all(isinstance(qubit, cirq.LineQubit) for qubit in qubits):
             raise TypeError(
                 "All qubits were not of type cirq.LineQubit, instead were "
                 f"{set(type(qubit) for qubit in qubits)}"
             )
         self.qubits = frozenset(qubits)
-        self.gateset = aqt_target_gateset.AQTTargetGateset()
 
         graph = nx.Graph()
         graph.add_edges_from([(a, b) for a in qubits for b in qubits if a != b], directed=False)
-        self._metadata = cirq.DeviceMetadata(self.qubits, graph)
+        self._metadata = aqt_device_metadata.AQTDeviceMetadata(
+            qubits=self.qubits,
+            measurement_duration=measurement_duration,
+            twoq_gates_duration=twoq_gates_duration,
+            oneq_gates_duration=oneq_gates_duration,
+        )
 
     @property
-    def metadata(self) -> cirq.DeviceMetadata:
+    def metadata(self) -> aqt_device_metadata.AQTDeviceMetadata:
         return self._metadata
 
-    def decompose_circuit(self, circuit: cirq.Circuit) -> cirq.Circuit:
-        return cirq.optimize_for_target_gateset(circuit, gateset=self.gateset)
-
-    def duration_of(self, operation):
-        if isinstance(operation.gate, cirq.XXPowGate):
-            return self._twoq_gates_duration
-        if isinstance(
-            operation.gate, (cirq.XPowGate, cirq.YPowGate, cirq.ZPowGate, cirq.PhasedXPowGate)
-        ):
-            return self._oneq_gates_duration
-        if isinstance(operation.gate, cirq.MeasurementGate):
-            return self._measurement_duration
-        raise ValueError(f'Unsupported gate type: {operation!r}')
-
     def validate_gate(self, gate: cirq.Gate):
-        if gate not in self.gateset:
+        if gate not in self.metadata.gateset:
             raise ValueError(f'Unsupported gate type: {gate!r}')
 
     def validate_operation(self, operation):
@@ -304,12 +292,7 @@ def at(self, position: int) -> Optional[cirq.LineQubit]:
         return q if q in self.qubits else None
 
     def _value_equality_values_(self) -> Any:
-        return (
-            self._measurement_duration,
-            self._twoq_gates_duration,
-            self._oneq_gates_duration,
-            self.qubits,
-        )
+        return (self.metadata, self.qubits)
 
     def __str__(self) -> str:
         diagram = cirq.TextDiagramDrawer()
@@ -324,9 +307,9 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return (
             f'cirq_aqt.aqt_device.AQTDevice('
-            f'measurement_duration={self._measurement_duration!r}, '
-            f'twoq_gates_duration={self._twoq_gates_duration!r}, '
-            f'oneq_gates_duration={self._oneq_gates_duration!r}, '
+            f'measurement_duration={self.metadata.measurement_duration!r}, '
+            f'twoq_gates_duration={self.metadata.twoq_gates_duration!r}, '
+            f'oneq_gates_duration={self.metadata.oneq_gates_duration!r}, '
             f'qubits={sorted(self.qubits)!r}'
             f')'
         )
diff --git a/cirq-aqt/cirq_aqt/aqt_device_metadata.py b/cirq-aqt/cirq_aqt/aqt_device_metadata.py
new file mode 100644
--- /dev/null
+++ b/cirq-aqt/cirq_aqt/aqt_device_metadata.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.
+
+
+"""DeviceMetadata for ion trap device with mutually linked qubits placed on a line.
+"""
+
+from typing import Any, Iterable, Mapping
+
+import networkx as nx
+
+import cirq
+from cirq_aqt import aqt_target_gateset
+
+
+class AQTDeviceMetadata(cirq.DeviceMetadata):
+    """Hardware metadata for ion trap device with all-connected qubits placed on a line."""
+
+    def __init__(
+        self,
+        qubits: Iterable['cirq.LineQubit'],
+        measurement_duration: 'cirq.DURATION_LIKE',
+        twoq_gates_duration: 'cirq.DURATION_LIKE',
+        oneq_gates_duration: 'cirq.DURATION_LIKE',
+    ):
+        """Create metadata object for AQTDevice.
+
+        Args:
+            qubits: Iterable of `cirq.LineQubit`s that exist on the device.
+            measurement_duration: The maximum duration of a measurement.
+            twoq_gates_duration: The maximum duration of a two qubit operation.
+            oneq_gates_duration: The maximum duration of a single qubit operation.
+        """
+
+        graph = nx.Graph()
+        graph.add_edges_from([(a, b) for a in qubits for b in qubits if a != b], directed=False)
+        super().__init__(qubits, graph)
+        self._gateset = aqt_target_gateset.AQTTargetGateset()
+        self._measurement_duration = cirq.Duration(measurement_duration)
+        self._twoq_gates_duration = cirq.Duration(twoq_gates_duration)
+        self._oneq_gates_duration = cirq.Duration(oneq_gates_duration)
+        self._gate_durations = {
+            cirq.GateFamily(cirq.MeasurementGate): self._measurement_duration,
+            cirq.GateFamily(cirq.XXPowGate): self._twoq_gates_duration,
+            cirq.GateFamily(cirq.XPowGate): self._oneq_gates_duration,
+            cirq.GateFamily(cirq.YPowGate): self._oneq_gates_duration,
+            cirq.GateFamily(cirq.ZPowGate): self._oneq_gates_duration,
+            cirq.GateFamily(cirq.PhasedXPowGate): self._oneq_gates_duration,
+        }
+        assert not self._gateset.gates.symmetric_difference(self._gate_durations.keys()), (
+            "AQTDeviceMetadata.gate_durations must have the same Gates " "as AQTTargetGateset."
+        )
+
+    @property
+    def gateset(self) -> 'cirq.Gateset':
+        """Returns the `cirq.Gateset` of supported gates on this device."""
+        return self._gateset
+
+    @property
+    def gate_durations(self) -> Mapping['cirq.GateFamily', 'cirq.Duration']:
+        """Get a dictionary of supported gate families and their gate operation durations.
+
+        Use `duration_of` to obtain duration of a specific `cirq.GateOperation` instance.
+        """
+        return self._gate_durations
+
+    @property
+    def measurement_duration(self) -> 'cirq.DURATION_LIKE':
+        """Return the maximum duration of the measurement operation."""
+        return self._measurement_duration
+
+    @property
+    def oneq_gates_duration(self) -> 'cirq.DURATION_LIKE':
+        """Return the maximum duration of an operation on one-qubit gates."""
+        return self._oneq_gates_duration
+
+    @property
+    def twoq_gates_duration(self) -> 'cirq.DURATION_LIKE':
+        """Return the maximum duration of an operation on two-qubit gates."""
+        return self._twoq_gates_duration
+
+    def duration_of(self, operation: 'cirq.Operation') -> 'cirq.DURATION_LIKE':
+        """Return the maximum duration of the specifed gate operation.
+
+        Args:
+            operation: The `cirq.Operation` for which to determine its duration.
+
+        Raises:
+            ValueError: if the operation has an unsupported gate type.
+        """
+        for gate_family, duration in self.gate_durations.items():
+            if operation in gate_family:
+                return duration
+        else:
+            raise ValueError(f'Unsupported gate type: {operation!r}')
+
+    def _value_equality_values_(self) -> Any:
+        return (
+            self._measurement_duration,
+            self._twoq_gates_duration,
+            self._oneq_gates_duration,
+            self.qubit_set,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f'cirq_aqt.aqt_device_metadata.AQTDeviceMetadata('
+            f'qubits={sorted(self.qubit_set)!r}, '
+            f'measurement_duration={self.measurement_duration!r}, '
+            f'twoq_gates_duration={self.twoq_gates_duration!r}, '
+            f'oneq_gates_duration={self.oneq_gates_duration!r}'
+            f')'
+        )


Test patch:
diff --git a/cirq-aqt/cirq_aqt/aqt_device_metadata_test.py b/cirq-aqt/cirq_aqt/aqt_device_metadata_test.py
new file mode 100644
--- /dev/null
+++ b/cirq-aqt/cirq_aqt/aqt_device_metadata_test.py
@@ -0,0 +1,63 @@
+# 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.
+
+"""Tests for AQTDeviceMetadata."""
+
+from typing import List
+
+import pytest
+
+import cirq
+from cirq_aqt.aqt_device_metadata import AQTDeviceMetadata
+from cirq_aqt.aqt_target_gateset import AQTTargetGateset
+
+
+@pytest.fixture
+def qubits() -> List[cirq.LineQubit]:
+    return cirq.LineQubit.range(5)
+
+
+@pytest.fixture
+def metadata(qubits) -> AQTDeviceMetadata:
+    return AQTDeviceMetadata(
+        qubits=qubits,
+        measurement_duration=cirq.Duration(millis=100),
+        twoq_gates_duration=cirq.Duration(millis=200),
+        oneq_gates_duration=cirq.Duration(millis=10),
+    )
+
+
+def test_aqtdevice_metadata(metadata, qubits):
+    assert metadata.qubit_set == frozenset(qubits)
+    assert set(qubits) == set(metadata.nx_graph.nodes())
+    edges = metadata.nx_graph.edges()
+    assert len(edges) == 10
+    assert all(q0 != q1 for q0, q1 in edges)
+    assert AQTTargetGateset() == metadata.gateset
+    assert len(metadata.gate_durations) == 6
+
+
+def test_aqtdevice_duration_of(metadata, qubits):
+    q0, q1 = qubits[:2]
+    ms = cirq.Duration(millis=1)
+    assert metadata.duration_of(cirq.Z(q0)) == 10 * ms
+    assert metadata.duration_of(cirq.measure(q0)) == 100 * ms
+    assert metadata.duration_of(cirq.measure(q0, q1)) == 100 * ms
+    assert metadata.duration_of(cirq.XX(q0, q1)) == 200 * ms
+    with pytest.raises(ValueError, match="Unsupported gate type"):
+        metadata.duration_of(cirq.I(q0))
+
+
+def test_repr(metadata):
+    cirq.testing.assert_equivalent_repr(metadata, setup_code='import cirq\nimport cirq_aqt\n')
diff --git a/cirq-aqt/cirq_aqt/aqt_device_test.py b/cirq-aqt/cirq_aqt/aqt_device_test.py
--- a/cirq-aqt/cirq_aqt/aqt_device_test.py
+++ b/cirq-aqt/cirq_aqt/aqt_device_test.py
@@ -13,22 +13,27 @@
 # limitations under the License.
 
 from datetime import timedelta
+from typing import List
 
 import pytest
 
 import cirq
-import cirq_aqt
-import cirq_aqt.aqt_device as cad
-import sympy.utilities.matchpy_connector
-
-
-def aqt_device(chain_length: int, use_timedelta=False) -> cad.AQTDevice:
-    ms = 1000 * cirq.Duration(nanos=1) if not use_timedelta else timedelta(microseconds=1)
-    return cad.AQTDevice(
-        measurement_duration=100 * ms,  # type: ignore
-        twoq_gates_duration=200 * ms,  # type: ignore
-        oneq_gates_duration=10 * ms,  # type: ignore
-        qubits=cirq.LineQubit.range(chain_length),
+from cirq_aqt import aqt_device, aqt_device_metadata
+
+
+@pytest.fixture
+def qubits() -> List[cirq.LineQubit]:
+    return cirq.LineQubit.range(3)
+
+
+@pytest.fixture
+def device(qubits) -> aqt_device.AQTDevice:
+    ms = cirq.Duration(millis=1)
+    return aqt_device.AQTDevice(
+        measurement_duration=100 * ms,
+        twoq_gates_duration=200 * ms,
+        oneq_gates_duration=10 * ms,
+        qubits=qubits,
     )
 
 
@@ -41,151 +46,112 @@ def qubits(self):
         raise NotImplementedError()
 
 
-def test_init():
-    d = aqt_device(3)
-    ms = 1000 * cirq.Duration(nanos=1)
-    q0 = cirq.LineQubit(0)
-    q1 = cirq.LineQubit(1)
-    q2 = cirq.LineQubit(2)
-
-    assert d.qubits == {q0, q1, q2}
-    assert d.duration_of(cirq.Z(q0)) == 10 * ms
-    assert d.duration_of(cirq.measure(q0)) == 100 * ms
-    assert d.duration_of(cirq.measure(q0, q1)) == 100 * ms
-    assert d.duration_of(cirq.XX(q0, q1)) == 200 * ms
-    with pytest.raises(ValueError, match="Unsupported gate type"):
-        _ = d.duration_of(cirq.I(q0))
-
+def test_init_qubits(device, qubits):
+    ms = cirq.Duration(millis=1)
+    assert device.qubits == frozenset(qubits)
     with pytest.raises(TypeError, match="NamedQubit"):
-        _ = cad.AQTDevice(
-            measurement_duration=ms,
-            twoq_gates_duration=ms,
-            oneq_gates_duration=ms,
+        aqt_device.AQTDevice(
+            measurement_duration=100 * ms,
+            twoq_gates_duration=200 * ms,
+            oneq_gates_duration=10 * ms,
             qubits=[cirq.LineQubit(0), cirq.NamedQubit("a")],
         )
 
 
-def test_metadata():
-    d = aqt_device(3)
-    assert d.metadata.qubit_set == frozenset(
-        {cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2)}
+@pytest.mark.parametrize('ms', [cirq.Duration(millis=1), timedelta(milliseconds=1)])
+def test_init_durations(ms, qubits):
+    dev = aqt_device.AQTDevice(
+        qubits=qubits,
+        measurement_duration=100 * ms,
+        twoq_gates_duration=200 * ms,
+        oneq_gates_duration=10 * ms,
     )
-    assert len(d.metadata.nx_graph.edges()) == 3
-
+    assert dev.metadata.twoq_gates_duration == cirq.Duration(millis=200)
+    assert dev.metadata.oneq_gates_duration == cirq.Duration(millis=10)
+    assert dev.metadata.measurement_duration == cirq.Duration(millis=100)
 
-def test_init_timedelta():
-    d = aqt_device(3, use_timedelta=True)
-    ms = 1000 * cirq.Duration(nanos=1)
-    q0 = cirq.LineQubit(0)
-    q1 = cirq.LineQubit(1)
-    q2 = cirq.LineQubit(2)
 
-    assert d.qubits == {q0, q1, q2}
-    assert d.duration_of(cirq.Z(q0)) == 10 * ms
-    assert d.duration_of(cirq.measure(q0)) == 100 * ms
-    assert d.duration_of(cirq.measure(q0, q1)) == 100 * ms
-    assert d.duration_of(cirq.XX(q0, q1)) == 200 * ms
-    with pytest.raises(ValueError):
-        _ = d.duration_of(cirq.testing.SingleQubitGate().on(q0))
+def test_metadata(device, qubits):
+    assert isinstance(device.metadata, aqt_device_metadata.AQTDeviceMetadata)
+    assert device.metadata.qubit_set == frozenset(qubits)
 
 
-def test_repr():
-    d = aqt_device(3)
-    assert repr(d) == (
+def test_repr(device):
+    assert repr(device) == (
         "cirq_aqt.aqt_device.AQTDevice("
-        "measurement_duration=cirq.Duration(micros=100), "
-        "twoq_gates_duration=cirq.Duration(micros=200), "
-        "oneq_gates_duration=cirq.Duration(micros=10), "
+        "measurement_duration=cirq.Duration(millis=100), "
+        "twoq_gates_duration=cirq.Duration(millis=200), "
+        "oneq_gates_duration=cirq.Duration(millis=10), "
         "qubits=[cirq.LineQubit(0), cirq.LineQubit(1), "
         "cirq.LineQubit(2)])"
     )
-    device = cirq_aqt.aqt_device.get_aqt_device(5)
     cirq.testing.assert_equivalent_repr(device, setup_code='import cirq\nimport cirq_aqt\n')
 
 
-def test_validate_measurement_non_adjacent_qubits_ok():
-    d = aqt_device(3)
-
-    d.validate_operation(
+def test_validate_measurement_non_adjacent_qubits_ok(device):
+    device.validate_operation(
         cirq.GateOperation(cirq.MeasurementGate(2, 'key'), (cirq.LineQubit(0), cirq.LineQubit(1)))
     )
 
 
-def test_validate_operation_existing_qubits():
-    d = aqt_device(3)
-
-    d.validate_operation(cirq.GateOperation(cirq.XX, (cirq.LineQubit(0), cirq.LineQubit(1))))
-    d.validate_operation(cirq.Z(cirq.LineQubit(0)))
-    d.validate_operation(
+def test_validate_operation_existing_qubits(device):
+    device.validate_operation(cirq.GateOperation(cirq.XX, (cirq.LineQubit(0), cirq.LineQubit(1))))
+    device.validate_operation(cirq.Z(cirq.LineQubit(0)))
+    device.validate_operation(
         cirq.PhasedXPowGate(phase_exponent=0.75, exponent=0.25, global_shift=0.1).on(
             cirq.LineQubit(1)
         )
     )
 
     with pytest.raises(ValueError):
-        d.validate_operation(cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(-1)))
+        device.validate_operation(cirq.CZ(cirq.LineQubit(0), cirq.LineQubit(-1)))
     with pytest.raises(ValueError):
-        d.validate_operation(cirq.Z(cirq.LineQubit(-1)))
+        device.validate_operation(cirq.Z(cirq.LineQubit(-1)))
     with pytest.raises(ValueError):
-        d.validate_operation(cirq.CZ(cirq.LineQubit(1), cirq.LineQubit(1)))
+        device.validate_operation(cirq.CZ(cirq.LineQubit(1), cirq.LineQubit(1)))
     with pytest.raises(ValueError):
-        d.validate_operation(cirq.X(cirq.NamedQubit("q1")))
-
+        device.validate_operation(cirq.X(cirq.NamedQubit("q1")))
 
-def test_validate_operation_supported_gate():
-    d = aqt_device(3)
 
+def test_validate_operation_supported_gate(device):
     class MyGate(cirq.Gate):
         def num_qubits(self):
             return 1
 
-    d.validate_operation(cirq.GateOperation(cirq.Z, [cirq.LineQubit(0)]))
+    device.validate_operation(cirq.GateOperation(cirq.Z, [cirq.LineQubit(0)]))
 
     assert MyGate().num_qubits() == 1
     with pytest.raises(ValueError):
-        d.validate_operation(cirq.GateOperation(MyGate(), [cirq.LineQubit(0)]))
+        device.validate_operation(cirq.GateOperation(MyGate(), [cirq.LineQubit(0)]))
     with pytest.raises(ValueError):
-        d.validate_operation(NotImplementedOperation())
+        device.validate_operation(NotImplementedOperation())
 
 
-def test_aqt_device_eq():
+def test_aqt_device_eq(device):
     eq = cirq.testing.EqualsTester()
-    eq.make_equality_group(lambda: aqt_device(3))
-    eq.make_equality_group(lambda: aqt_device(4))
+    eq.make_equality_group(lambda: device)
 
 
-def test_validate_circuit_repeat_measurement_keys():
-    d = aqt_device(3)
-
+def test_validate_circuit_repeat_measurement_keys(device):
     circuit = cirq.Circuit()
     circuit.append(
         [cirq.measure(cirq.LineQubit(0), key='a'), cirq.measure(cirq.LineQubit(1), key='a')]
     )
 
     with pytest.raises(ValueError, match='Measurement key a repeated'):
-        d.validate_circuit(circuit)
-
-
-def test_aqt_device_str():
-    assert str(aqt_device(3)) == "q(0)───q(1)───q(2)"
+        device.validate_circuit(circuit)
 
 
-def test_aqt_device_pretty_repr():
-    cirq.testing.assert_repr_pretty(aqt_device(3), "q(0)───q(1)───q(2)")
-    cirq.testing.assert_repr_pretty(aqt_device(3), "AQTDevice(...)", cycle=True)
+def test_aqt_device_str(device):
+    assert str(device) == "q(0)───q(1)───q(2)"
 
 
-def test_at():
-    d = aqt_device(3)
-    assert d.at(-1) is None
-    assert d.at(0) == cirq.LineQubit(0)
-    assert d.at(2) == cirq.LineQubit(2)
+def test_aqt_device_pretty_repr(device):
+    cirq.testing.assert_repr_pretty(device, "q(0)───q(1)───q(2)")
+    cirq.testing.assert_repr_pretty(device, "AQTDevice(...)", cycle=True)
 
 
-def test_decompose_parameterized_gates():
-    theta = sympy.Symbol("theta")
-    q = cirq.LineQubit.range(2)
-    circuit = cirq.Circuit(cirq.H(q[0]) ** theta, cirq.XX(*q) ** theta)
-    d = aqt_device(3)
-    assert d.gateset.validate(d.decompose_circuit(circuit))
-    assert d.gateset._decompose_two_qubit_operation(cirq.CZ(*q) ** theta, 0) is NotImplemented
+def test_at(device):
+    assert device.at(-1) is None
+    assert device.at(0) == cirq.LineQubit(0)
+    assert device.at(2) == cirq.LineQubit(2)