Update coverage for channel_fidelity

tests/test_channel_metrics/test_channel_fidelity.py

@@ -1,38 +1,36 @@
 """Tests for channel_fidelity."""
 import numpy as np
+import pytest
 
 from toqito.channel_metrics import channel_fidelity
 from toqito.channels import dephasing, depolarizing
 
-# def test_channel_fidelity_same_channel():
-#     """The fidelity of identical channels should yield 1."""
-#     choi_1 = dephasing(4)
-#     choi_2 = dephasing(4)
-#     np.testing.assert_equal(np.isclose(channel_fidelity(choi_1, choi_2), 1, atol=1e-3), True)
+dephasing_channel = dephasing(4)
+depolarizing_channel = depolarizing(4)
 
 
-# def test_channel_fidelity_different_channel():
-#     """Calculate the channel fidelity of different channels."""
-#     choi_1 = dephasing(4)
-#     choi_2 = depolarizing(4)
-#     np.testing.assert_equal(np.isclose(channel_fidelity(choi_1, choi_2), 1 / 2, atol=1e-3), True)
+def test_channel_fidelity_same_channel():
+    """The fidelity of identical channels should yield 1."""
+    assert np.isclose(channel_fidelity(dephasing_channel, dephasing_channel), 1, atol=1e-3)
+
+
+def test_channel_fidelity_different_channel():
+    """Calculate the channel fidelity of different channels."""
+    assert np.isclose(channel_fidelity(dephasing_channel, depolarizing_channel), 1 / 2, atol=1e-3)
 
 
 def test_channel_fidelity_inconsistent_dims():
     """Inconsistent dimensions between Choi matrices."""
-    with np.testing.assert_raises(ValueError):
-        choi_1 = depolarizing(4)
-        choi_2 = dephasing(2)
-        channel_fidelity(choi_1, choi_2)
+    with pytest.raises(
+        ValueError, match="The Choi matrices provided should be of equal dimension."
+    ):
+        small_dim_depolarizing_channel = depolarizing(2)
+        channel_fidelity(depolarizing_channel, small_dim_depolarizing_channel)
 
 
 def test_channel_fidelity_non_square():
     """Non-square inputs for channel fidelity."""
-    with np.testing.assert_raises(ValueError):
+    with pytest.raises(ValueError, match="The Choi matrix provided must be square."):
         choi_1 = np.array([[1, 2, 3], [4, 5, 6]])
         choi_2 = np.array([[1, 2, 3], [4, 5, 6]])
         channel_fidelity(choi_1, choi_2)
-
-
-if __name__ == "__main__":
-    np.testing.run_module_suite()

toqito/channel_metrics/channel_fidelity.py

@@ -2,7 +2,7 @@
 import cvxpy
 import numpy as np
 
-from picos import partial_trace
+from toqito.channels import partial_trace
 
 
 def channel_fidelity(choi_1: np.ndarray, choi_2: np.ndarray) -> float: