diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py index 7ce1c63e..606a269c 100644 --- a/tensorcircuit/circuit.py +++ b/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 diff --git a/tensorcircuit/quditcircuit.py b/tensorcircuit/quditcircuit.py index 3b235efe..a5c72afc 100644 --- a/tensorcircuit/quditcircuit.py +++ b/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, + ) diff --git a/tests/test_quditcircuit.py b/tests/test_quditcircuit.py index e2118986..3ab116a6 100644 --- a/tests/test_quditcircuit.py +++ b/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)