stv.py

from fractions import Fraction
from collections import defaultdict
import math
import hypothesis
import hypothesis.strategies as st
from hypothesis.errors import NoSuchExample


class Ambiguous(Exception):
    pass


def droop_quota(voters, winners):
    return Fraction(math.floor(voters / (winners + 1)) + 1)


def hare_quota(voters, winners):
    return Fraction(voters, winners)


def initial_weight_votes(votes):
    vote_weights = defaultdict(lambda: Fraction(0))
    for v in votes:
        vote_weights[tuple(v)] += Fraction(1)
    return vote_weights


def reweight_votes(original_vote_weights, candidates):
    vote_weights = defaultdict(lambda: Fraction(0))
    for v, w in original_vote_weights.items():
        if w > 0:
            vote_weights[tuple(w for w in v if w in candidates)] += w
    return vote_weights


def stv(votes, winners, quota_type=droop_quota, wright_restart=False):
    candidates = sorted(set(votes[0]))

    quota = quota_type(len(votes), winners)

    for v in votes:
        assert sorted(v) == candidates

    candidates = set(candidates)
    
    elected = set()
    disqualified = set()

    vote_weights = initial_weight_votes(votes)
        
    while len(elected) < winners:
        print("Beginning round")
        assert not (elected & disqualified)
        if len(elected) + len(disqualified) == len(candidates):
            raise Ambiguous()
        eligible = candidates - elected - disqualified
        vote_weights = reweight_votes(vote_weights, eligible)
        if not vote_weights:
            raise Ambiguous()

        current_scores = {c: Fraction(0) for c in eligible}
        voters = defaultdict(lambda: [])

        for v, w in vote_weights.items():
            c = v[0]
            assert c in eligible
            voters[c].append(v)
            current_scores[c] += w

        print("Scores:", current_scores)

        prev_total = sum(vote_weights.values())

        n_winners = 0
        for candidate, score in current_scores.items():
            if score >= quota:
                print("Elected:", candidate)
                n_winners += 1
                elected.add(candidate)
                excess = score - quota
                reweight = excess / score
                assert 0 <= reweight < 1
                for v in voters[candidate]:
                    old_weight = vote_weights[v]
                    vote_weights[v] *= reweight
                    print("Reweight:", v, old_weight, vote_weights[v])

        assert n_winners * quota + sum(vote_weights.values()) == prev_total

        if n_winners == 0:
            losing_score = min(current_scores.values())
            print("Losing score:", losing_score)
            losers = [
                c for c, v in current_scores.items() if v == losing_score]
            assert losers
            if len(losers) > 1:
                raise Ambiguous()
            print("Dropped out:", losers[0])
            disqualified.add(losers[0])
            if wright_restart:
                print("Restart")
                elected = set()
                vote_weights = initial_weight_votes(votes)
    return elected

@st.composite
def elections(draw, n_candidates=None):
    if n_candidates is None:
        n_candidates = draw(st.integers(3, 10))
    n_winners = draw(st.integers(1, n_candidates - 1))
    votes = draw(
        st.lists(
            st.permutations(range(n_candidates)),
            average_size=20, min_size=2),)
    return votes, n_winners


def distinguish(opt1, opt2):
    def is_distinguishing(election):
        votes, winners = election
        try:
            return stv(votes, winners, **opt1) != stv(votes, winners, **opt2)
        except Ambiguous:
            return False

    best = hypothesis.find(elections(10), is_distinguishing)

    hi = len(best[0][0])
    lo = 2
    while lo + 1 < hi:
        mid = (lo + hi) // 2
        print(lo, mid, hi)
        try:
            best = hypothesis.find(elections(mid), is_distinguishing)
            hi = mid
        except NoSuchExample:
            lo = mid
    return best 


@hypothesis.given(elections(), st.sampled_from([droop_quota, hare_quota]),
    st.booleans())
def test_works_for_all(election, quota_type, wright_restart):
    try:
        stv(*(election + (quota_type, wright_restart)))
    except Ambiguous:
        pass