diff --git a/src/votekit/pref_profile/pref_profile.py b/src/votekit/pref_profile/pref_profile.py index 9756e925..69d0ce3c 100644 --- a/src/votekit/pref_profile/pref_profile.py +++ b/src/votekit/pref_profile/pref_profile.py @@ -26,6 +26,8 @@ _validate_score_csv_format, ) from votekit.pref_profile.utils import ( + _subtract_rank_profiles, + _subtract_score_profiles, _sum_rank_profiles, _sum_score_profiles, convert_row_to_rank_ballot, @@ -703,7 +705,7 @@ def _translate_df_ranking_values( ranking_cols = [col for col in df.columns if col.startswith("Ranking_")] translated_df = df.copy() translated_df[ranking_cols] = translated_df[ranking_cols].map( - lambda ranking: (candidate_mapping[ranking]) + lambda ranking: candidate_mapping[ranking] ) return translated_df @@ -799,6 +801,16 @@ def __add__(self, other) -> RankProfile: """ return _sum_rank_profiles([self, other]) + def __sub__(self, other) -> RankProfile: + """ + Subtract the ballot weights of another profile, matched by identical + rankings. Voter sets of the minuend are retained. + + Raises: + ValueError: A ballot weight would become negative. + """ + return _subtract_rank_profiles(self, other) + def group_ballots(self) -> RankProfile: """ Groups ballots by rankings and updates weights. Retains voter sets, but @@ -1467,6 +1479,16 @@ def __add__(self, other): """ return _sum_score_profiles([self, other]) + def __sub__(self, other): + """ + Subtract the ballot weights of another profile, matched by identical + scores. Voter sets of the minuend are retained. + + Raises: + ValueError: A ballot weight would become negative. + """ + return _subtract_score_profiles(self, other) + def group_ballots(self) -> ScoreProfile: """ Groups ballots by scores and updates weights. Retains voter sets, but diff --git a/src/votekit/pref_profile/utils.py b/src/votekit/pref_profile/utils.py index 0902f534..efb84dd0 100644 --- a/src/votekit/pref_profile/utils.py +++ b/src/votekit/pref_profile/utils.py @@ -574,3 +574,159 @@ def sum_profiles(profiles: Sequence[PreferenceProfile]) -> PreferenceProfile: f"Cannot sum profiles of type {type(profiles[0]).__name__}. " "List can only contain RankProfiles or ScoreProfiles." ) + + +def subtract_profiles( + minuend: PreferenceProfile, subtrahend: PreferenceProfile +) -> PreferenceProfile: + """ + Subtract the ballot weights of one profile from another, matched by + identical ballot rankings (or scores). + + Profiles with the same ballot rankings but differing weights have their + weights subtracted, leaving the voter set of the minuend untouched. + + Args: + minuend (PreferenceProfile): The profile being subtracted from. + subtrahend (PreferenceProfile): The profile whose weights are subtracted. + + Returns: + PreferenceProfile: A new profile with the subtracted ballot weights. + + Raises: + TypeError: Both profiles must be of the same type, RankProfile or ScoreProfile. + ValueError: A ballot weight would become negative. + """ + + from votekit.pref_profile.pref_profile import RankProfile, ScoreProfile + + if type(minuend) is not type(subtrahend): + raise TypeError( + f"Both profiles must be of the same type. " + f"Got {type(minuend).__name__} and {type(subtrahend).__name__}." + ) + + if isinstance(minuend, RankProfile): + return _subtract_rank_profiles(minuend, subtrahend) + + if isinstance(minuend, ScoreProfile): + return _subtract_score_profiles(minuend, subtrahend) + + raise TypeError( + f"Cannot subtract profiles of type {type(minuend).__name__}. " + "Only RankProfiles or ScoreProfiles are supported." + ) + + +def _subtract_rank_profiles(minuend: "RankProfile", subtrahend: "RankProfile") -> "RankProfile": + """Helper function for subtract_profiles that subtracts RankProfiles.""" + + from votekit.pref_profile.pref_profile import RankProfile + + candidates = list(set().union(*[set(profile.candidates) for profile in [minuend, subtrahend]])) + max_ranking_length = max( + [ + profile.max_ranking_length + for profile in [minuend, subtrahend] + if profile.max_ranking_length is not None + ] + ) + + # Pad both profiles to the same ranking length + padded_dfs = [] + for profile in [minuend, subtrahend]: + assert profile.max_ranking_length is not None + curr_df = profile.df.copy() + for i in range(profile.max_ranking_length, max_ranking_length): + curr_df.insert( + len(curr_df.columns), + f"Ranking_{i + 1}", + pd.Series([frozenset("~")] * len(curr_df), dtype=object, index=curr_df.index), + ) + padded_dfs.append(curr_df) + + minuend_df, subtrahend_df = padded_dfs + ranking_cols = [f"Ranking_{i + 1}" for i in range(max_ranking_length)] + + # Align by ranking and subtract weights + subtrahend_weights = subtrahend_df.groupby(ranking_cols, dropna=False)["Weight"].sum() + minuend_grouped = minuend_df.groupby(ranking_cols, dropna=False) + new_rows = [] + for key, group in minuend_grouped: + key = key[0] if isinstance(key, tuple) and len(key) == 1 else key + weight = group["Weight"].sum() - subtrahend_weights.get(key, 0) + if weight < 0: + raise ValueError( + f"Cannot subtract profiles: ballot weight would become negative for ranking {key}." + ) + voter_set = set().union(*group["Voter Set"]) + ranking = { + col: key[i] if i < len(key) else frozenset("~") for i, col in enumerate(ranking_cols) + } + new_rows.append( + { + **ranking, + "Weight": weight, + "Voter Set": voter_set, + } + ) + + # Ballots present only in the subtrahend do not appear in the result + if new_rows: + new_df = pd.DataFrame(new_rows)[ranking_cols + ["Weight", "Voter Set"]] + else: + new_df = pd.DataFrame(columns=ranking_cols + ["Weight", "Voter Set"], dtype=object) + new_df.index.name = "Ballot Index" + + return RankProfile( + candidates=candidates, + df=new_df, + max_ranking_length=max_ranking_length, + ) + + +def _subtract_score_profiles(minuend: "ScoreProfile", subtrahend: "ScoreProfile") -> "ScoreProfile": + """Helper function for subtract_profiles that subtracts ScoreProfiles.""" + + from votekit.pref_profile.pref_profile import ScoreProfile + + total_cand = set().union(*[set(profile.candidates) for profile in [minuend, subtrahend]]) + minuend_df = minuend.df.copy() + subtrahend_df = subtrahend.df.copy() + for df in (minuend_df, subtrahend_df): + for cand in total_cand - set(df.columns) - {"Weight", "Voter Set"}: + df[cand] = [np.nan] * len(df) + + cand_cols = sort_candidates_pseudo_lexicographically(total_cand) + # NaN keys never compare equal, so fill missing scores with a sentinel + # before grouping so identical ballots match across profiles. + sentinel = -np.inf + subtrahend_filled = subtrahend_df.copy() + subtrahend_filled[cand_cols] = subtrahend_filled[cand_cols].fillna(sentinel) + minuend_filled = minuend_df.copy() + minuend_filled[cand_cols] = minuend_filled[cand_cols].fillna(sentinel) + + subtrahend_weights = subtrahend_filled.groupby(cand_cols, dropna=False)["Weight"].sum() + minuend_grouped = minuend_filled.groupby(cand_cols, dropna=False) + new_rows = [] + for key, group in minuend_grouped: + key = key[0] if isinstance(key, tuple) and len(key) == 1 else key + weight = group["Weight"].sum() - subtrahend_weights.get(key, 0) + if weight < 0: + raise ValueError( + f"Cannot subtract profiles: ballot weight would become negative for scores {key}." + ) + voter_set = set().union(*group["Voter Set"]) + scores = {cand: (np.nan if v == sentinel else v) for cand, v in zip(cand_cols, key)} + new_rows.append({**scores, "Weight": weight, "Voter Set": voter_set}) + + if new_rows: + new_df = pd.DataFrame(new_rows)[cand_cols + ["Weight", "Voter Set"]] + else: + new_df = pd.DataFrame(columns=cand_cols + ["Weight", "Voter Set"], dtype=object) + new_df.index.name = "Ballot Index" + + return ScoreProfile( + candidates=sort_candidates_pseudo_lexicographically(total_cand), + df=new_df, + ) diff --git a/tests/pref_profile/utils/test_subtract_profiles.py b/tests/pref_profile/utils/test_subtract_profiles.py new file mode 100644 index 00000000..f326e46c --- /dev/null +++ b/tests/pref_profile/utils/test_subtract_profiles.py @@ -0,0 +1,163 @@ +import pytest + +from votekit.ballot import RankBallot, ScoreBallot +from votekit.pref_profile import RankProfile, ScoreProfile +from votekit.pref_profile.utils import subtract_profiles + + +def test_subtract_rank_profiles_weights(): + minuend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}, {"C"}), weight=5), + RankBallot(ranking=({"A", "B"}, frozenset(), {"D"}), weight=3, voter_set={"Chris"}), + ], + candidates=["A", "B", "C", "D"], + max_ranking_length=3, + ) + subtrahend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}, {"C"}), weight=2), + ], + candidates=["A", "B", "C", "D"], + max_ranking_length=3, + ) + result = minuend - subtrahend + weights = {b.ranking: b.weight for b in result.ballots} + assert weights[(frozenset({"A"}), frozenset({"B"}), frozenset({"C"}))] == 3 + assert weights[(frozenset({"A", "B"}), frozenset(), frozenset({"D"}))] == 3 + + +def test_subtract_rank_profiles_retains_voter_set(): + minuend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=4, voter_set={"Chris", "Sam"}), + ], + candidates=["A", "B"], + max_ranking_length=2, + ) + subtrahend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=1), + ], + candidates=["A", "B"], + max_ranking_length=2, + ) + result = minuend - subtrahend + assert result.ballots[0].voter_set == {"Chris", "Sam"} + + +def test_subtract_rank_profiles_negative_weight_raises(): + minuend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=1), + ], + candidates=["A", "B"], + max_ranking_length=2, + ) + subtrahend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=3), + ], + candidates=["A", "B"], + max_ranking_length=2, + ) + with pytest.raises(ValueError, match="negative"): + minuend - subtrahend + + +def test_subtract_rank_profiles_different_lengths(): + minuend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=5), + ], + candidates=["A", "B", "C"], + max_ranking_length=3, + ) + subtrahend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=2), + ], + candidates=["A", "B", "C"], + max_ranking_length=2, + ) + result = minuend - subtrahend + assert result.ballots[0].weight == 3 + assert result.max_ranking_length == 3 + + +def test_subtract_score_profiles_weights(): + minuend = ScoreProfile( + ballots=[ + ScoreBallot(scores={"A": 2, "B": 2}, weight=4), + ScoreBallot(scores={"A": 1, "C": 3}, weight=2), + ], + candidates=["A", "B", "C"], + ) + subtrahend = ScoreProfile( + ballots=[ + ScoreBallot(scores={"A": 2, "B": 2}, weight=1), + ], + candidates=["A", "B", "C"], + ) + result = minuend - subtrahend + weights = {tuple(sorted(b.scores.items())): b.weight for b in result.ballots} + assert weights[(("A", 2), ("B", 2))] == 3 + assert weights[(("A", 1), ("C", 3))] == 2 + + +def test_subtract_score_profiles_negative_weight_raises(): + minuend = ScoreProfile( + ballots=[ + ScoreBallot(scores={"A": 2}, weight=1), + ], + candidates=["A"], + ) + subtrahend = ScoreProfile( + ballots=[ + ScoreBallot(scores={"A": 2}, weight=3), + ], + candidates=["A"], + ) + with pytest.raises(ValueError, match="negative"): + minuend - subtrahend + + +def test_subtract_profiles_mixed_types_raises(): + score_profile = ScoreProfile( + ballots=[ScoreBallot(scores={"A": 2}, weight=2)], + candidates=["A"], + ) + rank_profile = RankProfile( + ballots=[RankBallot(ranking=({"A"},), weight=2)], + candidates=["A"], + max_ranking_length=1, + ) + with pytest.raises(TypeError, match="same type"): + subtract_profiles(rank_profile, score_profile) + + +def test_subtract_profiles_function(): + minuend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=5), + ], + candidates=["A", "B"], + max_ranking_length=2, + ) + subtrahend = RankProfile( + ballots=[ + RankBallot(ranking=({"A"}, {"B"}), weight=2), + ], + candidates=["A", "B"], + max_ranking_length=2, + ) + result = subtract_profiles(minuend, subtrahend) + assert result.ballots[0].weight == 3 + + +def test_subtract_profiles_unknown_type_raises(): + class FakeProfile: + pass + + with pytest.raises(TypeError, match="Cannot subtract"): + subtract_profiles(FakeProfile(), FakeProfile()) # type: ignore[arg-type]