How can I implement partial transpose on a variable in Picos (Python, trying to solve an SDP)?

Question

I try to optimise a quantity via an SDP. I optimise over all PPT measurement operators and hence have the constraints $\Pi_k^{T_B} \succeq 0$ (PPT) for my measurement operators.

The part of the code where I define the SDP and constraints is:

[...]

# Defining the SDP
p = pic.Problem()
pic.new_param("rho_0", rho_0)
pic.new_param("rho_1", rho_1)
pic.new_param("rho_2", rho_2)

# Measurement operators
P_0 = p.add_variable("P_0", (9,9), "hermitian")
P_1 = p.add_variable("P_1", (9,9), "hermitian")
P_2 = p.add_variable("P_2", (9,9), "hermitian")
P_inc = p.add_variable("P_inc", (9,9), "hermitian")

# Partial transposes
p.add_constraint(PT_B(P_0, 3, 3) >> 0)
p.add_constraint(PT_B(P_1, 3, 3) >> 0)
p.add_constraint(PT_B(P_2, 3, 3) >> 0)
p.add_constraint(PT_B(P_inc, 3, 3) >> 0)

[...]

where PT_B is a function implementing the partial transpose, i.e.

def PT_B(M, d1, d2):
    """
    Partial Transpose map of M
    Input: Matrix M, Dimension d1 of subsystem 1, Dimension d2 of subsystem 2
    Output: Partial Transpose M_TB
    """
    assert M.shape == (d1 * d2, d1 * d2)

    # Reshape into 4 tensor
    M = M.reshape(d1, d2, d1, d2)

    # Transpose 2nd system
    M = M.transpose((0, 3, 2, 1))

    # Reshape back into a density matrix
    return M.reshape(d1 * d2, d1 * d2)

However, Picos doesn't let me manipulate the variable expressions. I also tried to implement the partial transpose by setting the elements of a new variable by hand to the elements of $\Pi_k^{T_B}$. But then Picos tells me that I can't slice variable expressions.

The error message for the above SDP code snippet I get is:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-38-36083f3fb159> in <module>
     12 
     13 # Partial transposes
---> 14 p.add_constraint(PT_B(P_0, 3, 3) >> 0)
     15 p.add_constraint(PT_B(P_1, 3, 3) >> 0)
     16 p.add_constraint(PT_B(P_2, 3, 3) >> 0)

<ipython-input-34-2327f6ca6850> in PT_B(M, d1, d2)
      5     Output: Partial Transpose M_TB
      6     """
----> 7     assert M.shape == (d1 * d2, d1 * d2)
      8 
      9     # Reshape into 4 tensor

AttributeError: 'Variable' object has no attribute 'shape'

If I try to force the variable expression into an np.array:

[...]

# Partial transposes
p.add_constraint(PT_B(np.array(P_0), 3, 3) >> 0)
p.add_constraint(PT_B(np.array(P_1), 3, 3) >> 0)
p.add_constraint(PT_B(np.array(P_2), 3, 3) >> 0)
p.add_constraint(PT_B(np.array(P_inc), 3, 3) >> 0)

[...]

Then I get the error message:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-39-c6e549dceb23> in <module>
     12 
     13 # Partial transposes
---> 14 p.add_constraint(PT_B(np.array(P_0), 3, 3) >> 0)
     15 p.add_constraint(PT_B(P_1, 3, 3) >> 0)
     16 p.add_constraint(PT_B(P_2, 3, 3) >> 0)

/usr/local/lib/python3.7/site-packages/picos/expressions.py in __getitem__(self, index)
   3973         JJ = rangeT
   3974         VV = [1.] * nsz
-> 3975         newfacs = {self: spmatrix(VV, II, JJ, (nsz, sz))}
   3976         if not self.constant is None:
   3977             newcons = self.constant[rangeT]

/usr/local/lib/python3.7/site-packages/picos/tools.py in spmatrix(*args, **kwargs)
   2187     """
   2188     try:
-> 2189         return cvx.spmatrix(*args, **kwargs)
   2190     except TypeError as error:
   2191         # CVXOPT does not like NumPy's int64 scalar type for indices, so attempt

TypeError: dimension too small

Does anyone have an idea how to implement the PPT constraint in the SDP in Python?

Mathematica doesn't allow complex matrix SDPs, and Picos and CVXPY give me a hard time implementing the PPT constraint.

Thanks for any suggestions!