Kraus part

tensorcircuit/circuit.py

@@ -475,7 +475,7 @@ def step_function(x: Tensor) -> Tensor:
             raise ValueError("no `get_gate_from_index` implementation is provided")
         g = get_gate_from_index(r, kraus)
         g = backend.reshape(g, [self._d for _ in range(sites * 2)])
-        self.any(*index, unitary=g, name=name)  # type: ignore
+        self.any(*index, unitary=g, name=name, dim=self._d)  # type: ignore
         return r
 
     def _general_kraus_tf(
@@ -600,9 +600,13 @@ def calculate_kraus_p(i: int) -> Tensor:
             for w, k in zip(prob, kraus_tensor)
         ]
         pick = self.unitary_kraus(
-            new_kraus, *index, prob=prob, status=status, name=name
+            new_kraus,
+            *index,
+            prob=prob,
+            status=status,
+            name=name,
         )
-        if with_prob is False:
+        if not with_prob:
             return pick
         else:
             return pick, prob
@@ -633,7 +637,11 @@ def general_kraus(
         :type status: Optional[float], optional
         """
         return self._general_kraus_2(
-            kraus, *index, status=status, with_prob=with_prob, name=name
+            kraus,
+            *index,
+            status=status,
+            with_prob=with_prob,
+            name=name,
         )
 
     apply_general_kraus = general_kraus

tensorcircuit/quditcircuit.py

@@ -144,7 +144,7 @@ def _apply_gate(self, *indices: int, name: str, **kwargs: Any) -> None:
         else:
             raise ValueError(f"Unsupported gate/arity: {name} on {len(indices)} qudits")
 
-    def any(self, *indices: int, unitary: Tensor, name: str = "any") -> None:
+    def any(self, *indices: int, unitary: Tensor, name: Optional[str] = None) -> None:
         """
         Apply an arbitrary unitary on one or two qudits.
 
@@ -155,6 +155,7 @@ def any(self, *indices: int, unitary: Tensor, name: str = "any") -> None:
         :param name: Optional label stored in the circuit history.
         :type name: str
         """
+        name = "any" if name is None else name
         self._circ.unitary(*indices, unitary=unitary, name=name, dim=self._d)  # type: ignore
 
     unitary = any
@@ -668,3 +669,65 @@ def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]:
         :rtype: List[Gate]
         """
         return self._circ.amplitude_before(l)
+
+    def general_kraus(
+        self,
+        kraus: Sequence[Gate],
+        *index: int,
+        status: Optional[float] = None,
+        with_prob: bool = False,
+        name: Optional[str] = None,
+    ) -> Tensor:
+        """
+        Monte Carlo trajectory simulation of general Kraus channel whose Kraus operators cannot be
+        amplified to unitary operators. For unitary operators composed Kraus channel, :py:meth:`unitary_kraus`
+        is much faster.
+
+        This function is jittable in theory. But only jax+GPU combination is recommended for jit
+        since the graph building time is too long for other backend options; though the running
+        time of the function is very fast for every case.
+
+        :param kraus: A list of ``tn.Node`` for Kraus operators.
+        :type kraus: Sequence[Gate]
+        :param index: The qubits index that Kraus channel is applied on.
+        :type index: int
+        :param status: Random tensor uniformly between 0 or 1, defaults to be None,
+            when the random number will be generated automatically
+        :type status: Optional[float], optional
+        """
+        return self._circ.general_kraus(
+            kraus,
+            *index,
+            status=status,
+            with_prob=with_prob,
+            name=name,
+        )
+
+    def unitary_kraus(
+        self,
+        kraus: Sequence[Gate],
+        *index: int,
+        prob: Optional[Sequence[float]] = None,
+        status: Optional[float] = None,
+        name: Optional[str] = None,
+    ) -> Tensor:
+        """
+        Apply unitary gates in ``kraus`` randomly based on corresponding ``prob``.
+        If ``prob`` is ``None``, this is reduced to kraus channel language.
+
+        :param kraus: List of ``tc.gates.Gate`` or just Tensors
+        :type kraus: Sequence[Gate]
+        :param prob: prob list with the same size as ``kraus``, defaults to None
+        :type prob: Optional[Sequence[float]], optional
+        :param status: random seed between 0 to 1, defaults to None
+        :type status: Optional[float], optional
+        :return: shape [] int dtype tensor indicates which kraus gate is actually applied
+        :rtype: Tensor
+        """
+        return self._circ.unitary_kraus(
+            kraus,
+            *index,
+            prob=prob,
+            status=status,
+            name=name,
+        )

tests/test_quditcircuit.py

@@ -518,3 +518,82 @@ def test_qudit_mutual_information_product_vs_entangled(backend):
     rho_A = qu.reduced_density_matrix(c2.state(), cut=[1], dim=d)
     SA = qu.entropy(rho_A)
     np.testing.assert_allclose(SA, np.log(d), rtol=1e-6, atol=1e-7)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_unitary_kraus_qutrit_single(backend):
+    r"""
+    Qutrit (d=3) deterministic unitary-kraus selection on a single site.
+    Case A: prob=[0,1] -> pick X (shift), |0\rangle -> |1\rangle.
+    Case B: prob=[1,0] -> pick I, state remains |0\rangle.
+    """
+    d = 3
+
+    # Identity and qutrit shift X (|k\rangle -> |k+1 mod 3)\rangle
+    I = tc.quditgates.i_matrix_func(d)
+    X = tc.quditgates.x_matrix_func(d)
+
+    # Case A: choose X branch deterministically
+    c = tc.QuditCircuit(1, dim=d)
+    idx = c.unitary_kraus([I, X], 0, prob=[0.0, 1.0])
+    assert idx == 1
+    np.testing.assert_allclose(c.amplitude("0"), 0.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c.amplitude("1"), 1.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c.amplitude("2"), 0.0 + 0j, atol=1e-6)
+
+    # Case B: choose I branch deterministically
+    c2 = tc.QuditCircuit(1, dim=d)
+    idx2 = c2.unitary_kraus([I, X], 0, prob=[1.0, 0.0])
+    assert idx2 == 0
+    np.testing.assert_allclose(c2.amplitude("0"), 1.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c2.amplitude("1"), 0.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c2.amplitude("2"), 0.0 + 0j, atol=1e-6)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_general_kraus_qutrit_single(backend):
+    r"""
+    Qutrit (d=3) tests for general_kraus on a single site (Part B only).
+
+    True general Kraus with normalization and `with_prob=True`:
+      K0 = sqrt(p) * I, K1 = sqrt(1-p) * X
+      (K0^\dagger K0 + K1^\dagger K1 = I)
+      `status` controls which branch is sampled.
+    """
+    d = 3
+
+    # Identity and qutrit shift X (|k> -> |k+1 mod 3)
+    I = tc.quditgates.i_matrix_func(d)
+    X = tc.quditgates.x_matrix_func(d)
+
+    p = 0.7
+    K0 = np.sqrt(p) * I
+    K1 = np.sqrt(1.0 - p) * X
+
+    # ---- completeness check in numpy space (works for all backends) ----
+    np.testing.assert_allclose(
+        tc.backend.transpose(tc.backend.conj(K0)) @ K0
+        + tc.backend.transpose(tc.backend.conj(K1)) @ K1,
+        I,
+        atol=1e-6,
+    )
+
+    # ---- Case B1: status small -> pick K0 with prob ~ p; state remains |0\rangle ----
+    c3 = tc.QuditCircuit(1, dim=d)
+    idx3, prob3 = c3.general_kraus([K0, K1], 0, status=0.2, with_prob=True)
+    assert idx3 == 0
+    np.testing.assert_allclose(np.array(prob3), np.array([p, 1 - p]), atol=1e-6)
+    np.testing.assert_allclose(np.array(prob3)[idx3], p, atol=1e-6)
+    np.testing.assert_allclose(c3.amplitude("0"), 1.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c3.amplitude("1"), 0.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c3.amplitude("2"), 0.0 + 0j, atol=1e-6)
+
+    # ---- Case B2: status large -> pick K1 with prob ~ (1-p); state becomes |1\rangle ----
+    c4 = tc.QuditCircuit(1, dim=d)
+    idx4, prob4 = c4.general_kraus([K0, K1], 0, status=0.95, with_prob=True)
+    assert idx4 == 1
+    np.testing.assert_allclose(np.array(prob4), np.array([p, 1 - p]), atol=1e-6)
+    np.testing.assert_allclose(np.array(prob4)[idx4], 1.0 - p, atol=1e-6)
+    np.testing.assert_allclose(c4.amplitude("0"), 0.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c4.amplitude("1"), 1.0 + 0j, atol=1e-6)
+    np.testing.assert_allclose(c4.amplitude("2"), 0.0 + 0j, atol=1e-6)