def _rho1(basis_states):
    """State ρ₁ from the "3states" functional"""
    d = len(basis_states)  # dimension of logical subspace
    return sum(
        [
            (2 * (d - i) / (d * (d + 1))) * psi * psi.dag()
            for (i, psi) in enumerate(basis_states)
            # note that i is 0-based, unlike in the paper
        ]
    )