Improved sample_gates.py implementation and unitary_protocol tests. Also added docstrings

Author: tanujkhattar

cirq-core/cirq/protocols/apply_unitary_protocol.py

@@ -134,10 +134,31 @@ def default(
         return ApplyUnitaryArgs(state, np.empty_like(state), range(num_qubits))
 
     @classmethod
-    def for_unitary(cls, qid_shapes: Tuple[int, ...]) -> 'ApplyUnitaryArgs':
-        state = qis.eye_tensor(qid_shapes, dtype=np.complex128)
-        buffer = np.empty_like(state)
-        return ApplyUnitaryArgs(state, buffer, range(len(qid_shapes)))
+    def for_unitary(
+        cls, num_qubits: Optional[int] = None, *, qid_shape: Optional[Tuple[int, ...]] = None
+    ) -> 'ApplyUnitaryArgs':
+        """A default instance corresponding to an identity matrix.
+
+        Specify exactly one argument.
+
+        Args:
+            num_qubits: The number of qubits to make space for in the state.
+            qid_shape: A tuple representing the number of quantum levels of each
+                qubit the identity matrix applies to. `qid_shape` is (2, 2, 2) for
+                a three-qubit identity operation tensor.
+
+        Raises:
+            TypeError: If exactly neither `num_qubits` or `qid_shape` is provided or
+                both are provided.
+        """
+        if (num_qubits is None) == (qid_shape is None):
+            raise TypeError('Specify exactly one of num_qubits or qid_shape.')
+        if num_qubits is not None:
+            qid_shape = (2,) * num_qubits
+        qid_shape = cast(Tuple[int, ...], qid_shape)  # Satisfy mypy
+        num_qubits = len(qid_shape)
+        state = qis.eye_tensor(qid_shape, dtype=np.complex128)
+        return ApplyUnitaryArgs(state, np.empty_like(state), range(num_qubits))
 
     def with_axes_transposed_to_start(self) -> 'ApplyUnitaryArgs':
         """Returns a transposed view of the same arguments.
@@ -471,7 +492,7 @@ def _strat_apply_unitary_from_decompose(val: Any, args: ApplyUnitaryArgs) -> Opt
     ordered_qubits = ancilla + tuple(qubits)
     all_qid_shapes = qid_shape_protocol.qid_shape(ordered_qubits)
     result = apply_unitaries(
-        operations, ordered_qubits, ApplyUnitaryArgs.for_unitary(all_qid_shapes), None
+        operations, ordered_qubits, ApplyUnitaryArgs.for_unitary(qid_shape=all_qid_shapes), None
     )
     if result is None or result is NotImplemented:
         return result

cirq-core/cirq/protocols/unitary_protocol.py

@@ -161,7 +161,7 @@ def _strat_unitary_from_apply_unitary(val: Any) -> Optional[np.ndarray]:
         return NotImplemented
 
     # Apply unitary effect to an identity matrix.
-    result = method(ApplyUnitaryArgs.for_unitary(val_qid_shape))
+    result = method(ApplyUnitaryArgs.for_unitary(qid_shape=val_qid_shape))
 
     if result is NotImplemented or result is None:
         return result
@@ -185,7 +185,7 @@ def _strat_unitary_from_decompose(val: Any) -> Optional[np.ndarray]:
 
     # Apply sub-operations' unitary effects to an identity matrix.
     result = apply_unitaries(
-        operations, ordered_qubits, ApplyUnitaryArgs.for_unitary(val_qid_shape), None
+        operations, ordered_qubits, ApplyUnitaryArgs.for_unitary(qid_shape=val_qid_shape), None
     )
 
     # Package result.

cirq-core/cirq/protocols/unitary_protocol_test.py

@@ -189,29 +189,40 @@ def test_has_unitary():
     assert not cirq.has_unitary(FullyImplemented(False))
 
 
-@pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 10))
-def test_decompose_gate_that_allocates_qubits(theta: float):
+def _test_gate_that_allocates_qubits(gate):
     from cirq.protocols.unitary_protocol import _strat_unitary_from_decompose
 
-    gate = testing.GateThatAllocatesAQubit(theta)
-    np.testing.assert_allclose(
-        cast(np.ndarray, _strat_unitary_from_decompose(gate)), gate.target_unitary()
-    )
-    np.testing.assert_allclose(
-        cast(np.ndarray, _strat_unitary_from_decompose(gate(a))), gate.target_unitary()
-    )
+    op = gate.on(*cirq.LineQubit.range(cirq.num_qubits(gate)))
+    moment = cirq.Moment(op)
+    circuit = cirq.FrozenCircuit(op)
+    circuit_op = cirq.CircuitOperation(circuit)
+    for val in [gate, op, moment, circuit, circuit_op]:
+        unitary_from_strat = _strat_unitary_from_decompose(val)
+        assert unitary_from_strat is not None
+        np.testing.assert_allclose(unitary_from_strat, gate.narrow_unitary())
 
 
 @pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 10))
-@pytest.mark.parametrize('n', [*range(1, 6)])
-def test_recusive_decomposition(n: int, theta: float):
-    from cirq.protocols.unitary_protocol import _strat_unitary_from_decompose
+@pytest.mark.parametrize('phase_state', [0, 1])
+@pytest.mark.parametrize('target_bitsize', [1, 2, 3])
+@pytest.mark.parametrize('ancilla_bitsize', [1, 4])
+def test_decompose_gate_that_allocates_clean_qubits(
+    theta: float, phase_state: int, target_bitsize: int, ancilla_bitsize: int
+):
 
-    g1 = testing.GateThatDecomposesIntoNGates(n, cirq.H, theta)
-    g2 = testing.GateThatDecomposesIntoNGates(n, g1, theta)
-    np.testing.assert_allclose(
-        cast(np.ndarray, _strat_unitary_from_decompose(g2)), g2.target_unitary()
-    )
+    gate = testing.PhaseUsingCleanAncilla(theta, phase_state, target_bitsize, ancilla_bitsize)
+    _test_gate_that_allocates_qubits(gate)
+
+
+@pytest.mark.parametrize('phase_state', [0, 1])
+@pytest.mark.parametrize('target_bitsize', [1, 2, 3])
+@pytest.mark.parametrize('ancilla_bitsize', [1, 4])
+def test_decompose_gate_that_allocates_dirty_qubits(
+    phase_state: int, target_bitsize: int, ancilla_bitsize: int
+):
+
+    gate = testing.PhaseUsingDirtyAncilla(phase_state, target_bitsize, ancilla_bitsize)
+    _test_gate_that_allocates_qubits(gate)
 
 
 def test_decompose_and_get_unitary():
@@ -227,15 +238,6 @@ def test_decompose_and_get_unitary():
     np.testing.assert_allclose(_strat_unitary_from_decompose(DummyComposite()), np.eye(1))
     np.testing.assert_allclose(_strat_unitary_from_decompose(OtherComposite()), m2)
 
-    np.testing.assert_allclose(
-        _strat_unitary_from_decompose(testing.GateThatAllocatesTwoQubits()),
-        testing.GateThatAllocatesTwoQubits.target_unitary(),
-    )
-    np.testing.assert_allclose(
-        _strat_unitary_from_decompose(testing.GateThatAllocatesTwoQubits().on(a, b)),
-        testing.GateThatAllocatesTwoQubits.target_unitary(),
-    )
-
 
 def test_decomposed_has_unitary():
     # Gates

cirq-core/cirq/testing/__init__.py

@@ -108,8 +108,4 @@
 
 from cirq.testing.sample_circuits import nonoptimal_toffoli_circuit
 
-from cirq.testing.sample_gates import (
-    GateThatAllocatesAQubit,
-    GateThatAllocatesTwoQubits,
-    GateThatDecomposesIntoNGates,
-)
+from cirq.testing.sample_gates import PhaseUsingCleanAncilla, PhaseUsingDirtyAncilla

cirq-core/cirq/testing/sample_gates.py

@@ -11,75 +11,69 @@
 # 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 functools
-import numpy as np
-from cirq import ops
-
-
-class GateThatAllocatesAQubit(ops.Gate):
-    r"""A gate that applies $Z^\theta$ indirectly through a clean ancilla."""
+import dataclasses
 
-    def __init__(self, theta: float) -> None:
-        super().__init__()
-        self._theta = theta
+import cirq
+import numpy as np
+from cirq import ops, qis
 
-    def _num_qubits_(self):
-        return 1
 
-    def _decompose_(self, q):
-        anc = ops.NamedQubit("anc")
-        yield ops.CX(*q, anc)
-        yield (ops.Z**self._theta)(anc)
-        yield ops.CX(*q, anc)
+def _matrix_for_phasing_state(num_qubits, phase_state, phase):
+    matrix = qis.eye_tensor((2,) * num_qubits, dtype=np.complex128)
+    matrix = matrix.reshape((2**num_qubits, 2**num_qubits))
+    matrix[phase_state, phase_state] = phase
+    print(num_qubits, phase_state, phase)
+    print(matrix)
+    return matrix
 
-    def target_unitary(self) -> np.ndarray:
-        return np.array([[1, 0], [0, (-1 + 0j) ** self._theta]])
 
+@dataclasses.dataclass(frozen=True)
+class PhaseUsingCleanAncilla(ops.Gate):
+    r"""Phases the state $|phase_state>$ by $\exp(1j * \pi * \theta)$ using one clean ancilla."""
 
-class GateThatAllocatesTwoQubits(ops.Gate):
-    r"""A gate that applies $-j Z \otimes Z$ indirectly through two ancillas."""
+    theta: float
+    phase_state: int = 1
+    target_bitsize: int = 1
+    ancilla_bitsize: int = 1
 
     def _num_qubits_(self):
-        return 2
+        return self.target_bitsize
 
-    def _decompose_(self, qs):
-        q0, q1 = qs
-        anc = ops.NamedQubit.range(2, prefix='two_ancillas_')
+    def _decompose_(self, qubits):
+        anc = ops.NamedQubit.range(self.ancilla_bitsize, prefix="anc")
+        cv = [int(x) for x in f'{self.phase_state:0{self.target_bitsize}b}']
+        cnot_ladder = [cirq.CNOT(anc[i - 1], anc[i]) for i in range(1, self.ancilla_bitsize)]
 
-        yield ops.X(anc[0])
-        yield ops.CX(q0, anc[0])
-        yield (ops.Y)(anc[0])
-        yield ops.CX(q0, anc[0])
+        yield ops.X(anc[0]).controlled_by(*qubits, control_values=cv)
+        yield [cnot_ladder, ops.Z(anc[-1]) ** self.theta, reversed(cnot_ladder)]
+        yield ops.X(anc[0]).controlled_by(*qubits, control_values=cv)
 
-        yield ops.CX(q1, anc[1])
-        yield (ops.Z)(anc[1])
-        yield ops.CX(q1, anc[1])
+    def narrow_unitary(self) -> np.ndarray:
+        """Narrowed unitary corresponding to the unitary effect applied on target qubits."""
+        phase = np.exp(1j * np.pi * self.theta)
+        return _matrix_for_phasing_state(self.target_bitsize, self.phase_state, phase)
 
-    @classmethod
-    def target_unitary(cls) -> np.ndarray:
-        # Unitary = -j Z \otimes Z
-        return np.array([[-1j, 0, 0, 0], [0, 1j, 0, 0], [0, 0, 1j, 0], [0, 0, 0, -1j]])
 
+@dataclasses.dataclass(frozen=True)
+class PhaseUsingDirtyAncilla(ops.Gate):
+    r"""Phases the state $|phase_state>$ by -1 using one dirty ancilla."""
 
-class GateThatDecomposesIntoNGates(ops.Gate):
-    r"""Applies $(Z^\theta)^{\otimes_n}$ on work qubits and `subgate` on $n$ borrowable ancillas."""
+    phase_state: int = 1
+    target_bitsize: int = 1
+    ancilla_bitsize: int = 1
 
-    def __init__(self, n: int, subgate: ops.Gate, theta: float) -> None:
-        super().__init__()
-        self._n = n
-        self._subgate = subgate
-        self._name = str(subgate)
-        self._theta = theta
-
-    def _num_qubits_(self) -> int:
-        return self._n
-
-    def _decompose_(self, qs):
-        ancilla = ops.NamedQubit.range(self._n, prefix=self._name)
-        yield self._subgate.on_each(ancilla)
-        yield (ops.Z**self._theta).on_each(qs)
-        yield self._subgate.on_each(ancilla)
-
-    def target_unitary(self) -> np.ndarray:
-        U = np.array([[1, 0], [0, (-1 + 0j) ** self._theta]])
-        return functools.reduce(np.kron, [U] * self._n)
+    def _num_qubits_(self):
+        return self.target_bitsize
+
+    def _decompose_(self, qubits):
+        anc = ops.NamedQubit.range(self.ancilla_bitsize, prefix="anc")
+        cv = [int(x) for x in f'{self.phase_state:0{self.target_bitsize}b}']
+        cnot_ladder = [cirq.CNOT(anc[i - 1], anc[i]) for i in range(1, self.ancilla_bitsize)]
+        yield ops.X(anc[0]).controlled_by(*qubits, control_values=cv)
+        yield [cnot_ladder, ops.Z(anc[-1]), reversed(cnot_ladder)]
+        yield ops.X(anc[0]).controlled_by(*qubits, control_values=cv)
+        yield [cnot_ladder, ops.Z(anc[-1]), reversed(cnot_ladder)]
+
+    def narrow_unitary(self) -> np.ndarray:
+        """Narrowed unitary corresponding to the unitary effect applied on target qubits."""
+        return _matrix_for_phasing_state(self.target_bitsize, self.phase_state, -1)

cirq-core/cirq/testing/sample_gates_test.py

@@ -11,48 +11,49 @@
 # 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 functools
 import pytest
 
 import numpy as np
 from cirq.testing import sample_gates
-from cirq import protocols, ops
+import cirq
 
 
 @pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 20))
-def test_GateThatAllocatesAQubit(theta: float):
-    g = sample_gates.GateThatAllocatesAQubit(theta)
-
-    want = np.array([[1, 0], [0, (-1 + 0j) ** theta]], dtype=np.complex128)
-    # test unitary
-    np.testing.assert_allclose(g.target_unitary(), want)
-
-    # test decomposition
-    np.testing.assert_allclose(protocols.unitary(g), g.target_unitary())
-
-
-def test_GateThatAllocatesTwoQubits():
-    g = sample_gates.GateThatAllocatesTwoQubits()
-
-    Z = np.array([[1, 0], [0, -1]])
-    want = -1j * np.kron(Z, Z)
-    # test unitary
-    np.testing.assert_allclose(g.target_unitary(), want)
-
-    # test decomposition
-    np.testing.assert_allclose(protocols.unitary(g), g.target_unitary())
-
-
-@pytest.mark.parametrize('n', [*range(1, 6)])
-@pytest.mark.parametrize('subgate', [ops.Z, ops.X, ops.Y, ops.T])
-@pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 5))
-def test_GateThatDecomposesIntoNGates(n: int, subgate: ops.Gate, theta: float):
-    g = sample_gates.GateThatDecomposesIntoNGates(n, subgate, theta)
-
-    U = np.array([[1, 0], [0, (-1 + 0j) ** theta]], dtype=np.complex128)
-    want = functools.reduce(np.kron, [U] * n)
-    # test unitary
-    np.testing.assert_allclose(g.target_unitary(), want)
-
-    # test decomposition
-    np.testing.assert_allclose(protocols.unitary(g), g.target_unitary())
+def test_phase_using_clean_ancilla(theta: float):
+    g = sample_gates.PhaseUsingCleanAncilla(theta)
+    q = cirq.LineQubit(0)
+    qubit_order = cirq.QubitOrder.explicit([q], fallback=cirq.QubitOrder.DEFAULT)
+    decomposed_unitary = cirq.Circuit(cirq.decompose_once(g.on(q))).unitary(qubit_order=qubit_order)
+    phase = np.exp(1j * np.pi * theta)
+    np.testing.assert_allclose(g.narrow_unitary(), np.array([[1, 0], [0, phase]]))
+    np.testing.assert_allclose(
+        decomposed_unitary,
+        # fmt: off
+        np.array(
+            [
+                [1 , 0    , 0    , 0],
+                [0 , phase, 0    , 0],
+                [0 , 0    , phase, 0],
+                [0 , 0    , 0    , 1],
+            ]
+        ),
+        # fmt: on
+    )
+
+
+@pytest.mark.parametrize(
+    'target_bitsize, phase_state', [(1, 0), (1, 1), (2, 0), (2, 1), (2, 2), (2, 3)]
+)
+@pytest.mark.parametrize('ancilla_bitsize', [1, 4])
+def test_phase_using_dirty_ancilla(target_bitsize, phase_state, ancilla_bitsize):
+    g = sample_gates.PhaseUsingDirtyAncilla(phase_state, target_bitsize, ancilla_bitsize)
+    q = cirq.LineQubit.range(target_bitsize)
+    qubit_order = cirq.QubitOrder.explicit(q, fallback=cirq.QubitOrder.DEFAULT)
+    decomposed_circuit = cirq.Circuit(cirq.decompose_once(g.on(*q)))
+    decomposed_unitary = decomposed_circuit.unitary(qubit_order=qubit_order)
+    phase_matrix = np.eye(2**target_bitsize)
+    phase_matrix[phase_state, phase_state] = -1
+    np.testing.assert_allclose(g.narrow_unitary(), phase_matrix)
+    np.testing.assert_allclose(
+        decomposed_unitary, np.kron(phase_matrix, np.eye(2**ancilla_bitsize)), atol=1e-5
+    )