Repository: quantumlib/Cirq

Issue:
Add efficient 2Q state preparation methods for Sqrt Iswaps and CZs as Optimizers
**Is your feature request related to a use case or problem? Please describe.**
Now that we have a `StatePreparationChannel` in Cirq (#4482), we should add efficient two qubit state preparation methods as optimizers for sqrt iswap and CZs (these are already being used internally by the physics team). 

These methods would also be useful as a reference for folks interested in adding more gateset specific efficient state preparation algorithms to Cirq (eg: https://github.com/quantumlib/Cirq/pull/4482#issuecomment-920800836)


**What is the urgency from your perspective for this issue? Is it blocking important work?**
P2 - we should do it in the next couple of quarters 

<!-- [optional] additional comment / context -->


Discussion:

Resolving patch:
diff --git a/cirq-core/cirq/__init__.py b/cirq-core/cirq/__init__.py
--- a/cirq-core/cirq/__init__.py
+++ b/cirq-core/cirq/__init__.py
@@ -344,6 +344,8 @@
     MergeInteractions,
     MergeInteractionsToSqrtIswap,
     MergeSingleQubitGates,
+    prepare_two_qubit_state_using_cz,
+    prepare_two_qubit_state_using_sqrt_iswap,
     single_qubit_matrix_to_gates,
     single_qubit_matrix_to_pauli_rotations,
     single_qubit_matrix_to_phased_x_z,
diff --git a/cirq-core/cirq/optimizers/__init__.py b/cirq-core/cirq/optimizers/__init__.py
--- a/cirq-core/cirq/optimizers/__init__.py
+++ b/cirq-core/cirq/optimizers/__init__.py
@@ -72,6 +72,11 @@
     MergeSingleQubitGates,
 )
 
+from cirq.optimizers.two_qubit_state_preparation import (
+    prepare_two_qubit_state_using_cz,
+    prepare_two_qubit_state_using_sqrt_iswap,
+)
+
 from cirq.optimizers.decompositions import (
     is_negligible_turn,
     single_qubit_matrix_to_gates,
@@ -106,6 +111,8 @@
     two_qubit_matrix_to_operations,
     two_qubit_matrix_to_diagonal_and_operations,
 )
+
+
 from cirq.optimizers.two_qubit_to_sqrt_iswap import (
     two_qubit_matrix_to_sqrt_iswap_operations,
 )
diff --git a/cirq-core/cirq/optimizers/two_qubit_state_preparation.py b/cirq-core/cirq/optimizers/two_qubit_state_preparation.py
new file mode 100644
--- /dev/null
+++ b/cirq-core/cirq/optimizers/two_qubit_state_preparation.py
@@ -0,0 +1,102 @@
+# 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.
+
+"""Utility methods for efficiently preparing two qubit states."""
+
+from typing import List, TYPE_CHECKING
+import numpy as np
+
+from cirq import ops, qis, circuits
+from cirq.optimizers import decompositions
+
+if TYPE_CHECKING:
+    import cirq
+
+
+def _1q_matrices_to_ops(g0, g1, q0, q1, include_identity=False):
+    ret = []
+    for g, q in zip(map(decompositions.single_qubit_matrix_to_phxz, [g0, g1]), [q0, q1]):
+        if g is not None:
+            ret.append(g.on(q))
+        elif include_identity:
+            ret.append(ops.I.on(q))
+    return ret
+
+
+def prepare_two_qubit_state_using_sqrt_iswap(
+    q0: 'cirq.Qid',
+    q1: 'cirq.Qid',
+    state: 'cirq.STATE_VECTOR_LIKE',
+    *,
+    use_sqrt_iswap_inv: bool = True,
+) -> List['cirq.Operation']:
+    """Prepares the given 2q state from |00> using at-most 1 √iSWAP gate + single qubit rotations.
+
+    Entangled states are prepared using exactly 1 √iSWAP gate while product states are prepared
+    using only single qubit rotations (0 √iSWAP gates)
+
+    Args:
+        q0: The first qubit being operated on.
+        q1: The other qubit being operated on.
+        state: 4x1 matrix representing two qubit state vector, ordered as 00, 01, 10, 11.
+        use_sqrt_iswap_inv: If True, uses `cirq.SQRT_ISWAP_INV` instead of `cirq.SQRT_ISWAP`.
+
+    Returns:
+        List of operations (at-most 1 √iSWAP + single qubit rotations) preparing `state` from |00>.
+    """
+    state = qis.to_valid_state_vector(state, num_qubits=2)
+    state = state / np.linalg.norm(state)
+    u, s, vh = np.linalg.svd(state.reshape(2, 2))
+    if np.isclose(s[0], 1):
+        # Product state can be prepare with just single qubit unitaries.
+        return _1q_matrices_to_ops(u, vh.T, q0, q1, True)
+    alpha = np.arccos(np.sqrt(np.clip(1 - s[0] * 2 * s[1], 0, 1)))
+    sqrt_iswap_gate = ops.SQRT_ISWAP_INV if use_sqrt_iswap_inv else ops.SQRT_ISWAP
+    op_list = [ops.ry(2 * alpha).on(q0), sqrt_iswap_gate.on(q0, q1)]
+    intermediate_state = circuits.Circuit(op_list).final_state_vector()
+    u_iSWAP, _, vh_iSWAP = np.linalg.svd(intermediate_state.reshape(2, 2))
+    return op_list + _1q_matrices_to_ops(
+        np.dot(u, np.linalg.inv(u_iSWAP)), np.dot(vh.T, np.linalg.inv(vh_iSWAP.T)), q0, q1
+    )
+
+
+def prepare_two_qubit_state_using_cz(
+    q0: 'cirq.Qid', q1: 'cirq.Qid', state: 'cirq.STATE_VECTOR_LIKE'
+) -> List['cirq.Operation']:
+    """Prepares the given 2q state from |00> using at-most 1 CZ gate + single qubit rotations.
+
+    Entangled states are prepared using exactly 1 CZ gate while product states are prepared
+    using only single qubit rotations (0 CZ gates)
+
+    Args:
+        q0: The first qubit being operated on.
+        q1: The other qubit being operated on.
+        state: 4x1 matrix representing two qubit state vector, ordered as 00, 01, 10, 11.
+
+    Returns:
+        List of operations (at-most 1 CZ + single qubit rotations) preparing `state` from |00>.
+    """
+    state = qis.to_valid_state_vector(state, num_qubits=2)
+    state = state / np.linalg.norm(state)
+    u, s, vh = np.linalg.svd(state.reshape(2, 2))
+    if np.isclose(s[0], 1):
+        # Product state can be prepare with just single qubit unitaries.
+        return _1q_matrices_to_ops(u, vh.T, q0, q1, True)
+    alpha = np.arccos(np.clip(s[0], 0, 1))
+    op_list = [ops.ry(2 * alpha).on(q0), ops.H.on(q1), ops.CZ.on(q0, q1)]
+    intermediate_state = circuits.Circuit(op_list).final_state_vector()
+    u_CZ, _, vh_CZ = np.linalg.svd(intermediate_state.reshape(2, 2))
+    return op_list + _1q_matrices_to_ops(
+        np.dot(u, np.linalg.inv(u_CZ)), np.dot(vh.T, np.linalg.inv(vh_CZ.T)), q0, q1
+    )


Test patch:
diff --git a/cirq-core/cirq/optimizers/two_qubit_state_preparation_test.py b/cirq-core/cirq/optimizers/two_qubit_state_preparation_test.py
new file mode 100644
--- /dev/null
+++ b/cirq-core/cirq/optimizers/two_qubit_state_preparation_test.py
@@ -0,0 +1,87 @@
+# 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.
+
+"""Tests for efficient two qubit state preparation methods."""
+
+import copy
+
+import pytest
+import numpy as np
+import cirq
+
+
+def random_state(seed: float):
+    return cirq.testing.random_superposition(4, random_state=seed)
+
+
+def states_with_phases(st: np.ndarray):
+    """Returns several states similar to st with modified global phases."""
+    st = np.array(st, dtype="complex64")
+    yield st
+    phases = [np.exp(1j * np.pi / 6), -1j, 1j, -1, np.exp(-1j * np.pi / 28)]
+    random = np.random.RandomState(1)
+    for _ in range(3):
+        curr_st = copy.deepcopy(st)
+        cirq.to_valid_state_vector(curr_st, num_qubits=2)
+        for i in range(4):
+            phase = random.choice(phases)
+            curr_st[i] *= phase
+        yield curr_st
+
+
+STATES_TO_PREPARE = [
+    *states_with_phases(np.array([1, 0, 0, 0])),
+    *states_with_phases(np.array([0, 1, 0, 0])),
+    *states_with_phases(np.array([0, 0, 1, 0])),
+    *states_with_phases(np.array([0, 0, 0, 1])),
+    *states_with_phases(np.array([1 / np.sqrt(2), 0, 0, 1 / np.sqrt(2)])),
+    *states_with_phases(np.array([1 / np.sqrt(2), 0, 0, -1 / np.sqrt(2)])),
+    *states_with_phases(np.array([0, 1 / np.sqrt(2), 1 / np.sqrt(2), 0])),
+    *states_with_phases(np.array([0, 1 / np.sqrt(2), -1 / np.sqrt(2), 0])),
+    *states_with_phases(random_state(97154)),
+    *states_with_phases(random_state(45375)),
+    *states_with_phases(random_state(78061)),
+    *states_with_phases(random_state(61474)),
+    *states_with_phases(random_state(22897)),
+]
+
+
+@pytest.mark.parametrize("state", STATES_TO_PREPARE)
+def test_prepare_two_qubit_state_using_cz(state):
+    state = cirq.to_valid_state_vector(state, num_qubits=2)
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.prepare_two_qubit_state_using_cz(*q, state))
+    ops_cz = [*circuit.findall_operations(lambda op: op.gate == cirq.CZ)]
+    ops_2q = [*circuit.findall_operations(lambda op: cirq.num_qubits(op) > 1)]
+    assert ops_cz == ops_2q
+    assert len(ops_cz) <= 1
+    assert cirq.allclose_up_to_global_phase(circuit.final_state_vector(), state)
+
+
+@pytest.mark.parametrize("state", STATES_TO_PREPARE)
+@pytest.mark.parametrize("use_sqrt_iswap_inv", [True, False])
+def test_prepare_two_qubit_state_using_sqrt_iswap(state, use_sqrt_iswap_inv):
+    state = cirq.to_valid_state_vector(state, num_qubits=2)
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.prepare_two_qubit_state_using_sqrt_iswap(
+            *q, state, use_sqrt_iswap_inv=use_sqrt_iswap_inv
+        )
+    )
+    sqrt_iswap_gate = cirq.SQRT_ISWAP_INV if use_sqrt_iswap_inv else cirq.SQRT_ISWAP
+    ops_iswap = [*circuit.findall_operations(lambda op: op.gate == sqrt_iswap_gate)]
+    ops_2q = [*circuit.findall_operations(lambda op: cirq.num_qubits(op) > 1)]
+    assert ops_iswap == ops_2q
+    assert len(ops_iswap) <= 1
+    assert cirq.allclose_up_to_global_phase(circuit.final_state_vector(), state)