Add TwoQubitDiagonalGate (#2891)

* Add TwoQubitDiagonalGate

format

Respond to reviews

revetr

revert again

ignore global phase

lint

Fixed edits for the Add TwoQubitDiagonalGate.

Change the docstring for the two_qubit_diagonal_gate.py

* Pull request review edits.

* GitHub merge fail fixes.

* Update two_qubit_diagonal_gate.py

Co-authored-by: Victory Omole <vtomole2@gmail.com>

cirq/__init__.py

@@ -255,6 +255,7 @@
     ThreeQubitDiagonalGate,
     TOFFOLI,
     transform_op_tree,
+    TwoQubitDiagonalGate,
     TwoQubitGate,
     WaitGate,
     X,

cirq/ops/__init__.py

@@ -220,5 +220,8 @@
     TOFFOLI,
 )
 
+from cirq.ops.two_qubit_diagonal_gate import (
+    TwoQubitDiagonalGate,)
+
 from cirq.ops.wait_gate import (
     WaitGate,)

cirq/ops/two_qubit_diagonal_gate.py

@@ -0,0 +1,104 @@
+# Copyright 2020 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.
+"""Creates the gate instrance for a two qubit diagonal gate.
+
+The gate is used to create a 4x4 matrix with the diagonal elements
+passed as a list.
+"""
+
+from typing import Any, Tuple, List, TYPE_CHECKING
+import numpy as np
+import sympy
+
+from cirq import protocols, value
+from cirq._compat import proper_repr
+from cirq.ops import gate_features
+
+if TYPE_CHECKING:
+    # pylint: disable=unused-import
+    import cirq
+
+
+@value.value_equality()
+class TwoQubitDiagonalGate(gate_features.TwoQubitGate):
+    """A gate given by a diagonal 4\\times 4 matrix."""
+
+    def __init__(self, diag_angles_radians: List[value.TParamVal]) -> None:
+        r"""A two qubit gate with only diagonal elements.
+
+        This gate's off-diagonal elements are zero and it's on diagonal
+        elements are all phases.
+
+        Args:
+            diag_angles_radians: The list of angles on the diagonal in radians.
+                If these values are $(x_0, x_1, \ldots , x_3)$ then the unitary
+                has diagonal values $(e^{i x_0}, e^{i x_1}, \ldots, e^{i x_3})$.
+        """
+        self._diag_angles_radians: Tuple[value.TParamVal, ...] = tuple(
+            diag_angles_radians)
+
+    def _is_parameterized_(self) -> bool:
+        return any(
+            isinstance(angle, sympy.Basic)
+            for angle in self._diag_angles_radians)
+
+    def _has_unitary_(self) -> bool:
+        return not self._is_parameterized_()
+
+    def _unitary_(self) -> np.ndarray:
+        if self._is_parameterized_():
+            return None
+        return np.diag(
+            [np.exp(1j * angle) for angle in self._diag_angles_radians])
+
+    def _apply_unitary_(self, args: 'protocols.ApplyUnitaryArgs') -> np.ndarray:
+        if self._is_parameterized_():
+            return NotImplemented
+        for index, angle in enumerate(self._diag_angles_radians):
+            subspace_index = args.subspace_index(big_endian_bits_int=index)
+            args.target_tensor[subspace_index] *= np.exp(1j * angle)
+        return args.target_tensor
+
+    def _resolve_parameters_(self, param_resolver: 'cirq.ParamResolver'
+                            ) -> 'TwoQubitDiagonalGate':
+        return TwoQubitDiagonalGate(
+            protocols.resolve_parameters(self._diag_angles_radians,
+                                         param_resolver))
+
+    def _circuit_diagram_info_(self, args: 'cirq.CircuitDiagramInfoArgs'
+                              ) -> 'cirq.CircuitDiagramInfo':
+        rounded_angles = np.array(self._diag_angles_radians)
+        if args.precision is not None:
+            rounded_angles = rounded_angles.round(args.precision)
+        diag_str = 'diag({})'.format(', '.join(
+            proper_repr(angle) for angle in rounded_angles))
+        return protocols.CircuitDiagramInfo((diag_str, '#2'))
+
+    def __pow__(self, exponent: Any) -> 'TwoQubitDiagonalGate':
+        if not isinstance(exponent, (int, float, sympy.Basic)):
+            return NotImplemented
+        angles = []
+        for angle in self._diag_angles_radians:
+            mulAngle = protocols.mul(angle, exponent, NotImplemented)
+            if mulAngle == NotImplemented:
+                return NotImplemented
+            angles.append(mulAngle)
+        return TwoQubitDiagonalGate(angles)
+
+    def _value_equality_values_(self) -> Any:
+        return tuple(self._diag_angles_radians)
+
+    def __repr__(self) -> str:
+        return 'cirq.TwoQubitDiagonalGate([{}])'.format(','.join(
+            proper_repr(angle) for angle in self._diag_angles_radians))

cirq/ops/two_qubit_diagonal_gate_test.py

