Update Neutral Atoms to Transformers (quantumlib#5311)

* Update Neutral Atoms to Transformers

Create NeutralAtomGateset that works as a compilation target
gateset.  Point ConvertToNeutralAtomGateset towards
optimize_for_target_gateset instead.

Author: dstrain115

cirq-core/cirq/neutral_atoms/__init__.py

@@ -21,3 +21,5 @@
     is_native_neutral_atom_gate,
     is_native_neutral_atom_op,
 )
+
+from cirq.neutral_atoms.neutral_atom_gateset import NeutralAtomGateset

cirq-core/cirq/neutral_atoms/convert_to_neutral_atom_gates.py

@@ -13,15 +13,19 @@
 # limitations under the License.
 from typing import List, Optional, TYPE_CHECKING
 
-from cirq import ops, protocols
+from cirq import ops, protocols, transformers
+from cirq._compat import deprecated_class
 from cirq.circuits.optimization_pass import PointOptimizationSummary, PointOptimizer
-from cirq.neutral_atoms import neutral_atom_devices
-from cirq import transformers
+from cirq.neutral_atoms.neutral_atom_gateset import NeutralAtomGateset
 
 if TYPE_CHECKING:
     import cirq
 
 
+@deprecated_class(
+    deadline='v0.16',
+    fix='Use cirq.optimize_for_target_gateset(circuit, gateset=NeutralAtomGateset()).',
+)
 class ConvertToNeutralAtomGates(PointOptimizer):
     """Attempts to convert gates into native Atom gates.
 
@@ -48,7 +52,7 @@ def __init__(self, ignore_failures=False) -> None:
         """
         super().__init__()
         self.ignore_failures = ignore_failures
-        self.gateset = neutral_atom_devices.neutral_atom_gateset()
+        self.gateset = NeutralAtomGateset()
 
     def _convert_one(self, op: ops.Operation) -> ops.OP_TREE:
         # Known matrix?
