Setting a constraint on the partial transpose of a variable matrix in CVXPY

Question

I'm calculating an SDP problem in Python using CVXPY and I want to set the constraint that not only my variable matrix is positive semidefinite (psd) but also its partial transpose over a certain axis is psd. I don't know how to set this requirement. My code would look as follows

#Set the variable matrix
P_0 = cp.Variable((d,d), symmetric=True)

Now i would like to define something like

def PT(d1, d2, rho):
    """Return rho^{T_1}, the partial trace of the density operator rho on C^{d_1} \ot C^{d_2} along the first system."""
    assert rho.shape == (d1 * d2, d1 * d2)

    # reshape into a 4-tensor (2 ket indices and 2 bra indices)
    rho = rho.reshape(d1, d2, d1, d2)

    # transpose the first subsystem
    rho = rho.transpose((0,3,2,1))

    # reshape back into a density operator
    return rho.reshape(d1 * d2, d1 * d2)

And then set the requirement that PT(3,3, P_0) >> 0 i.e. that it is psd. But this is not allowed in cvxpy. Also I could define a new matrix for my specific case like

P_0_tp = [[P_0[[0,0]], P_0[[1, 0]], P_0[[2, 0]], P_0[[0, 3]], P_0[[1, 3]], P_0[[2, 3]], 
  P_0[[0, 6]], P_0[[1, 6]], P_0[[2, 6]]], [P_0[[0, 1]], P_0[[1, 1]], P_0[[2, 1]], 
  P_0[[0, 4]], P_0[[1, 4]], P_0[[2, 4]], P_0[[0, 7]], P_0[[1, 7]], 
  P_0[[2, 7]]], [P_0[[0, 2]], P_0[[1, 2]], P_0[[2, 2]], P_0[[0, 5]], P_0[[1, 5]], 
  P_0[[2, 5]], P_0[[0, 8]], P_0[[1, 8]], P_0[[2, 8]]], [P_0[[3, 0]], P_0[[4, 0]], 
  P_0[[5, 0]], P_0[[3, 3]], P_0[[4, 3]], P_0[[5, 3]], P_0[[3, 6]], P_0[[4, 6]], 
  P_0[[5, 6]]], [P_0[[3, 1]], P_0[[4, 1]], P_0[[5, 1]], P_0[[3, 4]], P_0[[4, 4]], 
  P_0[[5, 4]], P_0[[3, 7]], P_0[[4, 7]], P_0[[5, 7]]], [P_0[[3, 2]], P_0[[4, 2]], 
  P_0[[5, 2]], P_0[[3, 5]], P_0[[4, 5]], P_0[[5, 5]], P_0[[3, 8]], P_0[[4, 8]], 
  P_0[[5, 8]]], [P_0[[6, 0]], P_0[[7, 0]], P_0[[8, 0]], P_0[[6, 3]], P_0[[7, 3]], 
  P_0[[8, 3]], P_0[[6, 6]], P_0[[7, 6]], P_0[[8, 6]]], [P_0[[6, 1]], P_0[[7, 1]], 
  P_0[[8, 1]], P_0[[6, 4]], P_0[[7, 4]], P_0[[8, 4]], P_0[[6, 7]], P_0[[7, 7]], 
  P_0[[8, 7]]], [P_0[[6, 2]], P_0[[7, 2]], P_0[[8, 2]], P_0[[6, 5]], P_0[[7, 5]], 
  P_0[[8, 5]], P_0[[6, 8]], P_0[[7, 8]], P_0[[8, 8]]]]

which is a 9x9 matrix that is now partially transposed in the second 3rd dimension. But how can I set this to be a variable in cvxpy?

Thanks in advance,

Answer 1

I had the same problem as you today: I wanted to create a new matrix that consisted of combinations of variables I had defined previously. I found an answer to my question in this answer. The problem with your second approach, the array P_0_tp, is that it is not made using cvxpy operations. You can construct it using for example cvxpy.vstack and cvxpy.hstack, or other available functions as you can find here (it seems there is a reshape function too...).

In the question I pose here, you can see how I ended up implementing it. I copy it here for completeness:

import cvxpy as cp

X = cp.Variable((3,3), PSD=True)

row_1 = cp.hstack((0, 1, X[0,0]))
row_2 = cp.hstack((1, 0, X[1,2]))
row_3 = cp.hstack((X[0,0], X[1,2], 0))

W = cp.vstack((row_1, row_2, row_3))
constraint = [W >> 0]

As you can see, I did not define W to be a variable, but now it is a cvxpy object:

In [1] W
Out[1]: Expression(AFFINE, UNKNOWN, (3, 3))