tensorcircuit-nightly 1.0.2.dev20250108__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/densitymatrix.py

@@ -17,7 +17,7 @@ from .channels import kraus_to_super_gate
 from .circuit import Circuit
 from .cons import backend, contractor, dtypestr
 from .basecircuit import BaseCircuit
-from .quantum import QuOperator
+from .quantum import QuOperator, _infer_num_sites
 
 Gate = gates.Gate
 Tensor = Any
@@ -35,9 +35,11 @@ class DMCircuit(BaseCircuit):
         dminputs: Optional[Tensor] = None,
         mpo_dminputs: Optional[QuOperator] = None,
         split: Optional[Dict[str, Any]] = None,
+        dim: Optional[int] = None,
     ) -> None:
         """
         The density matrix simulator based on tensornetwork engine.
+        Do not use this class with d!=2 directly
 
         :param nqubits: Number of qubits
         :type nqubits: int
@@ -55,6 +57,7 @@ class DMCircuit(BaseCircuit):
             ``max_singular_values`` and ``max_truncation_err``.
         :type split: Optional[Dict[str, Any]]
         """
+        self._d = 2 if dim is None else dim
         if not empty:
             if (
                 (inputs is None)
@@ -73,9 +76,9 @@ class DMCircuit(BaseCircuit):
                 inputs = backend.cast(inputs, dtype=dtypestr)
                 inputs = backend.reshape(inputs, [-1])
                 N = inputs.shape[0]
-                n = int(np.log(N) / np.log(2))
+                n = _infer_num_sites(N, self._d)
                 assert n == nqubits
-                inputs = backend.reshape(inputs, [2 for _ in range(n)])
+                inputs = backend.reshape(inputs, [self._d for _ in range(n)])
                 inputs_gate = Gate(inputs)
                 self._nodes = [inputs_gate]
                 self.coloring_nodes(self._nodes)
@@ -94,7 +97,9 @@ class DMCircuit(BaseCircuit):
             elif dminputs is not None:
                 dminputs = backend.convert_to_tensor(dminputs)
                 dminputs = backend.cast(dminputs, dtype=dtypestr)
-                dminputs = backend.reshape(dminputs, [2 for _ in range(2 * nqubits)])
+                dminputs = backend.reshape(
+                    dminputs, [self._d for _ in range(2 * nqubits)]
+                )
                 dminputs_gate = Gate(dminputs)
                 nodes = [dminputs_gate]
                 self._front = [dminputs_gate.get_edge(i) for i in range(2 * nqubits)]
@@ -173,14 +178,39 @@ class DMCircuit(BaseCircuit):
         self._nodes = [t]
 
     @staticmethod
-    def check_kraus(kraus: Sequence[Gate]) -> bool:  # TODO(@refraction-ray)
+    def check_kraus(kraus: Sequence[Gate]) -> bool:
+        r"""
+        Check if Kraus operators satisfy the completeness relation:
+        :math:`\sum_i K_i^\dagger 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))]
@@ -192,7 +222,7 @@ class DMCircuit(BaseCircuit):
             dd = dmc.densitymatrix()
             circuits.append(dd)
         tensor = reduce(add, circuits)
-        tensor = backend.reshape(tensor, [2 for _ in range(2 * self._nqubits)])
+        tensor = backend.reshaped(tensor, d=self._d)
         self._nodes = [Gate(tensor)]
         dangling = [e for e in self._nodes[0]]
         self._front = dangling
@@ -202,7 +232,7 @@ class DMCircuit(BaseCircuit):
 
     @staticmethod
     def apply_general_kraus_delayed(
-        krausf: Callable[..., Sequence[Gate]]
+        krausf: Callable[..., Sequence[Gate]],
     ) -> Callable[..., None]:
         def apply(self: "DMCircuit", *index: int, **vars: float) -> None:
             for key in ["status", "name"]:
@@ -230,7 +260,7 @@ class DMCircuit(BaseCircuit):
             t = contractor(nodes, output_edge_order=d_edges)
         else:
             t = nodes[0]
-        dm = backend.reshape(t.tensor, shape=[2**self._nqubits, 2**self._nqubits])
+        dm = backend.reshapem(t.tensor)
         if check:
             self.check_density_matrix(dm)
         return dm
@@ -249,7 +279,7 @@ class DMCircuit(BaseCircuit):
         dm = self.densitymatrix()
         e, v = backend.eigh(dm)
         np.testing.assert_allclose(
-            e[:-1], backend.zeros([2**self._nqubits - 1]), atol=1e-5
+            e[:-1], backend.zeros([self._d**self._nqubits - 1]), atol=1e-5
         )
         return v[:, -1]
 
@@ -272,7 +302,7 @@ class DMCircuit(BaseCircuit):
         reuse: bool = True,
         noise_conf: Optional[Any] = None,
         status: Optional[Tensor] = None,
-        **kws: Any
+        **kws: Any,
     ) -> tn.Node.tensor:
         """
         Compute the expectation of corresponding operators.
@@ -341,7 +371,6 @@ class DMCircuit2(DMCircuit):
             )
             for k in kraus
         ]
-        self.check_kraus(kraus)
         if not isinstance(
             index[0], int
         ):  # try best to be compatible with DMCircuit interface
@@ -351,7 +380,7 @@ class DMCircuit2(DMCircuit):
         #     index = [index[0] for _ in range(len(kraus))]
         super_op = kraus_to_super_gate(kraus)
         nlegs = 4 * len(index)
-        super_op = backend.reshape(super_op, [2 for _ in range(nlegs)])
+        super_op = backend.reshape(super_op, [self._d for _ in range(nlegs)])
         super_op = Gate(super_op)
         o2i = int(nlegs / 2)
         r2l = int(nlegs / 4)
@@ -367,7 +396,7 @@ class DMCircuit2(DMCircuit):
 
     @staticmethod
     def apply_general_kraus_delayed(
-        krausf: Callable[..., Sequence[Gate]]
+        krausf: Callable[..., Sequence[Gate]],
     ) -> Callable[..., None]:
         def apply(self: "DMCircuit2", *index: int, **vars: float) -> None:
             for key in ["status", "name"]: