Add is_cptp predicate (quantumlib#4365)

As requested in quantumlib#4194. Can be used for quantumlib#2271.

This predicate is meant to be invoked when constructing a channel to verify that the provided Kraus operators actually describe a valid quantum channel. Recommendations for cleaner `is_cptp` behavior or additional test cases are welcome.

cirq/__init__.py

@@ -130,6 +130,7 @@
     dot,
     expand_matrix_in_orthogonal_basis,
     hilbert_schmidt_inner_product,
+    is_cptp,
     is_diagonal,
     is_hermitian,
     is_normal,

cirq/linalg/__init__.py

@@ -61,6 +61,7 @@
 
 from cirq.linalg.predicates import (
     allclose_up_to_global_phase,
+    is_cptp,
     is_diagonal,
     is_hermitian,
     is_normal,

cirq/linalg/predicates.py

@@ -149,6 +149,21 @@ def is_normal(matrix: np.ndarray, *, rtol: float = 1e-5, atol: float = 1e-8) ->
     return matrix_commutes(matrix, matrix.T.conj(), rtol=rtol, atol=atol)
 
 
+def is_cptp(*, kraus_ops: Sequence[np.ndarray], rtol: float = 1e-5, atol: float = 1e-8):
+    """Determines if a channel is completely positive trace preserving (CPTP).
+
+    A channel composed of Kraus operators K[0:n] is a CPTP map if the sum of
+    the products `adjoint(K[i]) * K[i])` is equal to 1.
+
+    Args:
+        kraus_ops: The Kraus operators of the channel to check.
+        rtol: The relative tolerance on equality.
+        atol: The absolute tolerance on equality.
+    """
+    sum_ndarray = cast(np.ndarray, sum(matrix.T.conj() @ matrix for matrix in kraus_ops))
+    return np.allclose(sum_ndarray, np.eye(*sum_ndarray.shape), rtol=rtol, atol=atol)
+
+
 def matrix_commutes(
     m1: np.ndarray, m2: np.ndarray, *, rtol: float = 1e-5, atol: float = 1e-8
 ) -> bool:

cirq/linalg/predicates_test.py

@@ -292,6 +292,53 @@ def test_is_normal_tolerance():
     assert not cirq.is_normal(np.array([[0, 0.5, 0], [0, 0, 0.6], [0, 0, 0]]), atol=atol)
 
 
+def test_is_cptp():
+    rt2 = np.sqrt(0.5)
+    # Amplitude damping with gamma=0.5.
+    assert cirq.is_cptp(kraus_ops=[np.array([[1, 0], [0, rt2]]), np.array([[0, rt2], [0, 0]])])
+    # Depolarizing channel with p=0.75.
+    assert cirq.is_cptp(
+        kraus_ops=[
+            np.array([[1, 0], [0, 1]]) * 0.5,
+            np.array([[0, 1], [1, 0]]) * 0.5,
+            np.array([[0, -1j], [1j, 0]]) * 0.5,
+            np.array([[1, 0], [0, -1]]) * 0.5,
+        ]
+    )
+
+    assert not cirq.is_cptp(kraus_ops=[np.array([[1, 0], [0, 1]]), np.array([[0, 1], [0, 0]])])
+    assert not cirq.is_cptp(
+        kraus_ops=[
+            np.array([[1, 0], [0, 1]]),
+            np.array([[0, 1], [1, 0]]),
+            np.array([[0, -1j], [1j, 0]]),
+            np.array([[1, 0], [0, -1]]),
+        ]
+    )
+
+    # Makes 4 2x2 kraus ops.
+    one_qubit_u = cirq.testing.random_unitary(8)
+    one_qubit_kraus = np.reshape(one_qubit_u[:, :2], (-1, 2, 2))
+    assert cirq.is_cptp(kraus_ops=one_qubit_kraus)
+
+    # Makes 16 4x4 kraus ops.
+    two_qubit_u = cirq.testing.random_unitary(64)
+    two_qubit_kraus = np.reshape(two_qubit_u[:, :4], (-1, 4, 4))
+    assert cirq.is_cptp(kraus_ops=two_qubit_kraus)
+
+
+def test_is_cptp_tolerance():
+    rt2_ish = np.sqrt(0.5) - 0.01
+    atol = 0.25
+    # Moderately-incorrect amplitude damping with gamma=0.5.
+    assert cirq.is_cptp(
+        kraus_ops=[np.array([[1, 0], [0, rt2_ish]]), np.array([[0, rt2_ish], [0, 0]])], atol=atol
+    )
+    assert not cirq.is_cptp(
+        kraus_ops=[np.array([[1, 0], [0, rt2_ish]]), np.array([[0, rt2_ish], [0, 0]])], atol=1e-8
+    )
+
+
 def test_commutes():
     assert matrix_commutes(np.empty((0, 0)), np.empty((0, 0)))
     assert not matrix_commutes(np.empty((1, 0)), np.empty((0, 1)))