add check_kraus

Author: refraction-ray

tensorcircuit/densitymatrix.py

@@ -173,14 +173,39 @@ def _contract(self) -> None:
         self._nodes = [t]
 
     @staticmethod
-    def check_kraus(kraus: Sequence[Gate]) -> bool:  # TODO(@refraction-ray)
+    def check_kraus(kraus: Sequence[Gate]) -> bool:
+        """
+        Check if Kraus operators satisfy the completeness relation:
+        sum_i K_i^† K_i = I
+
+        :param kraus: Sequence of Kraus operators
+        :type kraus: Sequence[Gate]
+        :return: True if completeness relation is satisfied within tolerance
+        :rtype: bool
+        :raises ValueError: If completeness relation is not satisfied
+        """
+        # Convert all operators to tensors and get their shape
+        k_tensors = [k.tensor if isinstance(k, tn.Node) else k for k in kraus]
+        shape = k_tensors[0].shape
+        dim = shape[0]  # Assuming square matrices
+
+        # Sum K_i K_i^†
+        sum_kk = backend.zeros([dim, dim], dtype=k_tensors[0].dtype)
+        for k in k_tensors:
+            kdag = backend.conj(backend.transpose(k))
+            sum_kk = sum_kk + backend.matmul(kdag, k)
+
+        # Check if sum equals identity within tolerance
+        identity = backend.eye(dim, dtype=sum_kk.dtype)
+        if not np.allclose(sum_kk, identity, atol=1e-5):
+            raise ValueError("Kraus operators do not satisfy completeness relation")
+
         return True
 
     def apply_general_kraus(
         self, kraus: Sequence[Gate], index: Sequence[Tuple[int, ...]], **kws: Any
     ) -> None:
         # note the API difference for index arg between DM and DM2
-        self.check_kraus(kraus)
         assert len(kraus) == len(index) or len(index) == 1
         if len(index) == 1:
             index = [index[0] for _ in range(len(kraus))]
@@ -341,7 +366,6 @@ def apply_general_kraus(
             )
             for k in kraus
         ]
-        self.check_kraus(kraus)
         if not isinstance(
             index[0], int
         ):  # try best to be compatible with DMCircuit interface

tests/test_channels.py

@@ -363,3 +363,47 @@ def test_readout_mitigate(backend):
     readout_error.append([0.7, 0.9])  # readout error of qubit 2
 
     mitigate_readout(nqubit, c, readout_error)
+
+
+def test_valid_kraus_operators():
+    # Test with phase damping channel - should pass
+    kraus = phasedampingchannel(0.3)
+    assert tc.DMCircuit.check_kraus(kraus) is True
+
+    # Test with depolarizing channel - should pass
+    kraus = depolarizingchannel(0.1, 0.1, 0.1)
+    assert tc.DMCircuit.check_kraus(kraus) is True
+
+    # Test with reset channel - should pass
+    kraus = resetchannel()
+    assert tc.DMCircuit.check_kraus(kraus) is True
+
+    # Test with amplitude damping channel - should pass
+    kraus = amplitudedampingchannel(0.2, 0.3)
+    assert tc.DMCircuit.check_kraus(kraus) is True
+
+
+def test_invalid_kraus_operators():
+    # Create invalid Kraus operators that don't sum to identity
+    invalid_kraus = [
+        tc.gates.Gate(np.array([[1.0, 0], [0, 1.0]])),
+        tc.gates.Gate(np.array([[1.0, 0], [0, 1.0]])),
+    ]
+
+    with pytest.raises(ValueError):
+        tc.DMCircuit.check_kraus(invalid_kraus)
+
+
+def test_single_kraus_operator():
+    # Test with single unitary operator (should pass)
+    kraus = [tc.gates.Gate(np.array([[1.0, 0], [0, 1.0]]))]
+    assert tc.DMCircuit.check_kraus(kraus) is True
+
+
+def test_non_square_matrices():
+    # Test with non-square matrices (should raise error during matmul)
+    invalid_kraus = [
+        tc.gates.Gate(np.array([[1.0, 0]])),
+    ]
+    with pytest.raises(ValueError):
+        tc.DMCircuit.check_kraus(invalid_kraus)