[3 qubit decomposition] Extract diagonal (#3472)

Adds `cirq.two_qubit_matrix_to_diagonal_and_operations` that extracts a diagonal if it can from a 3CNOT unitary to a diagonal and a 2 CNOT unitary. 

Another part of #2873.

cirq/__init__.py

@@ -318,6 +318,7 @@
     stratified_circuit,
     SynchronizeTerminalMeasurements,
     two_qubit_matrix_to_operations,
+    two_qubit_matrix_to_diagonal_and_operations,
 )
 
 from cirq.qis import (

cirq/linalg/__init__.py

@@ -29,6 +29,7 @@
     axis_angle,
     AxisAngleDecomposition,
     deconstruct_single_qubit_matrix_into_angles,
+    extract_right_diag,
     kak_canonicalize_vector,
     kak_decomposition,
     kak_vector,

cirq/linalg/decompositions.py

@@ -15,14 +15,14 @@
 
 """Utility methods for breaking matrices into useful pieces."""
 
+import cmath
+import math
 from typing import (Any, Callable, Iterable, List, Optional, Sequence, Set,
                     Tuple, TYPE_CHECKING, TypeVar, Union)
 
-import math
-import cmath
+import matplotlib.pyplot as plt
 import numpy as np
 import scipy
-import matplotlib.pyplot as plt
 
 from cirq import value, protocols
 from cirq._compat import proper_repr
@@ -1045,3 +1045,24 @@ def _gamma(u: np.ndarray) -> np.ndarray:
         u @ yy @ u.T @ yy, where yy = Y ⊗ Y
     """
     return u @ YY @ u.T @ YY
+
+
+def extract_right_diag(u: np.ndarray) -> np.ndarray:
+    """Extract a diagonal unitary from a 3-CNOT two-qubit unitary.
+
+    Returns a 2-CNOT unitary D that is diagonal, so that U @ D needs only
+    two CNOT gates in case the original unitary is a 3-CNOT unitary.
+
+    See Proposition V.2 in Minimal Universal Two-Qubit CNOT-based Circuits.
+    https://arxiv.org/abs/quant-ph/0308033
+
+    Args:
+        u: three-CNOT two-qubit unitary
+    Returns:
+        diagonal extracted from U
+    """
+    t = _gamma(transformations.to_special(u).T).diagonal()
+    k = np.real(t[0] + t[3] - t[1] - t[2])
+    psi = np.arctan2(np.imag(np.sum(t)), k)
+    f = np.exp(1j * psi)
+    return np.diag([1, f, f, 1])

cirq/linalg/decompositions_test.py

@@ -748,7 +748,7 @@ def test_kak_decompose(unitary: np.ndarray):
 def test_num_two_qubit_gates_required():
     for i in range(4):
         assert cirq.num_cnots_required(
-            _two_qubit_circuit_with_cnots(i).unitary()) == i
+            cirq.testing.random_two_qubit_circuit_with_czs(i).unitary()) == i
 
     assert cirq.num_cnots_required(np.eye(4)) == 0
 
@@ -758,23 +758,19 @@ def test_num_two_qubit_gates_required_invalid():
         cirq.num_cnots_required(np.array([[1]]))
 
 
-def _two_qubit_circuit_with_cnots(num_cnots=3, a=None, b=None):
-    random.seed(32123)
-    if a is None or b is None:
-        a, b = cirq.LineQubit.range(2)
-
-    def random_one_qubit_gate():
-        return cirq.PhasedXPowGate(phase_exponent=random.random(),
-                                   exponent=random.random())
-
-    def one_cz():
-        return [
-            cirq.CZ.on(a, b),
-            random_one_qubit_gate().on(a),
-            random_one_qubit_gate().on(b),
-        ]
-
-    return cirq.Circuit([
-        random_one_qubit_gate().on(a),
-        random_one_qubit_gate().on(b), [one_cz() for _ in range(num_cnots)]
+@pytest.mark.parametrize(
+    "U",
+    [
+        cirq.testing.random_two_qubit_circuit_with_czs(3).unitary(),
+        # an example where gamma(special(u))=I, so the denominator becomes 0
+        1 / np.sqrt(2) * np.array(
+            [[(1 - 1j) * 2 / np.sqrt(5), 0, 0,
+              (1 - 1j) * 1 / np.sqrt(5)], [0, 0, 1 - 1j, 0], [0, 1 - 1j, 0, 0],
+             [-(1 - 1j) * 1 / np.sqrt(5), 0, 0, (1 - 1j) * 2 / np.sqrt(5)]],
+            dtype=np.complex128)
     ])
+def test_extract_right_diag(U):
+    assert cirq.num_cnots_required(U) == 3
+    diag = cirq.linalg.extract_right_diag(U)
+    assert cirq.is_diagonal(diag)
+    assert cirq.num_cnots_required(U @ diag) == 2

cirq/optimizers/__init__.py

@@ -64,7 +64,9 @@
     SynchronizeTerminalMeasurements,)
 
 from cirq.optimizers.two_qubit_decompositions import (
-    two_qubit_matrix_to_operations,)
+    two_qubit_matrix_to_operations,
+    two_qubit_matrix_to_diagonal_and_operations,
+)
 
 from cirq.optimizers.two_qubit_to_fsim import (
     decompose_two_qubit_interaction_into_four_fsim_gates,

cirq/optimizers/two_qubit_decompositions.py

@@ -18,6 +18,9 @@
 
 import numpy as np
 
+from cirq.linalg import predicates
+from cirq.linalg.decompositions import num_cnots_required, extract_right_diag
+
 from cirq import ops, linalg, protocols, circuits
 from cirq.optimizers import (
     decompositions,
@@ -61,6 +64,56 @@ def two_qubit_matrix_to_operations(
     return operations
 
 
+def two_qubit_matrix_to_diagonal_and_operations(
+        q0: 'cirq.Qid',
+        q1: 'cirq.Qid',
+        mat: np.ndarray,
+        allow_partial_czs: bool = False,
+        atol: float = 1e-8,
+        clean_operations: bool = True
+) -> Tuple[np.ndarray, List['cirq.Operation']]:
+    """Decomposes a 2-qubit unitary to a diagonal and the remaining operations.
+
+    For a 2-qubit unitary V, return ops, a list of operations and
+    D diagonal unitary, so that:
+        V = cirq.Circuit(ops) @ D
+
+    Args:
+        q0: The first qubit being operated on.
+        q1: The other qubit being operated on.
+        mat: the input unitary
+        allow_partial_czs: Enables the use of Partial-CZ gates.
+        atol: A limit on the amount of absolute error introduced by the
+            construction.
+        clean_operations: Enables optimizing resulting operation list by
+            merging operations and ejecting phased Paulis and Z operations.
+    Returns:
+        tuple(ops,D): operations `ops`, and the diagonal `D`
+    """
+    if predicates.is_diagonal(mat, atol=atol):
+        return mat, []
+
+    if num_cnots_required(mat) == 3:
+        right_diag = extract_right_diag(mat)
+        two_cnot_unitary = mat @ right_diag
+        # note that this implies that two_cnot_unitary @ d = mat
+        return right_diag.conj().T, two_qubit_matrix_to_operations(
+            q0,
+            q1,
+            two_cnot_unitary,
+            allow_partial_czs=allow_partial_czs,
+            atol=atol,
+            clean_operations=clean_operations)
+
+    return np.eye(4), two_qubit_matrix_to_operations(
+        q0,
+        q1,
+        mat,
+        allow_partial_czs=allow_partial_czs,
+        atol=atol,
+        clean_operations=clean_operations)
+
+
 def _xx_interaction_via_full_czs(q0: 'cirq.Qid', q1: 'cirq.Qid', x: float):
     a = x * -2 / np.pi
     yield ops.H(q1)

cirq/optimizers/two_qubit_decompositions_test.py

@@ -21,8 +21,9 @@
 import cirq
 from cirq import value
 from cirq.optimizers.two_qubit_decompositions import (
-    _parity_interaction, _is_trivial_angle
-)
+    _parity_interaction, _is_trivial_angle,
+    two_qubit_matrix_to_diagonal_and_operations)
+from cirq.testing import random_two_qubit_circuit_with_czs
 
 
 @pytest.mark.parametrize('rad,expected', (lambda err, largeErr: [
@@ -262,3 +263,20 @@ def test_kak_decomposition_depth_partial_cz():
     c = cirq.Circuit(operations_with_part)
     # 1 CP, 1+1 PhasedX, 1 Z
     assert len(c) <= 4
+
+
+@pytest.mark.parametrize("v", [
+    cirq.unitary(random_two_qubit_circuit_with_czs(3)),
+    cirq.unitary(random_two_qubit_circuit_with_czs(2)),
+    np.diag(np.exp(1j * np.pi * np.random.random(4))),
+])
+def test_decompose_to_diagonal_and_circuit(v):
+    b, c = cirq.LineQubit.range(2)
+    diagonal, ops = two_qubit_matrix_to_diagonal_and_operations(b, c, v)
+    assert cirq.is_diagonal(diagonal)
+    combined_circuit = cirq.Circuit(cirq.MatrixGate(diagonal)(b, c), ops)
+    circuit_unitary = combined_circuit.unitary(
+        qubits_that_should_be_present=[b, c])
+    cirq.testing.assert_allclose_up_to_global_phase(circuit_unitary,
+                                                    v,
+                                                    atol=1e-14)

cirq/testing/__init__.py

@@ -89,6 +89,7 @@
 from cirq.testing.random_circuit import (
     DEFAULT_GATE_DOMAIN,
     random_circuit,
+    random_two_qubit_circuit_with_czs,
 )
 
 from cirq.testing.sample_circuits import (

cirq/testing/random_circuit.py

@@ -15,6 +15,7 @@
 from typing import List, Union, Sequence, Dict, Optional, TYPE_CHECKING
 
 from cirq import ops, value
+from cirq.ops import Qid
 from cirq.circuits import Circuit
 from cirq._doc import document
 
@@ -119,3 +120,42 @@ def random_circuit(qubits: Union[Sequence[ops.Qid], int],
         moments.append(ops.Moment(operations))
 
     return Circuit(moments)
+
+
+def random_two_qubit_circuit_with_czs(
+        num_czs: int = 3,
+        q0: Qid = None,
+        q1: Qid = None,
+        random_state: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None) -> Circuit:
+    """Creates a random two qubit circuit with the given number of CNOTs.
+
+    The resulting circuit will have `num_cnots` number of CNOTs that will be
+    surrounded by random `PhasedXPowGate` instances on both qubits.
+
+    Args:
+         num_czs: the number of CNOTs to be guaranteed in the circuit
+         q0: the first qubit the circuit should operate on
+         q1: the second qubit the circuit should operate on
+         random_state: an optional random seed
+    Returns:
+         the random two qubit circuit
+    """
+    prng = value.parse_random_state(random_state)
+    q0 = ops.NamedQubit('q0') if q0 is None else q0
+    q1 = ops.NamedQubit('q1') if q1 is None else q1
+
+    def random_one_qubit_gate():
+        return ops.PhasedXPowGate(phase_exponent=prng.random(),
+                                  exponent=prng.random())
+
+    def one_cz():
+        return [
+            ops.CZ.on(q0, q1),
+            random_one_qubit_gate().on(q0),
+            random_one_qubit_gate().on(q1),
+        ]
+
+    return Circuit([
+        random_one_qubit_gate().on(q0),
+        random_one_qubit_gate().on(q1), [one_cz() for _ in range(num_czs)]
+    ])

cirq/testing/random_circuit_test.py

@@ -136,3 +136,62 @@ def test_random_circuit_reproducible_between_runs():
                   └──┘
     """
     cirq.testing.assert_has_diagram(circuit, expected_diagram)
+
+
+def test_random_two_qubit_circuit_with_czs():
+    num_czs = lambda circuit: len([
+        o for o in circuit.all_operations() if isinstance(
+            o.gate, cirq.CZPowGate)
+    ])
+
+    c = cirq.testing.random_two_qubit_circuit_with_czs()
+    assert num_czs(c) == 3
+    assert {cirq.NamedQubit('q0'), cirq.NamedQubit('q1')} == c.all_qubits()
+    assert all(
+        isinstance(op.gate, cirq.PhasedXPowGate) for op in c[0].operations)
+    assert c[0].qubits == c.all_qubits()
+
+    c = cirq.testing.random_two_qubit_circuit_with_czs(num_czs=0)
+    assert num_czs(c) == 0
+    assert {cirq.NamedQubit('q0'), cirq.NamedQubit('q1')} == c.all_qubits()
+    assert all(
+        isinstance(op.gate, cirq.PhasedXPowGate) for op in c[0].operations)
+    assert c[0].qubits == c.all_qubits()
+
+    a, b = cirq.LineQubit.range(2)
+    c = cirq.testing.random_two_qubit_circuit_with_czs(num_czs=1, q1=b)
+    assert num_czs(c) == 1
+    assert {b, cirq.NamedQubit('q0')} == c.all_qubits()
+    assert all(
+        isinstance(op.gate, cirq.PhasedXPowGate) for op in c[0].operations)
+    assert c[0].qubits == c.all_qubits()
+
+    c = cirq.testing.random_two_qubit_circuit_with_czs(num_czs=2, q0=a)
+    assert num_czs(c) == 2
+    assert {a, cirq.NamedQubit('q1')} == c.all_qubits()
+    assert all(
+        isinstance(op.gate, cirq.PhasedXPowGate) for op in c[0].operations)
+    assert c[0].qubits == c.all_qubits()
+
+    c = cirq.testing.random_two_qubit_circuit_with_czs(num_czs=3, q0=a, q1=b)
+    assert num_czs(c) == 3
+    assert c.all_qubits() == {a, b}
+    assert all(
+        isinstance(op.gate, cirq.PhasedXPowGate) for op in c[0].operations)
+    assert c[0].qubits == c.all_qubits()
+
+    seed = 77
+
+    c1 = cirq.testing.random_two_qubit_circuit_with_czs(num_czs=4,
+                                                        q0=a,
+                                                        q1=b,
+                                                        random_state=seed)
+    assert num_czs(c1) == 4
+    assert c1.all_qubits() == {a, b}
+
+    c2 = cirq.testing.random_two_qubit_circuit_with_czs(num_czs=4,
+                                                        q0=a,
+                                                        q1=b,
+                                                        random_state=seed)
+
+    assert c1 == c2

rtd_docs/api.rst

@@ -424,6 +424,7 @@ Classes and methods for rewriting circuits.
     cirq.single_qubit_matrix_to_phxz
     cirq.single_qubit_op_to_framed_phase_form
     cirq.stratified_circuit
+    cirq.two_qubit_matrix_to_diagonal_and_operations
     cirq.two_qubit_matrix_to_operations
     cirq.ConvertToCzAndSingleGates
     cirq.DropEmptyMoments
@@ -631,6 +632,7 @@ operation.
     cirq.testing.random_special_orthogonal
     cirq.testing.random_special_unitary
     cirq.testing.random_superposition
+    cirq.testing.random_two_qubit_circuit_with_czs
     cirq.testing.random_unitary
     cirq.testing.EqualsTester
     cirq.testing.NoIdentifierQubit