Add explicit validation checks and tests for partial_trace

Author: aman-coder03

toqito/matrix_ops/tests/test_partial_trace.py

@@ -646,9 +646,9 @@ def test_is_trace_prserving(input_mat, expected_result, sys_arg, dim_arg):
 @pytest.mark.parametrize(
     "sys_value",
     [
-        2,  # int out-of-bounds
-        [2],  # list out-of-bounds
-        [-1],  # negative index
+        2,      # int out-of-bounds
+        [2],    # list out-of-bounds
+        [-1],   # negative index
     ],
 )
 def test_sys_out_of_bounds(sys_value):
@@ -661,8 +661,8 @@ def test_sys_out_of_bounds(sys_value):
 @pytest.mark.parametrize(
     "dim_value",
     [
-        [2],  # len == 1 branch
-        [2, 2],  # len > 1 branch
+        [2],       # len == 1 branch
+        [2, 2],    # len > 1 branch
     ],
 )
 def test_dim_list_branches(dim_value):
@@ -672,29 +672,27 @@ def test_dim_list_branches(dim_value):
     assert result.shape == (2, 2)
 
 
-def test_partial_trace_non_square_matrix():
-    """Test that non-square matrices raise ValueError."""
-    mat = np.ones((2, 3))
-    with pytest.raises(ValueError, match="square"):
-        partial_trace(mat)
-
-
-def test_partial_trace_dim_product_mismatch():
-    """Test that mismatched dim product raises ValueError."""
-    mat = np.eye(4)
-    with pytest.raises(ValueError, match="Product of `dim`"):
-        partial_trace(mat, [0], [2, 3])
-
+@pytest.mark.parametrize(
+    "input_mat, sys_value, dim_value, error_msg",
+    [
+        # Non-square matrix
+        (np.ones((2, 3)), None, None, "square"),
 
-def test_dim_none_non_perfect_square():
-    """Test that non-perfect-square matrices with dim=None raise ValueError."""
-    mat = np.ones((6, 6))  # sqrt(6) not integer
-    with pytest.raises(ValueError, match="Cannot infer subsystem dimensions directly"):
-        partial_trace(mat, [0], None)
+        # Dim product mismatch
+        (np.eye(4), [0], [2, 3], "Product of `dim`"),
 
+        # dim=None and non-perfect-square
+        (np.ones((6, 6)), [0], None, "Cannot infer subsystem dimensions directly"),
 
-def test_dim_invalid_type_branch():
-    """Test that invalid dim type raises ValueError."""
-    mat = np.eye(4)
-    with pytest.raises(ValueError):
-        partial_trace(mat, [0], (2, 2))  # tuple instead of list
+        # Invalid dim type
+        (np.eye(4), [0], (2, 2), None),
+    ],
+)
+def test_partial_trace_invalid_inputs(input_mat, sys_value, dim_value, error_msg):
+    """Test various invalid parameter combinations."""
+    if error_msg:
+        with pytest.raises(ValueError, match=error_msg):
+            partial_trace(input_mat, sys_value, dim_value)
+    else:
+        with pytest.raises(ValueError):
+            partial_trace(input_mat, sys_value, dim_value)