@@ -91,8 +95,8 @@ def optimization_at(
 
 
 def is_native_neutral_atom_op(operation: ops.Operation) -> bool:
-    return operation in neutral_atom_devices.neutral_atom_gateset()
+    return operation in NeutralAtomGateset()
 
 
 def is_native_neutral_atom_gate(gate: ops.Gate) -> bool:
-    return gate in neutral_atom_devices.neutral_atom_gateset()
+    return gate in NeutralAtomGateset()

cirq-core/cirq/neutral_atoms/convert_to_neutral_atom_gates_test.py

@@ -16,37 +16,58 @@
 import pytest
 
 import cirq
-from cirq import ops
+
+
+Q = cirq.LineQubit.range(3)
+
+
+@pytest.mark.parametrize(
+    'expected',
+    (
+        cirq.Circuit(cirq.X.on(Q[0])),
+        cirq.Circuit(cirq.Y.on(Q[0])),
+        cirq.Circuit(cirq.ParallelGate(cirq.X, 3).on(*Q)),
+        cirq.Circuit(cirq.CNOT.on(Q[0], Q[1])),
+    ),
+)
+def test_gates_preserved(expected: cirq.Circuit):
+    actual = cirq.optimize_for_target_gateset(
+        expected, gateset=cirq.neutral_atoms.NeutralAtomGateset()
+    )
+    assert actual == expected
 
 
 def test_coverage():
-    q = cirq.LineQubit.range(3)
-    g = cirq.testing.ThreeQubitGate()
-
-    class FakeOperation(ops.Operation):
-        def __init__(self, gate, qubits):
-            self._gate = gate
-            self._qubits = qubits
-
-        @property
-        def qubits(self):
-            return self._qubits
-
-        def with_qubits(self, *new_qubits):
-            return FakeOperation(self._gate, new_qubits)
-
-    op = FakeOperation(g, q).with_qubits(*q)
-    circuit_ops = [cirq.Y(q[0]), cirq.ParallelGate(cirq.X, 3).on(*q)]
-    c = cirq.Circuit(circuit_ops)
-    cirq.neutral_atoms.ConvertToNeutralAtomGates().optimize_circuit(c)
-    assert c == cirq.Circuit(circuit_ops)
-    assert cirq.neutral_atoms.ConvertToNeutralAtomGates().convert(cirq.X.on(q[0])) == [
-        cirq.X.on(q[0])
-    ]
-    with pytest.raises(TypeError, match="Don't know how to work with"):
-        cirq.neutral_atoms.ConvertToNeutralAtomGates().convert(op)
-    assert not cirq.neutral_atoms.is_native_neutral_atom_op(op)
-    assert not cirq.neutral_atoms.is_native_neutral_atom_gate(g)
+    with cirq.testing.assert_deprecated(
+        "Use cirq.optimize_for_target_gateset", deadline='v0.16', count=5
+    ):
+        q = cirq.LineQubit.range(3)
+        g = cirq.testing.ThreeQubitGate()
+
+        class FakeOperation(cirq.Operation):
+            def __init__(self, gate, qubits):
+                self._gate = gate
+                self._qubits = qubits
+
+            @property
+            def qubits(self):
+                return self._qubits
+
+            def with_qubits(self, *new_qubits):
+                return FakeOperation(self._gate, new_qubits)
+
+        op = FakeOperation(g, q).with_qubits(*q)
+        circuit_ops = [cirq.Y(q[0]), cirq.ParallelGate(cirq.X, 3).on(*q)]
+        c = cirq.Circuit(circuit_ops)
+        cirq.neutral_atoms.ConvertToNeutralAtomGates().optimize_circuit(c)
+        assert c == cirq.Circuit(circuit_ops)
+        assert cirq.neutral_atoms.ConvertToNeutralAtomGates().convert(cirq.X.on(q[0])) == [
+            cirq.X.on(q[0])
+        ]
+        with pytest.raises(TypeError, match="Don't know how to work with"):
+            cirq.neutral_atoms.ConvertToNeutralAtomGates().convert(op)
+        assert not cirq.neutral_atoms.is_native_neutral_atom_op(op)
+        assert not cirq.neutral_atoms.is_native_neutral_atom_gate(g)
 
 
 def test_avoids_decompose_fallback_when_matrix_available_single_qubit():
@@ -60,8 +81,13 @@ def _decompose_(self, qubits):
 
     q = cirq.GridQubit(0, 0)
     c = cirq.Circuit(OtherX().on(q), OtherOtherX().on(q))
-    cirq.neutral_atoms.ConvertToNeutralAtomGates().optimize_circuit(c)
-    cirq.testing.assert_has_diagram(c, '(0, 0): ───PhX(1)───PhX(1)───')
+    converted = cirq.optimize_for_target_gateset(c, gateset=cirq.neutral_atoms.NeutralAtomGateset())
+    cirq.testing.assert_has_diagram(converted, '(0, 0): ───PhX(1)───PhX(1)───')
+    with cirq.testing.assert_deprecated(
+        "Use cirq.optimize_for_target_gateset", deadline='v0.16', count=2
+    ):
+        cirq.neutral_atoms.ConvertToNeutralAtomGates().optimize_circuit(c)
+        cirq.testing.assert_has_diagram(c, '(0, 0): ───PhX(1)───PhX(1)───')
 
 
 def test_avoids_decompose_fallback_when_matrix_available_two_qubit():
@@ -76,12 +102,15 @@ def _decompose_(self, qubits):
     q00 = cirq.GridQubit(0, 0)
     q01 = cirq.GridQubit(0, 1)
     c = cirq.Circuit(OtherCZ().on(q00, q01), OtherOtherCZ().on(q00, q01))
-    cirq.neutral_atoms.ConvertToNeutralAtomGates().optimize_circuit(c)
-    cirq.testing.assert_has_diagram(
-        c,
-        """
+    expected_diagram = """
 (0, 0): ───@───@───
            │   │
 (0, 1): ───@───@───
-""",
-    )
+"""
+    converted = cirq.optimize_for_target_gateset(c, gateset=cirq.neutral_atoms.NeutralAtomGateset())
+    cirq.testing.assert_has_diagram(converted, expected_diagram)
+    with cirq.testing.assert_deprecated(
+        "Use cirq.optimize_for_target_gateset", deadline='v0.16', count=2
+    ):
+        cirq.neutral_atoms.ConvertToNeutralAtomGates().optimize_circuit(c)
+        cirq.testing.assert_has_diagram(c, expected_diagram)

cirq-core/cirq/neutral_atoms/neutral_atom_devices.py

@@ -20,7 +20,8 @@
 from cirq.devices.grid_qubit import GridQubit
 from cirq.ops import raw_types
 from cirq.value import Duration
-from cirq.neutral_atoms import convert_to_neutral_atom_gates
+from cirq.neutral_atoms.convert_to_neutral_atom_gates import ConvertToNeutralAtomGates
+from cirq.neutral_atoms.neutral_atom_gateset import NeutralAtomGateset
 
 if TYPE_CHECKING:
     import cirq
@@ -31,22 +32,6 @@ def _subgate_if_parallel_gate(gate: 'cirq.Gate') -> 'cirq.Gate':
     return gate.sub_gate if isinstance(gate, ops.ParallelGate) else gate
 
 
-def neutral_atom_gateset(max_parallel_z=None, max_parallel_xy=None):
-    return ops.Gateset(
-        ops.AnyIntegerPowerGateFamily(ops.CNotPowGate),
-        ops.AnyIntegerPowerGateFamily(ops.CCNotPowGate),
-        ops.AnyIntegerPowerGateFamily(ops.CZPowGate),
-        ops.AnyIntegerPowerGateFamily(ops.CCZPowGate),
-        ops.ParallelGateFamily(ops.ZPowGate, max_parallel_allowed=max_parallel_z),
-        ops.ParallelGateFamily(ops.XPowGate, max_parallel_allowed=max_parallel_xy),
-        ops.ParallelGateFamily(ops.YPowGate, max_parallel_allowed=max_parallel_xy),
-        ops.ParallelGateFamily(ops.PhasedXPowGate, max_parallel_allowed=max_parallel_xy),
-        ops.MeasurementGate,
-        ops.IdentityGate,
-        unroll_circuit_op=False,
-    )
-
-
 @value.value_equality
 class NeutralAtomDevice(devices.Device):
     """A device with qubits placed on a grid."""
@@ -107,7 +92,7 @@ def __init__(
             ops.AnyIntegerPowerGateFamily(ops.CCZPowGate),
             unroll_circuit_op=False,
         )
-        self.gateset = neutral_atom_gateset(max_parallel_z, max_parallel_xy)
+        self.gateset = NeutralAtomGateset(max_parallel_z, max_parallel_xy)
         for q in qubits:
             if not isinstance(q, GridQubit):
                 raise ValueError(f'Unsupported qubit type: {q!r}')
@@ -133,7 +118,7 @@ def qubit_list(self):
         deadline='v0.15',
     )
     def decompose_operation(self, operation: ops.Operation) -> ops.OP_TREE:
-        return convert_to_neutral_atom_gates.ConvertToNeutralAtomGates().convert(operation)
+        return ConvertToNeutralAtomGates().convert(operation)
 
     def duration_of(self, operation: ops.Operation):
         """Provides the duration of the given operation on this device.

cirq-core/cirq/neutral_atoms/neutral_atom_devices_test.py

@@ -114,7 +114,7 @@ def test_init_errors():
 
 def test_decompose_error_deprecated():
     d = square_device(2, 2, holes=[cirq.GridQubit(1, 1)])
-    with cirq.testing.assert_deprecated('ConvertToNeutralAtomGates', deadline='v0.15'):
+    with cirq.testing.assert_deprecated('ConvertToNeutralAtomGates', deadline='v0.15', count=2):
         for op in d.decompose_operation((cirq.CCZ**1.5).on(*(d.qubit_list()))):
             d.validate_operation(op)
 

cirq-core/cirq/neutral_atoms/neutral_atom_gateset.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.
+from typing import List, Optional, TYPE_CHECKING
+
+from cirq import ops, protocols, transformers
+from cirq.protocols.decompose_protocol import DecomposeResult
+
+
+if TYPE_CHECKING:
+    import cirq
+
+
+class NeutralAtomGateset(transformers.CompilationTargetGateset):
+    """A Compilation target intended for neutral atom devices.
+
+    This gateset supports CNOT, CCNOT (TOFFOLI) gates, CZ,
+    CCZ gates, as well as single qubits gates that can be used
+    in a parallel fashion.  The maximum amount of parallelism
+    can be set by arguments.
+
+    This compilation gateset decomposes operations into CZ
+    because CZ gates are the highest fidelity two qubit gates
+    for neutral atoms.
+
+    Args:
+        max_parallel_z: The maximum amount of parallelism for
+            Z gates.
+        max_parallel_xy: The maximum amount of parallelism for
+            X, Y and PhasedXPow gates.
+    """
+
+    def __init__(self, max_parallel_z: Optional[int] = None, max_parallel_xy: Optional[int] = None):
+        super().__init__(
+            ops.AnyIntegerPowerGateFamily(ops.CNotPowGate),
+            ops.AnyIntegerPowerGateFamily(ops.CCNotPowGate),
+            ops.AnyIntegerPowerGateFamily(ops.CZPowGate),
+            ops.AnyIntegerPowerGateFamily(ops.CCZPowGate),
+            ops.ParallelGateFamily(ops.ZPowGate, max_parallel_allowed=max_parallel_z),
+            ops.ParallelGateFamily(ops.XPowGate, max_parallel_allowed=max_parallel_xy),
+            ops.ParallelGateFamily(ops.YPowGate, max_parallel_allowed=max_parallel_xy),
+            ops.ParallelGateFamily(ops.PhasedXPowGate, max_parallel_allowed=max_parallel_xy),
+            ops.MeasurementGate,
+            ops.IdentityGate,
+            unroll_circuit_op=False,
+        )
+
+    def num_qubits(self) -> int:
+        """Maximum number of qubits on which a gate from this gateset can act upon."""
+        return 2
+
+    def decompose_to_target_gateset(self, op: 'cirq.Operation', moment_idx: int) -> DecomposeResult:
+        """Method to rewrite the given operation using gates from this gateset.
+        Args:
+            op: `cirq.Operation` to be rewritten using gates from this gateset.
+            moment_idx: Moment index where the given operation `op` occurs in a circuit.
+        Returns:
+            - An equivalent `cirq.OP_TREE` implementing `op` using gates from this gateset.
+            - `None` or `NotImplemented` if does not know how to decompose `op`.
+        """
+        # Known matrix?
+        mat = protocols.unitary(op, None) if len(op.qubits) <= 2 else None
+        if mat is not None and len(op.qubits) == 1:
+            gates = transformers.single_qubit_matrix_to_phased_x_z(mat)
+            return [g.on(op.qubits[0]) for g in gates]
+        if mat is not None and len(op.qubits) == 2:
+            return transformers.two_qubit_matrix_to_cz_operations(
+                op.qubits[0], op.qubits[1], mat, allow_partial_czs=False, clean_operations=True
+            )
+
+        return NotImplemented
+
+    @property
+    def preprocess_transformers(self) -> List['cirq.TRANSFORMER']:
+        return []
+
+    @property
+    def postprocess_transformers(self) -> List['cirq.TRANSFORMER']:
+        return []

cirq-core/cirq/neutral_atoms/neutral_atom_gateset_test.py

@@ -0,0 +1,66 @@
+# 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
+import cirq.neutral_atoms.neutral_atom_gateset as nag
+
+Q = cirq.LineQubit.range(10)
+
+
+@pytest.mark.parametrize(
+    'op,is_in_gateset',
+    (
+        (cirq.X.on(Q[0]), True),
+        (cirq.Y.on(Q[0]), True),
+        (cirq.I.on(Q[0]), True),
+        (cirq.measure(*Q), True),
+        (cirq.ParallelGate(cirq.X, 3).on(*Q[:3]), True),
+        (cirq.ParallelGate(cirq.X, 4).on(*Q[:4]), False),
+        (cirq.ParallelGate(cirq.X, 10).on(*Q), False),
+        (cirq.ParallelGate(cirq.Y, 3).on(*Q[:3]), True),
+        (cirq.ParallelGate(cirq.Y, 4).on(*Q[:4]), False),
+        (cirq.ParallelGate(cirq.Y, 10).on(*Q), False),
+        (cirq.ParallelGate(cirq.Z, 3).on(*Q[:3]), True),
+        (cirq.ParallelGate(cirq.Z, 4).on(*Q[:4]), True),
+        (cirq.ParallelGate(cirq.Z, 5).on(*Q[:5]), False),
+        (cirq.ParallelGate(cirq.Z, 10).on(*Q), False),
+        (
+            cirq.ParallelGate(cirq.PhasedXPowGate(exponent=0.5, phase_exponent=0.25), 3).on(*Q[:3]),
+            True,
+        ),
+        (
+            cirq.ParallelGate(cirq.PhasedXPowGate(exponent=0.5, phase_exponent=0.25), 4).on(*Q[:4]),
+            False,
+        ),
+        (cirq.CNOT.on(Q[0], Q[1]), True),
+        ((cirq.CNOT**0.5).on(Q[0], Q[1]), False),
+        (cirq.CZ.on(Q[0], Q[1]), True),
+        ((cirq.CZ**0.5).on(Q[0], Q[1]), False),
+        (cirq.CCZ.on(Q[0], Q[1], Q[2]), True),
+        ((cirq.CCZ**0.5).on(Q[0], Q[1], Q[2]), False),
+        ((cirq.TOFFOLI**0.5).on(Q[0], Q[1], Q[2]), False),
+    ),
+)
+def test_gateset(op: cirq.Operation, is_in_gateset: bool):
+    gateset = nag.NeutralAtomGateset(max_parallel_z=4, max_parallel_xy=3)
+    assert gateset.validate(op) == is_in_gateset
+    converted_ops = cirq.optimize_for_target_gateset(cirq.Circuit(op), gateset=gateset)
+    if is_in_gateset:
+        assert converted_ops == cirq.Circuit(op)
+        assert gateset.validate(converted_ops)
+
+
+def test_gateset_qubits():
+    assert nag.NeutralAtomGateset(max_parallel_z=4, max_parallel_xy=3).num_qubits() == 2