Feat/remove ballots with cand

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

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()
 
@@ -314,16 +318,16 @@ 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.
 
     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():
@@ -380,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,
@@ -395,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,
     )
@@ -405,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.
@@ -426,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(
@@ -503,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]
@@ -513,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,
@@ -521,7 +515,93 @@ 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,
     )
+
+
+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'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
+            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[ballots_to_remove].index)
+
+    if remove_empty_ballots:
+        mask = cleaned_df[ranking_cols].map(lambda x: x == frozenset({"~"})).all(axis=1)
+        cleaned_df = cleaned_df[~mask]
+
+    if remove_zero_weight_ballots:
+        cleaned_df = cleaned_df[cleaned_df["Weight"] > 0]
+
+    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=set(),
+        no_rank_altr_idxs=removed_ballot_idxs,
+        nonempty_altr_idxs=set(),
+        unaltr_idxs=set(unaltered_idxs),
+    )

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}

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}

tests/cleaning/rank_profiles/test_remove_ballots_with_cand_rank_profile.py

@@ -0,0 +1,79 @@
+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 == 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():
+    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 == 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():
+    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 == 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():
+    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)

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()