Modify apply_unitary and unitary to take in numpy unitary matrix. (#7419)

Addressing issue: #7050

TESTED=Ran following tests for files affected.
```
check/pytest-changed-files
check/pytest
check/format-incremental --apply
check/pylint-changed-files
```

Author: iamsusiep

cirq-core/cirq/protocols/apply_mixture_protocol.py

@@ -22,7 +22,6 @@
 import numpy as np
 from typing_extensions import Protocol
 
-from cirq import linalg
 from cirq._doc import doc_private
 from cirq.protocols import qid_shape_protocol
 from cirq.protocols.apply_unitary_protocol import apply_unitary, ApplyUnitaryArgs
@@ -332,32 +331,6 @@ def _apply_unitary_strat(
     return right_result
 
 
-def _apply_unitary_from_matrix_strat(
-    val: np.ndarray, args: ApplyMixtureArgs, is_density_matrix: bool
-) -> np.ndarray | None:
-    """Used to enact mixture tuples that are given as (probability, np.ndarray)
-
-    If `val` does not support `apply_unitary` returns None.
-    """
-    qid_shape = tuple(args.target_tensor.shape[i] for i in args.left_axes)
-    matrix_tensor = np.reshape(val.astype(args.target_tensor.dtype), qid_shape * 2)
-    linalg.targeted_left_multiply(
-        matrix_tensor, args.target_tensor, args.left_axes, out=args.auxiliary_buffer0
-    )
-
-    if not is_density_matrix:
-        return args.auxiliary_buffer0
-    # No need to transpose as we are acting on the tensor
-    # representation of matrix, so transpose is done for us.
-    linalg.targeted_left_multiply(
-        np.conjugate(matrix_tensor),
-        args.auxiliary_buffer0,
-        cast(tuple[int], args.right_axes),
-        out=args.target_tensor,
-    )
-    return args.target_tensor
-
-
 def _apply_mixture_from_mixture_strat(
     val: Any, args: ApplyMixtureArgs, is_density_matrix: bool
 ) -> np.ndarray | None:
@@ -373,8 +346,6 @@ def _apply_mixture_from_mixture_strat(
     for prob, op in prob_mix:
         np.copyto(dst=args.target_tensor, src=args.auxiliary_buffer1)
         right_result = _apply_unitary_strat(op, args, is_density_matrix)
-        if right_result is None:
-            right_result = _apply_unitary_from_matrix_strat(op, args, is_density_matrix)
 
         args.out_buffer += prob * right_result
 

cirq-core/cirq/protocols/apply_unitary_protocol.py

@@ -469,15 +469,20 @@ def _apply_unitary_from_matrix(matrix: np.ndarray, unitary_value: Any, args: App
 def _strat_apply_unitary_from_unitary(
     unitary_value: Any, args: ApplyUnitaryArgs
 ) -> np.ndarray | None:
-    # Check for magic method.
-    method = getattr(unitary_value, '_unitary_', None)
-    if method is None:
-        return NotImplemented
-
-    # Attempt to get the unitary matrix.
-    matrix = method()
-    if matrix is NotImplemented or matrix is None:
-        return matrix
+    if isinstance(unitary_value, np.ndarray):
+        matrix = unitary_value
+        if not linalg.is_unitary(matrix):
+            return None
+    else:
+        # Check for magic method.
+        method = getattr(unitary_value, '_unitary_', None)
+        if method is None:
+            return NotImplemented
+
+        # Attempt to get the unitary matrix.
+        matrix = method()
+        if matrix is NotImplemented or matrix is None:
+            return matrix
 
     return _apply_unitary_from_matrix(matrix, unitary_value, args)
 

cirq-core/cirq/protocols/apply_unitary_protocol_test.py

@@ -56,12 +56,13 @@ def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs) -> np.ndarray:
             args.target_tensor[one] *= -1
             return args.target_tensor
 
-    fails = [NoUnitaryEffect(), HasApplyReturnsNotImplemented()]
+    fails = [NoUnitaryEffect(), HasApplyReturnsNotImplemented(), m * 2]
     passes = [
         HasUnitary(),
         HasApplyReturnsNotImplementedButHasUnitary(),
         HasApplyOutputInBuffer(),
         HasApplyMutateInline(),
+        m,
     ]
 
     def make_input():

cirq-core/cirq/protocols/has_unitary_protocol.py

@@ -19,7 +19,7 @@
 import numpy as np
 from typing_extensions import Protocol
 
-from cirq import qis
+from cirq import linalg, qis
 from cirq._doc import doc_private
 from cirq.protocols import qid_shape_protocol
 from cirq.protocols.apply_unitary_protocol import ApplyUnitaryArgs
@@ -112,6 +112,8 @@ def has_unitary(val: Any, *, allow_decompose: bool = True) -> bool:
 
 def _strat_has_unitary_from_has_unitary(val: Any) -> bool | None:
     """Attempts to infer a value's unitary-ness via its _has_unitary_ method."""
+    if isinstance(val, np.ndarray):
+        return linalg.is_unitary(val)
     if hasattr(val, '_has_unitary_'):
         result = val._has_unitary_()
         if result is NotImplemented:

cirq-core/cirq/protocols/has_unitary_protocol_test.py

@@ -61,10 +61,13 @@ class Yes:
         def _unitary_(self):
             return np.array([[1]])
 
+    m = np.diag([1, -1])
     assert not cirq.has_unitary(No1())
     assert not cirq.has_unitary(No2())
+    assert not cirq.has_unitary(m * 2)
     assert cirq.has_unitary(Yes())
     assert cirq.has_unitary(Yes(), allow_decompose=False)
+    assert cirq.has_unitary(m)
 
 
 def test_via_apply_unitary() -> None:

cirq-core/cirq/protocols/kraus_protocol.py

@@ -148,9 +148,7 @@ def kraus(
     mixture_getter = getattr(val, '_mixture_', None)
     mixture_result = NotImplemented if mixture_getter is None else mixture_getter()
     if mixture_result is not NotImplemented and mixture_result is not None:
-        return tuple(
-            np.sqrt(p) * (u if isinstance(u, np.ndarray) else unitary(u)) for p, u in mixture_result
-        )
+        return tuple(np.sqrt(p) * unitary(u) for p, u in mixture_result)
 
     unitary_getter = getattr(val, '_unitary_', None)
     unitary_result = NotImplemented if unitary_getter is None else unitary_getter()

cirq-core/cirq/protocols/mixture_protocol.py

@@ -94,7 +94,7 @@ def mixture(
     mixture_getter = getattr(val, '_mixture_', None)
     result = NotImplemented if mixture_getter is None else mixture_getter()
     if result is not NotImplemented and result is not None:
-        return tuple((p, u if isinstance(u, np.ndarray) else unitary(u)) for p, u in result)
+        return tuple((p, unitary(u)) for p, u in result)
 
     unitary_getter = getattr(val, '_unitary_', None)
     result = NotImplemented if unitary_getter is None else unitary_getter()

cirq-core/cirq/protocols/mixture_protocol_test.py

@@ -19,8 +19,8 @@
 
 import cirq
 
-a = np.array([1])
-b = np.array([1j])
+a = np.array([[1]])
+b = np.array([[1j]])
 
 
 class NoMethod:

cirq-core/cirq/protocols/pauli_expansion_protocol_test.py

@@ -54,7 +54,7 @@ def _unitary_(self) -> np.ndarray:
 
 
 @pytest.mark.parametrize(
-    'val', (NoMethod(), ReturnsNotImplemented(), HasQuditUnitary(), 123, np.eye(2), object(), cirq)
+    'val', (NoMethod(), ReturnsNotImplemented(), HasQuditUnitary(), 123, object(), cirq)
 )
 def test_raises_no_pauli_expansion(val) -> None:
     assert cirq.pauli_expansion(val, default=None) is None

cirq-core/cirq/protocols/unitary_protocol.py

@@ -20,6 +20,7 @@
 import numpy as np
 from typing_extensions import Protocol
 
+from cirq import linalg
 from cirq._doc import doc_private
 from cirq.protocols import qid_shape_protocol
 from cirq.protocols.apply_unitary_protocol import apply_unitaries, ApplyUnitaryArgs
@@ -84,6 +85,7 @@ def unitary(
 
     The matrix is determined by any one of the following techniques:
 
+    - If the value is a numpy array, it is returned directly.
     - The value has a `_unitary_` method that returns something besides None or
         NotImplemented. The matrix is whatever the method returned.
     - The value has a `_decompose_` method that returns a list of operations,
@@ -110,7 +112,13 @@ def unitary(
     Raises:
         TypeError: `val` doesn't have a unitary effect and no default value was
             specified.
+        ValueError: `val` is a numpy array that is not unitary.
     """
+    if isinstance(val, np.ndarray):
+        if not linalg.is_unitary(val):
+            raise ValueError("The provided numpy array is not unitary.")
+        return val
+
     strats = [
         _strat_unitary_from_unitary,
         _strat_unitary_from_apply_unitary,

cirq-core/cirq/protocols/unitary_protocol_test.py

@@ -161,6 +161,15 @@ def test_unitary():
         _ = cirq.unitary(ReturnsNotImplemented())
     assert cirq.unitary(ReturnsMatrix()) is m1
 
+    # Test that numpy arrays are handled directly
+    test_matrix = np.array([[1, 0], [0, 1]])
+    assert cirq.unitary(test_matrix, NotImplemented) is test_matrix
+
+    # Test that non-unitary numpy arrays raise ValueError
+    non_unitary_matrix = np.array([[1, 1], [0, 1]])
+    with pytest.raises(ValueError, match="The provided numpy array is not unitary"):
+        _ = cirq.unitary(non_unitary_matrix)
+
     assert cirq.unitary(NoMethod(), None) is None
     assert cirq.unitary(ReturnsNotImplemented(), None) is None
     assert cirq.unitary(ReturnsMatrix(), None) is m1

cirq-google/cirq_google/transformers/analytical_decompositions/two_qubit_to_sycamore_test.py

@@ -86,12 +86,12 @@ def test_known_two_qubit_op_decomposition(op, theta_range):
         cirq.FSimGate(0.25, 0.85).on(*_QUBITS),
         cirq.XX(*_QUBITS),
         cirq.YY(*_QUBITS),
-        *[cirq.testing.random_unitary(4, random_state=1234) for _ in range(10)],
+        cirq.MatrixGate(cirq.testing.random_unitary(4)).on(*_QUBITS),
     ],
 )
 def test_unknown_two_qubit_op_decomposition(op):
     assert cg.known_2q_op_to_sycamore_operations(op) is None
-    if cirq.has_unitary(op) and cirq.num_qubits(op) == 2:
+    if not cirq.is_parameterized(op) and cirq.num_qubits(op) == 2:
         matrix_2q_circuit = cirq.Circuit(
             cg.two_qubit_matrix_to_sycamore_operations(_QUBITS[0], _QUBITS[1], cirq.unitary(op))
         )