Extract stateless IRV tabulation helpers

src/vse_sim/methods/irv.py

@@ -2,6 +2,88 @@
 from ..voter_models import DeterministicModel, Voter  # noqa: F401
 
 
+def build_preference_schedule(ballots):
+    """Count identical candidate rankings."""
+    preferences = {}
+    for ballot in ballots:
+        ranking = tuple(ballot)
+        preferences[ranking] = preferences.get(ranking, 0) + 1
+    return preferences
+
+
+def eliminate_candidate(preferences, candidate_to_eliminate):
+    """Return a schedule with one candidate removed from every ranking."""
+    if not isinstance(candidate_to_eliminate, CandidateWithCount):
+        return preferences
+
+    updated_preferences = {}
+    for ranking, votes in preferences.items():
+        updated_ranking = tuple(
+            candidate
+            for candidate in ranking
+            if candidate != candidate_to_eliminate.candidate
+        )
+        if updated_ranking:
+            updated_preferences[updated_ranking] = (
+                updated_preferences.get(updated_ranking, 0) + votes
+            )
+    return updated_preferences
+
+
+def candidate_votes(preference_schedule):
+    """Return active candidates ordered from most to fewest first choices."""
+    candidates = {}
+    for ranking, votes in preference_schedule.items():
+        candidate = ranking[0]
+        if candidate in candidates:
+            candidates[candidate].votes += votes
+        else:
+            candidates[candidate] = CandidateWithCount(candidate, votes)
+
+    # VSE needs a complete ranking even for candidates with no active first
+    # choices.
+    alternates = []
+    tracked_alternates = set()
+    for ranking in preference_schedule:
+        for alternate in ranking[1:]:
+            if alternate not in candidates and alternate not in tracked_alternates:
+                alternates.append(CandidateWithCount(alternate, 0))
+                tracked_alternates.add(alternate)
+
+    active = sorted(
+        candidates.values(),
+        key=lambda candidate: (candidate.votes, candidate.candidate),
+        reverse=True,
+    )
+    return active + alternates
+
+
+def least_candidate(vote_ranking, keep=None):
+    """Return the lowest-ranked candidate not present in ``keep``."""
+    keep = () if keep is None else keep
+    for candidate in reversed(vote_ranking):
+        if candidate.candidate not in keep:
+            return candidate
+    return None
+
+
+def rank_vector_to_preference(ballot):
+    """Return candidate IDs in descending preference order from a rank vector."""
+    return sorted(
+        range(len(ballot)),
+        key=lambda candidate: ballot[candidate],
+        reverse=True,
+    )
+
+
+def finish_order_to_results(finish_order):
+    """Convert winner-first finish order to high-is-better candidate scores."""
+    results = [-1] * len(finish_order)
+    for score, candidate in enumerate(reversed(finish_order)):
+        results[candidate] = score
+    return results
+
+
 class Irv(Method):
     """Implement Instant-Runoff Voting over complete ranked ballots.
 
@@ -13,67 +95,12 @@ class Irv(Method):
 
     stratTargetFor = Method.stratTarget3
 
-    def buildPreferenceSchedule(self, ballots):
-        """Gets a dictionary of the form {ranking as tuple, vote count}."""
-
-        prefs = {}
-        for b in ballots:
-            key = tuple(b)
-            if key in prefs:
-                prefs[key] += 1
-            else:
-                prefs[key] = 1
-        return prefs
-
-    def eliminateCandidate(self, inputPrefs, toEliminate):
-        """Gets a dictionary of the form {ranking as tuple, vote count} with toEliminate removed."""
-
-        if not isinstance(toEliminate, CandidateWithCount):
-            return inputPrefs
-
-        prefs = {}
-        for ranking, votes in inputPrefs.items():
-            newranking = [
-                candidate
-                for candidate in ranking
-                if candidate != toEliminate.candidate
-            ]
-
-            if not newranking:
-                continue
-            newkey = tuple(newranking)
-            if newkey in prefs:
-                prefs[newkey] += votes
-            else:
-                prefs[newkey] = votes
-        return prefs
-
-    def candidateVotes(self, prefSchedule):
-        """Gets a list of CandidateWithCount, from highest to lowest."""
-        candidates = {}
-        for ranking, votes in prefSchedule.items():
-            candidate = ranking[0]
-            if candidate in candidates:
-                candidates[candidate].votes += votes
-            else:
-                candidates[candidate] = CandidateWithCount(candidate, votes)
-
-        # Simply for VSE which requires ranking of non-winners; in real election we don't really
-        # care
-        alternates = []
-        trackedalt = set()
-        for ranking, _votes in prefSchedule.items():
-            for alternate in ranking[1:]:
-                if (alternate not in candidates) and alternate not in trackedalt:
-                    alternates.append(CandidateWithCount(alternate, 0))
-                    trackedalt.add(alternate)
-
-        return sorted(candidates.values(), key=lambda c: (c.votes, c.candidate), reverse = True) + alternates
-
-    def getLeast(self, voteRanking, keep = {}):
-        for candidate in reversed(voteRanking):
-            if candidate.candidate not in keep:
-                return candidate
+    buildPreferenceSchedule = staticmethod(build_preference_schedule)
+    eliminateCandidate = staticmethod(eliminate_candidate)
+    candidateVotes = staticmethod(candidate_votes)
+    getLeast = staticmethod(least_candidate)
+    rankVectorToPreference = staticmethod(rank_vector_to_preference)
+    finishOrderToResults = staticmethod(finish_order_to_results)
 
     def runIrv(self, remaining, ncand):
         """IRV results."""
@@ -85,21 +112,6 @@ def runIrv(self, remaining, ncand):
             remaining = self.eliminateCandidate(remaining, toEliminate)
         return results
 
-    @staticmethod
-    def rankVectorToPreference(ballot):
-        """Return candidate IDs in descending preference order from a rank vector."""
-        return sorted(range(len(ballot)), key=lambda candidate: ballot[candidate],
-                      reverse=True)
-
-    @staticmethod
-    def finishOrderToResults(finishOrder):
-        """Convert winner-first finish order to high-is-better candidate scores."""
-        ncand = len(finishOrder)
-        results = [-1] * ncand
-        for score, candidate in enumerate(reversed(finishOrder)):
-            results[candidate] = score
-        return results
-
     def results(self, ballots, **kwargs):
         """IRV results.
 

