# Source code for qutip.partial_transpose

    __all__ = ['partial_transpose']

    import numpy as np
    import scipy.sparse as sp

    from . import (
        Qobj, state_index_number, state_number_index, state_number_enumerate,
    )
    from .core.dimensions import flatten

    def partial_transpose(rho, mask, method='dense'):
        """
        Return the partial transpose of a Qobj instance `rho`,
        where `mask` is an array/list with length that equals
        the number of components of `rho` (that is, the length of
        `rho.dims[0]`), and the values in `mask` indicates whether
        or not the corresponding subsystem is to be transposed.
        The elements in `mask` can be boolean or integers `0` or `1`,
        where `True`/`1` indicates that the corresponding subsystem
        should be tranposed.

        Parameters
        ----------

        rho : :class:`.Qobj`
            A density matrix.

        mask : *list* / *array*
            A mask that selects which subsystems should be transposed.

        method : str {"dense", "sparse"}, default: "dense"
            Choice of method. The "sparse" implementation can be faster for
            large and sparse systems (hundreds of quantum states).

        Returns
        -------

        rho_pr: :class:`.Qobj`
            A density matrix with the selected subsystems transposed.

        """
        rho._dims._require_pure_dims("partial transpose")

        mask = [int(i) for i in mask]

        if method == 'sparse':
            return _partial_transpose_sparse(rho, mask)
        else:
            return _partial_transpose_dense(rho, mask)

    def _partial_transpose_dense(rho, mask):
        """
        Based on Jonas' implementation using numpy.
        Very fast for dense problems.
        """
        nsys = len(mask)
        pt_dims = np.arange(2 * nsys).reshape([2, nsys]).T
        pt_idx = np.concatenate([[pt_dims[n, mask[n]] for n in range(nsys)],
                                [pt_dims[n, 1 - mask[n]] for n in range(nsys)]])

        data = (rho.full()
                .reshape(flatten(rho.dims))
                .transpose(pt_idx)
                .reshape(rho.shape))
        return Qobj(data, dims=rho.dims)

    def _partial_transpose_sparse(rho, mask):
        """
        Implement the partial transpose using the CSR sparse matrix.
        """

        data = sp.lil_matrix((rho.shape[0], rho.shape[1]), dtype=complex)
        rho_data = rho.to("CSR").data.as_scipy()

        for m in range(len(rho_data.indptr) - 1):

            n1 = rho_data.indptr[m]
            n2 = rho_data.indptr[m + 1]

            psi_A = state_index_number(rho.dims[0], m)

            for idx, n in enumerate(rho_data.indices[n1:n2]):

                psi_B = state_index_number(rho.dims[1], n)

                m_pt = state_number_index(
                    rho.dims[1], np.choose(mask, [psi_A, psi_B]))
                n_pt = state_number_index(
                    rho.dims[0], np.choose(mask, [psi_B, psi_A]))

                data[m_pt, n_pt] = rho_data.data[n1 + idx]

        return Qobj(data.tocsr(), dims=rho.dims)

    def _partial_transpose_reference(rho, mask):
        """
        This is a reference implementation that explicitly loops over
        all states and performs the transpose. It's slow but easy to
        understand and useful for testing.
        """

        A_pt = np.zeros(rho.shape, dtype=complex)

        for psi_A in state_number_enumerate(rho.dims[0]):
            m = state_number_index(rho.dims[0], psi_A)

            for psi_B in state_number_enumerate(rho.dims[1]):
                n = state_number_index(rho.dims[1], psi_B)

                m_pt = state_number_index(
                    rho.dims[1], np.choose(mask, [psi_A, psi_B]))
                n_pt = state_number_index(
                    rho.dims[0], np.choose(mask, [psi_B, psi_A]))

                A_pt[m_pt, n_pt] = rho.to("CSR").data.as_scipy()[m, n]

        return Qobj(A_pt, dims=rho.dims)