Implemented State Preparation Gate (#4482)

Implements a new gate to reset on all channels and prepare an arbitrary state over n-qubits by returning the Kraus operator for the same. Associated tests included.

Closes #4119.

cirq-core/cirq/__init__.py

@@ -266,6 +266,7 @@
     PhasedXPowGate,
     PhasedXZGate,
     PhaseFlipChannel,
+    StatePreparationChannel,
     ProjectorString,
     ProjectorSum,
     RandomGateChannel,

cirq-core/cirq/json_resolver_cache.py

@@ -123,6 +123,7 @@ def two_qubit_matrix_gate(matrix):
         'PhasedISwapPowGate': cirq.PhasedISwapPowGate,
         'PhasedXPowGate': cirq.PhasedXPowGate,
         'PhasedXZGate': cirq.PhasedXZGate,
+        'StatePreparationChannel': cirq.StatePreparationChannel,
         'ProjectorString': cirq.ProjectorString,
         'ProjectorSum': cirq.ProjectorSum,
         'RandomGateChannel': cirq.RandomGateChannel,

cirq-core/cirq/ops/__init__.py

@@ -299,3 +299,5 @@
     wait,
     WaitGate,
 )
+
+from cirq.ops.state_preparation_channel import StatePreparationChannel

cirq-core/cirq/ops/state_preparation_channel.py

@@ -0,0 +1,120 @@
+# Copyright 2021 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.
+
+"""Quantum gates to prepare a given target state."""
+
+from typing import Any, Dict, Tuple, TYPE_CHECKING
+
+import numpy as np
+
+from cirq import protocols
+from cirq.ops import raw_types
+from cirq._compat import proper_repr
+
+if TYPE_CHECKING:
+    import cirq
+
+
+class StatePreparationChannel(raw_types.Gate):
+    """A channel which prepares any state provided as the state vector on it's target qubits."""
+
+    def __init__(self, target_state: np.ndarray, name: str = "StatePreparation") -> None:
+        """Initializes a State Preparation channel.
+
+        Args:
+            target_state: The state vector that this gate should prepare.
+            name: the name of the gate
+
+        Raises:
+            ValueError: if the array is not 1D, or does not have 2**n elements for some integer n.
+        """
+        if len(target_state.shape) != 1:
+            raise ValueError('`target_state` must be a 1d numpy array.')
+
+        n = int(np.round(np.log2(target_state.shape[0] or 1)))
+        if 2 ** n != target_state.shape[0]:
+            raise ValueError(f'Matrix width ({target_state.shape[0]}) is not a power of 2')
+
+        self._state = target_state.astype(np.complex128) / np.linalg.norm(target_state)
+        self._num_qubits = n
+        self._name = name
+        self._qid_shape = (2,) * n
+
+    @staticmethod
+    def _has_unitary_() -> bool:
+        """Checks and returns if the gate has a unitary representation.
+        It doesn't, since the resetting of the channels is a non-unitary operations,
+        it involves measurement."""
+        return False
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        """Converts the gate object into a serializable dictionary"""
+        return {
+            'cirq_type': self.__class__.__name__,
+            'target_state': self._state.tolist(),
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, target_state, **kwargs):
+        """Recreates the channel object from it's serialized form
+
+        Args:
+            target_state: the state to prepare using this channel
+            kwargs: other keyword arguments, ignored
+        """
+        return cls(target_state=np.array(target_state))
+
+    def _num_qubits_(self):
+        return self._num_qubits
+
+    def _qid_shape_(self) -> Tuple[int, ...]:
+        return self._qid_shape
+
+    def _circuit_diagram_info_(
+        self, _args: 'cirq.CircuitDiagramInfoArgs'
+    ) -> 'cirq.CircuitDiagramInfo':
+        """Returns the information required to draw out the circuit diagram for this channel."""
+        symbols = (
+            [self._name]
+            if self._num_qubits == 1
+            else [f'{self._name}[{i+1}]' for i in range(0, self._num_qubits)]
+        )
+        return protocols.CircuitDiagramInfo(wire_symbols=symbols)
+
+    @staticmethod
+    def _has_kraus_():
+        return True
+
+    def _kraus_(self):
+        """Returns the Kraus operator for this gate
+        The Kraus Operator is |Psi><i| for all |i>, where |Psi> is the target state.
+        This allows is to take any input state to the target state.
+        The operator satisfies the completeness relation Sum(E^ E) = I.
+        """
+        operator = np.zeros(shape=(2 ** self._num_qubits,) * 3, dtype=np.complex128)
+        for i in range(len(operator)):
+            operator[i, :, i] = self._state
+        return operator
+
+    def __repr__(self) -> str:
+        return f'cirq.StatePreparationChannel({proper_repr(self._state)})'
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, StatePreparationChannel):
+            return False
+        return np.allclose(self.state, other.state)
+
+    @property
+    def state(self):
+        return self._state

cirq-core/cirq/ops/state_preparation_channel_test.py

