Repository: quantumlib/Cirq

Issue:
Implement ThreeQubitDiagonalGate
A three-qubit diagonal gate is a gate whose unitary matrix is of the form `np.diag([a, b, c, d, e, f, g, h])`. Such gates can be decomposed and simulated efficiently, and this decomposition/simulation code would be useful in several places (e.g. when defining the `CXXYY` gate in openfermion-cirq).

The `ThreeQubitDiagonalGate` class should have the following features:

- An `__init__` method taking the rotation angle for each element along diagonal. (The *angles*, not the complex values the angles produces. Also store the angles, not the complex values. This makes the behavior of `__pow__` better.)
- A `def _unitary_(self)` method that returns its matrix.
- A `def _circuit_diagram_info_(self, args)` method that uses symbols like `diag(a, b, ...)----#2----#3`.
- A `def default_decompose(self)` method that looks almost identical to the CCZ decomposition, except that the various T/phasing gates are replaced with gates that depend on the diagonal components of the matrix.
    - For reference, this is the CCZ code:

        ```python
        def default_decompose(self, qubits):
            """An adjacency-respecting decomposition.
    
            0: ───p───@──────────────@───────@──────────@──────────
                      │              │       │          │
            1: ───p───X───@───p^-1───X───@───X──────@───X──────@───
                          │              │          │          │
            2: ───p───────X───p──────────X───p^-1───X───p^-1───X───
    
            where p = T**self._exponent
            """
            a, b, c = qubits
    
            # Hacky magic: avoid the non-adjacent edge.
            if hasattr(b, 'is_adjacent'):
                if not b.is_adjacent(a):
                    b, c = c, b
                elif not b.is_adjacent(c):
                    a, b = b, a
    
            p = common_gates.T**self._exponent
            sweep_abc = [common_gates.CNOT(a, b),
                         common_gates.CNOT(b, c)]
    
            yield p(a), p(b), p(c)
            yield sweep_abc
            yield p(b)**-1, p(c)
            yield sweep_abc
            yield p(c)**-1
            yield sweep_abc
            yield p(c)**-1
            yield sweep_abc
            ```

    Each of the `p`s in the code controls the phasing of four of the eight diagonal entries in the matrix. By varying all 7 you can implement any diagonal matrix (up to global phase). You will have to write code to solve the actual_diagonal-to-phasing_angles problem.

- A `def _apply_unitary_to_tensor_` method that multiplies the appropriate entries by the appropriate phase values from the matrix.
     - For reference, here is how the CCZ would be implemented:
        ```python
        def _apply_unitary_to_tensor_(self,
                                      target_tensor: np.ndarray,
                                      available_buffer: np.ndarray,
                                      axes: Sequence[int],
                                      ) -> np.ndarray:
            ooo = linalg.slice_for_qubits_equal_to(axes, 0b111)
            target_tensor[ooo] *= np.exp(1j * self.exponent * np.pi)
            return target_tensor
        ```
    The code should avoid doing multiplications where the phase values equal 1.

- A `__pow__` method to extrapolate the phasing effect.
- `__eq__`/`__ne__`/`__hash__`




Discussion:
@Strilanc what did you mean by the `#2---#3` in `def _circuit_diagram_info_(self, args)`?
It's a way of indicating the order of the other qubits in the operation. First qubit gets the name, others get numbers. This matters for operations that depend on the qubit order.

```
----op----
    |     
----#3----
    |     
----#2----
```

Resolving patch:
diff --git a/cirq/__init__.py b/cirq/__init__.py
--- a/cirq/__init__.py
+++ b/cirq/__init__.py
@@ -183,6 +183,7 @@
     SwapPowGate,
     T,
     ThreeQubitGate,
+    ThreeQubitDiagonalGate,
     TOFFOLI,
     transform_op_tree,
     TwoQubitGate,
diff --git a/cirq/ops/__init__.py b/cirq/ops/__init__.py
--- a/cirq/ops/__init__.py
+++ b/cirq/ops/__init__.py
@@ -183,5 +183,6 @@
     CSWAP,
     CSwapGate,
     FREDKIN,
+    ThreeQubitDiagonalGate,
     TOFFOLI,
 )
diff --git a/cirq/ops/three_qubit_gates.py b/cirq/ops/three_qubit_gates.py
--- a/cirq/ops/three_qubit_gates.py
+++ b/cirq/ops/three_qubit_gates.py
@@ -14,9 +14,10 @@
 
 """Common quantum gates that target three qubits."""
 
-from typing import Optional, Tuple
+from typing import Any, List, Optional, Tuple, TYPE_CHECKING, Union
 
 import numpy as np
