fix(channel_ops): accept wrapped CP Kraus forms in kraus_to_choi

toqito/channel_ops/kraus_to_choi.py

@@ -3,6 +3,7 @@
 import numpy as np
 
 from toqito.channel_ops import partial_channel
+from toqito.channel_props._kraus_format import normalize_kraus
 from toqito.channel_props.channel_dim import channel_dim
 from toqito.states import max_entangled
 
@@ -68,12 +69,10 @@ def kraus_to_choi(kraus_ops: list[np.ndarray] | list[list[np.ndarray]], sys: int
         # Fast path for the default convention (apply the channel to the second subsystem of the unnormalized
         # maximally entangled state). The Choi matrix is then sum_i vec(A_i) vec(B_i)^dagger, where vec stacks the
         # columns of an operator. This avoids building the entangled state and the dense Kronecker products that the
-        # general partial_channel path uses. For a flat completely-positive list [K_1, ...], B_i = A_i.
-        if isinstance(kraus_ops[0], np.ndarray):
-            left = right = np.stack([np.asarray(k, dtype=complex) for k in kraus_ops])
-        else:
-            left = np.stack([np.asarray(pair[0], dtype=complex) for pair in kraus_ops])
-            right = np.stack([np.asarray(pair[1], dtype=complex) for pair in kraus_ops])
+        # general partial_channel path uses. For the completely-positive forms, B_i = A_i.
+        left_ops, right_ops, is_cp = normalize_kraus(kraus_ops)
+        left = np.stack([np.asarray(k, dtype=complex) for k in left_ops])
+        right = left if is_cp else np.stack([np.asarray(k, dtype=complex) for k in right_ops])
         # Column-major vectorization of each operator: transpose to (r, d_in, d_out) then flatten the last two axes.
         v_left = left.transpose(0, 2, 1).reshape(left.shape[0], -1)
         v_right = right.transpose(0, 2, 1).reshape(right.shape[0], -1)

toqito/channel_ops/tests/test_kraus_to_choi.py

@@ -166,3 +166,16 @@ def test_kraus_to_choi_general_sys_branch():
     herm = (choi + choi.conj().T) / 2
     np.testing.assert_allclose(choi, herm)
     assert np.linalg.eigvalsh(herm).min() > -1e-10
+
+
+@pytest.mark.parametrize(
+    "wrap", [lambda ops: [[k] for k in ops], lambda ops: [ops]], ids=["one_per_wrapper", "single_wrapper"]
+)
+def test_kraus_to_choi_wrapped_cp_formats(wrap):
+    """The wrapped CP forms [[K1], [K2], ...] and [[K1, K2, K3]] match the flat list."""
+    kraus_ops = [
+        np.sqrt(1 / 3) * np.eye(2),
+        np.sqrt(1 / 3) * np.array([[0, 1], [1, 0]]),
+        np.sqrt(1 / 3) * np.diag([1, -1]),
+    ]
+    np.testing.assert_allclose(kraus_to_choi(wrap(kraus_ops)), kraus_to_choi(kraus_ops))