1# Copyright 2018 The Cirq Developers
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14from typing import List, Optional, TYPE_CHECKING
15
16from cirq import ops, protocols
17from cirq.circuits.optimization_pass import (
18    PointOptimizationSummary,
19    PointOptimizer,
20)
21from cirq.neutral_atoms import neutral_atom_devices
22from cirq import optimizers
23
24if TYPE_CHECKING:
25    import cirq
26
27
28class ConvertToNeutralAtomGates(PointOptimizer):
29    """Attempts to convert gates into native Atom gates.
30
31    First, checks if the given operation is already a native neutral atom
32    operation.
33
34    Second, checks if the operation has a known unitary. If so, and the gate
35        is a 1-qubit or 2-qubit gate, then performs circuit synthesis of the
36        operation. The 2-qubit gates are decomposed using CZ gates because
37        CZ gates are the highest fidelity 2-qubit gates for neutral atoms.
38
39    Third, attempts to `cirq.decompose` to the operation.
40
41    Fourth, if ignore_failures is set, gives up and returns the gate unchanged.
42        Otherwise raises a TypeError.
43    """
44
45    def __init__(self, ignore_failures=False) -> None:
46        """Inits ConvertToNeutralAtomGates.
47
48        Args:
49            ignore_failures: If set, gates that fail to convert are forwarded
50                unchanged. If not set, conversion failures raise a TypeError.
51        """
52        super().__init__()
53        self.ignore_failures = ignore_failures
54        self.gateset = neutral_atom_devices.neutral_atom_gateset()
55
56    def _convert_one(self, op: ops.Operation) -> ops.OP_TREE:
57        # Known matrix?
58        mat = protocols.unitary(op, None) if len(op.qubits) <= 2 else None
59        if mat is not None and len(op.qubits) == 1:
60            gates = optimizers.single_qubit_matrix_to_phased_x_z(mat)
61            return [g.on(op.qubits[0]) for g in gates]
62        if mat is not None and len(op.qubits) == 2:
63            return optimizers.two_qubit_matrix_to_operations(
64                op.qubits[0], op.qubits[1], mat, allow_partial_czs=False, clean_operations=True
65            )
66
67        return NotImplemented
68
69    def convert(self, op: ops.Operation) -> List[ops.Operation]:
70        def on_stuck_raise(bad):
71            return TypeError(
72                "Don't know how to work with {!r}. "
73                "It isn't a native atom operation, "
74                "a 1 or 2 qubit gate with a known unitary, "
75                "or composite.".format(bad)
76            )
77
78        return protocols.decompose(
79            op,
80            keep=self.gateset._validate_operation,
81            intercepting_decomposer=self._convert_one,
82            on_stuck_raise=None if self.ignore_failures else on_stuck_raise,
83        )
84
85    def optimization_at(
86        self, circuit: 'cirq.Circuit', index: int, op: 'cirq.Operation'
87    ) -> Optional['cirq.PointOptimizationSummary']:
88        converted = self.convert(op)
89        if len(converted) == 1 and converted[0] is op:
90            return None
91        return PointOptimizationSummary(
92            clear_span=1, new_operations=converted, clear_qubits=op.qubits
93        )
94
95
96def is_native_neutral_atom_op(operation: ops.Operation) -> bool:
97    return operation in neutral_atom_devices.neutral_atom_gateset()
98
99
100def is_native_neutral_atom_gate(gate: ops.Gate) -> bool:
101    return gate in neutral_atom_devices.neutral_atom_gateset()
102