+import sympy
 
 from cirq import linalg, protocols, value
 from cirq._compat import proper_repr
@@ -29,6 +30,9 @@
     op_tree,
     raw_types,
 )
+if TYPE_CHECKING:
+    # pylint: disable=unused-import
+    import cirq
 
 
 class CCZPowGate(eigen_gate.EigenGate,
@@ -150,6 +154,136 @@ def __str__(self) -> str:
         return 'CCZ**{}'.format(self._exponent)
 
 
+@value.value_equality()
+class ThreeQubitDiagonalGate(gate_features.ThreeQubitGate):
+    """A gate given by a diagonal 8x8 matrix."""
+
+    def __init__(self,
+                 diag_angles_radians: List[Union[float, sympy.Basic]]) -> None:
+        """
+
+        Args:
+            diag_angles_radians: The list of angles on the diagonal in radians.
+
+        """
+        self._diag_angles_radians = diag_angles_radians \
+        # type: List[Union[float, sympy.Basic]]
+
+    def _is_parameterized_(self):
+        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 NotImplemented
+        return np.diag(
+            [np.exp(1j * angle) for angle in self._diag_angles_radians])
+
+    def _resolve_parameters_(self, param_resolver: 'cirq.ParamResolver'
+                            ) -> 'ThreeQubitDiagonalGate':
+        return ThreeQubitDiagonalGate([
+            protocols.resolve_parameters(angle, param_resolver)
+            for angle in self._diag_angles_radians
+        ])
+
+    def _circuit_diagram_info_(self, args: protocols.CircuitDiagramInfoArgs
+                              ) -> protocols.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', '#3'))
+
+    def __pow__(self, exponent: Any) -> 'ThreeQubitDiagonalGate':
+        if not isinstance(exponent, (int, float, sympy.Basic)):
+            return NotImplemented
+        return ThreeQubitDiagonalGate([
+            protocols.mul(angle, exponent, NotImplemented)
+            for angle in self._diag_angles_radians
+        ])
+
+    def _decompose_(self, qubits):
+        """An adjacency-respecting decomposition.
+
+        0: ───p_0───@──────────────@───────@──────────@──────────
+                    │              │       │          │
+        1: ───p_1───X───@───p_3────X───@───X──────@───X──────@───
+                        │              │          │          │
+        2: ───p_2───────X───p_4────────X───p_5────X───p_6────X───
+
+        where p_i = T**(4*x_i) and x_i solve the system of equations
+                    [0, 0, 1, 0, 1, 1, 1][x_0]   [r_1]
+                    [0, 1, 0, 1, 1, 0, 1][x_1]   [r_2]
+                    [0, 1, 1, 1, 0, 1, 0][x_2]   [r_3]
+                    [1, 0, 0, 1, 1, 1, 0][x_3] = [r_4]
+                    [1, 0, 1, 1, 0, 0, 1][x_4]   [r_5]
+                    [1, 1, 0, 0, 0, 1, 1][x_5]   [r_6]
+                    [1, 1, 1, 0, 1, 0, 0][x_6]   [r_7]
+        where r_i is self._diag_angles_radians[i].
+
+        The above system was created by equating the composition of the gates
+        in the circuit diagram to np.diag(self._diag_angles) (shifted by a
+        global phase of np.exp(-1j * self._diag_angles[0])).
+        """
+
+        a, b, c = qubits
+        if hasattr(b, 'is_adjacent'):
+            if not b.is_adjacent(a):
+                b, c = c, b
+            elif not b.is_adjacent(c):
+                a, b = b, a
+        sweep_abc = [common_gates.CNOT(a, b), common_gates.CNOT(b, c)]
+        phase_matrix_inverse = .25 * np.array(
+            [[-1, -1, -1, 1, 1, 1, 1], [-1, 1, 1, -1, -1, 1, 1],
+             [1, -1, 1, -1, 1, -1, 1], [-1, 1, 1, 1, 1, -1, -1],
+             [1, 1, -1, 1, -1, -1, 1], [1, -1, 1, 1, -1, 1, -1],
+             [1, 1, -1, -1, 1, 1, -1]])
+        shifted_angles_tail = [
+            angle - self._diag_angles_radians[0]
+            for angle in self._diag_angles_radians[1:]
+        ]
+        phase_solutions = phase_matrix_inverse.dot(shifted_angles_tail)
+        p_gates = [
+            pauli_gates.Z**(solution / np.pi) for solution in phase_solutions
+        ]
+
+        return [
+            p_gates[0](a),
+            p_gates[1](b),
+            p_gates[2](c),
+            sweep_abc,
+            p_gates[3](b),
+            p_gates[4](c),
+            sweep_abc,
+            p_gates[5](c),
+            sweep_abc,
+            p_gates[6](c),
+            sweep_abc,
+        ]
+
+    def _apply_unitary_(self, args: protocols.ApplyUnitaryArgs) -> np.ndarray:
+        if self._is_parameterized_():
+            return NotImplemented
+        for index, angle in enumerate(self._diag_angles_radians):
+            little_endian_index = 4 * (index & 1) + 2 * ((index >> 1) & 1) + (
+                (index >> 2) & 1)
+            subspace_index = args.subspace_index(little_endian_index)
+            args.target_tensor[subspace_index] *= np.exp(1j * angle)
+        return args.target_tensor
+
+    def _value_equality_values_(self):
+        return tuple(self._diag_angles_radians)
+
+    def __repr__(self) -> str:
+        return 'cirq.ThreeQubitDiagonalGate([{}])'.format(','.join(
+            proper_repr(angle) for angle in self._diag_angles_radians))
+
+
 class CCXPowGate(eigen_gate.EigenGate,
                  gate_features.ThreeQubitGate,
                  gate_features.InterchangeableQubitsGate):


Test patch:
diff --git a/cirq/ops/three_qubit_gates_test.py b/cirq/ops/three_qubit_gates_test.py
--- a/cirq/ops/three_qubit_gates_test.py
+++ b/cirq/ops/three_qubit_gates_test.py
@@ -15,6 +15,7 @@
 import itertools
 import numpy as np
 import pytest
+import sympy
 
 import cirq
 
@@ -32,6 +33,10 @@ def test_eigen_gates_consistent_protocols(eigen_gate_type):
 def test_consistent_protocols():
     cirq.testing.assert_implements_consistent_protocols(cirq.CSWAP)
 
+    cirq.testing.assert_implements_consistent_protocols(
+        cirq.ThreeQubitDiagonalGate([2, 3, 5, 7, 11, 13, 17, 19]),
+        ignoring_global_phase=True)
+
 
 def test_init():
     assert (cirq.CCZ**0.5).exponent == 0.5
@@ -87,6 +92,13 @@ def test_unitary():
         [0, 0, 0, 0, 0, 0, 0, 1],
     ]), atol=1e-8)
 