@@ -0,0 +1,95 @@
+# Copyright 2020 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 numpy as np
+import pytest
+import sympy
+
+import cirq
+
+
+@pytest.mark.parametrize(
+    'gate', ((cirq.TwoQubitDiagonalGate(
+        [2, 3, 5, 7]), cirq.TwoQubitDiagonalGate([0, 0, 0, 0]),
+              cirq.TwoQubitDiagonalGate([2, 3, 5, sympy.Symbol('a')]),
+              cirq.TwoQubitDiagonalGate([0.34, 0.12, 0, 0.96]))))
+def test_consistent_protocols(gate):
+    cirq.testing.assert_implements_consistent_protocols(gate)
+
+
+def test_unitary():
+    diagonal_angles = [2, 3, 5, 7]
+    assert cirq.has_unitary(cirq.TwoQubitDiagonalGate(diagonal_angles))
+    np.testing.assert_allclose(
+        cirq.unitary(cirq.TwoQubitDiagonalGate(diagonal_angles)),
+        np.diag([np.exp(1j * angle) for angle in diagonal_angles]),
+        atol=1e-8)
+
+
+def test_diagram():
+    a, b = cirq.LineQubit.range(2)
+
+    diagonal_circuit = cirq.Circuit(
+        cirq.TwoQubitDiagonalGate([2, 3, 5, 7])(a, b))
+    cirq.testing.assert_has_diagram(
+        diagonal_circuit, """
+0: ───diag(2, 3, 5, 7)───
+      │
+1: ───#2─────────────────
+""")
+    cirq.testing.assert_has_diagram(diagonal_circuit,
+                                    """
+0: ---diag(2, 3, 5, 7)---
+      |
+1: ---#2-----------------
+""",
+                                    use_unicode_characters=False)
+
+
+def test_diagonal_exponent():
+    diagonal_angles = [2, 3, 5, 7]
+    diagonal_gate = cirq.TwoQubitDiagonalGate(diagonal_angles)
+
+    sqrt_diagonal_gate = diagonal_gate**.5
+
+    expected_angles = [prime / 2 for prime in diagonal_angles]
+    assert cirq.approx_eq(sqrt_diagonal_gate,
+                          cirq.TwoQubitDiagonalGate(expected_angles))
+
+    assert cirq.pow(cirq.TwoQubitDiagonalGate(diagonal_angles), "test",
+                    None) is None
+
+
+def test_protocols_mul_not_implemented():
+    diagonal_angles = [2, 3, None, 7]
+    diagonal_gate = cirq.TwoQubitDiagonalGate(diagonal_angles)
+    with pytest.raises(TypeError):
+        cirq.protocols.pow(diagonal_gate, 3)
+
+
+def test_resolve():
+    diagonal_angles = [2, 3, 5, 7]
+    diagonal_gate = cirq.TwoQubitDiagonalGate(
+        diagonal_angles[:2] +
+        [sympy.Symbol('a'), sympy.Symbol('b')])
+    assert cirq.is_parameterized(diagonal_gate)
+
+    diagonal_gate = cirq.resolve_parameters(diagonal_gate, {'a': 5})
+    assert diagonal_gate == cirq.TwoQubitDiagonalGate(diagonal_angles[:3] +
+                                                      [sympy.Symbol('b')])
+    assert cirq.is_parameterized(diagonal_gate)
+
+    diagonal_gate = cirq.resolve_parameters(diagonal_gate, {'b': 7})
+    assert diagonal_gate == cirq.TwoQubitDiagonalGate(diagonal_angles)
+    assert not cirq.is_parameterized(diagonal_gate)

cirq/protocols/json_serialization.py

@@ -143,6 +143,7 @@ def two_qubit_matrix_gate(matrix):
                 'ThreeDGridQubit': cirq.pasqal.ThreeDGridQubit,
                 'TrialResult': cirq.TrialResult,
                 'TwoQubitMatrixGate': two_qubit_matrix_gate,
+                'TwoQubitDiagonalGate': cirq.TwoQubitDiagonalGate,
                 '_UnconstrainedDevice':
                 cirq.devices.unconstrained_device._UnconstrainedDevice,
                 'WaitGate': cirq.WaitGate,

cirq/protocols/json_serialization_test.py

@@ -315,6 +315,7 @@ def test_mutually_exclusive_blacklist():
     'TextDiagramDrawer',
     'ThreeQubitDiagonalGate',
     'Timestamp',
+    'TwoQubitDiagonalGate',
     'UnitSweep',
     'WaveFunctionSimulatorState',
     'WaveFunctionTrialResult',

docs/api.rst

@@ -97,6 +97,7 @@ Unitary effects that can be applied to one or more qubits.
     cirq.TaggedOperation
     cirq.ThreeQubitDiagonalGate
     cirq.ThreeQubitGate
+    cirq.TwoQubitDiagonalGate
     cirq.TwoQubitGate
     cirq.WaitGate
     cirq.XPowGate