diff --git a/src/votekit/cleaning/__init__.py b/src/votekit/cleaning/__init__.py index 0e4e7edb..dce60578 100644 --- a/src/votekit/cleaning/__init__.py +++ b/src/votekit/cleaning/__init__.py @@ -8,6 +8,7 @@ condense_rank_profile, remove_and_condense_rank_profile, remove_cand_rank_profile, + remove_rank_ballots_with_cands, remove_repeat_cands_rank_profile, ) from .score_ballots_cleaning import remove_cand_score_ballot @@ -17,6 +18,7 @@ "clean_rank_profile", "remove_repeat_cands_rank_profile", "remove_cand_rank_profile", + "remove_rank_ballots_with_cands", "condense_rank_profile", "remove_and_condense_rank_profile", "remove_cand_rank_ballot", diff --git a/src/votekit/cleaning/rank_profiles_cleaning.py b/src/votekit/cleaning/rank_profiles_cleaning.py index 396f2057..7624bc8d 100644 --- a/src/votekit/cleaning/rank_profiles_cleaning.py +++ b/src/votekit/cleaning/rank_profiles_cleaning.py @@ -288,6 +288,61 @@ def remove_cand_rank_profile( ) +def remove_rank_ballots_with_cands( + removed: Candidate | CandidateList, + profile: RankProfile, + remove_zero_weight_ballots: bool = True, +) -> CleanedRankProfile: + """ + Remove ballots containing any specified candidate or marker from a ranked profile. + + Args: + removed (Candidate | list[Candidate] | list[str] | list[int]): + Candidate, marker, or list of candidates and markers to filter on. + profile (RankProfile): Profile to remove matching ballots from. + remove_zero_weight_ballots (bool, optional): Whether or not to remove ballots that have no + weight. Defaults to True. + + Returns: + CleanedRankProfile: A cleaned ``RankProfile``. + + Raises: + ProfileError: Profile must only contain ranked ballots. + """ + if not isinstance(profile, RankProfile): + raise ProfileError("Profile must be a RankProfile.") + + if isinstance(removed, Candidate): + removed = [removed] + + removed_set = set(removed) + + assert profile.max_ranking_length is not None + ranking_cols = [f"Ranking_{i}" for i in range(1, profile.max_ranking_length + 1)] + contains_removed = profile.df[ranking_cols].map( + lambda cands: isinstance(cands, frozenset) and not cands.isdisjoint(removed_set) + ) + matching_idxs = set(profile.df.index[contains_removed.any(axis=1)]) + + cleaned_df = profile.df.drop(index=matching_idxs) + if remove_zero_weight_ballots: + cleaned_df = cleaned_df[cleaned_df["Weight"] > 0] + + unaltr_idxs = set(profile.df.index) - matching_idxs + + return CleanedRankProfile( + df=cleaned_df, + candidates=profile.candidates, + max_ranking_length=profile.max_ranking_length, + parent_profile=profile, + df_index_column=list(cleaned_df.index), + no_wt_altr_idxs=set(), + no_rank_altr_idxs=matching_idxs, + nonempty_altr_idxs=set(), + unaltr_idxs=unaltr_idxs, + ) + + def condense_ranking_row( ranking_tup: tuple, ) -> tuple: diff --git a/tests/cleaning/rank_profiles/test_remove_rank_ballots_with_cands.py b/tests/cleaning/rank_profiles/test_remove_rank_ballots_with_cands.py new file mode 100644 index 00000000..ebad4746 --- /dev/null +++ b/tests/cleaning/rank_profiles/test_remove_rank_ballots_with_cands.py @@ -0,0 +1,65 @@ +import pytest + +from votekit.ballot import RankBallot, ScoreBallot +from votekit.cleaning import remove_rank_ballots_with_cands +from votekit.pref_profile import CleanedRankProfile, ProfileError, RankProfile, ScoreProfile + +profile = RankProfile( + ballots=[ + RankBallot(ranking=[{"A"}, {"B"}], weight=1), + RankBallot(ranking=[{"A"}, {"overvote"}, {"C"}], weight=2), + RankBallot(ranking=[{"C"}, {"A"}], weight=3), + RankBallot(ranking=[{"undervote", "B"}, {"C"}], weight=4), + RankBallot(ranking=[{"A"}], weight=0), + ] +) + + +def test_remove_rank_ballots_with_candidate(): + cleaned_profile = remove_rank_ballots_with_cands("overvote", profile) + + assert isinstance(cleaned_profile, CleanedRankProfile) + assert cleaned_profile.parent_profile == profile + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"A"}, {"B"}], weight=1), + RankBallot(ranking=[{"C"}, {"A"}], weight=3), + RankBallot(ranking=[{"undervote", "B"}, {"C"}], weight=4), + ) + assert cleaned_profile.candidates == profile.candidates + assert cleaned_profile.no_wt_altr_idxs == set() + assert cleaned_profile.no_rank_altr_idxs == {1} + assert cleaned_profile.nonempty_altr_idxs == set() + assert cleaned_profile.unaltr_idxs == {0, 2, 3, 4} + + +def test_remove_rank_ballots_with_multiple_candidates_and_ties(): + cleaned_profile = remove_rank_ballots_with_cands(["overvote", "undervote"], profile) + + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"A"}, {"B"}], weight=1), + RankBallot(ranking=[{"C"}, {"A"}], weight=3), + ) + assert cleaned_profile.no_rank_altr_idxs == {1, 3} + assert cleaned_profile.unaltr_idxs == {0, 2, 4} + + +def test_remove_rank_ballots_with_candidate_retains_zero_weight_when_requested(): + cleaned_profile = remove_rank_ballots_with_cands( + "overvote", + profile, + remove_zero_weight_ballots=False, + ) + + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"A"}, {"B"}], weight=1), + RankBallot(ranking=[{"C"}, {"A"}], weight=3), + RankBallot(ranking=[{"undervote", "B"}, {"C"}], weight=4), + RankBallot(ranking=[{"A"}], weight=0), + ) + + +def test_remove_rank_ballots_with_cands_requires_rank_profile(): + score_profile = ScoreProfile(ballots=[ScoreBallot(scores={"A": 1})]) + + with pytest.raises(ProfileError, match="Profile must be a RankProfile."): + remove_rank_ballots_with_cands("A", score_profile) # type: ignore[arg-type]