Refactor: Partial trace and partial transpose

[vprusso]

The problem:

The present implementation of partial_trace and partial_transpose is not particularly favorable for a few different reasons:

• Counterintuitively, the first index is "1" and opposed to "0". The violates the conventions used elsewhere and is a result of these functions originally being converted over from MATLAB

• The functions themselves are quite convoluted, and hence, more likely to induce error. While the test suite for these functions is fairly thorough, it is still not ideal.

The solution:

The solution is fairly straightforward:

• Replace the partial_transpose and partial_trace functions with the corresponding equivalent functions provided from the picos module (that is, picos.partial_transpose and picos.partial_trace, respectively.

• The other places in the code that make use of either partial_trace or partial_tranpose will necessarily need to be updated along with any mention of usage in the docs/tutorials.

[ErikQQY]

Hi there! I am interested in this issue, it would be great if you can assign this issue to me!😄

[vprusso]

Thank you for your interest @ErikQQY ! You have been assigned!

[ErikQQY]

@vprusso I just refactored partial_trace and partial_transpose, but as for those input matrix is Variable matrix generated from cvxpy, partial_trace and partial_transpose in picos can't handle them😢, how do we handle these usage?

[ErikQQY]

For example:

toqito/toqito/state_opt/optimal_clone.py

Line 176 in 835fcd3

partial_trace(x_var, sys, dim) == np.identity(2**num_reps),

toqito/toqito/nonlocal_games/quantum_hedging.py

Line 131 in 835fcd3

partial_trace(x_var, self._sys, self._dim) == np.identity(2**self._num_reps)

toqito/toqito/state_opt/ppt_distinguishability.py

Line 180 in 835fcd3

constraints.append(partial_transpose(meas[i], sys_list, dim_list) >> 0)

[vprusso]

@vprusso I just refactored partial_trace and partial_transpose, but as for those input matrix is Variable matrix generated from cvxpy, partial_trace and partial_transpose in picos can't handle them😢, how do we handle these usage?

Hi @ErikQQY . Thanks for the work thus far and for the question. Two options come to mind:

1. The more involved solution would be to rewrite the cvxpy SDPs to equivalent ones in picos. However, I think that's beyond the scope of this issue and is probably overkill for what we need here.
2. A more prudent approach may be to leverage something like this:

    # If the input matrix is a CVX variable for an SDP, we convert it to a numpy array,
    # perform the partial trace, and convert it back to a CVX variable.
    if isinstance(input_mat, Variable):
        rho_np = expr_as_np_array(input_mat)
        traced_rho = partial_trace(rho_np, sys, dim)
        traced_rho = np_array_as_expr(traced_rho)
        return traced_rho

Note that the above snippet of code is used in both the partial_trace and partial_transpose functions that presently reside in toqito. The general idea is that if the type of the input variable is a cvxpy.Variable then this helper does some gymnastics to convert it into something that the partial_* function can operate on and then converts it back into a cvxpy.Variable` expression.

I believe the same general approach could also be applied to the portions of the code that you pointed to.

Does that make sense?

[ErikQQY]

Yeah, I think the prudent approach is appliable in this case, I think we have two choices:

1. Directly use the above snippet of code in where we use partial_trace and partial_transpose with cvxpy.Variable. (But I think this would be a little messy)
2. Instead of directly delete our partial_trace and partial_transpose, we can still have our own partial_trace and partial_transpose, but we made a wrapper for picos.partial_trace and picos.partial_transpose:

def partial_trace(mat, sys, dim):
    if isinstance(input_mat, Variable):
        rho_np = expr_as_np_array(input_mat)
        traced_rho = partial_trace(rho_np, sys, dim)
        traced_rho = np_array_as_expr(traced_rho)
        return traced_rho
    picos.partial_trace(mat, sys, dim)

[vprusso]

Yep, I think the approach you outlined in (2) sounds like the way to go!

[ErikQQY]

@vprusso I complete the partial_trace refactor, but as for partial_transpose, the usage of partial_transpose in toqito is quite different from picos, for example, partial_transpose in toqito is capable of handling non square matrices, while picos can't:

rho = np.kron(np.eye(2, 3), np.ones((2, 3)))
rho = np.kron(rho, np.eye(2, 3))
dim = np.array([[2, 2, 2], [3, 3, 3]])
res = partial_transpose(rho, sys=2, dim=dim)

It confuses me a lot🤔

[vprusso]

Hmm, you know, I'm trying to imagine a scenario for when the partial transpose function would ever need to be applied to a non-square matrix. My initial thought is to make the assumption that the input matrix must be square.

As a question, modulo the test_partial_transpose_non_square_matrix and test_partial_transpose_non_square_matrix_2 tests in tests/test_partial_transpose.py, does anything else in toqito require non-square matrices?

[ErikQQY]

Yeah, I think the partial_transpose should be applied to a square matrix, but as a matter of fact, there is indeed a test for non-square matrices in toqito:

toqito/tests/test_channels/test_partial_transpose.py

Lines 711 to 722 in 1da8ad2

	def test_partial_transpose_non_square():
	"""Test partial transpose on non square matrices ."""
	rho = np.kron(np.eye(2, 3), np.ones((2, 3)))
	rho = np.kron(rho, np.eye(2, 3))
	dim = np.array([[2, 2, 2], [3, 3, 3]])
	res = partial_transpose(rho, sys=2, dim=dim)
	expected = np.kron(np.eye(2, 3), np.ones((3, 2)))
	expected = np.kron(expected, np.eye(2, 3))
	np.testing.assert_equal(np.allclose(res, expected), True)

[vprusso]

Yes, I did indeed see the test for the non-square matrices, however, assuming that the partial transpose is only applicable to square matrices (modulo these tests) I think we can remove these test cases.

I did not happen to see any other places where applying a partial transpose to a non-square matrix is needed so I think we can:

• Replace the partial_transpose function with the picos.partial_transpose function
• Deal with the cvx.Variable case the same way that you did with the partial trace
• Remove these two unit tests on non-square inputs for the partial transpose
• Ensure all other test cases of the partial transpose pass and that any usages of the partial transpose are properly updated.
• The docs (for both partial transpose and partial trace) should be updated to reflect this change.

The above is probably not an exhaustive list, but I think it covers the bulk of it. Does that make sense?

[ErikQQY]

Yeah, this make sense to me!

I got another question here:

toqito/toqito/channels/realignment.py

Lines 80 to 84 in 150aa3a

	dim_x = np.array([[dim[0][1], dim[0][0]], [dim[1][0], dim[1][1]]])
	dim_y = np.array([[dim[1][0], dim[0][0]], [dim[0][1], dim[1][1]]])
	x_tmp = swap(input_mat, [1, 2], dim, True)
	y_tmp = partial_transpose(x_tmp, sys=1, dim=dim_x)

An MWE can be reproduce using:

a=np.arange(1, 17).reshape(4,4)
partial_transpose(a, 1,  [[2, 2], [2, 2]])

As far as I know, picos doesn't have this feature(passing dim as [[dim1, dim1], [dim2, dim2]]), how we handle this by replacing with picos? Do we do the same as we did for those cvxpy variables?

[vprusso]

Hmm. That is a bit of a hurdle. Okay, reassess. What if we kept the partial_trace and partial_tranpose functions as they presently are in toqito but we modified them so that the index that they take is an offset of 0 (instead of starting at 1 as it does now).

There may indeed be some extra functionality here in this that actually makes sense to keep, and the ability to perform this operation on non-square matrices might be useful.

Thoughts?

[ErikQQY]

I think that would be the right choice! PR updated!