ensure partial_trace raises ValueError for invalid parameters

[aman-coder03]

Description

closes #1388

adds explicit validation checks to partial_trace to ensure that invalid parameters raise clear and consistent ValueError.
specifically, this PR ensures:
input_mat must be square.
np.prod(dim) must match the matrix dimension.
sys indices must be within valid subsystem bounds.

previously, some invalid inputs could result in NumPy IndexError or unintended behavior instead of a clear ValueError.
corresponding unit tests have been added to maintain coverage for the new validation logic.

Changes

• added explicit validation for non-square matrices
• added validation for dim product mismatch
• added validation for out-of-bounds sys indices
• added unit tests covering new validation branches
• maintained full coverage for modified logic

Checklist

• Use ruff for errors related to code style and formatting.
• Verify all previous and newly added unit tests pass in pytest.
• Check the documentation build does not lead to any failures. Sphinx build can be checked locally for any failures related to your PR

[aman-coder03]

@vprusso please have look and suggest if some changes need to be made!

[purva-thakre]

As discussed in #1385, we need a better plan for invalid inputs. I am in favor of closing this PR as well but I will let @vprusso decide.

[vprusso]

There is an is_square function in matrix_props/ that could be used instead.

[aman-coder03]

thanks @vprusso

[aman-coder03]

if i am importing is_square on the top of the file there is circular import error happening, so i am planning to insert it inside the function so that circular import error is solved.

[aman-coder03]

but doing so is causing the error in ruff check

[aman-coder03]

maybe i should use this logic only instead of importing?
if input_mat.ndim != 2 or input_mat.shape[0] != input_mat.shape[1]

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.3%. Comparing base (be6586b) to head (fd64fe5).
⚠️ Report is 6 commits behind head on master.

Additional details and impacted files

@@          Coverage Diff           @@
##           master   #1391   +/-   ##
======================================
  Coverage    98.3%   98.3%           
======================================
  Files         202     202           
  Lines        5212    5223   +11     
  Branches     1196    1201    +5     
======================================
+ Hits         5124    5135   +11     
  Misses         46      46           
  Partials       42      42           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[aman-coder03]

@vprusso codecov locally is showing 100% now after updating the PR and rebasing

[vprusso]

Clean PR -- the validation improvements are exactly the kind of defensive checks this function needed. A few minor notes:

Minor

1. Negative sys indices are now rejected

Previously, sys=-1 would silently work (Python negative indexing → last subsystem). The PR now rejects it. This is the right call for this domain (negative subsystem indices are almost certainly a bug), but worth noting as a behavior change.

2. Error message for 1D input

If someone passes a 1D array, the check input_mat.ndim != 2 triggers with "input_mat must be a square matrix." Technically correct but "must be a 2D square matrix" would be slightly more helpful. Very minor.

3. Pre-existing test file issue (not your PR)

Line 14 of the existing test file has a double comma in the parametrize string:

"input_mat, , sys_arg, dim_arg, msg",

Not caused by this PR, just FYI.

Looks good to merge once CI passes.

[aman-coder03]

i have updated the error message to "must be a 2D square matrix" and fixed the double comma in the parametrize string @vprusso

[aman-coder03]

@vprusso can you please have a look at this PR now?

[vprusso]

This looks good overall — the validation logic is solid and the tests are comprehensive. A few minor items:

1. Changed behavior for non-square matrix error message

The test on line 32 changed the expected error from "Cannot infer subsystem dimensions directly" to "input_mat must be a 2D square matrix". This is a behavior change — previously a 3x4 matrix would fall through to the dim-inference logic and fail there. The new explicit check is better, but just noting it's a subtle change in error message for existing users.

2. The dim type broadening is good

Accepting tuple and np.ndarray in addition to list (line 140) is a nice improvement — users often pass numpy arrays or tuples for dimensions.

3. The prod_dim != n check

Good addition at line 149. Previously this would silently produce wrong results for mismatched dims.

4. Negative index handling

The sys < 0 check is correct — numpy would wrap negative indices silently, which could produce unexpected results. Good catch.

Approve with minor suggestion

Consider adding the dim type check (else: raise ValueError) to the docstring's Raises section so users know what input types are accepted. Otherwise this is ready to merge.

[aman-coder03]

added all ValueError conditions to the Raises section in the docstring, it should be ready to merge! @vprusso

[aman-coder03]