@@ -0,0 +1,128 @@
+# Copyright 2021 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 cirq
+import pytest
+
+
+@pytest.mark.parametrize(
+    'state',
+    np.array(
+        [
+            [1, 0, 0, 0],
+            [1, 0, 0, 1],
+            [3, 5, 2, 7],
+            [0.7823, 0.12323, 0.4312, 0.12321],
+            [23, 43, 12, 19],
+            [1j, 0, 0, 0],
+            [1j, 0, 0, 1j],
+            [1j, -1j, -1j, 1j],
+            [1 + 1j, 0, 0, 0],
+            [1 + 1j, 0, 1 + 1j, 0],
+            [3 + 1j, 5 + 8j, 21, 0.85j],
+        ]
+    ),
+)
+def test_state_prep_channel_kraus(state):
+    qubits = cirq.LineQubit.range(2)
+    gate = cirq.StatePreparationChannel(state)(qubits[0], qubits[1])
+    state = state / np.linalg.norm(state)
+    np.testing.assert_almost_equal(
+        cirq.kraus(gate),
+        (
+            np.array([state, np.zeros(4), np.zeros(4), np.zeros(4)]).T,
+            np.array([np.zeros(4), state, np.zeros(4), np.zeros(4)]).T,
+            np.array([np.zeros(4), np.zeros(4), state, np.zeros(4)]).T,
+            np.array([np.zeros(4), np.zeros(4), np.zeros(4), state]).T,
+        ),
+    )
+
+
+def test_state_prep_channel_kraus_small():
+    gate = cirq.StatePreparationChannel(np.array([0.0, 1.0]))(cirq.LineQubit(0))
+    np.testing.assert_almost_equal(
+        cirq.kraus(gate), (np.array([[0.0, 0.0], [1.0, 0.0]]), np.array([[0.0, 0.0], [0.0, 1.0]]))
+    )
+    assert cirq.has_kraus(gate)
+    assert not cirq.has_mixture(gate)
+
+    gate = cirq.StatePreparationChannel(np.array([1.0, 0.0]))(cirq.LineQubit(0))
+    np.testing.assert_almost_equal(
+        cirq.kraus(gate), (np.array([[1.0, 0.0], [0.0, 0.0]]), np.array([[0.0, 1.0], [0.0, 0.0]]))
+    )
+    assert cirq.has_kraus(gate)
+    assert not cirq.has_mixture(gate)
+
+
+def test_state_prep_gate_printing():
+    circuit = cirq.Circuit()
+    qubits = cirq.LineQubit.range(2)
+    gate = cirq.StatePreparationChannel(np.array([1, 0, 0, 1]) / np.sqrt(2))
+    circuit.append(cirq.H(qubits[0]))
+    circuit.append(cirq.CNOT(qubits[0], qubits[1]))
+    circuit.append(gate(qubits[0], qubits[1]))
+    cirq.testing.assert_has_diagram(
+        circuit,
+        """
+0: ───H───@───StatePreparation[1]───
+          │   │
+1: ───────X───StatePreparation[2]───
+""",
+    )
+
+
+@pytest.mark.parametrize('name', ['Prep', 'S'])
+def test_state_prep_gate_printing_with_name(name):
+    circuit = cirq.Circuit()
+    qubits = cirq.LineQubit.range(2)
+    gate = cirq.StatePreparationChannel(np.array([1, 0, 0, 1]) / np.sqrt(2), name=name)
+    circuit.append(cirq.H(qubits[0]))
+    circuit.append(cirq.CNOT(qubits[0], qubits[1]))
+    circuit.append(gate(qubits[0], qubits[1]))
+    cirq.testing.assert_has_diagram(
+        circuit,
+        f"""
+0: ───H───@───{name}[1]───
+          │   │
+1: ───────X───{name}[2]───
+""",
+    )
+
+
+def test_gate_params():
+    state = np.array([1, 0, 0, 0], dtype=np.complex64)
+    gate = cirq.StatePreparationChannel(state)
+    assert gate.num_qubits() == 2
+    assert not gate._has_unitary_()
+    assert gate._has_kraus_()
+    assert (
+        repr(gate)
+        == 'cirq.StatePreparationChannel(np.array([(1+0j), 0j, 0j, 0j], dtype=np.complex128))'
+    )
+
+
+def test_gate_error_handling():
+    with pytest.raises(ValueError, match='`target_state` must be a 1d numpy array.'):
+        cirq.StatePreparationChannel(np.eye(2))
+    with pytest.raises(ValueError, match=f'Matrix width \\(5\\) is not a power of 2'):
+        cirq.StatePreparationChannel(np.ones(shape=5))
+
+
+def test_equality_of_gates():
+    state = np.array([1, 0, 0, 0], dtype=np.complex64)
+    gate_1 = cirq.StatePreparationChannel(state)
+    gate_2 = cirq.StatePreparationChannel(state)
+    assert gate_1 == gate_2, "Equal state not leading to same gate"
+    assert not gate_1 == state, "Incompatible objects shouldn't be equal"

cirq-core/cirq/protocols/json_test_data/StatePreparationChannel.json

@@ -0,0 +1,25 @@
+{
+  "cirq_type": "StatePreparationChannel",
+  "target_state": [
+    {
+      "cirq_type": "complex",
+      "real": 1.0,
+      "imag": 0.0
+    },
+    {
+      "cirq_type": "complex",
+      "real": 0.0,
+      "imag": 0.0
+    },
+    {
+      "cirq_type": "complex",
+      "real": 0.0,
+      "imag": 0.0
+    },
+    {
+      "cirq_type": "complex",
+      "real": 0.0,
+      "imag": 0.0
+    }
+  ]
+}

cirq-core/cirq/protocols/json_test_data/StatePreparationChannel.repr

@@ -0,0 +1 @@
+cirq.StatePreparationChannel(np.array([(1+0j), 0j, 0j, 0j], dtype=np.complex128))