Partial Trace

function idx = mrg(i,k,b)
% mixes the bits of i with the bits of k, taking bits from i if b is set to
% 1, and from k otherwise
idx = 0;
ss = size(b);

for i=ss:-1:1
    if b(i)
        idx = idx + bitshift( bitand(i,1), i-ss);
        i = bitshift(i,-1);
    else
        idx = idx + bitshift( bitand(k,1), i-ss);
        k = bitshift(k,-1);
    end
end

end


% assumes par is a cvx matrix of size 2^b * 2^b
% b is an array containing which qubits are traced out
function res=prtrace(par, b)

% number of bits in the remaining states
cnt = 0;
n = size(b);
for i=1:n
 cnt = cnt + b(i);
end

SZ = bitshift(1,n);
sz = bitshift(1,cnt);
rem = bitshift(1,n-cnt);


res = par;

for i = sz+1:SZ
 res(sz+1,:)=[];
end

for i = sz+1:SZ
 res(:,sz+1)=[];
end

for i=1:sz
 for j=1:sz
     res(i,j) = 0;
 end
end

for i=1:sz
 for j=1:sz
     for k=1:rem
         res(i,j) = res(i,j) + res(1+mrg(i-1,k-1,b), 1 + mrg(j-1,k-1,b));
     end
 end
end

end