gistfile1.py

def transform_state_vector(rho, U):
    N = len(U)
    if rho.shape[-1] == N:
        # rho is in Hilbert space
        pass
    elif rho.shape[-1] == N ** 2:
        # rho is in Liouville space
        U = np.kron(U, U)
    else:
        raise ValueError('basis transformation incompatible with '
                         'operator dimensions')
    return np.einsum('ij,...j->...i', U.T.conj(), rho)