From aeceeb681e58965119719ebcd192c0adf52693ca Mon Sep 17 00:00:00 2001 From: Grace Gibson Date: Fri, 14 Aug 2026 10:29:51 -0500 Subject: [PATCH 1/4] add remove_ballots_with_cand fxn and tests --- src/votekit/cleaning/__init__.py | 2 + .../cleaning/rank_profiles_cleaning.py | 91 +++++++++++++++++++ ...t_remove_ballots_with_cand_rank_profile.py | 76 ++++++++++++++++ 3 files changed, 169 insertions(+) create mode 100644 tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py diff --git a/src/votekit/cleaning/__init__.py b/src/votekit/cleaning/__init__.py index 0e4e7edb..be2064e0 100644 --- a/src/votekit/cleaning/__init__.py +++ b/src/votekit/cleaning/__init__.py @@ -7,6 +7,7 @@ clean_rank_profile, condense_rank_profile, remove_and_condense_rank_profile, + remove_ballots_with_cand_rank_profile, remove_cand_rank_profile, remove_repeat_cands_rank_profile, ) @@ -25,6 +26,7 @@ "remove_cand_score_ballot", "clean_score_profile", "remove_cand_score_profile", + "remove_ballots_with_cand_rank_profile", ] # Patch __module__ on every exported symbol so that Sphinx autodoc displays diff --git a/src/votekit/cleaning/rank_profiles_cleaning.py b/src/votekit/cleaning/rank_profiles_cleaning.py index 396f2057..95edd5d1 100644 --- a/src/votekit/cleaning/rank_profiles_cleaning.py +++ b/src/votekit/cleaning/rank_profiles_cleaning.py @@ -525,3 +525,94 @@ def remove_and_condense_rank_profile( nonempty_altr_idxs=new_nonempty_altr_idxs, unaltr_idxs=new_unaltr_idxs, ) + + +def remove_ballots_with_cand_rank_profile( + removed: Candidate | list[Candidate], + profile: RankProfile, + remove_empty_ballots: bool = True, + remove_zero_weight_ballots: bool = True, + retain_original_candidate_list: bool = False, +) -> CleanedRankProfile: + """ + Given a ranked profile, remove the ballots that contain the given candidate(s). + + A removed ballot is considered altered for both weight and ranking. No ballots are altered in + place, only removed. + + Args: + removed (Candidate | list[Candidate]): Candidate or list of candidates to remove their + ballots. Candidates can be strings, integers, or mix of both. + profile (RankProfile): Profile to remove ballots from. + remove_empty_ballots (bool, optional): Whether or not to remove ballots that have no + ranking or scores as a result of cleaning. Defaults to True. + remove_zero_weight_ballots (bool, optional): Whether or not to remove ballots that have no + weight as a result of cleaning. Defaults to True. + retain_original_candidate_list (bool, optional): Whether or not to retain the original list + of candidates. Defaults to False. + + Returns: + CleanedRankProfile: A cleaned ``RankProfile`` with ballots containing the specified + candidate(s) removed. + + Raises: + ProfileError: Profile must contain ranked ballots. + TypeError: Candidates to be removed must be strings or integers. A boolean or float + candidate of the same value as an integer candidate would result in removing the ballots + with that integer candidate. + """ + if not isinstance(profile, RankProfile): + raise ProfileError("Profile must be a RankProfile.") + + if isinstance(removed, Candidate) and not isinstance(removed, bool): + removed = [removed] + elif isinstance(removed, list): + if any(not isinstance(cand, (str, int)) or isinstance(cand, bool) for cand in removed): + raise TypeError("Candidates must be strings or integers within removed.") + else: + raise TypeError("removed must be a str/int candidate or a list of candidates.") + + cand_ids = [] + for cand in removed: + cand_ids.extend( + [ + profile.candidate_id_map[cand_set] + for cand_set in profile.candidate_id_map + if cand in cand_set + ] + ) + ranking_cols = [f"Ranking_{i}" for i in range(1, profile.max_ranking_length + 1)] + ballots_to_remove = profile._df[ranking_cols].isin(cand_ids).any(axis=1) + cleaned_df = profile.df[~ballots_to_remove] + removed_ballot_idxs = set(profile.df.index[ballots_to_remove]) + + empty_ballot_idxs: set[int] = set() + if remove_empty_ballots: + mask = cleaned_df[ranking_cols].map(lambda x: x == frozenset({"~"})).all(axis=1) + empty_ballot_idxs = set(cleaned_df.index[mask]) + cleaned_df = cleaned_df[~mask] + + zero_weight_ballot_idxs: set[int] = set() + if remove_zero_weight_ballots: + mask = cleaned_df["Weight"] > 0 + zero_weight_ballot_idxs = set(cleaned_df.index[~mask]) + cleaned_df = cleaned_df[mask] + + candidates = ( + profile.candidates + if retain_original_candidate_list + else tuple(set(profile.candidates) - set(removed)) + ) + + unaltered_idxs = list(cleaned_df.index) + return CleanedRankProfile( + df=cleaned_df, + candidates=candidates, + max_ranking_length=profile.max_ranking_length, + parent_profile=profile, + df_index_column=unaltered_idxs, + no_wt_altr_idxs=removed_ballot_idxs | zero_weight_ballot_idxs, + no_rank_altr_idxs=removed_ballot_idxs | empty_ballot_idxs, + nonempty_altr_idxs=set(), + unaltr_idxs=set(unaltered_idxs), + ) diff --git a/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py b/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py new file mode 100644 index 00000000..76ddc0e4 --- /dev/null +++ b/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py @@ -0,0 +1,76 @@ +import pytest + +from votekit.ballot import RankBallot +from votekit.cleaning import remove_ballots_with_cand_rank_profile +from votekit.pref_profile import CleanedRankProfile, RankProfile + +profile_no_ties = RankProfile( + ballots=[ + RankBallot(ranking=[{"A"}, {"B"}, {"C"}], weight=1), + RankBallot(ranking=[{"B"}, {"C"}], weight=1 / 2), + RankBallot(ranking=[{"C"}], weight=3), + RankBallot(ranking=[{"B", "C", "A"}], weight=3), + ], + max_ranking_length=3, +) + +profile_with_ties = RankProfile( + ballots=[ + RankBallot(ranking=[{"A", "B"}, {"C"}], weight=1), + RankBallot(ranking=[{"B", "C"}], weight=1 / 2), + RankBallot(ranking=[{"C"}], weight=3), + RankBallot(ranking=[{"B", "C", "A"}], weight=3), + ], + max_ranking_length=3, +) + + +def test_remove_ballots_with_cand_rank_profile(): + cleaned_profile = remove_ballots_with_cand_rank_profile("A", profile_no_ties) + assert isinstance(cleaned_profile, CleanedRankProfile) + assert cleaned_profile.parent_profile == profile_no_ties + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"B"}, {"C"}], weight=1 / 2), + RankBallot(ranking=[{"C"}], weight=3), + ) + assert cleaned_profile != profile_no_ties + assert cleaned_profile.no_wt_altr_idxs == {0, 3} + assert cleaned_profile.no_rank_altr_idxs == {0, 3} + assert cleaned_profile.nonempty_altr_idxs == set() + assert cleaned_profile.unaltr_idxs == {1, 2} + + +def test_remove_ballots_with_cand_rank_profile_with_ties(): + cleaned_profile = remove_ballots_with_cand_rank_profile("A", profile_with_ties) + assert isinstance(cleaned_profile, CleanedRankProfile) + assert cleaned_profile.parent_profile == profile_with_ties + assert cleaned_profile.ballots == ( + RankBallot(ranking=[{"B", "C"}], weight=1 / 2), + RankBallot(ranking=[{"C"}], weight=3), + ) + assert cleaned_profile != profile_with_ties + assert cleaned_profile.no_wt_altr_idxs == {0, 3} + assert cleaned_profile.no_rank_altr_idxs == {0, 3} + assert cleaned_profile.nonempty_altr_idxs == set() + assert cleaned_profile.unaltr_idxs == {1, 2} + + +def test_remove_ballots_with_mult_cands(): + cleaned_profile = remove_ballots_with_cand_rank_profile(["A", "B"], profile_no_ties) + assert isinstance(cleaned_profile, CleanedRankProfile) + assert cleaned_profile.parent_profile == profile_no_ties + assert cleaned_profile.ballots == (RankBallot(ranking=[{"C"}], weight=3),) + assert cleaned_profile != profile_no_ties + assert cleaned_profile.no_wt_altr_idxs == {0, 1, 3} + assert cleaned_profile.no_rank_altr_idxs == {0, 1, 3} + assert cleaned_profile.nonempty_altr_idxs == set() + assert cleaned_profile.unaltr_idxs == {2} + + +def test_remove_ballots_with_invalid_cand_type(): + with pytest.raises(TypeError, match="Candidates must be strings or integers within removed."): + remove_ballots_with_cand_rank_profile([1.0], profile_no_ties) # type: ignore[arg-type] + with pytest.raises( + TypeError, match="removed must be a str/int candidate or a list of candidates." + ): + remove_ballots_with_cand_rank_profile(True, profile_no_ties) From 44d7d537d6369b9d453eecce29d5657f89408a08 Mon Sep 17 00:00:00 2001 From: Grace Gibson Date: Mon, 17 Aug 2026 06:38:05 -0500 Subject: [PATCH 2/4] fix removed ballots to be no_rank_altr_idxs --- src/votekit/cleaning/rank_profiles_cleaning.py | 17 ++++++----------- ...est_remove_ballots_with_cand_rank_profile.py | 9 ++++++--- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/src/votekit/cleaning/rank_profiles_cleaning.py b/src/votekit/cleaning/rank_profiles_cleaning.py index 95edd5d1..08726189 100644 --- a/src/votekit/cleaning/rank_profiles_cleaning.py +++ b/src/votekit/cleaning/rank_profiles_cleaning.py @@ -537,8 +537,8 @@ def remove_ballots_with_cand_rank_profile( """ Given a ranked profile, remove the ballots that contain the given candidate(s). - A removed ballot is considered altered for both weight and ranking. No ballots are altered in - place, only removed. + A removed ballot's ranking is considered empty after cleaning and recorded in the + ``no_rank_altr_idxs`` of the returned ``CleanedRankProfile``. Args: removed (Candidate | list[Candidate]): Candidate or list of candidates to remove their @@ -584,19 +584,14 @@ def remove_ballots_with_cand_rank_profile( ranking_cols = [f"Ranking_{i}" for i in range(1, profile.max_ranking_length + 1)] ballots_to_remove = profile._df[ranking_cols].isin(cand_ids).any(axis=1) cleaned_df = profile.df[~ballots_to_remove] - removed_ballot_idxs = set(profile.df.index[ballots_to_remove]) + removed_ballot_idxs = set(profile.df[ballots_to_remove].index) - empty_ballot_idxs: set[int] = set() if remove_empty_ballots: mask = cleaned_df[ranking_cols].map(lambda x: x == frozenset({"~"})).all(axis=1) - empty_ballot_idxs = set(cleaned_df.index[mask]) cleaned_df = cleaned_df[~mask] - zero_weight_ballot_idxs: set[int] = set() if remove_zero_weight_ballots: - mask = cleaned_df["Weight"] > 0 - zero_weight_ballot_idxs = set(cleaned_df.index[~mask]) - cleaned_df = cleaned_df[mask] + cleaned_df = cleaned_df[cleaned_df["Weight"] > 0] candidates = ( profile.candidates @@ -611,8 +606,8 @@ def remove_ballots_with_cand_rank_profile( max_ranking_length=profile.max_ranking_length, parent_profile=profile, df_index_column=unaltered_idxs, - no_wt_altr_idxs=removed_ballot_idxs | zero_weight_ballot_idxs, - no_rank_altr_idxs=removed_ballot_idxs | empty_ballot_idxs, + no_wt_altr_idxs=set(), + no_rank_altr_idxs=removed_ballot_idxs, nonempty_altr_idxs=set(), unaltr_idxs=set(unaltered_idxs), ) diff --git a/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py b/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py index 76ddc0e4..b75cdf32 100644 --- a/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py +++ b/tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py @@ -34,10 +34,11 @@ def test_remove_ballots_with_cand_rank_profile(): RankBallot(ranking=[{"C"}], weight=3), ) assert cleaned_profile != profile_no_ties - assert cleaned_profile.no_wt_altr_idxs == {0, 3} + assert cleaned_profile.no_wt_altr_idxs == set() assert cleaned_profile.no_rank_altr_idxs == {0, 3} assert cleaned_profile.nonempty_altr_idxs == set() assert cleaned_profile.unaltr_idxs == {1, 2} + assert cleaned_profile.candidates == ("B", "C") def test_remove_ballots_with_cand_rank_profile_with_ties(): @@ -49,10 +50,11 @@ def test_remove_ballots_with_cand_rank_profile_with_ties(): RankBallot(ranking=[{"C"}], weight=3), ) assert cleaned_profile != profile_with_ties - assert cleaned_profile.no_wt_altr_idxs == {0, 3} + assert cleaned_profile.no_wt_altr_idxs == set() assert cleaned_profile.no_rank_altr_idxs == {0, 3} assert cleaned_profile.nonempty_altr_idxs == set() assert cleaned_profile.unaltr_idxs == {1, 2} + assert cleaned_profile.candidates == ("B", "C") def test_remove_ballots_with_mult_cands(): @@ -61,10 +63,11 @@ def test_remove_ballots_with_mult_cands(): assert cleaned_profile.parent_profile == profile_no_ties assert cleaned_profile.ballots == (RankBallot(ranking=[{"C"}], weight=3),) assert cleaned_profile != profile_no_ties - assert cleaned_profile.no_wt_altr_idxs == {0, 1, 3} + assert cleaned_profile.no_wt_altr_idxs == set() assert cleaned_profile.no_rank_altr_idxs == {0, 1, 3} assert cleaned_profile.nonempty_altr_idxs == set() assert cleaned_profile.unaltr_idxs == {2} + assert cleaned_profile.candidates == ("C",) def test_remove_ballots_with_invalid_cand_type(): From 22253f02127409fa04232f1286daaceca0e28e05 Mon Sep 17 00:00:00 2001 From: Grace Gibson Date: Mon, 17 Aug 2026 07:15:04 -0500 Subject: [PATCH 3/4] exclude pre-existing empty ballots from no_rank_altr_idxs, add emptyset to no_rank_altr_idxs check --- src/votekit/cleaning/rank_profiles_cleaning.py | 6 +++++- .../rank_profiles/test_clean_ranked_profile.py | 14 ++++++++------ .../test_remove_cand_ranked_profile.py | 8 ++++---- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/votekit/cleaning/rank_profiles_cleaning.py b/src/votekit/cleaning/rank_profiles_cleaning.py index 08726189..cfb0c98b 100644 --- a/src/votekit/cleaning/rank_profiles_cleaning.py +++ b/src/votekit/cleaning/rank_profiles_cleaning.py @@ -40,10 +40,14 @@ def _iterate_and_clean_ranking_tuples( cleaned_df[ranking_cols] = pd.DataFrame(cleaned_rows, index=cleaned_df.index) tilde = frozenset({"~"}) + empty = frozenset() idxs = cleaned_df.index unaltr_idxs = {idx for idx, (o, c) in zip(idxs, zip(orig_rows, cleaned_rows)) if o == c} - no_rank_altr_idxs = {idx for idx, c in zip(idxs, cleaned_rows) if all(x == tilde for x in c)} + no_rank_altr_idxs = { + idx for idx, c in zip(idxs, cleaned_rows) if all(x == tilde or x == empty for x in c) + } + no_rank_altr_idxs = no_rank_altr_idxs - unaltr_idxs nonempty_altr_idxs = set(idxs) - unaltr_idxs - no_rank_altr_idxs no_wt_altr_idxs: set[int] = set() diff --git a/tests/cleaning/rank_profiles/test_clean_ranked_profile.py b/tests/cleaning/rank_profiles/test_clean_ranked_profile.py index d7e18d01..57b3b688 100644 --- a/tests/cleaning/rank_profiles/test_clean_ranked_profile.py +++ b/tests/cleaning/rank_profiles/test_clean_ranked_profile.py @@ -9,6 +9,7 @@ RankBallot(ranking=[{"C"}, {"B"}, {"A"}], weight=3), RankBallot(ranking=({"A"},)), RankBallot(ranking=({"B"},), weight=0), + RankBallot(ranking=(), weight=2), ] ) @@ -30,9 +31,9 @@ def test_clean_profile_with_defaults(): assert adj_profile != profile assert adj_profile.no_wt_altr_idxs == set() - assert adj_profile.no_rank_altr_idxs == set() - assert adj_profile.nonempty_altr_idxs == {0, 1, 2, 3} - assert adj_profile.unaltr_idxs == {4} + assert adj_profile.no_rank_altr_idxs == {3} + assert adj_profile.nonempty_altr_idxs == {0, 1, 2} + assert adj_profile.unaltr_idxs == {4, 5} def test_clean_profile_change_defaults(): @@ -59,12 +60,13 @@ def test_clean_profile_change_defaults(): ), RankBallot(ranking=(frozenset(),)), RankBallot(ranking=({"B"},), weight=0), + RankBallot(weight=2), ) ) assert adj_profile.candidates == profile.candidates assert adj_profile.max_ranking_length == 3 assert adj_profile.no_wt_altr_idxs == set() - assert adj_profile.no_rank_altr_idxs == set() - assert adj_profile.nonempty_altr_idxs == {0, 1, 2, 3} - assert adj_profile.unaltr_idxs == {4} + assert adj_profile.no_rank_altr_idxs == {3} + assert adj_profile.nonempty_altr_idxs == {0, 1, 2} + assert adj_profile.unaltr_idxs == {4, 5} diff --git a/tests/cleaning/rank_profiles/test_remove_cand_ranked_profile.py b/tests/cleaning/rank_profiles/test_remove_cand_ranked_profile.py index 6a96a057..381b55b4 100644 --- a/tests/cleaning/rank_profiles/test_remove_cand_ranked_profile.py +++ b/tests/cleaning/rank_profiles/test_remove_cand_ranked_profile.py @@ -55,8 +55,8 @@ def test_remove_mult_cands(): ) assert cleaned_profile != profile_no_ties assert cleaned_profile.no_wt_altr_idxs == set() - assert cleaned_profile.no_rank_altr_idxs == set() - assert cleaned_profile.nonempty_altr_idxs == {0, 1, 2} + assert cleaned_profile.no_rank_altr_idxs == {0} + assert cleaned_profile.nonempty_altr_idxs == {1, 2} assert cleaned_profile.unaltr_idxs == set() @@ -76,6 +76,6 @@ def test_remove_cand_with_ties(): ) assert cleaned_profile != profile_with_ties assert cleaned_profile.no_wt_altr_idxs == set() - assert cleaned_profile.no_rank_altr_idxs == set() - assert cleaned_profile.nonempty_altr_idxs == {0, 1, 2} + assert cleaned_profile.no_rank_altr_idxs == {0} + assert cleaned_profile.nonempty_altr_idxs == {1, 2} assert cleaned_profile.unaltr_idxs == set() From 272e4c9f4dc94deb39a916190124aa062944057f Mon Sep 17 00:00:00 2001 From: Grace Gibson Date: Wed, 19 Aug 2026 09:52:14 -0500 Subject: [PATCH 4/4] fix _is_equiv fxns rules and include no_rank_altr_idxs in check --- .../cleaning/rank_profiles_cleaning.py | 34 +++++++------------ .../test_condense_ranked_profile.py | 8 ++--- 2 files changed, 16 insertions(+), 26 deletions(-) diff --git a/src/votekit/cleaning/rank_profiles_cleaning.py b/src/votekit/cleaning/rank_profiles_cleaning.py index cfb0c98b..88d74bfe 100644 --- a/src/votekit/cleaning/rank_profiles_cleaning.py +++ b/src/votekit/cleaning/rank_profiles_cleaning.py @@ -318,7 +318,7 @@ def _is_equiv_to_condensed(ranking: pd.Series) -> bool: """ Returns True if the given ranking is equivalent to its condensed form. It is equivalent if the rankings are identical, or if the original ranking only has trailing empty frozensets - in its ranking after some listed candidate. + or tilde frozensets in its ranking after some listed candidate. Args: ranking (pd.Series): Ranking to check. @@ -326,8 +326,8 @@ def _is_equiv_to_condensed(ranking: pd.Series) -> bool: Returns: bool: True if the given ranking is equivalent to its condensed form. """ - if all(cs == frozenset() for cs in ranking): - return False + if all(cs != frozenset() for cs in ranking): + return True for i, cand_set in enumerate(ranking): if cand_set != frozenset(): @@ -384,13 +384,14 @@ def condense_rank_profile( additional_unaltr_idxs = set( [ i - for i in condensed_profile.nonempty_altr_idxs + for i in (condensed_profile.nonempty_altr_idxs | condensed_profile.no_rank_altr_idxs) if _is_equiv_to_condensed(ranking_df.loc[i]) # type: ignore[arg-type] ] ) new_unaltr_idxs = condensed_profile.unaltr_idxs | additional_unaltr_idxs new_nonempty_altr_idxs = condensed_profile.nonempty_altr_idxs.difference(additional_unaltr_idxs) + new_no_rank_altr_idxs = condensed_profile.no_rank_altr_idxs.difference(additional_unaltr_idxs) return CleanedRankProfile( df=condensed_profile.df, @@ -399,7 +400,7 @@ def condense_rank_profile( parent_profile=profile, df_index_column=condensed_profile.df_index_column, no_wt_altr_idxs=condensed_profile.no_wt_altr_idxs, - no_rank_altr_idxs=condensed_profile.no_rank_altr_idxs, + no_rank_altr_idxs=new_no_rank_altr_idxs, nonempty_altr_idxs=new_nonempty_altr_idxs, unaltr_idxs=new_unaltr_idxs, ) @@ -409,8 +410,8 @@ def _is_equiv_for_remove_and_condense(removed: CandidateList, ranking: pd.Series """ Returns True if the given ranking is equivalent to its removed and condensed form. It is equivalent if the ranking has no candidate in the removed list and either no empty - frozensets or only trailing ones. If its has internal empty frozensets or any candidate - in the removed list, it is not equivalent. + frozensets or only trailing ones. Tilde frozensets can also be trailing. If its has internal + empty frozensets or any candidate in the removed list, it is not equivalent. Args: removed (list[Candidate] | list[str] | list[int]): Candidates to be removed. @@ -430,19 +431,7 @@ def _is_equiv_for_remove_and_condense(removed: CandidateList, ranking: pd.Series ): return False - if all(c_set != frozenset() for c_set in ranking): - return True - - for i, cand_set in enumerate(ranking): - if cand_set != frozenset(): - continue - - if all(cs == frozenset() for cs in ranking[i:]): - return True - - return False - - return True + return _is_equiv_to_condensed(ranking) def remove_and_condense_rank_profile( @@ -507,7 +496,7 @@ def remove_and_condense_rank_profile( additional_unaltr_idxs = set( [ i - for i in cleaned_profile.nonempty_altr_idxs + for i in (cleaned_profile.nonempty_altr_idxs | cleaned_profile.no_rank_altr_idxs) if _is_equiv_for_remove_and_condense( removed, ranking_df.loc[i], # type: ignore[arg-type] @@ -517,6 +506,7 @@ def remove_and_condense_rank_profile( new_unaltr_idxs = cleaned_profile.unaltr_idxs | additional_unaltr_idxs new_nonempty_altr_idxs = cleaned_profile.nonempty_altr_idxs.difference(additional_unaltr_idxs) + new_no_rank_altr_idxs = cleaned_profile.no_rank_altr_idxs.difference(additional_unaltr_idxs) return CleanedRankProfile( df=cleaned_profile.df, @@ -525,7 +515,7 @@ def remove_and_condense_rank_profile( parent_profile=cleaned_profile.parent_profile, df_index_column=cleaned_profile.df_index_column, no_wt_altr_idxs=cleaned_profile.no_wt_altr_idxs, - no_rank_altr_idxs=cleaned_profile.no_rank_altr_idxs, + no_rank_altr_idxs=new_no_rank_altr_idxs, nonempty_altr_idxs=new_nonempty_altr_idxs, unaltr_idxs=new_unaltr_idxs, ) diff --git a/tests/cleaning/rank_profiles/test_condense_ranked_profile.py b/tests/cleaning/rank_profiles/test_condense_ranked_profile.py index 4e1d75b2..2b268b18 100644 --- a/tests/cleaning/rank_profiles/test_condense_ranked_profile.py +++ b/tests/cleaning/rank_profiles/test_condense_ranked_profile.py @@ -22,9 +22,9 @@ def test_condense_profile(): ) assert cleaned_profile != profile assert cleaned_profile.no_wt_altr_idxs == set() - assert cleaned_profile.no_rank_altr_idxs == {2} + assert cleaned_profile.no_rank_altr_idxs == set() assert cleaned_profile.nonempty_altr_idxs == {0} - assert cleaned_profile.unaltr_idxs == {1} + assert cleaned_profile.unaltr_idxs == {1, 2} def test_condense_profile_idempotent(): @@ -54,5 +54,5 @@ def test_condense_profile_equivalence(): cleaned = condense_rank_profile(profile) assert cleaned.nonempty_altr_idxs == {0} - assert cleaned.no_rank_altr_idxs == {2} - assert cleaned.unaltr_idxs == {1} + assert cleaned.no_rank_altr_idxs == set() + assert cleaned.unaltr_idxs == {1, 2}