tests/test_regressions.py

@@ -6,7 +6,7 @@
 import pytest
 
 from scripts.recalculate_irv_pages import recalculate
-from vse_sim.core import SideTally
+from vse_sim.core import CandidateWithCount, SideTally
 from vse_sim.diagnostics import TRACE, setDebug, trace
 from vse_sim.methods import (
     Borda,
@@ -18,6 +18,12 @@
     Schulze,
     Score,
 )
+from vse_sim.methods.irv import (
+    build_preference_schedule,
+    candidate_votes,
+    eliminate_candidate,
+    least_candidate,
+)
 from vse_sim.simulation import CsvBatch, seedRandomGenerators
 from vse_sim.strategies import ProbChooser, beHon, beStrat
 from vse_sim.voter_models import Electorate, Voter
@@ -130,6 +136,23 @@ def test_irv_results_keep_simulator_score_contract_for_strategy():
     assert polls[0][0] == 0
 
 
+def test_irv_tabulation_helpers_are_stateless():
+    schedule = build_preference_schedule(
+        [[0, 1, 2], [0, 1, 2], [1, 0, 2]]
+    )
+    assert schedule == {(0, 1, 2): 2, (1, 0, 2): 1}
+
+    votes = candidate_votes(schedule)
+    assert [(candidate.candidate, candidate.votes) for candidate in votes] == [
+        (0, 2),
+        (1, 1),
+        (2, 0),
+    ]
+    assert least_candidate(votes).candidate == 2
+    assert least_candidate(votes, keep={2}).candidate == 1
+    assert eliminate_candidate(schedule, CandidateWithCount(1)) == {(0, 2): 3}
+
+
 @pytest.mark.parametrize(
     "probabilities",
     [