Calculate the purity of a mixed quantum state.

    Purity, defined as tr(rho^2), has an upper bound of 1 for a pure state,
    and a lower bound of 1/D (where D is the Hilbert space dimension) for a
    competently mixed state.

    Two closely related measures are the linear entropy, 1- purity, and the
    participation ratio, 1/purity.

def purity(rho: Density) -> bk.BKTensor:
    """
    Calculate the purity of a mixed quantum state.

    Purity, defined as tr(rho^2), has an upper bound of 1 for a pure state,
    and a lower bound of 1/D (where D is the Hilbert space dimension) for a
    competently mixed state.

    Two closely related measures are the linear entropy, 1- purity, and the
    participation ratio, 1/purity.
    """
    tensor = rho.tensor
    N = rho.qubit_nb
    matrix = bk.reshape(tensor, [2**N, 2**N])
    return bk.trace(bk.matmul(matrix, matrix))