diff --git a/src/votekit/cleaning/__init__.py b/src/votekit/cleaning/__init__.py index 0e4e7edb..ad0761e7 100644 --- a/src/votekit/cleaning/__init__.py +++ b/src/votekit/cleaning/__init__.py @@ -9,6 +9,7 @@ remove_and_condense_rank_profile, remove_cand_rank_profile, remove_repeat_cands_rank_profile, + truncate_rank_profile, ) from .score_ballots_cleaning import remove_cand_score_ballot from .score_profiles_cleaning import clean_score_profile, remove_cand_score_profile @@ -22,6 +23,7 @@ "remove_cand_rank_ballot", "condense_rank_ballot", "remove_repeat_cands_rank_ballot", + "truncate_rank_profile", "remove_cand_score_ballot", "clean_score_profile", "remove_cand_score_profile", diff --git a/src/votekit/cleaning/rank_profiles_cleaning.py b/src/votekit/cleaning/rank_profiles_cleaning.py index 396f2057..08f5c3db 100644 --- a/src/votekit/cleaning/rank_profiles_cleaning.py +++ b/src/votekit/cleaning/rank_profiles_cleaning.py @@ -288,6 +288,83 @@ def remove_cand_rank_profile( ) +def truncate_ranking_row( + removed: Candidate | CandidateList, + ranking_tup: tuple[frozenset, ...], +) -> tuple[frozenset, ...]: + """ + Truncate a ranking at the first position containing a specified candidate or marker. + + The matching position and every position below it are replaced with trailing ``~`` + placeholders so that the ranking keeps its original width in a profile dataframe. + + Args: + removed (Candidate | list[Candidate]): Candidate or list of candidates or markers at + which to truncate. + ranking_tup (tuple): Ranking to truncate. + + Returns: + tuple: Ranking truncated at the first matching position. + """ + if isinstance(removed, Candidate): + removed = [removed] + + removed_set = set(removed) + out: list[frozenset] = [] + + for cand_set in ranking_tup: + if cand_set.isdisjoint(removed_set): + out.append(cand_set) + continue + + out.extend([frozenset("~")] * (len(ranking_tup) - len(out))) + break + + return tuple(out) + + +def truncate_rank_profile( + removed: Candidate | CandidateList, + profile: RankProfile, + remove_empty_ballots: bool = True, + remove_zero_weight_ballots: bool = True, + retain_original_candidate_list: bool = True, +) -> CleanedRankProfile: + """ + Truncate ranked ballots at the first position containing a specified candidate or marker. + + This is useful for cleaning CVR data where values such as ``"overvote"`` or ``"undervote"`` + terminate the meaningful portion of a ballot. The matching position and all lower-ranked + positions are removed. Ballots without a matching value are retained unchanged. + + Args: + removed (Candidate | list[Candidate]): Candidate, marker, or list of candidates and + markers at which to truncate. + profile (RankProfile): Profile to truncate. + remove_empty_ballots (bool, optional): Whether or not to remove ballots with no ranking + after truncation. Defaults to True. + remove_zero_weight_ballots (bool, optional): Whether or not to remove zero-weight ballots. + Defaults to True. + retain_original_candidate_list (bool, optional): Whether or not to retain the original + candidate list. Defaults to True. + + Returns: + CleanedRankProfile: A cleaned ``RankProfile``. + + Raises: + ProfileError: Profile must only contain ranked ballots. + """ + cleaned_profile = clean_rank_profile( + profile, + partial(truncate_ranking_row, removed), + remove_empty_ballots, + remove_zero_weight_ballots, + retain_original_candidate_list, + ) + + return cleaned_profile + + def condense_ranking_row( ranking_tup: tuple, ) -> tuple: diff --git a/tests/cleaning/rank_profiles/test_truncate_ranked_profile.py b/tests/cleaning/rank_profiles/test_truncate_ranked_profile.py new file mode 100644 index 00000000..903f4f92 --- /dev/null +++ b/tests/cleaning/rank_profiles/test_truncate_ranked_profile.py @@ -0,0 +1,52 @@ +import pytest + +from votekit.ballot import RankBallot, ScoreBallot +from votekit.cleaning import truncate_rank_profile +from votekit.pref_profile import CleanedRankProfile, ProfileError, RankProfile, ScoreProfile + +profile = RankProfile( + ballots=[ + RankBallot(ranking=[{"A"}, {"overvote"}, {"B"}, {"C"}], weight=1), + RankBallot(ranking=[{"A"}, {"B"}, {"C"}], weight=2), + RankBallot(ranking=[{"undervote"}, {"C"}], weight=3), + RankBallot(ranking=[{"A"}, {"B"}], weight=0), + ] +) + + +def test_truncate_rank_profile_at_candidate_or_marker(): + cleaned_profile = truncate_rank_profile(["overvote", "undervote"], profile) + + assert isinstance(cleaned_profile, CleanedRankProfile) + assert cleaned_profile.parent_profile == profile + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"A"}], weight=1), + RankBallot(ranking=[{"A"}, {"B"}, {"C"}], weight=2), + ) + assert cleaned_profile.no_rank_altr_idxs == {2} + assert cleaned_profile.nonempty_altr_idxs == {0} + assert cleaned_profile.unaltr_idxs == {1, 3} + assert cleaned_profile.no_wt_altr_idxs == set() + + +def test_truncate_rank_profile_can_retain_empty_and_zero_weight_ballots(): + cleaned_profile = truncate_rank_profile( + "overvote", + profile, + remove_empty_ballots=False, + remove_zero_weight_ballots=False, + ) + + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"A"}], weight=1), + RankBallot(ranking=[{"A"}, {"B"}, {"C"}], weight=2), + RankBallot(ranking=[{"undervote"}, {"C"}], weight=3), + RankBallot(ranking=[{"A"}, {"B"}], weight=0), + ) + + +def test_truncate_rank_profile_requires_rank_profile(): + score_profile = ScoreProfile(ballots=[ScoreBallot(scores={"A": 1})]) + + with pytest.raises(ProfileError, match="Profile must be a RankProfile."): + truncate_rank_profile("overvote", score_profile) # type: ignore[arg-type]