@vprusso i have addressed all the feedback and the PR is ready for a final review and merge at your convenience. thank you!

[vprusso]

Thanks for tackling #1388 — the validation logic itself is a clear improvement (the early square check in particular replaces a misleading "Cannot infer subsystem dimensions" error with an accurate one). A few things need to be addressed before this can merge.

Blocking

1. The docstring was duplicated into two formats. The PR inserts a new Sphinx-RST block (Examples, .. jupyter-execute::, .. math::, :code:, :raises:, :param:, .. footbibliography::) from lines 32–126 while leaving the original mkdocs-style docstring (Args: / Returns: / Examples: with ```python exec="1" ``` fences) intact below it. toqito's docs are built with mkdocs + mkdocstrings + markdown-exec, not Sphinx — none of the RST directives will render, and help(partial_trace) now prints the description twice. The docstring rewrite is also outside the scope of #1388. Please remove the added RST block entirely and keep only the validation changes in this PR. If you want to overhaul the docstring format, that should be a separate PR/discussion.

2. Orphan text at lines 127–128 (subsystems are given by the vector \dim`…) is a leftover fragment from the duplication and will appear as garbled floating text in help()`. Will go away when the RST block is removed.

3. Trailing newline was removed from partial_trace.py (\ No newline at end of file at the end of the diff). This will trip ruff's W292. Please add it back.

Non-blocking

4. Tuple inconsistency between dim and sys. dim now accepts tuple, but sys still only accepts int | list | np.ndarray. Either accept tuple for both, or neither — and update the type annotations (sys: int | list[int] | None, dim: int | list[int] | np.ndarray | None) to match whichever you choose.

5. Scope creep in the test file. test_partial_trace.py also fixes an unrelated pre-existing typo in the parametrize id string ("input_mat, , sys_arg, dim_arg, msg" → "input_mat, sys_arg, dim_arg, msg"). That's a fine fix, but it should be called out in the PR description so reviewers aren't surprised.

6. test_dim_list_branches only asserts .shape. Both parametrized cases ([2] and [2, 2]) will produce a 2×2 result regardless of whether the partial trace is computed correctly. Consider comparing against a known expected matrix so the test actually guards correctness of those branches.

7. Consider validating the CVXPY Variable path. The new square check runs after the isinstance(input_mat, Variable) recursion, so it's exercised for Variable inputs too — but there's no test for that path. A one-line test passing a non-square Variable and asserting ValueError would close the loop.

8. Nit (l. 271): any(s < 0 or s >= num_sys for s in sys) — after converting to np.asarray you can do this vectorized (np.any((arr < 0) | (arr >= num_sys))), but the current form is fine.

[vprusso]

This entire RST block (roughly lines 32–126) is Sphinx syntax — .. jupyter-execute::, .. math::, :code:, .. footbibliography::, :raises:, :param:, :return:. toqito's docs are rendered by mkdocs + mkdocstrings + markdown-exec, not Sphinx, so none of this will render, and it duplicates the original docstring that still follows below starting at the Args: section (line 130). help(partial_trace) will also print the description twice.

Please remove this RST block entirely — the docstring rewrite is out of scope for #1388 (which is just about validation). If you want to update the docstring later, that should be a separate PR.

[vprusso]

Orphan sentence left over from the duplication — this is a continuation of the opening paragraph at line 28–30 that got stranded between the new :return: line and the original Args: section. Will be cleaned up automatically once the RST block above is removed.

[vprusso]

Good extension — accepting tuple and np.ndarray is more consistent with the type annotation. Two small things:

1. The function annotation at line 14 declares dim: int | list[int] | np.ndarray | None — tuple isn't listed there. Please add it (or drop tuple from the runtime check).
2. For a numpy input, you now always call np.array(dim) even when it's already an ndarray. np.asarray(dim) avoids the unnecessary copy.

[vprusso]

dim was just extended to accept tuple, but sys still only accepts list/ndarray. A user passing partial_trace(mat, (0, 1), [2, 2]) will hit the else branch and get a ValueError — inconsistent with dim. Either accept tuple here too, or (preferably) standardize on one rule for both parameters and update the annotations to match.

[vprusso]

The file no longer ends with a newline (\ No newline at end of file in the diff). ruff's W292 will flag this — please re-add the trailing newline.

[aman-coder03]

@vprusso , addressed all feedback!!