Repository: vprusso/toqito

Issue:
Question/Bug: Not able to apply channel correctly
## Expected Behavior

I'm trying to apply a channel to a state, but the channel is not preserving the trace. Could I be using the functions incorrectly or is there a bug? Thanks.

```
from toqito import states, channels, channel_ops, matrix_props, channel_props
import numpy as np

ind = 0

# Density matrix of the first Bell pair
state = np.outer(states.bell(ind), states.bell(ind))
# Depolarizing channel with p=0.5
dep_chan = channels.depolarizing(4, param_p=0.5)
# This prints True
print(channel_props.is_quantum_channel(dep_chan))
d = channel_ops.apply_channel(state, dep_chan)
# This prints False, but I expect it to be True.
print(matrix_props.is_density(d))
```

## Specifications

- Version: toqito~=1.0.5
- Platform: Mac

Discussion:

Patch:
diff --git a/toqito/channel_ops/apply_channel.py b/toqito/channel_ops/apply_channel.py
index 1cfe4c52..90a03eb5 100644
--- a/toqito/channel_ops/apply_channel.py
+++ b/toqito/channel_ops/apply_channel.py
@@ -1,6 +1,7 @@
 """Apply channel to an operator."""
 from __future__ import annotations
 import numpy as np
+import itertools
 
 from toqito.matrix_ops import vec
 from toqito.perms import swap
@@ -101,23 +102,29 @@ def apply_channel(mat: np.ndarray, phi_op: np.ndarray | list[list[np.ndarray]])
     if isinstance(phi_op, list):
         s_phi_op = [len(phi_op), len(phi_op[0])]
 
-        # Map is completely positive.
-        if s_phi_op[1] == 1 or (s_phi_op[0] == 1 and s_phi_op[1] > 2):
-            for i in range(s_phi_op[0]):
-                phi_op[i][1] = phi_op[i][0].conj().T
-        else:
-            for i in range(s_phi_op[0]):
-                phi_op[i][1] = phi_op[i][1].conj().T
         phi_0_list = []
         phi_1_list = []
-        for i in range(s_phi_op[0]):
-            phi_0_list.append(phi_op[i][0])
-            phi_1_list.append(phi_op[i][1])
+
+        # Map is completely positive if input is given as:
+        # 1. [K1, K2, .. Kr]
+        # 2. [[K1], [K2], .. [Kr]]
+        # 3. [[K1, K2, .. Kr]] and r > 2
+        if isinstance(phi_op[0], np.ndarray):
+            phi_0_list = phi_op
+        elif s_phi_op[1] == 1 or (s_phi_op[0] == 1 and s_phi_op[1] > 2):
+            phi_0_list = list(itertools.chain(*phi_op))
+        else:
+            # Input is given as: [[A1, B1], [A2, B2], .. [Ar, Br]]
+            phi_0_list = [k_mat[0] for k_mat in phi_op]
+            phi_1_list = [k_mat[1].conj().T for k_mat in phi_op]
+
+        if not phi_1_list:
+            phi_1_list = [k_mat.conj().T for k_mat in phi_0_list]
 
         k_1 = np.concatenate(phi_0_list, axis=1)
         k_2 = np.concatenate(phi_1_list, axis=0)
 
-        a_mat = np.kron(np.identity(len(phi_op)), mat)
+        a_mat = np.kron(np.identity(len(phi_0_list)), mat)
         return k_1 @ a_mat @ k_2
 
     # The superoperator was given as a Choi matrix:
@@ -134,6 +141,7 @@ def apply_channel(mat: np.ndarray, phi_op: np.ndarray | list[list[np.ndarray]])
                 True,
             ).T,
             (int(phi_size[0] * np.prod(mat_size)), int(phi_size[1])),
