From 1ef0145a93481a73e44de89fa67ae183128f0f65 Mon Sep 17 00:00:00 2001 From: Weiguo Ma Date: Tue, 16 Sep 2025 14:09:09 +0800 Subject: [PATCH 1/6] add kraus functions to quditcircuit, untested, un-debug --- tensorcircuit/circuit.py | 23 +++++++++-- tensorcircuit/quditcircuit.py | 75 +++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 3 deletions(-) diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py index 7ce1c63e..610e7754 100644 --- a/tensorcircuit/circuit.py +++ b/tensorcircuit/circuit.py @@ -391,6 +391,7 @@ def unitary_kraus( prob: Optional[Sequence[float]] = None, status: Optional[float] = None, name: Optional[str] = None, + return_gate: bool = False, ) -> Tensor: """ Apply unitary gates in ``kraus`` randomly based on corresponding ``prob``. @@ -422,6 +423,7 @@ def index2gate(r: Tensor, kraus: Sequence[Tensor]) -> Tensor: status=status, get_gate_from_index=index2gate, name=name, + return_gate=return_gate, ) def _unitary_kraus_template( @@ -434,6 +436,7 @@ def _unitary_kraus_template( Callable[[Tensor, Sequence[Tensor]], Tensor] ] = None, name: Optional[str] = None, + return_gate: bool = False, ) -> Tensor: # DRY sites = len(index) kraus = [k.tensor if isinstance(k, tn.Node) else k for k in kraus] @@ -475,6 +478,8 @@ 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)]) + if return_gate: + return r, g self.any(*index, unitary=g, name=name) # type: ignore return r @@ -553,6 +558,7 @@ def _general_kraus_2( status: Optional[float] = None, with_prob: bool = False, name: Optional[str] = None, + return_gate: bool = False, ) -> Tensor: # the graph building time is frustratingly slow, several minutes # though running time is in terms of ms @@ -600,9 +606,14 @@ 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, + return_gate=return_gate, ) - if with_prob is False: + if not with_prob: return pick else: return pick, prob @@ -614,6 +625,7 @@ def general_kraus( status: Optional[float] = None, with_prob: bool = False, name: Optional[str] = None, + return_gate: bool = False, ) -> Tensor: """ Monte Carlo trajectory simulation of general Kraus channel whose Kraus operators cannot be @@ -633,7 +645,12 @@ 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, + return_gate=return_gate, ) apply_general_kraus = general_kraus diff --git a/tensorcircuit/quditcircuit.py b/tensorcircuit/quditcircuit.py index 3b235efe..807563f4 100644 --- a/tensorcircuit/quditcircuit.py +++ b/tensorcircuit/quditcircuit.py @@ -668,3 +668,78 @@ 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 + """ + res = self._circ.general_kraus( + kraus=kraus, + *index, + status=status, + with_prob=with_prob, + name=name, + return_gate=True, + ) + + if with_prob: + (pick, gate), prob = res + self.any(*index, unitary=gate, name=name) + return pick, prob + else: + pick, gate = res + self.any(*index, unitary=gate, name=name) + return pick + + 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 + """ + r, g = self._circ.unitary_kraus( + kraus=kraus, + *index, + prob=prob, + status=status, + name=name, + return_gate=True, + ) + self.any(*index, unitary=g, name=name) + return r From 4ace9caccaa604831218f82d559771acea557d6a Mon Sep 17 00:00:00 2001 From: Weiguo Ma Date: Tue, 16 Sep 2025 14:19:28 +0800 Subject: [PATCH 2/6] a bug fixed, and the functions are correct --- tensorcircuit/quditcircuit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorcircuit/quditcircuit.py b/tensorcircuit/quditcircuit.py index 807563f4..d98fad58 100644 --- a/tensorcircuit/quditcircuit.py +++ b/tensorcircuit/quditcircuit.py @@ -695,7 +695,7 @@ def general_kraus( :type status: Optional[float], optional """ res = self._circ.general_kraus( - kraus=kraus, + kraus, *index, status=status, with_prob=with_prob, @@ -734,7 +734,7 @@ def unitary_kraus( :rtype: Tensor """ r, g = self._circ.unitary_kraus( - kraus=kraus, + kraus, *index, prob=prob, status=status, From 230900b5e80906ffae9015a2660e1e781054aeac Mon Sep 17 00:00:00 2001 From: Weiguo Ma Date: Tue, 16 Sep 2025 14:19:55 +0800 Subject: [PATCH 3/6] add test for qudit kraus --- tests/test_quditcircuit.py | 40 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/test_quditcircuit.py b/tests/test_quditcircuit.py index e2118986..39ffb82b 100644 --- a/tests/test_quditcircuit.py +++ b/tests/test_quditcircuit.py @@ -518,3 +518,43 @@ 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> -> |1>. + Case B: prob=[1,0] -> pick I, state remains |0>. + """ + 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) + + # Case A: choose X branch deterministically + c = tc.QuditCircuit(1, dim=d) + idx = c.unitary_kraus([I, X], 0, prob=[0.0, 1.0]) + # idx_val = int(tc.backend.numpy(idx)) if hasattr(tc.backend, "numpy") else int(idx) + assert idx == 1 + + a0 = c.amplitude("0") + a1 = c.amplitude("1") + a2 = c.amplitude("2") + np.testing.assert_allclose(tc.backend.numpy(a1), 1.0 + 0j, atol=1e-6) + np.testing.assert_allclose(tc.backend.numpy(a0), 0.0 + 0j, atol=1e-6) + np.testing.assert_allclose(tc.backend.numpy(a2), 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]) + # idx2_val = int(tc.backend.numpy(idx2)) if hasattr(tc.backend, "numpy") else int(idx2) + assert idx2 == 0 + + b0 = c2.amplitude("0") + b1 = c2.amplitude("1") + b2 = c2.amplitude("2") + np.testing.assert_allclose(b0, 1.0 + 0j, atol=1e-6) + np.testing.assert_allclose(b1, 0.0 + 0j, atol=1e-6) + np.testing.assert_allclose(b2, 0.0 + 0j, atol=1e-6) \ No newline at end of file From f8cd7784c58ea522240c2ff9b8632d27fd2a0d10 Mon Sep 17 00:00:00 2001 From: Weiguo Ma Date: Tue, 16 Sep 2025 14:25:56 +0800 Subject: [PATCH 4/6] formatted codes. --- tensorcircuit/quditcircuit.py | 3 ++- tests/test_quditcircuit.py | 16 +++++++--------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/tensorcircuit/quditcircuit.py b/tensorcircuit/quditcircuit.py index d98fad58..014da0a4 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 diff --git a/tests/test_quditcircuit.py b/tests/test_quditcircuit.py index 39ffb82b..86f738d1 100644 --- a/tests/test_quditcircuit.py +++ b/tests/test_quditcircuit.py @@ -524,32 +524,30 @@ def test_qudit_mutual_information_product_vs_entangled(backend): 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> -> |1>. - Case B: prob=[1,0] -> pick I, state remains |0>. + 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> -> |k+1 mod 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]) - # idx_val = int(tc.backend.numpy(idx)) if hasattr(tc.backend, "numpy") else int(idx) assert idx == 1 a0 = c.amplitude("0") a1 = c.amplitude("1") a2 = c.amplitude("2") - np.testing.assert_allclose(tc.backend.numpy(a1), 1.0 + 0j, atol=1e-6) - np.testing.assert_allclose(tc.backend.numpy(a0), 0.0 + 0j, atol=1e-6) - np.testing.assert_allclose(tc.backend.numpy(a2), 0.0 + 0j, atol=1e-6) + np.testing.assert_allclose(a1, 1.0 + 0j, atol=1e-6) + np.testing.assert_allclose(a0, 0.0 + 0j, atol=1e-6) + np.testing.assert_allclose(a2, 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]) - # idx2_val = int(tc.backend.numpy(idx2)) if hasattr(tc.backend, "numpy") else int(idx2) assert idx2 == 0 b0 = c2.amplitude("0") @@ -557,4 +555,4 @@ def test_unitary_kraus_qutrit_single(backend): b2 = c2.amplitude("2") np.testing.assert_allclose(b0, 1.0 + 0j, atol=1e-6) np.testing.assert_allclose(b1, 0.0 + 0j, atol=1e-6) - np.testing.assert_allclose(b2, 0.0 + 0j, atol=1e-6) \ No newline at end of file + np.testing.assert_allclose(b2, 0.0 + 0j, atol=1e-6) From e56f7e897f0b07cf04e5ecc3dccd0da838391b59 Mon Sep 17 00:00:00 2001 From: Weiguo Ma Date: Tue, 16 Sep 2025 14:48:16 +0800 Subject: [PATCH 5/6] add a test to general_kraus --- tests/test_quditcircuit.py | 67 ++++++++++++++++++++++++++++++-------- 1 file changed, 54 insertions(+), 13 deletions(-) diff --git a/tests/test_quditcircuit.py b/tests/test_quditcircuit.py index 86f738d1..3ab116a6 100644 --- a/tests/test_quditcircuit.py +++ b/tests/test_quditcircuit.py @@ -537,22 +537,63 @@ def test_unitary_kraus_qutrit_single(backend): c = tc.QuditCircuit(1, dim=d) idx = c.unitary_kraus([I, X], 0, prob=[0.0, 1.0]) assert idx == 1 - - a0 = c.amplitude("0") - a1 = c.amplitude("1") - a2 = c.amplitude("2") - np.testing.assert_allclose(a1, 1.0 + 0j, atol=1e-6) - np.testing.assert_allclose(a0, 0.0 + 0j, atol=1e-6) - np.testing.assert_allclose(a2, 0.0 + 0j, atol=1e-6) + 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, + ) - b0 = c2.amplitude("0") - b1 = c2.amplitude("1") - b2 = c2.amplitude("2") - np.testing.assert_allclose(b0, 1.0 + 0j, atol=1e-6) - np.testing.assert_allclose(b1, 0.0 + 0j, atol=1e-6) - np.testing.assert_allclose(b2, 0.0 + 0j, 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) From c3cec245f37b6f4c5f2a1b5fb48874dd17f691cb Mon Sep 17 00:00:00 2001 From: Weiguo Ma Date: Tue, 16 Sep 2025 15:17:58 +0800 Subject: [PATCH 6/6] make everything easier. --- tensorcircuit/circuit.py | 11 +---------- tensorcircuit/quditcircuit.py | 17 ++--------------- 2 files changed, 3 insertions(+), 25 deletions(-) diff --git a/tensorcircuit/circuit.py b/tensorcircuit/circuit.py index 610e7754..606a269c 100644 --- a/tensorcircuit/circuit.py +++ b/tensorcircuit/circuit.py @@ -391,7 +391,6 @@ def unitary_kraus( prob: Optional[Sequence[float]] = None, status: Optional[float] = None, name: Optional[str] = None, - return_gate: bool = False, ) -> Tensor: """ Apply unitary gates in ``kraus`` randomly based on corresponding ``prob``. @@ -423,7 +422,6 @@ def index2gate(r: Tensor, kraus: Sequence[Tensor]) -> Tensor: status=status, get_gate_from_index=index2gate, name=name, - return_gate=return_gate, ) def _unitary_kraus_template( @@ -436,7 +434,6 @@ def _unitary_kraus_template( Callable[[Tensor, Sequence[Tensor]], Tensor] ] = None, name: Optional[str] = None, - return_gate: bool = False, ) -> Tensor: # DRY sites = len(index) kraus = [k.tensor if isinstance(k, tn.Node) else k for k in kraus] @@ -478,9 +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)]) - if return_gate: - return r, g - 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( @@ -558,7 +553,6 @@ def _general_kraus_2( status: Optional[float] = None, with_prob: bool = False, name: Optional[str] = None, - return_gate: bool = False, ) -> Tensor: # the graph building time is frustratingly slow, several minutes # though running time is in terms of ms @@ -611,7 +605,6 @@ def calculate_kraus_p(i: int) -> Tensor: prob=prob, status=status, name=name, - return_gate=return_gate, ) if not with_prob: return pick @@ -625,7 +618,6 @@ def general_kraus( status: Optional[float] = None, with_prob: bool = False, name: Optional[str] = None, - return_gate: bool = False, ) -> Tensor: """ Monte Carlo trajectory simulation of general Kraus channel whose Kraus operators cannot be @@ -650,7 +642,6 @@ def general_kraus( status=status, with_prob=with_prob, name=name, - return_gate=return_gate, ) apply_general_kraus = general_kraus diff --git a/tensorcircuit/quditcircuit.py b/tensorcircuit/quditcircuit.py index 014da0a4..a5c72afc 100644 --- a/tensorcircuit/quditcircuit.py +++ b/tensorcircuit/quditcircuit.py @@ -695,24 +695,14 @@ def general_kraus( when the random number will be generated automatically :type status: Optional[float], optional """ - res = self._circ.general_kraus( + return self._circ.general_kraus( kraus, *index, status=status, with_prob=with_prob, name=name, - return_gate=True, ) - if with_prob: - (pick, gate), prob = res - self.any(*index, unitary=gate, name=name) - return pick, prob - else: - pick, gate = res - self.any(*index, unitary=gate, name=name) - return pick - def unitary_kraus( self, kraus: Sequence[Gate], @@ -734,13 +724,10 @@ def unitary_kraus( :return: shape [] int dtype tensor indicates which kraus gate is actually applied :rtype: Tensor """ - r, g = self._circ.unitary_kraus( + return self._circ.unitary_kraus( kraus, *index, prob=prob, status=status, name=name, - return_gate=True, ) - self.any(*index, unitary=g, name=name) - return r