Generalize Partial Transpose for arbitrary pair of axis

[EmilianoG-byte]

Problem description: Currently partial_transpose only allows to transpose two axis within the same subsystem. See docs. More general transposition among 2 arbitrary axis could be helpful (See context for example).

Solution: Instead of using the current "manual" function, we could np.swapaxis() to swap any two axis.

The input expr will still be a 2D matrix, and as long as the conditions expr.shape[0] != np.prod(dims) and expr.shape[0] != expr.shape[1] are satisfied, the reshaping can be done in between. The output will again be a 2D matrix of the original shape.

Describe alternatives you've considered More general solution: work with >2D dimensional arrays in CVXPY. More involved solution.

Additional context Example: in the study of Open Quantum Systems, changing from a Superoperator representation onto a Choi Matrix representation is done through a swap of axis (1,2 in row-order vectorization). Reference. Choi matrices are PSD so SDP can be used to solve a convex optimization involving Choi Matrices instead of Superoperators.

[EmilianoG-byte]

My proposed solution would be to change partial_transpose.py as such:

from typing import Optional, Tuple

import numpy as np
from cvxpy.atoms.atom import Atom

def partial_transpose(expr, dims: Tuple[int], axis: Optional[int] = 0):
    """
    Assumes :math:`\\texttt{expr} = X_1 \\otimes ... \\otimes X_n` is a 2D Kronecker
    product composed of :math:`n = \\texttt{len(dims)}` implicit subsystems.
    Letting :math:`k = \\texttt{axis}`, the returned expression is a
    *partial transpose* of :math:`\\texttt{expr}`, with the transpose applied to its
    :math:`k^{\\text{th}}` implicit subsystem:

    .. math::
        X_1 \\otimes ... \\otimes X_k^T \\otimes ... \\otimes X_n.

    Parameters
    ----------
    expr : :class:`~cvxpy.expressions.expression.Expression`
        The 2D expression to take the partial transpose of.
    dims : tuple of ints.
        A tuple of integers encoding the dimensions of each subsystem.
    axis : int
        The index of the subsystem to be transposed
        from the tensor product that defines expr.
    """
    expr = Atom.cast_to_const(expr)
    if expr.ndim != 2 or expr.shape[0] != expr.shape[1]:
        raise ValueError("Only supports square matrices.")
    if axis < 0 or axis >= len(dims):
        raise ValueError(
            f"Invalid axis argument, should be between 0 and {len(dims)}, got {axis}."
        )
    if expr.shape[0] != np.prod(dims):
        raise ValueError("Dimension of system doesn't correspond to dimension of subsystems.")

    return np.reshape(expr, dims*2).swapaxes(axis, axis + len(dims)).reshape([np.prod(dims)]*2)  

EDIT: I see now that using the numpy functions np.reshape and swap axes is probably not an option. I apologize for the confusion! I keep the code here in case it is still useful for checking the results.

[rileyjmurray]

Thanks for opening this issue, @EmilianoG-byte! A question for you: is this proposed API compatible with our current API? Since CVXPY is pretty widely used, we're reluctant to make breaking changes. That's not to say we won't support this; it just determines if we need a different function name than partial_transpose (perhaps implicit_tensor_transpose or swap_implicit_tensor_axis).

[EmilianoG-byte]

Hello, @rileyjmurray! You are totally right, I was sort of mixing two ideas I had in mind related to this issue.

1. I agree that a function swapping two axis from different subsystems should be a separate function. Since this is basically what swap_axis from numpy does (after some reshape of a 2D array), this could be an option for the name. if this new function accepts axis1 and axis2 as parameters, then partial_tranpose could be implemented as a particular case when axis2 = axis1 + len(dims). This is coherent with what is known in the literature as "partial transpose".

2. The code I sent in the previous comment was mostly a suggestion on re-implementing the current partial_transpose without changing its functionality. When trying to make a PR out of it and running the test suite, I realized that it is not as straight forward as I thought to just use the numpy functions within CVXPY. I just thought of finding ways to improve it, since in my current use case partial_transpose makes the compilation step take ages.