+    diagonal_angles = [2, 3, 5, 7, 11, 13, 17, 19]
+    assert cirq.has_unitary(cirq.ThreeQubitDiagonalGate(diagonal_angles))
+    np.testing.assert_allclose(
+        cirq.unitary(cirq.ThreeQubitDiagonalGate(diagonal_angles)),
+        np.diag([np.exp(1j * angle) for angle in diagonal_angles]),
+        atol=1e-8)
+
 
 def test_str():
     assert str(cirq.CCX) == 'TOFFOLI'
@@ -133,14 +145,12 @@ def test_eq():
 @pytest.mark.parametrize('op,max_two_cost', [
     (cirq.CCZ(*cirq.LineQubit.range(3)), 8),
     (cirq.CCX(*cirq.LineQubit.range(3)), 8),
-    (cirq.CCZ(cirq.LineQubit(0),
-              cirq.LineQubit(2),
-              cirq.LineQubit(1)), 8),
+    (cirq.CCZ(cirq.LineQubit(0), cirq.LineQubit(2), cirq.LineQubit(1)), 8),
     (cirq.CSWAP(*cirq.LineQubit.range(3)), 9),
     (cirq.CSWAP(*reversed(cirq.LineQubit.range(3))), 9),
-    (cirq.CSWAP(cirq.LineQubit(1),
-                cirq.LineQubit(0),
-                cirq.LineQubit(2)), 12),
+    (cirq.CSWAP(cirq.LineQubit(1), cirq.LineQubit(0), cirq.LineQubit(2)), 12),
+    (cirq.ThreeQubitDiagonalGate([2, 3, 5, 7, 11, 13, 17, 19])(
+        cirq.LineQubit(1), cirq.LineQubit(2), cirq.LineQubit(3)), 8),
 ])
 def test_decomposition_cost(op: cirq.Operation, max_two_cost: int):
     ops = tuple(
@@ -152,7 +162,10 @@ def test_decomposition_cost(op: cirq.Operation, max_two_cost: int):
 
 
 @pytest.mark.parametrize('gate', [
-    cirq.CCX, cirq.CSWAP, cirq.CCZ,
+    cirq.CCX,
+    cirq.CSWAP,
+    cirq.CCZ,
+    cirq.ThreeQubitDiagonalGate([2, 3, 5, 7, 11, 13, 17, 19]),
 ])
 def test_decomposition_respects_locality(gate):
     a = cirq.GridQubit(0, 0)
@@ -167,16 +180,13 @@ def test_decomposition_respects_locality(gate):
 
 def test_diagram():
     a, b, c, d = cirq.LineQubit.range(4)
-    circuit = cirq.Circuit.from_ops(
-        cirq.TOFFOLI(a, b, c),
-        cirq.TOFFOLI(a, b, c)**0.5,
-        cirq.CCX(a, c, b),
-        cirq.CCZ(a, d, b),
-        cirq.CCZ(a, d, b)**0.5,
-        cirq.CSWAP(a, c, d),
-        cirq.FREDKIN(a, b, c)
-    )
-    cirq.testing.assert_has_diagram(circuit, """
+    circuit = cirq.Circuit.from_ops(cirq.TOFFOLI(a, b, c),
+                                    cirq.TOFFOLI(a, b, c)**0.5,
+                                    cirq.CCX(a, c, b), cirq.CCZ(a, d, b),
+                                    cirq.CCZ(a, d, b)**0.5, cirq.CSWAP(a, c, d),
+                                    cirq.FREDKIN(a, b, c))
+    cirq.testing.assert_has_diagram(
+        circuit, """
 0: ───@───@───────@───@───@───────@───@───
       │   │       │   │   │       │   │
 1: ───@───@───────X───@───@───────┼───×───
@@ -185,7 +195,8 @@ def test_diagram():
                       │   │       │
 3: ───────────────────@───@^0.5───×───────
 """)
-    cirq.testing.assert_has_diagram(circuit, """
+    cirq.testing.assert_has_diagram(circuit,
+                                    """
 0: ---@---@-------@---@---@-------@------@------
       |   |       |   |   |       |      |
 1: ---@---@-------X---@---@-------|------swap---
@@ -193,4 +204,57 @@ def test_diagram():
 2: ---X---X^0.5---@---|---|-------swap---swap---
                       |   |       |
 3: -------------------@---@^0.5---swap----------
-""", use_unicode_characters=False)
+""",
+                                    use_unicode_characters=False)
+
+    diagonal_circuit = cirq.Circuit.from_ops(
+        cirq.ThreeQubitDiagonalGate([2, 3, 5, 7, 11, 13, 17, 19])(a, b, c))
+    cirq.testing.assert_has_diagram(
+        diagonal_circuit, """
+0: ───diag(2, 3, 5, 7, 11, 13, 17, 19)───
+      │
+1: ───#2─────────────────────────────────
+      │
+2: ───#3─────────────────────────────────
+""")
+    cirq.testing.assert_has_diagram(diagonal_circuit,
+                                    """
+0: ---diag(2, 3, 5, 7, 11, 13, 17, 19)---
+      |
+1: ---#2---------------------------------
+      |
+2: ---#3---------------------------------
+""",
+                                    use_unicode_characters=False)
+
+
+def test_diagonal_exponent():
+    diagonal_angles = [2, 3, 5, 7, 11, 13, 17, 19]
+    diagonal_gate = cirq.ThreeQubitDiagonalGate(diagonal_angles)
+
+    sqrt_diagonal_gate = diagonal_gate**.5
+
+    expected_angles = [prime / 2 for prime in diagonal_angles]
+    np.testing.assert_allclose(expected_angles,
+                               sqrt_diagonal_gate._diag_angles_radians,
+                               atol=1e-8)
+
+    assert cirq.pow(cirq.ThreeQubitDiagonalGate(diagonal_angles), "test",
+                    None) is None
+
+
+def test_resolve():
+    diagonal_angles = [2, 3, 5, 7, 11, 13, 17, 19]
+    diagonal_gate = cirq.ThreeQubitDiagonalGate(
+        diagonal_angles[:6] +
+        [sympy.Symbol('a'), sympy.Symbol('b')])
+    assert cirq.is_parameterized(diagonal_gate)
+
+    diagonal_gate = cirq.resolve_parameters(diagonal_gate, {'a': 17})
+    assert diagonal_gate == cirq.ThreeQubitDiagonalGate(diagonal_angles[:7] +
+                                                        [sympy.Symbol('b')])
+    assert cirq.is_parameterized(diagonal_gate)
+
+    diagonal_gate = cirq.resolve_parameters(diagonal_gate, {'b': 19})
+    assert diagonal_gate == cirq.ThreeQubitDiagonalGate(diagonal_angles)
+    assert not cirq.is_parameterized(diagonal_gate)