+            order="F"
         )
         return a_mat @ b_mat
     raise ValueError(
diff --git a/toqito/channel_ops/partial_channel.py b/toqito/channel_ops/partial_channel.py
index 580288f0..d3235a7f 100644
--- a/toqito/channel_ops/partial_channel.py
+++ b/toqito/channel_ops/partial_channel.py
@@ -1,6 +1,7 @@
 """Apply channel a subsystem of an operator."""
 from __future__ import annotations
 import numpy as np
+import itertools
 
 from toqito.channel_ops import apply_channel
 from toqito.states import max_entangled
@@ -93,11 +94,19 @@ def partial_channel(
     if isinstance(phi_map, list):
         # Compute the Kraus operators on the full system.
         s_phi_1, s_phi_2 = len(phi_map), len(phi_map[0])
-
-        # Map is completely positive.
-        if s_phi_2 == 1 or s_phi_1 == 1 and s_phi_2 > 2:
+        phi_list = []
+        # Map is completely positive if input is given as:
+        # 1. [K1, K2, .. Kr]
+        # 2. [[K1], [K2], .. [Kr]]
+        # 3. [[K1, K2, .. Kr]] and r > 2
+        if isinstance(phi_map[0], np.ndarray):
+            phi_list = phi_map
+        elif s_phi_2 == 1 or s_phi_1 == 1 and s_phi_2 > 2:
+            phi_list = list(itertools.chain(*phi_map))
+
+        if phi_list:
             phi = []
-            for m in phi_map:
+            for m in phi_list:
                 phi.append(
                     np.kron(
                         np.kron(np.identity(prod_dim_r1), m),
@@ -134,27 +143,27 @@ def partial_channel(
         dim = np.array(
             [
                 [
-                    prod_dim_r2,
-                    prod_dim_r2,
-                    int(dim[0, sys - 1]),
-                    int(dim_phi[0] / dim[0, sys - 1]),
                     prod_dim_r1,
                     prod_dim_r1,
+                    int(dim[0, sys - 1]),
+                    int(dim_phi[0] / dim[0, sys - 1]),
+                    prod_dim_r2,
+                    prod_dim_r2,
                 ],
                 [
-                    prod_dim_c2,
-                    prod_dim_c2,
-                    int(dim[1, sys - 1]),
-                    int(dim_phi[1] / dim[1, sys - 1]),
                     prod_dim_c1,
                     prod_dim_c1,
+                    int(dim[1, sys - 1]),
+                    int(dim_phi[1] / dim[1, sys - 1]),
+                    prod_dim_c2,
+                    prod_dim_c2,
                 ],
             ]
         )
-        psi_r1 = max_entangled(prod_dim_r2, False, False)
-        psi_c1 = max_entangled(prod_dim_c2, False, False)
-        psi_r2 = max_entangled(prod_dim_r1, False, False)
-        psi_c2 = max_entangled(prod_dim_c1, False, False)
+        psi_r1 = max_entangled(prod_dim_r1, False, False)
+        psi_c1 = max_entangled(prod_dim_c1, False, False)
+        psi_r2 = max_entangled(prod_dim_r2, False, False)
+        psi_c2 = max_entangled(prod_dim_c2, False, False)
 
         phi_map = permute_systems(
             np.kron(np.kron(psi_r1 * psi_c1.conj().T, phi_map), psi_r2 * psi_c2.conj().T),
diff --git a/toqito/channels/dephasing.py b/toqito/channels/dephasing.py
index 897a62e1..ec173aaf 100644
--- a/toqito/channels/dephasing.py
+++ b/toqito/channels/dephasing.py
@@ -66,10 +66,10 @@ def dephasing(dim: int, param_p: float = 0) -> np.ndarray:
     >>>     [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
     >>> )
     >>> apply_channel(test_input_mat, dephasing(4, 0.5))
-    [[17.5  0.   0.   0. ]
-     [ 0.  20.   0.   0. ]
-     [ 0.   0.  22.5  0. ]
-     [ 0.   0.   0.  25. ]]
+    [[ 1.   1.   1.5  2. ]
+     [ 2.5  6.   3.5  4. ]
+     [ 4.5  5.  11.   6. ]
+     [ 6.5  7.   7.5 16. ]]
 
     References
     ==========
diff --git a/toqito/channels/depolarizing.py b/toqito/channels/depolarizing.py
index e518a49e..7b56703f 100644
--- a/toqito/channels/depolarizing.py
+++ b/toqito/channels/depolarizing.py
@@ -42,10 +42,10 @@ def depolarizing(dim: int, param_p: float = 0) -> np.ndarray:
 
     .. math::
         \Phi(\rho) = \frac{1}{4} \begin{pmatrix}
-                                    \frac{1}{2} & 0 & 0 & \frac{1}{2} \\
-                                    0 & 0 & 0 & 0 \\
-                                    0 & 0 & 0 & 0 \\
-                                    \frac{1}{2} & 0 & 0 & \frac{1}{2}
+                                    1 & 0 & 0 & 0 \\
+                                    0 & 1 & 0 & 0 \\
+                                    0 & 0 & 1 & 0 \\
+                                    0 & 0 & 0 & 1
                                  \end{pmatrix}.
 
     This can be observed in :code:`toqito` as follows.
@@ -57,10 +57,10 @@ def depolarizing(dim: int, param_p: float = 0) -> np.ndarray:
     >>>     [[1 / 2, 0, 0, 1 / 2], [0, 0, 0, 0], [0, 0, 0, 0], [1 / 2, 0, 0, 1 / 2]]
     >>> )
     >>> apply_channel(test_input_mat, depolarizing(4))
-    [[0.125 0.    0.    0.125]
-     [0.    0.    0.    0.   ]
-     [0.    0.    0.    0.   ]
-     [0.125 0.    0.    0.125]]
+    [[0.25 0.   0.   0.  ]
+     [0.   0.25 0.   0.  ]
+     [0.   0.   0.25 0.  ]
+     [0.   0.   0.   0.25]]
 
     >>> from toqito.channel_ops import apply_channel
     >>> from toqito.channels import depolarizing
@@ -69,10 +69,10 @@ def depolarizing(dim: int, param_p: float = 0) -> np.ndarray:
     >>>     [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]
     >>> )
     >>> apply_channel(test_input_mat, depolarizing(4, 0.5))
-    [[17.125  0.25   0.375  0.5  ]
-     [ 0.625 17.75   0.875  1.   ]
-     [ 1.125  1.25  18.375  1.5  ]
-     [ 1.625  1.75   1.875 19.   ]]
+    [[ 4.75  1.    1.5   2.  ]
+     [ 2.5   7.25  3.5   4.  ]
+     [ 4.5   5.    9.75  6.  ]
+     [ 6.5   7.    7.5  12.25]]
 
 
     References


Test patch:
diff --git a/tests/test_channel_ops/test_apply_channel.py b/tests/test_channel_ops/test_apply_channel.py
index dd403672..13720d49 100644
--- a/tests/test_channel_ops/test_apply_channel.py
+++ b/tests/test_channel_ops/test_apply_channel.py
@@ -1,8 +1,10 @@
 """Tests for apply_channel."""
 import numpy as np
+import pytest
 
 from toqito.channel_ops import apply_channel
 from toqito.perms import swap_operator
+from toqito.matrices import pauli
 
 
 def test_apply_channel_choi():
@@ -21,6 +23,21 @@ def test_apply_channel_choi():
     bool_mat = np.isclose(res, expected_res)
     np.testing.assert_equal(np.all(bool_mat), True)
 
+def test_apply_channel_choi_non_square():
+    """
+    The swap operator is the Choi matrix of the transpose map.
+
+    The following test is a (non-ideal, but illustrative) way of computing
+    the transpose of a non square matrix.
+    """
+    test_input_mat = np.array([[0, 1], [2, 3], [4, 5]])
+
+    expected_res = np.array([[0, 2, 4], [1, 3, 5]])
+
+    res = apply_channel(test_input_mat, swap_operator([2, 3]))
+
+    bool_mat = np.isclose(res, expected_res)
+    np.testing.assert_equal(np.all(bool_mat), True)
 
 def test_apply_channel_kraus():
     """
@@ -45,6 +62,25 @@ def test_apply_channel_kraus():
     bool_mat = np.isclose(res, expected_res)
     np.testing.assert_equal(np.all(bool_mat), True)
 
+@pytest.mark.parametrize("nested", [1, 2, 3])
+def test_apply_channel_cpt_kraus(nested):
+    """
+    Apply Kraus map of single qubit depolarizing channel.
+    """
+    test_input_mat = np.array([[1, 0], [0, 0]])
+
+    expected_res = np.array([[0.5, 0], [0, 0.5]])
+
+    kraus = [0.5 * pauli(ind) for ind in range(4)]
+    if nested == 2:
+        kraus = [kraus]
+    elif nested == 3:
+        kraus = [[mat] for mat in kraus]
+
+    res = apply_channel(test_input_mat, kraus)
+
+    bool_mat = np.isclose(res, expected_res)
+    np.testing.assert_equal(np.all(bool_mat), True)
 
 def test_apply_channel_invalid_input():
     """Invalid input for apply map."""
diff --git a/tests/test_channel_ops/test_partial_channel.py b/tests/test_channel_ops/test_partial_channel.py
index 399129e7..3a374127 100644
--- a/tests/test_channel_ops/test_partial_channel.py
+++ b/tests/test_channel_ops/test_partial_channel.py
@@ -1,8 +1,10 @@
 """Tests for partial_channel."""
 import numpy as np
+import pytest
 
 from toqito.channel_ops import partial_channel
 from toqito.channels import depolarizing
+from toqito.matrices import pauli
 
 
 def test_partial_channel_depolarizing_first_system():
@@ -122,6 +124,34 @@ def test_partial_channel_dim_list():
     np.testing.assert_equal(np.all(bool_mat), True)
 
 
+@pytest.mark.parametrize("nested", [1, 2, 3])
+def test_partial_channel_cpt_kraus(nested):
+    """
+    Perform the partial map using the Kraus representation of 
+    the depolarizing channel.
+    """
+    rho = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]])
+    kraus = [0.5 * pauli(ind) for ind in range(4)]
+    if nested == 2:
+        kraus = [kraus]
+    elif nested == 3:
+        kraus = [[mat] for mat in kraus]
+
+    res = partial_channel(rho, kraus)
+
+    expected_res = np.array(
+        [
+            [3.5, 0.0, 5.5, 0.0],
+            [0.0, 3.5, 0.0, 5.5],
+            [11.5, 0.0, 13.5, 0.0],
+            [0.0, 11.5, 0.0, 13.5],
+        ]
+    )
+
+    bool_mat = np.isclose(expected_res, res)
+    np.testing.assert_equal(np.all(bool_mat), True)
+
+
 def test_partial_channel_non_square_matrix():
     """Matrix must be square."""
     with np.testing.assert_raises(ValueError):
diff --git a/tests/test_channel_props/test_is_unital.py b/tests/test_channel_props/test_is_unital.py
index 98fdf429..95510241 100644
--- a/tests/test_channel_props/test_is_unital.py
+++ b/tests/test_channel_props/test_is_unital.py
@@ -31,9 +31,9 @@ def test_is_unital_swap_operator_choi_true():
     np.testing.assert_equal(is_unital(swap_operator(3)), True)
 
 
-def test_is_unital_depolarizing_choi_false():
-    """Verify Choi matrix of the depolarizing map is not unital."""
-    np.testing.assert_equal(is_unital(depolarizing(4)), False)
+def test_is_unital_depolarizing_choi_true():
+    """Verify Choi matrix of the depolarizing map is unital."""
+    np.testing.assert_equal(is_unital(depolarizing(4)), True)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_channels/test_dephasing.py b/tests/test_channels/test_dephasing.py
index 1839555f..a43b5213 100644
--- a/tests/test_channels/test_dephasing.py
+++ b/tests/test_channels/test_dephasing.py
@@ -20,10 +20,10 @@ def test_dephasing_completely_dephasing():
 def test_dephasing_partially_dephasing():
     """The partially dephasing channel for `p = 0.5`."""
     test_input_mat = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]])
+    param_p = 0.5
 
-    expected_res = np.array([[17.5, 0, 0, 0], [0, 20, 0, 0], [0, 0, 22.5, 0], [0, 0, 0, 25]])
-
-    res = apply_channel(test_input_mat, dephasing(4, 0.5))
+    res = apply_channel(test_input_mat, dephasing(4, param_p))
+    expected_res = (1 - param_p) * np.diag(np.diag(test_input_mat)) + param_p * test_input_mat
 
     bool_mat = np.isclose(expected_res, res)
     np.testing.assert_equal(np.all(bool_mat), True)
diff --git a/tests/test_channels/test_depolarizing.py b/tests/test_channels/test_depolarizing.py
index 622f8da8..9b09eefd 100644
--- a/tests/test_channels/test_depolarizing.py
+++ b/tests/test_channels/test_depolarizing.py
@@ -11,9 +11,7 @@ def test_depolarizing_complete_depolarizing():
         [[1 / 2, 0, 0, 1 / 2], [0, 0, 0, 0], [0, 0, 0, 0], [1 / 2, 0, 0, 1 / 2]]
     )
 
-    expected_res = (
-        1 / 4 * np.array([[1 / 2, 0, 0, 1 / 2], [0, 0, 0, 0], [0, 0, 0, 0], [1 / 2, 0, 0, 1 / 2]])
-    )
+    expected_res = 1 / 4 * np.identity(4)
 
     res = apply_channel(test_input_mat, depolarizing(4))
 
@@ -24,17 +22,12 @@ def test_depolarizing_complete_depolarizing():
 def test_depolarizing_partially_depolarizing():
     """The partially depolarizing channel for `p = 0.5`."""
     test_input_mat = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]])
+    param_p = 0.5
 
-    expected_res = np.array(
-        [
-            [17.125, 0.25, 0.375, 0.5],
-            [0.625, 17.75, 0.875, 1],
-            [1.125, 1.25, 18.375, 1.5],
-            [1.625, 1.75, 1.875, 19],
-        ]
-    )
-
-    res = apply_channel(test_input_mat, depolarizing(4, 0.5))
+    res = apply_channel(test_input_mat, depolarizing(4, param_p))
+    expected_res = (1 - param_p) * np.trace(test_input_mat) * np.identity(
+        4
+    ) / 4 + param_p * test_input_mat
 
     bool_mat = np.isclose(expected_res, res)
     np.testing.assert_equal(np.all(bool_mat), True)