feat: add profile subtraction utility

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

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,
+    )

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]