Merge pull request #378 from mggg/enhancement/faster-mentions

Enhancement/faster mentions

Author: peterrrock2

src/votekit/utils/common_utils.py

@@ -4,7 +4,7 @@
 import random
 import warnings
 from itertools import permutations
-from typing import TYPE_CHECKING, Iterable, Literal, Optional, Sequence
+from typing import TYPE_CHECKING, Any, Iterable, Literal, Optional, Sequence, cast
 
 import numpy as np
 import pandas as pd
@@ -388,18 +388,23 @@ def first_place_votes(
     )
 
 
-def mentions(
+def _mentions_from_ballots(
     profile: RankProfile,
 ) -> dict[Candidate, float]:
     """
-    Calculates total mentions for all candidates in a ``RankProfile``.
+    Calculates total mentions for all candidates by iterating the profile's ballot list.
 
     Args:
         profile (RankProfile): RankProfile of ballots.
 
     Returns:
         dict[Candidate, float]:
             Dictionary mapping candidates to mention totals (values).
+            Candidates can be strings, integers, or mix of both.
+
+    Raises:
+        TypeError: If profile is not a RankProfile.
+        TypeError: If any ballot has no ranking.
     """
     from votekit.pref_profile import RankProfile
 
@@ -409,13 +414,109 @@ def mentions(
     for ballot in profile.ballots:
         if ballot.ranking is None:
             raise TypeError("Ballots must have rankings.")
-        else:
-            for s in ballot.ranking:
-                for cand in s:
-                    mentions[cand] += float(ballot.weight)
+        # a candidate mentioned only on zero-weight ballots is absent from profile.candidates,
+        # and the ballot contributes nothing anyway
+        if ballot.weight == 0:
+            continue
+        for s in ballot.ranking:
+            for cand in s:
+                mentions[cand] += float(ballot.weight)
     return mentions
 
 
+def _mentions_from_df(profile: RankProfile) -> dict[Candidate, float]:
+    """
+    Calculates total mentions for all candidates from the profile df, without materializing
+    ballots.
+
+    Args:
+        profile (RankProfile): RankProfile of ballots.
+
+    Returns:
+        dict[Candidate, float]:
+            Dictionary mapping candidates to mention totals (values).
+            Candidates can be strings, integers, or mix of both.
+
+    Raises:
+        TypeError: If profile is not a RankProfile.
+        TypeError: If any ballot has no ranking, i.e. its df row is entirely "~" cells.
+    """
+    from votekit.pref_profile import RankProfile
+
+    if not isinstance(profile, RankProfile):
+        raise TypeError("Profile must be of type RankProfile.")
+
+    assert profile.max_ranking_length is not None
+
+    ranking_cols = [f"Ranking_{i}" for i in range(1, profile.max_ranking_length + 1)]
+
+    tilde = frozenset({"~"})
+
+    # no ranking columns with ballots present means every ballot has ranking=None
+    if len(profile.df) and not ranking_cols:
+        raise TypeError("Ballots must have rankings.")
+
+    # positional index so duplicate df index labels cannot break the weight lookup below
+    rank_sets = cast(Any, profile.df[ranking_cols].reset_index(drop=True).stack())
+
+    mask = rank_sets.map(lambda s: isinstance(s, frozenset) and bool(s) and s != tilde)
+
+    # a row of only "~" cells is the df encoding of ranking=None; match the ballot path's
+    # error. A present-but-empty frozenset ranking is allowed in both paths.
+    is_tilde = rank_sets.map(lambda s: s == tilde)
+    if len(is_tilde) and is_tilde.groupby(level=0).all().any():
+        raise TypeError("Ballots must have rankings.")
+
+    rank_sets = rank_sets[mask]
+    exploded = rank_sets.explode()
+
+    if exploded.empty:
+        return {c: 0.0 for c in profile.candidates}
+
+    weights = profile.df["Weight"].to_numpy()[exploded.index.get_level_values(0)]
+
+    totals = pd.Series(weights).groupby(exploded.to_numpy(), sort=False).sum()
+
+    # float() normalizes np.float64 so both mentions paths return the same types
+    return {c: float(totals.get(c, 0.0)) for c in profile.candidates}
+
+
+def mentions(profile: RankProfile) -> dict[Candidate, float]:
+    """
+    Calculates total mentions for all candidates in a ``RankProfile``.
+
+    Every position that lists a candidate contributes the full ballot weight, so each member of
+    a tied position receives the full weight, and a candidate ranked in multiple positions of
+    one ballot is counted once per position.
+
+    Computes from the profile's ballot list when it is already materialized and from the
+    underlying df otherwise; both give identical results.
+
+    Args:
+        profile (RankProfile): RankProfile of ballots.
+
+    Returns:
+        dict[Candidate, float]:
+            Dictionary mapping candidates to mention totals (values).
+            Candidates can be strings, integers, or mix of both.
+
+    Raises:
+        TypeError: If profile is not a RankProfile.
+        TypeError: If any ballot has no ranking.
+    """
+    from votekit.pref_profile import RankProfile
+
+    if not isinstance(profile, RankProfile):
+        raise TypeError("Profile must be of type RankProfile.")
+
+    # NOTE: If the ballots are materialized, then iterating through the ballots directly tends
+    # to be faster than using the DataFrame.
+    if "ballots" in profile.__dict__:
+        return _mentions_from_ballots(profile)
+
+    return _mentions_from_df(profile)
+
+
 def borda_scores(
     profile: RankProfile,
     borda_max: Optional[int] = None,

tests/utils/test_common_utils.py

@@ -1,6 +1,10 @@
+import random
+from fractions import Fraction
 from itertools import permutations
+from pathlib import Path
 from typing import Literal, cast
 
+import pandas as pd
 import pytest
 
 from votekit.ballot import RankBallot, ScoreBallot
@@ -25,6 +29,9 @@
     tiebroken_ranking,
     validate_score_vector,
 )
+from votekit.utils.common_utils import _mentions_from_ballots, _mentions_from_df
+
+CSV_DIR = Path(__file__).resolve().parents[1] / "data" / "csv"
 
 profile_no_ties = RankProfile(
     ballots=(
@@ -42,6 +49,14 @@
     )
 )
 
+profile_with_duplicates = RankProfile(
+    ballots=(
+        RankBallot(ranking=tuple(map(frozenset, [{"A"}, {"B"}, {"B"}])), weight=1),
+        RankBallot(ranking=tuple(map(frozenset, [{"A"}, {"B"}, {"C"}])), weight=1 / 2),
+        RankBallot(ranking=tuple(map(frozenset, [{"B"}, {"B"}, {"B"}])), weight=3),
+    )
+)
+
 profile_with_missing = RankProfile(
     ballots=(
         RankBallot(ranking=tuple(map(frozenset, [{"A", "B"}, {"D"}])), weight=1),
@@ -246,13 +261,133 @@ def test_fpv_errors():
         first_place_votes(cast(RankProfile, ScoreProfile(ballots=(ScoreBallot(scores={"A": 3}),))))
 
 
-def test_mentions():
+def test_mentions_from_ballots():
+    correct = {"A": 9 / 2, "B": 9 / 2, "C": 7 / 2}
+    test = _mentions_from_ballots(profile_no_ties)
+    assert correct == test
+    assert isinstance(test["A"], float)
+
+
+def test_mentions_from_ballots_with_ties():
+    correct = {"A": 9 / 2, "B": 9 / 2, "C": 7 / 2}
+    test = _mentions_from_ballots(profile_with_ties)
+    assert correct == test
+    assert isinstance(test["A"], float)
+
+
+def test_mentions_from_ballots_with_duplicates():
+    correct = {"A": 3 / 2, "B": 23 / 2, "C": 1 / 2}
+    test = _mentions_from_ballots(profile_with_duplicates)
+    assert correct == test
+    assert isinstance(test["A"], float)
+
+
+def test_mentions_from_df():
+    correct = {"A": 9 / 2, "B": 9 / 2, "C": 7 / 2}
+    test = _mentions_from_df(profile_no_ties)
+    assert correct == test
+    assert isinstance(test["A"], float)
+
+
+def test_mentions_from_df_with_ties():
     correct = {"A": 9 / 2, "B": 9 / 2, "C": 7 / 2}
-    test = mentions(profile_no_ties)
+    test = _mentions_from_df(profile_with_ties)
+    assert correct == test
+    assert isinstance(test["A"], float)
+
+
+def test_mentions_from_df_with_duplicates():
+    correct = {"A": 3 / 2, "B": 23 / 2, "C": 1 / 2}
+    test = _mentions_from_df(profile_with_duplicates)
     assert correct == test
     assert isinstance(test["A"], float)
 
 
+@pytest.mark.slow
+def test_fast_and_slow_mentions_are_same():
+    profile = cast(RankProfile, RankProfile.from_csv(CSV_DIR / "albany_profile.csv"))
+    assert _mentions_from_ballots(profile) == _mentions_from_df(profile)
+
+
+def test_mentions_zero_weight_ballots():
+    # "A" appears only on a zero-weight ballot, so it is not in profile.candidates; both
+    # mentions paths must skip it rather than raise KeyError
+    profile = RankProfile(
+        ballots=(
+            RankBallot(ranking=tuple(map(frozenset, [{"A"}])), weight=0),
+            RankBallot(ranking=tuple(map(frozenset, [{"B"}])), weight=2),
+        )
+    )
+    correct = {"B": 2.0}
+    assert _mentions_from_ballots(profile) == correct
+    assert _mentions_from_df(profile) == correct
+
+
+def test_mentions_from_df_returns_python_floats():
+    assert all(type(v) is float for v in _mentions_from_df(profile_no_ties).values())
+
+
+def test_mentions_no_ranking_ballot_raises_in_both_paths():
+    # an unranked ballot materializes with ranking=None, so both paths must raise the same error
+    profile = RankProfile(ballots=(RankBallot(weight=2), RankBallot(ranking=({"A"},), weight=1)))
+    with pytest.raises(TypeError, match="Ballots must have rankings."):
+        _mentions_from_df(profile)
+    with pytest.raises(TypeError, match="Ballots must have rankings."):
+        _mentions_from_ballots(profile)
+
+
+def test_mentions_all_unranked_profile_raises_in_both_paths():
+    # every ballot has ranking=None, so the profile has no ranking columns at all
+    profile = RankProfile(ballots=(RankBallot(weight=2),))
+    with pytest.raises(TypeError, match="Ballots must have rankings."):
+        _mentions_from_df(profile)
+    with pytest.raises(TypeError, match="Ballots must have rankings."):
+        _mentions_from_ballots(profile)
+
+
+def test_mentions_empty_frozenset_ranking_consistent():
+    # a present-but-empty frozenset ranking is not ranking=None; neither path should raise
+    profile = RankProfile(ballots=(RankBallot(ranking=(frozenset(),), weight=1),))
+    assert _mentions_from_df(profile) == {}
+    assert _mentions_from_ballots(profile) == {}
+
+
+def test_mentions_from_df_duplicate_ballot_index():
+    # profiles built from a df keep the given index; duplicate labels must not break the df path
+    df = RankProfile(ballots=(RankBallot(ranking=({"A"}, {"B"}), weight=1),)).df
+    profile = RankProfile(df=pd.concat([df, df]), candidates=("A", "B"), max_ranking_length=2)
+    correct = {"A": 2.0, "B": 2.0}
+    assert _mentions_from_df(profile) == correct
+    assert _mentions_from_ballots(profile) == correct
+
+
+def _random_mentions_profile(rng: random.Random) -> RankProfile:
+    # samples every dimension that has produced a path divergence: ties, duplicate candidates,
+    # short ballots, zero and Fraction weights, mixed str/int candidates
+    cands = ["A", "B", 1, 2, "C"][: rng.randint(2, 5)]
+    ballots = []
+    for _ in range(rng.randint(1, 8)):
+        pool = [rng.choice(cands) for _ in range(rng.randint(1, len(cands)))]
+        ranking = []
+        while pool:
+            n = rng.randint(1, min(2, len(pool)))
+            ranking.append(frozenset(pool[:n]))
+            pool = pool[n:]
+        weight = rng.choice([0, 1, 3, 1 / 2, Fraction(1, 3)])
+        ballots.append(RankBallot(ranking=tuple(ranking), weight=weight))
+    # sometimes omit the explicit candidate list so profile.candidates can exclude candidates
+    # that appear only on zero-weight ballots
+    if rng.random() < 0.5:
+        return RankProfile(ballots=tuple(ballots), max_ranking_length=len(cands))
+    return RankProfile(ballots=tuple(ballots), candidates=cands, max_ranking_length=len(cands))
+
+
+@pytest.mark.parametrize("seed", range(25))
+def test_mentions_paths_stay_in_sync(seed):
+    profile = _random_mentions_profile(random.Random(seed))
+    assert _mentions_from_ballots(profile) == _mentions_from_df(profile)
+
+
 def test_mentions_errors():
     with pytest.raises(TypeError, match="Profile must be of type RankProfile"):
         mentions(cast(RankProfile, ScoreProfile(ballots=(ScoreBallot(scores={"A": 3}),))))