From 1ce3f5d031cb42bb71c3e5830614214f80eb9ab9 Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 15 Jul 2025 11:21:20 -0400 Subject: [PATCH 01/38] Initial commit for 3.3.0-stv-animation branch. Add animations module with empty STVAnimation class. --- src/votekit/animations.py | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 src/votekit/animations.py diff --git a/src/votekit/animations.py b/src/votekit/animations.py new file mode 100644 index 00000000..ed04eac9 --- /dev/null +++ b/src/votekit/animations.py @@ -0,0 +1,6 @@ +class STVAnimation(): + def __init__(self): + pass + + def render(self): + pass From 5a2c8a742a0456587603ece675cae1adfa4d860e Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 15 Jul 2025 15:45:31 -0400 Subject: [PATCH 02/38] index on 3.3.0-stv-animation: 1ce3f5d Initial commit for 3.3.0-stv-animation branch. Add animations module with empty STVAnimation class. --- src/votekit/animations.py | 472 +++++++++++++++++++++++++++++++++++++- 1 file changed, 468 insertions(+), 4 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index ed04eac9..6eb73955 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,6 +1,470 @@ +from copy import deepcopy +from manim import * + +class ElectionScene(Scene): + colors = [color.ManimColor(hex) for hex in [ + "#16DEBD", + "#163EDE", + "#9F34F6", + "#FF6F00", + "#8F560C", + "#E2AD00", + "#8AD412"]] + bar_color = LIGHT_GRAY + bar_opacity = 1 + win_bar_color = GREEN + eliminated_bar_color = RED + ghost_opacity = 0.3 + ticker_tape_height = 2 + offscreen_sentinel = '__offscreen__' + offscreen_candidate_color = GRAY + title_font_size = 48 + + def __init__(self, candidates, rounds, title=None): + Scene.__init__(self) + self.candidates = candidates + self.rounds = rounds + self.title = title + + + self.width = 8 + self.bar_height = 3.5 / len(self.candidates) + self.font_size = 3*40 / len(self.candidates) + self.bar_buffer_size = 1 / len(self.candidates) + self.strikethrough_width = self.font_size / 5 + self.max_support = 1.1 * max([round['quota'] for round in self.rounds]) + + self.quota_line = None + self.ticker_tape_line = None + + + def construct(self): + if self.title is not None: + self.__draw_title__(self.title) + self.__draw_initial_bars__() + self.__initialize_ticker_tape__() + + # Go round by round and animate the events + for round_number, round in enumerate(self.rounds): + print(f"Animating round: {round['event']}, {round['candidate']}") + + # Remove the candidate from the candidate list + if round['candidate'] == self.offscreen_sentinel: #If the eliminated candidate is offscreen + eliminated_candidate = None + else: + eliminated_candidate = self.candidates[round['candidate']] + self.candidates.pop(round['candidate']) + + self.wait(2) + + # Draw or move the quota line + self.__update_quota_line__(round['quota']) + + self.__ticker_animation_shift__(round) + self.__ticker_animation_highlight__(round) + + if round['event'] == 'elimination': + if eliminated_candidate is None: #Offscreen candidate eliminated + self.__animate_elimination_offscreen__(round) + else: #Onscreen candidate eliminated + self.__animate_elimination__(eliminated_candidate, round) + elif round['event'] == 'win': + self.__animate_win__(eliminated_candidate, round) + else: + raise Exception(f"Event type {round['event']} not recognized.") + + self.wait(2) + + def __draw_title__(self, message): + text = Tex( + r"{7cm}\centering " + message, + tex_environment="minipage" + ).scale_to_fit_width(10) + self.play(Create(text)) + self.wait(3) + self.play(Uncreate(text)) + + def __draw_initial_bars__(self): + # Sort candidates by starting first place votes + sorted_candidates = sorted(self.candidates.keys(), key=lambda x: self.candidates[x]['support'], reverse=True) + + # Create bars + for i, name in enumerate(sorted_candidates): + color = self.colors[i % len(self.colors)] + self.candidates[name]['color'] = color + self.candidates[name]['bars'] = [Rectangle( + width=self.__support_to_bar_width__(self.candidates[name]['support']), + height=self.bar_height, + color=self.bar_color, + fill_color=color, + fill_opacity=self.bar_opacity + )] + + self.candidates[sorted_candidates[0]]['bars'][0].to_edge(UP) #First candidate goes at the top + for i, name in enumerate(sorted_candidates[1:], start=1): + self.candidates[name]['bars'][0].next_to( + self.candidates[sorted_candidates[i-1]]['bars'][0], + DOWN, + buff=self.bar_buffer_size + ).align_to( + self.candidates[sorted_candidates[0]]['bars'][0], + LEFT + ) #The rest of the candidates go below + + # Draw a large black rectangle for the background so that the ticker tape vanishes behind it + background = Rectangle(width = self.camera.frame_width, + height = self.camera.frame_height, + fill_color = BLACK, + color = BLACK, + fill_opacity = 1).shift(UP * self.ticker_tape_height).set_z_index(-1) + + + # Create and place candidate names + for name in sorted_candidates: + candidate = self.candidates[name] + candidate['name_text'] = Text( + name, + font_size=self.font_size, + color=candidate['color'] + ).next_to(candidate['bars'][0], LEFT, buff=0.2) + + + + # Draw the bars and names + self.play(*[FadeIn(self.candidates[name]['name_text']) for name in sorted_candidates], FadeIn(background)) + self.play(*[Create(self.candidates[name]['bars'][0]) for name in sorted_candidates]) + + def __initialize_ticker_tape__(self): + line_length = self.camera.frame_width + ticker_line = Line(start = [-line_length/2, 0., 0.], + end = [line_length/2, 0., 0.]) + ticker_line.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height) + ticker_line.set_z_index(2) #Keep this line in front of the bars and the quota line + self.ticker_tape_line = ticker_line + + for i, round in enumerate(self.rounds): + new_message = Text( + round['message'], + font_size = 24, + color=DARK_GRAY) + round['ticker_text'] = new_message + if i == 0: + new_message.to_edge(DOWN,buff=0).shift(DOWN) + else: + new_message.next_to(self.rounds[i-1]['ticker_text'], DOWN) + new_message.set_z_index(-2) + + self.play(Create(ticker_line)) + self.play(*[Create(round['ticker_text']) for round in self.rounds]) + + def __ticker_animation_shift__(self, round): + shift_to_round = round['ticker_text'].animate.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height/3) + drag_other_messages = [ + MaintainPositionRelativeTo(other_round['ticker_text'], round['ticker_text']) + for other_round in self.rounds if not other_round == round + ] + self.play(shift_to_round, *drag_other_messages) + + def __ticker_animation_highlight__(self, round): + highlight_message = round['ticker_text'].animate.set_color(WHITE) + unhighlight_other_messages = [ + other_round['ticker_text'].animate.set_color(DARK_GRAY) + for other_round in self.rounds if not other_round == round + ] + self.play(highlight_message, *unhighlight_other_messages) + + + + def __update_quota_line__(self, quota): + some_candidate = list(self.candidates.values())[0] + + if not self.quota_line: + line_top = self.camera.frame_height / 2 + line_bottom = self.ticker_tape_line.get_top()[1] + self.quota_line = Line( + start = [0., line_top, 0.], + end = [0., line_bottom, 0.], + color=self.win_bar_color) + self.quota_line.align_to(some_candidate['bars'][0], LEFT) + self.quota_line.shift((self.width * quota/self.max_support) * RIGHT) + self.quota_line.set_z_index(1) # Keep the quota line in the front + + quota_text = Text("Threshold", font_size = 24).next_to(self.quota_line).align_to(self.quota_line, UP).shift(DOWN / 2) + self.play(Create(self.quota_line), Create(quota_text)) + self.wait(2) + self.play(Uncreate(quota_text)) + else: + self.play(self.quota_line.animate.align_to(some_candidate['bars'][0], LEFT).shift((self.width * quota/self.max_support) * RIGHT)) + + + def __animate_win__(self, from_candidate, round): + assert from_candidate is not None + + #Box the winner's name + winner_box = SurroundingRectangle(from_candidate['name_text'], color=GREEN, buff=0.1) + + #Save the winner's votes as green votes + winner_rectangle = Rectangle( + width=self.__support_to_bar_width__(round['quota']), + height=self.bar_height, + color=self.bar_color, + fill_color=self.win_bar_color, + fill_opacity=self.bar_opacity + ).align_to(from_candidate['bars'][0], LEFT).align_to(from_candidate['bars'][0], UP) + + # Create short bars that will represent the transferred votes + destinations = round['support_transferred'] + candidate_color = from_candidate['color'] + old_bars = from_candidate['bars'] + new_bars = [] + transformations = [] + for destination, votes in destinations.items(): + if votes <= 0: + continue + sub_bar = Rectangle( + width = self.__support_to_bar_width__(votes), + height=self.bar_height, + color=self.bar_color, + fill_color=candidate_color, + fill_opacity=self.bar_opacity) + self.candidates[destination]['support'] += votes + new_bars.append(sub_bar) + transformations.append( + new_bars[-1].animate.next_to( + self.candidates[destination]['bars'][-1], + RIGHT, + buff=0) + ) #The sub-bars will move to be next to the bars of their destination candidates + self.candidates[destination]['bars'].append(sub_bar) + # Create a final short bar representing the exhausted votes + exhausted_votes = from_candidate['support'] - round['quota'] - np.sum(list(destinations.values())) + exhausted_bar = Rectangle( + width = self.__support_to_bar_width__(exhausted_votes), + height=self.bar_height, + color=self.bar_color, + fill_color=candidate_color, + fill_opacity=self.bar_opacity) + assert self.quota_line is not None + exhausted_bar.align_to(self.quota_line, LEFT).align_to(old_bars[0], UP) + + # The short bars should start at the right side of the old bar + if len(new_bars) > 0: + new_bars[0].align_to(old_bars[-1], RIGHT).align_to(old_bars[-1],UP) + for i, sub_bar in enumerate(new_bars[1:], start=1): + sub_bar.next_to(new_bars[i-1], LEFT, buff=0) + + + # Animate the box around the candidate name and the message text + self.play(Create(winner_box)) + # Animate the splitting of the old bar into sub-bars + self.play(*[FadeOut(bar) for bar in old_bars], + *[FadeIn(bar) for bar in new_bars], + FadeIn(exhausted_bar), + FadeIn(winner_rectangle)) + # Animate moving the sub-bars to the destination bars, and the destruction of the exhausted votes + transformations.append(Uncreate(exhausted_bar)) + if len(transformations) > 0: + self.play(*transformations) + + + def __animate_elimination__(self, from_candidate, round): + destinations = round['support_transferred'] + + #Cross out the candidate name + cross = Line( + from_candidate['name_text'].get_left(), + from_candidate['name_text'].get_right(), + color=RED, + ) + cross.set_stroke(width=self.strikethrough_width) + self.play(Create(cross)) + + # Create short bars that will replace the candidate's current bars + candidate_color = from_candidate['color'] + old_bars = from_candidate['bars'] + new_bars = [] #The bits to be redistributed + transformations = [] + for destination, votes in destinations.items(): + if votes <= 0: + continue + sub_bar = Rectangle( + width = self.__support_to_bar_width__(votes), + height=self.bar_height, + color=self.bar_color, + fill_color=candidate_color, + fill_opacity=self.bar_opacity) + self.candidates[destination]['support'] += votes + new_bars.append(sub_bar) + transformations.append( + new_bars[-1].animate.next_to( + self.candidates[destination]['bars'][-1], + RIGHT, + buff=0) + ) #The sub-bars will move to be next to the bars of their destination candidates + self.candidates[destination]['bars'].append(sub_bar) + # Create a final short bar representing the exhausted votes + exhausted_votes = from_candidate['support'] - np.sum(list(destinations.values())) + exhausted_bar = Rectangle( + width = self.__support_to_bar_width__(exhausted_votes), + height=self.bar_height, + color=self.bar_color, + fill_color=candidate_color, + fill_opacity=self.bar_opacity) + exhausted_bar.align_to(old_bars[0], LEFT).align_to(old_bars[0], UP) + + if len(new_bars) > 0: + # The short bars should start in the same place as the old bars. Place them. + rightmost_old_bar = old_bars[-1] + new_bars[0].align_to(rightmost_old_bar, RIGHT).align_to(rightmost_old_bar,UP) + for i, sub_bar in enumerate(new_bars[1:], start=1): + sub_bar.next_to(new_bars[i-1], LEFT, buff=0) + + + # Animate the splitting of the old bar into sub-bars + self.play(*[bar.animate.set_fill(self.eliminated_bar_color).set_opacity(self.ghost_opacity) for bar in old_bars], + *[FadeIn(bar) for bar in new_bars], + FadeIn(exhausted_bar)) + # Animate the exhaustion of votes and moving the sub-bars to the destination bars + self.play(Uncreate(exhausted_bar), *transformations) + + + + def __animate_elimination_offscreen__(self, round): + destinations = round['support_transferred'] + + # Create short bars that will replace the candidate's current bars + new_bars = [] #The bits to be redistributed + transformations = [] + for destination, votes in destinations.items(): + if votes <= 0: + continue + sub_bar = Rectangle( + width = self.__support_to_bar_width__(votes), + height=self.bar_height, + color=self.bar_color, + fill_color=self.offscreen_candidate_color, + fill_opacity=self.bar_opacity, + ) + self.candidates[destination]['support'] += votes + new_bars.append(sub_bar) + transformations.append( + new_bars[-1].animate.next_to( + self.candidates[destination]['bars'][-1], + RIGHT, + buff=0) + ) #The sub-bars will move to be next to the bars of their destination candidates + self.candidates[destination]['bars'].append(sub_bar) + + for bar in new_bars: + bar.to_edge(DOWN).shift((self.bar_height + 2)*DOWN) + + # Animate the exhaustion of votes and moving the sub-bars to the destination bars + self.play(*transformations) + + + def __rescale_bars__(self): + #Re-scale the bars so they're not out of frame + self.max_support = max([candidate['support'] for candidate in self.candidates.values()]) + + transformations = [] + for candidate in self.candidates.values(): + old_bar = candidate['bar'] + new_bar = Rectangle( + width=self.__support_to_bar_width__(candidate['support']), + height=self.bar_height, + color=self.bar_color, + fill_color=candidate['color'], + fill_opacity=self.bar_opacity).next_to(candidate['name_text'], RIGHT) + bar_shortening_transformation = Transform( + old_bar, + new_bar, + ) + transformations.append(bar_shortening_transformation) + self.play(*[transformations]) + + def __support_to_bar_width__(self, support): + return self.width * support / self.max_support + + + class STVAnimation(): - def __init__(self): - pass + def __init__(self, election): + self.candidates = self.__make_candidate_dict__(election) + self.rounds = self.__make_event_list__(election) + + def __make_candidate_dict__(self, election): #this creates the candidate dictionary + viz_candidates = {} + for fset in election.get_remaining(0): + name, = fset + viz_candidates[name] = {'support': int(election.election_states[0].scores[name])} + return viz_candidates - def render(self): - pass + def __get_transferred_votes__(self, election, candidate, round_number): #this helps create the round dictionary, for transferred points + prev_state = election.election_states[round_number-1] + current_state = election.election_states[round_number] + prev_score = int(prev_state.scores[candidate]) + current_score = int(current_state.scores[candidate]) + return current_score - prev_score + + def __make_event_list__(self, election): + events = [] + for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 + elected_candidates = [] + for fset in election_round.elected: + name, = fset + elected_candidates.append(name) + for name in elected_candidates: #Create a new event for each candidate elected this round + event_type = 'win' + message = f'Round {round_number}: {name} Elected' + support_transferred = {} + if round_number < len(election): #If it's the last round, don't worry about the runoff + for (UH-OH) + + + def __make_round_list__(self, election): + rounds = [0]*(len(election)) + for election_round in range(1, len(election)+1): #I think we start at 1 because 0 is how people start? + if len(election.election_states[r].elected[0]) > 0: #in the case that someone is elected + event = 'win' + candidate = list(election.election_states[r].elected[0])[0] #Is it possible for more than one person to win at the same time? + message = f'Round {r}: {candidate} Elected' + support_transferred = {} + for party in runners: #aka going over all of the parties, then going through all of the voters in each party, to remove any candidates that have won and are now out of the election + if candidate in runners[party]: + runners[party].remove(candidate) + + if r != len(election): + for cand in runners[party]: + support_transferred[cand] = self.__get_transferred_votes__(election, cand, r) + else: + support_transferred = {} + rounds[r-1] = {'event': event, + 'candidate': candidate, + 'support_transferred': support_transferred, + 'quota': election.threshold, + 'message' : message} + + else: #in the case that someone was eliminated + event = 'elimination' + candidate = list(election.election_states[r].eliminated[0])[0] + message = f'Round {r}: {candidate} Eliminated' + support_transferred = {} + for party in runners: + if candidate in runners[party]: + runners[party].remove(candidate) + + if r!= len(election): + for cand in runners[party]: + support_transferred[cand] = self.__get_transferred_votes__(election, cand, r) + else: + support_transferred = {} + rounds[r-1] = {'event': event, + 'candidate': candidate, + 'support_transferred': support_transferred, + 'quota': election.threshold, + 'message': message} + return rounds + + def render(self, preview=False): + manimation = ElectionScene(self.candidates, self.rounds) + manimation.render(preview=preview) From 82a50ee819f371d6e66a6bb2e11b8c066da1d155 Mon Sep 17 00:00:00 2001 From: prismika Date: Wed, 16 Jul 2025 09:45:43 -0400 Subject: [PATCH 03/38] Collect animation code in votekit.animation module. --- src/votekit/animations.py | 121 ++++++++++++++++++-------------------- 1 file changed, 58 insertions(+), 63 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 6eb73955..cbbec71f 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -26,7 +26,6 @@ def __init__(self, candidates, rounds, title=None): self.rounds = rounds self.title = title - self.width = 8 self.bar_height = 3.5 / len(self.candidates) self.font_size = 3*40 / len(self.candidates) @@ -118,7 +117,6 @@ def __draw_initial_bars__(self): color = BLACK, fill_opacity = 1).shift(UP * self.ticker_tape_height).set_z_index(-1) - # Create and place candidate names for name in sorted_candidates: candidate = self.candidates[name] @@ -128,8 +126,6 @@ def __draw_initial_bars__(self): color=candidate['color'] ).next_to(candidate['bars'][0], LEFT, buff=0.2) - - # Draw the bars and names self.play(*[FadeIn(self.candidates[name]['name_text']) for name in sorted_candidates], FadeIn(background)) self.play(*[Create(self.candidates[name]['bars'][0]) for name in sorted_candidates]) @@ -174,11 +170,11 @@ def __ticker_animation_highlight__(self, round): self.play(highlight_message, *unhighlight_other_messages) - def __update_quota_line__(self, quota): some_candidate = list(self.candidates.values())[0] if not self.quota_line: + # If the quota line doesn't exist yet, draw it. line_top = self.camera.frame_height / 2 line_bottom = self.ticker_tape_line.get_top()[1] self.quota_line = Line( @@ -189,10 +185,8 @@ def __update_quota_line__(self, quota): self.quota_line.shift((self.width * quota/self.max_support) * RIGHT) self.quota_line.set_z_index(1) # Keep the quota line in the front - quota_text = Text("Threshold", font_size = 24).next_to(self.quota_line).align_to(self.quota_line, UP).shift(DOWN / 2) - self.play(Create(self.quota_line), Create(quota_text)) + self.play(Create(self.quota_line)) self.wait(2) - self.play(Uncreate(quota_text)) else: self.play(self.quota_line.animate.align_to(some_candidate['bars'][0], LEFT).shift((self.width * quota/self.max_support) * RIGHT)) @@ -321,7 +315,9 @@ def __animate_elimination__(self, from_candidate, round): # Animate the splitting of the old bar into sub-bars - self.play(*[bar.animate.set_fill(self.eliminated_bar_color).set_opacity(self.ghost_opacity) for bar in old_bars], + self.play(*[ + bar.animate.set_opacity(self.ghost_opacity) + for bar in old_bars], *[FadeIn(bar) for bar in new_bars], FadeIn(exhausted_bar)) # Animate the exhaustion of votes and moving the sub-bars to the destination bars @@ -388,18 +384,16 @@ def __support_to_bar_width__(self, support): class STVAnimation(): - def __init__(self, election): + def __init__(self, election, title = None): self.candidates = self.__make_candidate_dict__(election) self.rounds = self.__make_event_list__(election) + self.title = title def __make_candidate_dict__(self, election): #this creates the candidate dictionary - viz_candidates = {} - for fset in election.get_remaining(0): - name, = fset - viz_candidates[name] = {'support': int(election.election_states[0].scores[name])} + viz_candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} return viz_candidates - def __get_transferred_votes__(self, election, candidate, round_number): #this helps create the round dictionary, for transferred points + def __get_transferred_votes__(self, election, candidate, round_number): prev_state = election.election_states[round_number-1] current_state = election.election_states[round_number] prev_score = int(prev_state.scores[candidate]) @@ -409,62 +403,63 @@ def __get_transferred_votes__(self, election, candidate, round_number): #this he def __make_event_list__(self, election): events = [] for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 + + remaining_candidates = [] + for fset in election_round.remaining: + remaining_candidates += list(fset) + elected_candidates = [] for fset in election_round.elected: + if len(fset) == 0: + break name, = fset elected_candidates.append(name) - for name in elected_candidates: #Create a new event for each candidate elected this round + + eliminated_candidates = [] + for fset in election_round.eliminated: + if len(fset) == 0: + break + name, = fset + eliminated_candidates.append(name) + + for name in elected_candidates[:1]: #TODO This code currently skips all but the first elected candidate. It will take some work to deduce how support transferred if multiple candidates were elected at once. + #Create a new event for each candidate elected this round event_type = 'win' message = f'Round {round_number}: {name} Elected' support_transferred = {} - if round_number < len(election): #If it's the last round, don't worry about the runoff - for (UH-OH) - - - def __make_round_list__(self, election): - rounds = [0]*(len(election)) - for election_round in range(1, len(election)+1): #I think we start at 1 because 0 is how people start? - if len(election.election_states[r].elected[0]) > 0: #in the case that someone is elected - event = 'win' - candidate = list(election.election_states[r].elected[0])[0] #Is it possible for more than one person to win at the same time? - message = f'Round {r}: {candidate} Elected' - support_transferred = {} - for party in runners: #aka going over all of the parties, then going through all of the voters in each party, to remove any candidates that have won and are now out of the election - if candidate in runners[party]: - runners[party].remove(candidate) - - if r != len(election): - for cand in runners[party]: - support_transferred[cand] = self.__get_transferred_votes__(election, cand, r) - else: - support_transferred = {} - rounds[r-1] = {'event': event, - 'candidate': candidate, - 'support_transferred': support_transferred, - 'quota': election.threshold, - 'message' : message} - - else: #in the case that someone was eliminated - event = 'elimination' - candidate = list(election.election_states[r].eliminated[0])[0] - message = f'Round {r}: {candidate} Eliminated' + if round_number < len(election): #If it's the last round, don't worry about the transferred votes + for candidate in remaining_candidates: + if candidate == name: + continue + support_transferred[candidate] = self.__get_transferred_votes__(election, candidate, round_number) + events.append(dict( + event = event_type, + candidate = name, + support_transferred = support_transferred, + quota = election.threshold, + message = message + )) + + for name in eliminated_candidates[:1]: #TODO This code currently skips all but the first eliminated candidate. It will take some work to deduce how support transferred if multiple candidates were eliminated at once. + #Create a new event for each candidate eliminated this round + event_type = 'elimination' + message = f'Round {round_number}: {name} Eliminated' support_transferred = {} - for party in runners: - if candidate in runners[party]: - runners[party].remove(candidate) - - if r!= len(election): - for cand in runners[party]: - support_transferred[cand] = self.__get_transferred_votes__(election, cand, r) - else: - support_transferred = {} - rounds[r-1] = {'event': event, - 'candidate': candidate, - 'support_transferred': support_transferred, - 'quota': election.threshold, - 'message': message} - return rounds + + if round_number < len(election): #If it's the last round, don't worry about the transferred votes + for candidate in remaining_candidates: + if candidate == name: + continue + support_transferred[candidate] = self.__get_transferred_votes__(election, candidate, round_number) + events.append(dict( + event = event_type, + candidate = name, + support_transferred = support_transferred, + quota = election.threshold, + message = message + )) + return events def render(self, preview=False): - manimation = ElectionScene(self.candidates, self.rounds) + manimation = ElectionScene(self.candidates, self.rounds, title=self.title) manimation.render(preview=preview) From eab7683636668c448693eb76ddba931ec395c5ba Mon Sep 17 00:00:00 2001 From: prismika Date: Wed, 23 Jul 2025 14:56:00 -0400 Subject: [PATCH 04/38] Add support for multiple winners in one round. --- src/votekit/animations.py | 263 +++++++++++++++++++++++--------------- 1 file changed, 157 insertions(+), 106 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index cbbec71f..7dec9ba3 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,5 +1,12 @@ from copy import deepcopy from manim import * +from votekit.cleaning import condense_ballot_ranking, remove_cand_from_ballot +from votekit.utils import ballots_by_first_cand +from votekit.elections.election_types.ranking.stv import STV +from typing import Literal, List +from collections import defaultdict +import logging + class ElectionScene(Scene): colors = [color.ManimColor(hex) for hex in [ @@ -38,6 +45,8 @@ def __init__(self, candidates, rounds, title=None): def construct(self): + logging.getLogger("manim").setLevel(logging.WARNING) + if self.title is not None: self.__draw_title__(self.title) self.__draw_initial_bars__() @@ -45,14 +54,16 @@ def construct(self): # Go round by round and animate the events for round_number, round in enumerate(self.rounds): - print(f"Animating round: {round['event']}, {round['candidate']}") + print(f"Animating round: {round['event']}, {round['candidates']}") # Remove the candidate from the candidate list - if round['candidate'] == self.offscreen_sentinel: #If the eliminated candidate is offscreen - eliminated_candidate = None + if round['candidates'] == self.offscreen_sentinel: #If the eliminated candidate is offscreen + eliminated_candidates = None else: - eliminated_candidate = self.candidates[round['candidate']] - self.candidates.pop(round['candidate']) + eliminated_candidates = {} + for name in round['candidates']: + eliminated_candidates[name] = (self.candidates[name]) + self.candidates.pop(name) self.wait(2) @@ -63,12 +74,13 @@ def construct(self): self.__ticker_animation_highlight__(round) if round['event'] == 'elimination': - if eliminated_candidate is None: #Offscreen candidate eliminated + if eliminated_candidates is None: #Offscreen candidate eliminated self.__animate_elimination_offscreen__(round) else: #Onscreen candidate eliminated - self.__animate_elimination__(eliminated_candidate, round) + self.__animate_elimination__(eliminated_candidates, round) elif round['event'] == 'win': - self.__animate_win__(eliminated_candidate, round) + assert eliminated_candidates is not None + self.__animate_win__(eliminated_candidates, round) else: raise Exception(f"Event type {round['event']} not recognized.") @@ -191,77 +203,96 @@ def __update_quota_line__(self, quota): self.play(self.quota_line.animate.align_to(some_candidate['bars'][0], LEFT).shift((self.width * quota/self.max_support) * RIGHT)) - def __animate_win__(self, from_candidate, round): - assert from_candidate is not None + def __animate_win__(self, from_candidates : dict, round): + for c in from_candidates: + print("Candidate", c) + print(f"{round=}") - #Box the winner's name - winner_box = SurroundingRectangle(from_candidate['name_text'], color=GREEN, buff=0.1) - - #Save the winner's votes as green votes - winner_rectangle = Rectangle( - width=self.__support_to_bar_width__(round['quota']), - height=self.bar_height, - color=self.bar_color, - fill_color=self.win_bar_color, - fill_opacity=self.bar_opacity - ).align_to(from_candidate['bars'][0], LEFT).align_to(from_candidate['bars'][0], UP) + #Box the winners' names + winner_boxes = [SurroundingRectangle(from_candidate['name_text'], color=GREEN, buff=0.1) + for from_candidate in from_candidates.values()] - # Create short bars that will represent the transferred votes - destinations = round['support_transferred'] - candidate_color = from_candidate['color'] - old_bars = from_candidate['bars'] - new_bars = [] - transformations = [] - for destination, votes in destinations.items(): - if votes <= 0: - continue - sub_bar = Rectangle( - width = self.__support_to_bar_width__(votes), + + # Animate the box around the candidate name and the message text + self.play(*[Create(box) for box in winner_boxes]) + + # Create and animate the subdivision and redistribution of winners' leftover votes + for from_candidate_name, from_candidate in from_candidates.items(): + old_bars : List[Rectangle] = from_candidate['bars'] + new_bars : List[Rectangle] = [] + transformations : List[Animation] = [] + destinations = round['support_transferred'][from_candidate_name] + candidate_color = from_candidate['color'] + + winner_rectangle = Rectangle( + width=self.__support_to_bar_width__(round['quota']), + height=self.bar_height, + color=self.bar_color, + fill_color=self.win_bar_color, + fill_opacity=self.bar_opacity + ).align_to(from_candidate['bars'][0], LEFT).align_to(from_candidate['bars'][0], UP) + + #Create a sub-bar for each destination + for destination, votes in destinations.items(): + if votes <= 0: + continue + sub_bar = Rectangle( + width = self.__support_to_bar_width__(votes), + height=self.bar_height, + color=self.bar_color, + fill_color=candidate_color, + fill_opacity=self.bar_opacity) + # The first sub-bar should start at the right end of the eliminated candidate's stack. The rest should be arranged to the left of that one + if len(new_bars) == 0: + sub_bar.align_to(from_candidate['bars'][-1], RIGHT).align_to(from_candidate['bars'][-1], UP) + else: + sub_bar.next_to(new_bars[-1], LEFT, buff=0) + new_bars.append(sub_bar) + self.candidates[destination]['support'] += votes + + transformation = sub_bar.animate.next_to( + self.candidates[destination]['bars'][-1], + RIGHT, + buff=0 + ) + transformations.append( + transformation + ) #The sub-bars will move to be next to the bars of their destination candidates + + # Let the new sub-bar be owned by its destination candidate + self.candidates[destination]['bars'] += sub_bar + + + + # Create a final short bar representing the exhausted votes + exhausted_votes = from_candidate['support'] - round['quota'] - np.sum(list(destinations.values())) + exhausted_bar = Rectangle( + width = self.__support_to_bar_width__(exhausted_votes), height=self.bar_height, color=self.bar_color, fill_color=candidate_color, fill_opacity=self.bar_opacity) - self.candidates[destination]['support'] += votes - new_bars.append(sub_bar) - transformations.append( - new_bars[-1].animate.next_to( - self.candidates[destination]['bars'][-1], - RIGHT, - buff=0) - ) #The sub-bars will move to be next to the bars of their destination candidates - self.candidates[destination]['bars'].append(sub_bar) - # Create a final short bar representing the exhausted votes - exhausted_votes = from_candidate['support'] - round['quota'] - np.sum(list(destinations.values())) - exhausted_bar = Rectangle( - width = self.__support_to_bar_width__(exhausted_votes), - height=self.bar_height, - color=self.bar_color, - fill_color=candidate_color, - fill_opacity=self.bar_opacity) - assert self.quota_line is not None - exhausted_bar.align_to(self.quota_line, LEFT).align_to(old_bars[0], UP) - - # The short bars should start at the right side of the old bar - if len(new_bars) > 0: - new_bars[0].align_to(old_bars[-1], RIGHT).align_to(old_bars[-1],UP) - for i, sub_bar in enumerate(new_bars[1:], start=1): - sub_bar.next_to(new_bars[i-1], LEFT, buff=0) + assert self.quota_line is not None + exhausted_bar.align_to(self.quota_line, LEFT).align_to(from_candidate['bars'][-1], UP) + exhausted_bar.set_z_index(1) + + # Animate the splitting of the old bar into the new sub_bars + self.play( + *[FadeOut(bar) for bar in old_bars], + FadeIn(winner_rectangle), + *[FadeIn(bar) for bar in new_bars], + FadeIn(exhausted_bar) + ) + # Animate moving the sub-bars to the destination bars, and the destruction of the exhausted votes + if len(transformations) > 0: + self.play(*transformations, Uncreate(exhausted_bar)) - # Animate the box around the candidate name and the message text - self.play(Create(winner_box)) - # Animate the splitting of the old bar into sub-bars - self.play(*[FadeOut(bar) for bar in old_bars], - *[FadeIn(bar) for bar in new_bars], - FadeIn(exhausted_bar), - FadeIn(winner_rectangle)) - # Animate moving the sub-bars to the destination bars, and the destruction of the exhausted votes - transformations.append(Uncreate(exhausted_bar)) - if len(transformations) > 0: - self.play(*transformations) - def __animate_elimination__(self, from_candidate, round): + def __animate_elimination__(self, from_candidates, round): + print(from_candidates) + from_candidate = list(from_candidates.values())[0] destinations = round['support_transferred'] #Cross out the candidate name @@ -362,7 +393,7 @@ def __rescale_bars__(self): #Re-scale the bars so they're not out of frame self.max_support = max([candidate['support'] for candidate in self.candidates.values()]) - transformations = [] + transformations : List[Animation] = [] for candidate in self.candidates.values(): old_bar = candidate['bar'] new_bar = Rectangle( @@ -387,20 +418,52 @@ class STVAnimation(): def __init__(self, election, title = None): self.candidates = self.__make_candidate_dict__(election) self.rounds = self.__make_event_list__(election) + print(self.rounds) self.title = title def __make_candidate_dict__(self, election): #this creates the candidate dictionary viz_candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} return viz_candidates - def __get_transferred_votes__(self, election, candidate, round_number): - prev_state = election.election_states[round_number-1] + def __get_transferred_votes__(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']): + prev_profile, prev_state = election.get_step(round_number-1) current_state = election.election_states[round_number] - prev_score = int(prev_state.scores[candidate]) - current_score = int(current_state.scores[candidate]) - return current_score - prev_score - def __make_event_list__(self, election): + if event_type == 'elimination': + if len(from_candidates) > 1: + raise ValueError(f'Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported.') + from_candidate = from_candidates[0] + result_dict = {} + for to_candidate in [c for s in current_state.remaining for c in s]: + prev_score = int(prev_state.scores[to_candidate]) + current_score = int(current_state.scores[to_candidate]) + result_dict[to_candidate] = current_score - prev_score + return result_dict + + elif event_type == 'win': + ballots_by_fpv = ballots_by_first_cand(prev_profile) + transfers = {} + for from_candidate in from_candidates: + new_ballots = election.transfer( + from_candidate, + prev_state.scores[from_candidate], + ballots_by_fpv[from_candidate], + election.threshold + ) + clean_ballots = [ + condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) + for b in new_ballots + ] + transfer_weights_from_candidate = defaultdict(float) + for ballot in clean_ballots: + if ballot.ranking is not None: + to_candidate, = ballot.ranking[0] + transfer_weights_from_candidate[to_candidate] += ballot.weight + + transfers[from_candidate] = transfer_weights_from_candidate + return transfers + + def __make_event_list__(self, election : STV): events = [] for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 @@ -410,50 +473,38 @@ def __make_event_list__(self, election): elected_candidates = [] for fset in election_round.elected: - if len(fset) == 0: - break - name, = fset - elected_candidates.append(name) + if len(fset) > 0: + name, = fset + elected_candidates.append(name) eliminated_candidates = [] for fset in election_round.eliminated: - if len(fset) == 0: - break - name, = fset - eliminated_candidates.append(name) - - for name in elected_candidates[:1]: #TODO This code currently skips all but the first elected candidate. It will take some work to deduce how support transferred if multiple candidates were elected at once. - #Create a new event for each candidate elected this round + if len(fset) > 0: + name, = fset + eliminated_candidates.append(name) + + if len(elected_candidates) > 0: event_type = 'win' - message = f'Round {round_number}: {name} Elected' + message = f'Round {round_number}: {elected_candidates} Elected' support_transferred = {} - if round_number < len(election): #If it's the last round, don't worry about the transferred votes - for candidate in remaining_candidates: - if candidate == name: - continue - support_transferred[candidate] = self.__get_transferred_votes__(election, candidate, round_number) + if round_number == len(election): #If it's the last round, don't worry about the transferred votes + support_transferred = {cand : {} for cand in elected_candidates} + else: + support_transferred = self.__get_transferred_votes__(election, round_number, elected_candidates, 'win') events.append(dict( event = event_type, - candidate = name, + candidates = elected_candidates, support_transferred = support_transferred, quota = election.threshold, message = message )) - - for name in eliminated_candidates[:1]: #TODO This code currently skips all but the first eliminated candidate. It will take some work to deduce how support transferred if multiple candidates were eliminated at once. - #Create a new event for each candidate eliminated this round + elif len(eliminated_candidates) > 0: event_type = 'elimination' - message = f'Round {round_number}: {name} Eliminated' - support_transferred = {} - - if round_number < len(election): #If it's the last round, don't worry about the transferred votes - for candidate in remaining_candidates: - if candidate == name: - continue - support_transferred[candidate] = self.__get_transferred_votes__(election, candidate, round_number) + message = f'Round {round_number}: {eliminated_candidates} Eliminated' + support_transferred = self.__get_transferred_votes__(election, round_number, eliminated_candidates, 'elimination') events.append(dict( event = event_type, - candidate = name, + candidates = eliminated_candidates, support_transferred = support_transferred, quota = election.threshold, message = message From 927c5b1a85c801cd15ae5147750ebe48583f8afe Mon Sep 17 00:00:00 2001 From: prismika Date: Wed, 23 Jul 2025 16:06:29 -0400 Subject: [PATCH 05/38] Clean up the animation code. Specifically, - Remove double underscores - Remove wildcard imports - Change votekit imports to local imports - Minor changes --- src/votekit/animations.py | 127 +++++++++++++++++++------------------- 1 file changed, 65 insertions(+), 62 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 7dec9ba3..845fb0bb 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,15 +1,16 @@ from copy import deepcopy -from manim import * -from votekit.cleaning import condense_ballot_ranking, remove_cand_from_ballot -from votekit.utils import ballots_by_first_cand -from votekit.elections.election_types.ranking.stv import STV +import manim +from manim import Animation, Transform, Rectangle, SurroundingRectangle, Line, Create, Uncreate, FadeIn, FadeOut, Text, UP, DOWN, LEFT, RIGHT +from .cleaning import condense_ballot_ranking, remove_cand_from_ballot +from .utils import ballots_by_first_cand +from .elections.election_types.ranking.stv import STV from typing import Literal, List from collections import defaultdict import logging -class ElectionScene(Scene): - colors = [color.ManimColor(hex) for hex in [ +class ElectionScene(manim.Scene): + colors = [manim.color.ManimColor(hex) for hex in [ "#16DEBD", "#163EDE", "#9F34F6", @@ -17,18 +18,18 @@ class ElectionScene(Scene): "#8F560C", "#E2AD00", "#8AD412"]] - bar_color = LIGHT_GRAY + bar_color = manim.LIGHT_GRAY bar_opacity = 1 - win_bar_color = GREEN - eliminated_bar_color = RED + win_bar_color = manim.GREEN + eliminated_bar_color = manim.RED ghost_opacity = 0.3 ticker_tape_height = 2 offscreen_sentinel = '__offscreen__' - offscreen_candidate_color = GRAY + offscreen_candidate_color = manim.GRAY title_font_size = 48 def __init__(self, candidates, rounds, title=None): - Scene.__init__(self) + super().__init__() self.candidates = candidates self.rounds = rounds self.title = title @@ -36,8 +37,9 @@ def __init__(self, candidates, rounds, title=None): self.width = 8 self.bar_height = 3.5 / len(self.candidates) self.font_size = 3*40 / len(self.candidates) + self.bar_opacity = 1 self.bar_buffer_size = 1 / len(self.candidates) - self.strikethrough_width = self.font_size / 5 + self.strikethrough_thickness = self.font_size / 5 self.max_support = 1.1 * max([round['quota'] for round in self.rounds]) self.quota_line = None @@ -48,9 +50,9 @@ def construct(self): logging.getLogger("manim").setLevel(logging.WARNING) if self.title is not None: - self.__draw_title__(self.title) - self.__draw_initial_bars__() - self.__initialize_ticker_tape__() + self._draw_title(self.title) + self._draw_initial_bars() + self._initialize_ticker_tape() # Go round by round and animate the events for round_number, round in enumerate(self.rounds): @@ -68,26 +70,26 @@ def construct(self): self.wait(2) # Draw or move the quota line - self.__update_quota_line__(round['quota']) + self._update_quota_line(round['quota']) - self.__ticker_animation_shift__(round) - self.__ticker_animation_highlight__(round) + self._ticker_animation_shift(round) + self._ticker_animation_highlight(round) if round['event'] == 'elimination': if eliminated_candidates is None: #Offscreen candidate eliminated - self.__animate_elimination_offscreen__(round) + self._animate_elimination_offscreen(round) else: #Onscreen candidate eliminated - self.__animate_elimination__(eliminated_candidates, round) + self._animate_elimination(eliminated_candidates, round) elif round['event'] == 'win': assert eliminated_candidates is not None - self.__animate_win__(eliminated_candidates, round) + self._animate_win(eliminated_candidates, round) else: raise Exception(f"Event type {round['event']} not recognized.") self.wait(2) - def __draw_title__(self, message): - text = Tex( + def _draw_title(self, message): + text = manim.Tex( r"{7cm}\centering " + message, tex_environment="minipage" ).scale_to_fit_width(10) @@ -95,7 +97,7 @@ def __draw_title__(self, message): self.wait(3) self.play(Uncreate(text)) - def __draw_initial_bars__(self): + def _draw_initial_bars(self): # Sort candidates by starting first place votes sorted_candidates = sorted(self.candidates.keys(), key=lambda x: self.candidates[x]['support'], reverse=True) @@ -104,7 +106,7 @@ def __draw_initial_bars__(self): color = self.colors[i % len(self.colors)] self.candidates[name]['color'] = color self.candidates[name]['bars'] = [Rectangle( - width=self.__support_to_bar_width__(self.candidates[name]['support']), + width=self._support_to_bar_width(self.candidates[name]['support']), height=self.bar_height, color=self.bar_color, fill_color=color, @@ -125,8 +127,8 @@ def __draw_initial_bars__(self): # Draw a large black rectangle for the background so that the ticker tape vanishes behind it background = Rectangle(width = self.camera.frame_width, height = self.camera.frame_height, - fill_color = BLACK, - color = BLACK, + fill_color = manim.BLACK, + color = manim.BLACK, fill_opacity = 1).shift(UP * self.ticker_tape_height).set_z_index(-1) # Create and place candidate names @@ -142,7 +144,7 @@ def __draw_initial_bars__(self): self.play(*[FadeIn(self.candidates[name]['name_text']) for name in sorted_candidates], FadeIn(background)) self.play(*[Create(self.candidates[name]['bars'][0]) for name in sorted_candidates]) - def __initialize_ticker_tape__(self): + def _initialize_ticker_tape(self): line_length = self.camera.frame_width ticker_line = Line(start = [-line_length/2, 0., 0.], end = [line_length/2, 0., 0.]) @@ -154,7 +156,7 @@ def __initialize_ticker_tape__(self): new_message = Text( round['message'], font_size = 24, - color=DARK_GRAY) + color=manim.DARK_GRAY) round['ticker_text'] = new_message if i == 0: new_message.to_edge(DOWN,buff=0).shift(DOWN) @@ -165,30 +167,31 @@ def __initialize_ticker_tape__(self): self.play(Create(ticker_line)) self.play(*[Create(round['ticker_text']) for round in self.rounds]) - def __ticker_animation_shift__(self, round): + def _ticker_animation_shift(self, round): shift_to_round = round['ticker_text'].animate.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height/3) drag_other_messages = [ - MaintainPositionRelativeTo(other_round['ticker_text'], round['ticker_text']) + manim.MaintainPositionRelativeTo(other_round['ticker_text'], round['ticker_text']) for other_round in self.rounds if not other_round == round ] self.play(shift_to_round, *drag_other_messages) - def __ticker_animation_highlight__(self, round): - highlight_message = round['ticker_text'].animate.set_color(WHITE) + def _ticker_animation_highlight(self, round): + highlight_message = round['ticker_text'].animate.set_color(manim.WHITE) unhighlight_other_messages = [ - other_round['ticker_text'].animate.set_color(DARK_GRAY) + other_round['ticker_text'].animate.set_color(manim.DARK_GRAY) for other_round in self.rounds if not other_round == round ] self.play(highlight_message, *unhighlight_other_messages) - def __update_quota_line__(self, quota): + def _update_quota_line(self, quota): some_candidate = list(self.candidates.values())[0] if not self.quota_line: # If the quota line doesn't exist yet, draw it. - line_top = self.camera.frame_height / 2 + assert self.ticker_tape_line is not None line_bottom = self.ticker_tape_line.get_top()[1] + line_top = self.camera.frame_height / 2 self.quota_line = Line( start = [0., line_top, 0.], end = [0., line_bottom, 0.], @@ -203,13 +206,13 @@ def __update_quota_line__(self, quota): self.play(self.quota_line.animate.align_to(some_candidate['bars'][0], LEFT).shift((self.width * quota/self.max_support) * RIGHT)) - def __animate_win__(self, from_candidates : dict, round): + def _animate_win(self, from_candidates : dict, round): for c in from_candidates: print("Candidate", c) print(f"{round=}") #Box the winners' names - winner_boxes = [SurroundingRectangle(from_candidate['name_text'], color=GREEN, buff=0.1) + winner_boxes = [SurroundingRectangle(from_candidate['name_text'], color=manim.GREEN, buff=0.1) for from_candidate in from_candidates.values()] @@ -224,8 +227,8 @@ def __animate_win__(self, from_candidates : dict, round): destinations = round['support_transferred'][from_candidate_name] candidate_color = from_candidate['color'] - winner_rectangle = Rectangle( - width=self.__support_to_bar_width__(round['quota']), + winner_bar = Rectangle( + width=self._support_to_bar_width(round['quota']), height=self.bar_height, color=self.bar_color, fill_color=self.win_bar_color, @@ -237,7 +240,7 @@ def __animate_win__(self, from_candidates : dict, round): if votes <= 0: continue sub_bar = Rectangle( - width = self.__support_to_bar_width__(votes), + width = self._support_to_bar_width(votes), height=self.bar_height, color=self.bar_color, fill_color=candidate_color, @@ -265,9 +268,9 @@ def __animate_win__(self, from_candidates : dict, round): # Create a final short bar representing the exhausted votes - exhausted_votes = from_candidate['support'] - round['quota'] - np.sum(list(destinations.values())) + exhausted_votes = from_candidate['support'] - round['quota'] - sum(list(destinations.values())) exhausted_bar = Rectangle( - width = self.__support_to_bar_width__(exhausted_votes), + width = self._support_to_bar_width(exhausted_votes), height=self.bar_height, color=self.bar_color, fill_color=candidate_color, @@ -279,7 +282,7 @@ def __animate_win__(self, from_candidates : dict, round): # Animate the splitting of the old bar into the new sub_bars self.play( *[FadeOut(bar) for bar in old_bars], - FadeIn(winner_rectangle), + FadeIn(winner_bar), *[FadeIn(bar) for bar in new_bars], FadeIn(exhausted_bar) ) @@ -290,7 +293,7 @@ def __animate_win__(self, from_candidates : dict, round): - def __animate_elimination__(self, from_candidates, round): + def _animate_elimination(self, from_candidates, round): print(from_candidates) from_candidate = list(from_candidates.values())[0] destinations = round['support_transferred'] @@ -299,9 +302,9 @@ def __animate_elimination__(self, from_candidates, round): cross = Line( from_candidate['name_text'].get_left(), from_candidate['name_text'].get_right(), - color=RED, + color=manim.RED, ) - cross.set_stroke(width=self.strikethrough_width) + cross.set_stroke(width=self.strikethrough_thickness) self.play(Create(cross)) # Create short bars that will replace the candidate's current bars @@ -313,7 +316,7 @@ def __animate_elimination__(self, from_candidates, round): if votes <= 0: continue sub_bar = Rectangle( - width = self.__support_to_bar_width__(votes), + width = self._support_to_bar_width(votes), height=self.bar_height, color=self.bar_color, fill_color=candidate_color, @@ -328,9 +331,9 @@ def __animate_elimination__(self, from_candidates, round): ) #The sub-bars will move to be next to the bars of their destination candidates self.candidates[destination]['bars'].append(sub_bar) # Create a final short bar representing the exhausted votes - exhausted_votes = from_candidate['support'] - np.sum(list(destinations.values())) + exhausted_votes = from_candidate['support'] - sum(list(destinations.values())) exhausted_bar = Rectangle( - width = self.__support_to_bar_width__(exhausted_votes), + width = self._support_to_bar_width(exhausted_votes), height=self.bar_height, color=self.bar_color, fill_color=candidate_color, @@ -356,7 +359,7 @@ def __animate_elimination__(self, from_candidates, round): - def __animate_elimination_offscreen__(self, round): + def _animate_elimination_offscreen(self, round): destinations = round['support_transferred'] # Create short bars that will replace the candidate's current bars @@ -366,7 +369,7 @@ def __animate_elimination_offscreen__(self, round): if votes <= 0: continue sub_bar = Rectangle( - width = self.__support_to_bar_width__(votes), + width = self._support_to_bar_width(votes), height=self.bar_height, color=self.bar_color, fill_color=self.offscreen_candidate_color, @@ -389,7 +392,7 @@ def __animate_elimination_offscreen__(self, round): self.play(*transformations) - def __rescale_bars__(self): + def _rescale_bars(self): #Re-scale the bars so they're not out of frame self.max_support = max([candidate['support'] for candidate in self.candidates.values()]) @@ -397,7 +400,7 @@ def __rescale_bars__(self): for candidate in self.candidates.values(): old_bar = candidate['bar'] new_bar = Rectangle( - width=self.__support_to_bar_width__(candidate['support']), + width=self._support_to_bar_width(candidate['support']), height=self.bar_height, color=self.bar_color, fill_color=candidate['color'], @@ -409,23 +412,23 @@ def __rescale_bars__(self): transformations.append(bar_shortening_transformation) self.play(*[transformations]) - def __support_to_bar_width__(self, support): + def _support_to_bar_width(self, support): return self.width * support / self.max_support class STVAnimation(): def __init__(self, election, title = None): - self.candidates = self.__make_candidate_dict__(election) - self.rounds = self.__make_event_list__(election) + self.candidates = self._make_candidate_dict(election) + self.rounds = self._make_event_list(election) print(self.rounds) self.title = title - def __make_candidate_dict__(self, election): #this creates the candidate dictionary + def _make_candidate_dict(self, election): #this creates the candidate dictionary viz_candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} return viz_candidates - def __get_transferred_votes__(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']): + def _get_transferred_votes(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']): prev_profile, prev_state = election.get_step(round_number-1) current_state = election.election_states[round_number] @@ -463,7 +466,7 @@ def __get_transferred_votes__(self, election : STV, round_number : int, from_can transfers[from_candidate] = transfer_weights_from_candidate return transfers - def __make_event_list__(self, election : STV): + def _make_event_list(self, election : STV): events = [] for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 @@ -490,7 +493,7 @@ def __make_event_list__(self, election : STV): if round_number == len(election): #If it's the last round, don't worry about the transferred votes support_transferred = {cand : {} for cand in elected_candidates} else: - support_transferred = self.__get_transferred_votes__(election, round_number, elected_candidates, 'win') + support_transferred = self._get_transferred_votes(election, round_number, elected_candidates, 'win') events.append(dict( event = event_type, candidates = elected_candidates, @@ -501,7 +504,7 @@ def __make_event_list__(self, election : STV): elif len(eliminated_candidates) > 0: event_type = 'elimination' message = f'Round {round_number}: {eliminated_candidates} Eliminated' - support_transferred = self.__get_transferred_votes__(election, round_number, eliminated_candidates, 'elimination') + support_transferred = self._get_transferred_votes(election, round_number, eliminated_candidates, 'elimination') events.append(dict( event = event_type, candidates = eliminated_candidates, From 875be7b0eb87c28e1c0aebffd26458c767983843 Mon Sep 17 00:00:00 2001 From: prismika Date: Wed, 23 Jul 2025 17:32:45 -0400 Subject: [PATCH 06/38] Fix static typechecking errors. Add type hints and documentation to all functions. --- src/votekit/animations.py | 190 ++++++++++++++++++++++++++++++-------- 1 file changed, 150 insertions(+), 40 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 845fb0bb..da3234d3 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -10,6 +10,14 @@ class ElectionScene(manim.Scene): + """ + Class for Manim animation of an STV election. This class is instantiated by the class STVAnimation. It should not be instantiated directly. + + Args: + candidates (dict[str,dict]): A dictionary mapping each candidate to a dictionary of attributes of the candidate. + rounds (List[dict]): A list of dictionaries representing the rounds of the election. Each dictionary is a summary of the events in the round it represents. + title (str): A string to be displayed at the beginning of the animation as a title screen. If None, the animation will skip the title screen. + """ colors = [manim.color.ManimColor(hex) for hex in [ "#16DEBD", "#163EDE", @@ -28,7 +36,7 @@ class ElectionScene(manim.Scene): offscreen_candidate_color = manim.GRAY title_font_size = 48 - def __init__(self, candidates, rounds, title=None): + def __init__(self, candidates : dict[str,dict], rounds : List[dict], title : (str | None) = None): super().__init__() self.candidates = candidates self.rounds = rounds @@ -46,7 +54,12 @@ def __init__(self, candidates, rounds, title=None): self.ticker_tape_line = None - def construct(self): + def construct(self) -> None: + """ + Constructs the animation. + """ + + # Manim produces a lot of logging output. Set the logging level to WARNING. logging.getLogger("manim").setLevel(logging.WARNING) if self.title is not None: @@ -56,8 +69,6 @@ def construct(self): # Go round by round and animate the events for round_number, round in enumerate(self.rounds): - print(f"Animating round: {round['event']}, {round['candidates']}") - # Remove the candidate from the candidate list if round['candidates'] == self.offscreen_sentinel: #If the eliminated candidate is offscreen eliminated_candidates = None @@ -88,7 +99,14 @@ def construct(self): self.wait(2) - def _draw_title(self, message): + + def _draw_title(self, message : str) -> None: + """ + Draw the title screen. + + Args: + message (str): String that the title screen will display. + """ text = manim.Tex( r"{7cm}\centering " + message, tex_environment="minipage" @@ -97,7 +115,10 @@ def _draw_title(self, message): self.wait(3) self.play(Uncreate(text)) - def _draw_initial_bars(self): + def _draw_initial_bars(self) -> None: + """ + Instantiate and draw the names and bars for each candidate. + """ # Sort candidates by starting first place votes sorted_candidates = sorted(self.candidates.keys(), key=lambda x: self.candidates[x]['support'], reverse=True) @@ -125,8 +146,10 @@ def _draw_initial_bars(self): ) #The rest of the candidates go below # Draw a large black rectangle for the background so that the ticker tape vanishes behind it - background = Rectangle(width = self.camera.frame_width, - height = self.camera.frame_height, + frame_width = manim.config.frame_width + frame_height = manim.config.frame_height + background = Rectangle(width = frame_width, + height = frame_height, fill_color = manim.BLACK, color = manim.BLACK, fill_opacity = 1).shift(UP * self.ticker_tape_height).set_z_index(-1) @@ -144,10 +167,13 @@ def _draw_initial_bars(self): self.play(*[FadeIn(self.candidates[name]['name_text']) for name in sorted_candidates], FadeIn(background)) self.play(*[Create(self.candidates[name]['bars'][0]) for name in sorted_candidates]) - def _initialize_ticker_tape(self): - line_length = self.camera.frame_width - ticker_line = Line(start = [-line_length/2, 0., 0.], - end = [line_length/2, 0., 0.]) + def _initialize_ticker_tape(self) -> None: + """ + Instantiate and draw the ticker tape line and text. + """ + line_length = manim.config.frame_width + ticker_line = Line(start = LEFT * line_length/2, + end = RIGHT * line_length/2,) ticker_line.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height) ticker_line.set_z_index(2) #Keep this line in front of the bars and the quota line self.ticker_tape_line = ticker_line @@ -167,7 +193,13 @@ def _initialize_ticker_tape(self): self.play(Create(ticker_line)) self.play(*[Create(round['ticker_text']) for round in self.rounds]) - def _ticker_animation_shift(self, round): + def _ticker_animation_shift(self, round : dict) -> None: + """ + Animate the shifting of the ticker tape to display the message for a given round. + + Args: + round (dict): The round whose message will shift into view. + """ shift_to_round = round['ticker_text'].animate.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height/3) drag_other_messages = [ manim.MaintainPositionRelativeTo(other_round['ticker_text'], round['ticker_text']) @@ -175,7 +207,13 @@ def _ticker_animation_shift(self, round): ] self.play(shift_to_round, *drag_other_messages) - def _ticker_animation_highlight(self, round): + def _ticker_animation_highlight(self, round : dict) -> None: + """ + Play an animation graying out all ticker tape message but one. + + Args: + round (dict): The round whose message will be highlighted. + """ highlight_message = round['ticker_text'].animate.set_color(manim.WHITE) unhighlight_other_messages = [ other_round['ticker_text'].animate.set_color(manim.DARK_GRAY) @@ -184,17 +222,22 @@ def _ticker_animation_highlight(self, round): self.play(highlight_message, *unhighlight_other_messages) - def _update_quota_line(self, quota): - some_candidate = list(self.candidates.values())[0] + def _update_quota_line(self, quota : float) -> None: + """ + Update the position of the quota line to reflect the given quota. If no quota line exists, create it and animate its creation. + Args: + quota (float): The threshold number of votes necessary to be elected in the current round. + """ + some_candidate = list(self.candidates.values())[0] if not self.quota_line: # If the quota line doesn't exist yet, draw it. assert self.ticker_tape_line is not None line_bottom = self.ticker_tape_line.get_top()[1] - line_top = self.camera.frame_height / 2 + line_top = manim.config.frame_height / 2 self.quota_line = Line( - start = [0., line_top, 0.], - end = [0., line_bottom, 0.], + start = UP * line_top, + end = UP * line_bottom, color=self.win_bar_color) self.quota_line.align_to(some_candidate['bars'][0], LEFT) self.quota_line.shift((self.width * quota/self.max_support) * RIGHT) @@ -206,11 +249,14 @@ def _update_quota_line(self, quota): self.play(self.quota_line.animate.align_to(some_candidate['bars'][0], LEFT).shift((self.width * quota/self.max_support) * RIGHT)) - def _animate_win(self, from_candidates : dict, round): - for c in from_candidates: - print("Candidate", c) - print(f"{round=}") + def _animate_win(self, from_candidates : dict[str,dict], round : dict) -> None: + """ + Animate a round in which one or more candidates are elected. + Args: + from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates elected this round and the values are dictionaries recording the candidate's attributes. + round (dict): A dictionary recording the events of this round. + """ #Box the winners' names winner_boxes = [SurroundingRectangle(from_candidate['name_text'], color=manim.GREEN, buff=0.1) for from_candidate in from_candidates.values()] @@ -223,7 +269,7 @@ def _animate_win(self, from_candidates : dict, round): for from_candidate_name, from_candidate in from_candidates.items(): old_bars : List[Rectangle] = from_candidate['bars'] new_bars : List[Rectangle] = [] - transformations : List[Animation] = [] + transformations = [] destinations = round['support_transferred'][from_candidate_name] candidate_color = from_candidate['color'] @@ -293,8 +339,14 @@ def _animate_win(self, from_candidates : dict, round): - def _animate_elimination(self, from_candidates, round): - print(from_candidates) + def _animate_elimination(self, from_candidates : dict[str,dict], round : dict) -> None: + """ + Animate a round in which a candidate was eliminated. While the interface supports multiple candidate eliminations in one round for future extensibility, this function currently only supports elimination of one candidate at a time. The from_candidates argument should have exactly one entry. + + Args: + from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates eliminated this round and the values are dictionaries recording the candidate's attributes. This function currently only supports one elimination at a time, so from_candidates should be a dictionary with exactly one entry. + round (dict): A dictionary recording the events of this round. + """ from_candidate = list(from_candidates.values())[0] destinations = round['support_transferred'] @@ -359,7 +411,13 @@ def _animate_elimination(self, from_candidates, round): - def _animate_elimination_offscreen(self, round): + def _animate_elimination_offscreen(self, round : dict) -> None: + """ + Animate a round in which offscreen candidates were eliminated. + + Args: + round (dict): A dictionary recording the events of this round. + """ destinations = round['support_transferred'] # Create short bars that will replace the candidate's current bars @@ -392,11 +450,11 @@ def _animate_elimination_offscreen(self, round): self.play(*transformations) - def _rescale_bars(self): - #Re-scale the bars so they're not out of frame + def _rescale_bars(self) -> None: + """Re-scale the bars so they fit nicely in frame.""" self.max_support = max([candidate['support'] for candidate in self.candidates.values()]) - transformations : List[Animation] = [] + transformations : List[Transform] = [] for candidate in self.candidates.values(): old_bar = candidate['bar'] new_bar = Rectangle( @@ -410,25 +468,62 @@ def _rescale_bars(self): new_bar, ) transformations.append(bar_shortening_transformation) - self.play(*[transformations]) + self.play(*transformations) + + def _support_to_bar_width(self, support : float) -> float: + """ + Convert a number of votes to the width of a bar in manim coordinates representing that many votes. - def _support_to_bar_width(self, support): + Args: + support (float): A number of votes. + + Returns: + float: The width, in manim coordinates, of a bar representing the support. + """ return self.width * support / self.max_support class STVAnimation(): - def __init__(self, election, title = None): + """ + A class which creates round-by-round animations of STV elections. + + Args: + election (STV): An STV election to animate. + title (str): Text to be displayed at the beginning of the animation as a title screen. If None, the title screen will be skipped. + """ + def __init__(self, election : STV, title : (str | None) = None): + # Extract only the salient details from the election. self.candidates = self._make_candidate_dict(election) self.rounds = self._make_event_list(election) - print(self.rounds) self.title = title - def _make_candidate_dict(self, election): #this creates the candidate dictionary - viz_candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} - return viz_candidates + def _make_candidate_dict(self, election : STV) -> dict: + """ + Create the dictionary of candidates and relevant facts about each one. + + Args: + election (STV): An STV election from which to extract the candidates. + + Returns: + dict: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. + """ + candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} + return candidates - def _get_transferred_votes(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']): + def _get_transferred_votes(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']) -> dict[str, dict[str, float]]: + """ + Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. + + Args: + election (STV): The election. + round_number (int): The number of the round in question. + from_candidates (List[str]): A list of the names of the elected or eliminated candidates. + event_type (str): "win" if candidates were elected this round, "elimination" otherwise. + + Returns: + dict[str, dict[str, float]]: A nested dictionary. If d is the return value, c1 was a candidate eliminated this round, and c2 is a remaining candidate, then d[c1][c2] will be the total support transferred this round from c1 to c2. + """ prev_profile, prev_state = election.get_step(round_number-1) current_state = election.election_states[round_number] @@ -466,7 +561,16 @@ def _get_transferred_votes(self, election : STV, round_number : int, from_candid transfers[from_candidate] = transfer_weights_from_candidate return transfers - def _make_event_list(self, election : STV): + def _make_event_list(self, election : STV) -> List[dict]: + """ + Process an STV election into a condensed list of only the salient details from each round. + + Args: + election (STV): The STV election to process. + + Returns: + List[dict]: A list of dictionaries corresponding to the rounds of the election. Each dictionary records salient attributes of the corresponding round. + """ events = [] for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 @@ -514,6 +618,12 @@ def _make_event_list(self, election : STV): )) return events - def render(self, preview=False): + def render(self, preview : bool = False) -> None: + """ + Renders the STV animation using Manim. + + Args: + preview (bool): If true, display the result in a video player immediately upon completion. + """ manimation = ElectionScene(self.candidates, self.rounds, title=self.title) manimation.render(preview=preview) From 2337e2012de410fcf7393df4e4300014fda07dc2 Mon Sep 17 00:00:00 2001 From: prismika Date: Wed, 23 Jul 2025 17:36:55 -0400 Subject: [PATCH 07/38] Rearrange classes and functions for readability. --- src/votekit/animations.py | 294 +++++++++++++++++++------------------- 1 file changed, 148 insertions(+), 146 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index da3234d3..393f5973 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -9,6 +9,154 @@ import logging + +class STVAnimation(): + """ + A class which creates round-by-round animations of STV elections. + + Args: + election (STV): An STV election to animate. + title (str): Text to be displayed at the beginning of the animation as a title screen. If None, the title screen will be skipped. + """ + def __init__(self, election : STV, title : (str | None) = None): + # Extract only the salient details from the election. + self.candidates = self._make_candidate_dict(election) + self.rounds = self._make_event_list(election) + self.title = title + + def _make_candidate_dict(self, election : STV) -> dict: + """ + Create the dictionary of candidates and relevant facts about each one. + + Args: + election (STV): An STV election from which to extract the candidates. + + Returns: + dict: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. + """ + candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} + return candidates + + + def _make_event_list(self, election : STV) -> List[dict]: + """ + Process an STV election into a condensed list of only the salient details from each round. + + Args: + election (STV): The STV election to process. + + Returns: + List[dict]: A list of dictionaries corresponding to the rounds of the election. Each dictionary records salient attributes of the corresponding round. + """ + events = [] + for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 + + remaining_candidates = [] + for fset in election_round.remaining: + remaining_candidates += list(fset) + + elected_candidates = [] + for fset in election_round.elected: + if len(fset) > 0: + name, = fset + elected_candidates.append(name) + + eliminated_candidates = [] + for fset in election_round.eliminated: + if len(fset) > 0: + name, = fset + eliminated_candidates.append(name) + + if len(elected_candidates) > 0: + event_type = 'win' + message = f'Round {round_number}: {elected_candidates} Elected' + support_transferred = {} + if round_number == len(election): #If it's the last round, don't worry about the transferred votes + support_transferred = {cand : {} for cand in elected_candidates} + else: + support_transferred = self._get_transferred_votes(election, round_number, elected_candidates, 'win') + events.append(dict( + event = event_type, + candidates = elected_candidates, + support_transferred = support_transferred, + quota = election.threshold, + message = message + )) + elif len(eliminated_candidates) > 0: + event_type = 'elimination' + message = f'Round {round_number}: {eliminated_candidates} Eliminated' + support_transferred = self._get_transferred_votes(election, round_number, eliminated_candidates, 'elimination') + events.append(dict( + event = event_type, + candidates = eliminated_candidates, + support_transferred = support_transferred, + quota = election.threshold, + message = message + )) + return events + + + def _get_transferred_votes(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']) -> dict[str, dict[str, float]]: + """ + Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. + + Args: + election (STV): The election. + round_number (int): The number of the round in question. + from_candidates (List[str]): A list of the names of the elected or eliminated candidates. + event_type (str): "win" if candidates were elected this round, "elimination" otherwise. + + Returns: + dict[str, dict[str, float]]: A nested dictionary. If d is the return value, c1 was a candidate eliminated this round, and c2 is a remaining candidate, then d[c1][c2] will be the total support transferred this round from c1 to c2. + """ + prev_profile, prev_state = election.get_step(round_number-1) + current_state = election.election_states[round_number] + + if event_type == 'elimination': + if len(from_candidates) > 1: + raise ValueError(f'Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported.') + from_candidate = from_candidates[0] + result_dict = {} + for to_candidate in [c for s in current_state.remaining for c in s]: + prev_score = int(prev_state.scores[to_candidate]) + current_score = int(current_state.scores[to_candidate]) + result_dict[to_candidate] = current_score - prev_score + return result_dict + + elif event_type == 'win': + ballots_by_fpv = ballots_by_first_cand(prev_profile) + transfers = {} + for from_candidate in from_candidates: + new_ballots = election.transfer( + from_candidate, + prev_state.scores[from_candidate], + ballots_by_fpv[from_candidate], + election.threshold + ) + clean_ballots = [ + condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) + for b in new_ballots + ] + transfer_weights_from_candidate = defaultdict(float) + for ballot in clean_ballots: + if ballot.ranking is not None: + to_candidate, = ballot.ranking[0] + transfer_weights_from_candidate[to_candidate] += ballot.weight + + transfers[from_candidate] = transfer_weights_from_candidate + return transfers + + def render(self, preview : bool = False) -> None: + """ + Renders the STV animation using Manim. + + Args: + preview (bool): If true, display the result in a video player immediately upon completion. + """ + manimation = ElectionScene(self.candidates, self.rounds, title=self.title) + manimation.render(preview=preview) + + class ElectionScene(manim.Scene): """ Class for Manim animation of an STV election. This class is instantiated by the class STVAnimation. It should not be instantiated directly. @@ -481,149 +629,3 @@ def _support_to_bar_width(self, support : float) -> float: float: The width, in manim coordinates, of a bar representing the support. """ return self.width * support / self.max_support - - - -class STVAnimation(): - """ - A class which creates round-by-round animations of STV elections. - - Args: - election (STV): An STV election to animate. - title (str): Text to be displayed at the beginning of the animation as a title screen. If None, the title screen will be skipped. - """ - def __init__(self, election : STV, title : (str | None) = None): - # Extract only the salient details from the election. - self.candidates = self._make_candidate_dict(election) - self.rounds = self._make_event_list(election) - self.title = title - - def _make_candidate_dict(self, election : STV) -> dict: - """ - Create the dictionary of candidates and relevant facts about each one. - - Args: - election (STV): An STV election from which to extract the candidates. - - Returns: - dict: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. - """ - candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} - return candidates - - def _get_transferred_votes(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']) -> dict[str, dict[str, float]]: - """ - Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. - - Args: - election (STV): The election. - round_number (int): The number of the round in question. - from_candidates (List[str]): A list of the names of the elected or eliminated candidates. - event_type (str): "win" if candidates were elected this round, "elimination" otherwise. - - Returns: - dict[str, dict[str, float]]: A nested dictionary. If d is the return value, c1 was a candidate eliminated this round, and c2 is a remaining candidate, then d[c1][c2] will be the total support transferred this round from c1 to c2. - """ - prev_profile, prev_state = election.get_step(round_number-1) - current_state = election.election_states[round_number] - - if event_type == 'elimination': - if len(from_candidates) > 1: - raise ValueError(f'Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported.') - from_candidate = from_candidates[0] - result_dict = {} - for to_candidate in [c for s in current_state.remaining for c in s]: - prev_score = int(prev_state.scores[to_candidate]) - current_score = int(current_state.scores[to_candidate]) - result_dict[to_candidate] = current_score - prev_score - return result_dict - - elif event_type == 'win': - ballots_by_fpv = ballots_by_first_cand(prev_profile) - transfers = {} - for from_candidate in from_candidates: - new_ballots = election.transfer( - from_candidate, - prev_state.scores[from_candidate], - ballots_by_fpv[from_candidate], - election.threshold - ) - clean_ballots = [ - condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) - for b in new_ballots - ] - transfer_weights_from_candidate = defaultdict(float) - for ballot in clean_ballots: - if ballot.ranking is not None: - to_candidate, = ballot.ranking[0] - transfer_weights_from_candidate[to_candidate] += ballot.weight - - transfers[from_candidate] = transfer_weights_from_candidate - return transfers - - def _make_event_list(self, election : STV) -> List[dict]: - """ - Process an STV election into a condensed list of only the salient details from each round. - - Args: - election (STV): The STV election to process. - - Returns: - List[dict]: A list of dictionaries corresponding to the rounds of the election. Each dictionary records salient attributes of the corresponding round. - """ - events = [] - for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 - - remaining_candidates = [] - for fset in election_round.remaining: - remaining_candidates += list(fset) - - elected_candidates = [] - for fset in election_round.elected: - if len(fset) > 0: - name, = fset - elected_candidates.append(name) - - eliminated_candidates = [] - for fset in election_round.eliminated: - if len(fset) > 0: - name, = fset - eliminated_candidates.append(name) - - if len(elected_candidates) > 0: - event_type = 'win' - message = f'Round {round_number}: {elected_candidates} Elected' - support_transferred = {} - if round_number == len(election): #If it's the last round, don't worry about the transferred votes - support_transferred = {cand : {} for cand in elected_candidates} - else: - support_transferred = self._get_transferred_votes(election, round_number, elected_candidates, 'win') - events.append(dict( - event = event_type, - candidates = elected_candidates, - support_transferred = support_transferred, - quota = election.threshold, - message = message - )) - elif len(eliminated_candidates) > 0: - event_type = 'elimination' - message = f'Round {round_number}: {eliminated_candidates} Eliminated' - support_transferred = self._get_transferred_votes(election, round_number, eliminated_candidates, 'elimination') - events.append(dict( - event = event_type, - candidates = eliminated_candidates, - support_transferred = support_transferred, - quota = election.threshold, - message = message - )) - return events - - def render(self, preview : bool = False) -> None: - """ - Renders the STV animation using Manim. - - Args: - preview (bool): If true, display the result in a video player immediately upon completion. - """ - manimation = ElectionScene(self.candidates, self.rounds, title=self.title) - manimation.render(preview=preview) From 13785c46314adc8ebed776be9f1a94461ee40b0f Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 24 Jul 2025 16:18:27 -0400 Subject: [PATCH 08/38] Fix bug in which rendering an STVAnimation caused the saved state to mutate. --- src/votekit/animations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 393f5973..9b872668 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,6 +1,6 @@ from copy import deepcopy import manim -from manim import Animation, Transform, Rectangle, SurroundingRectangle, Line, Create, Uncreate, FadeIn, FadeOut, Text, UP, DOWN, LEFT, RIGHT +from manim import Transform, Rectangle, SurroundingRectangle, Line, Create, Uncreate, FadeIn, FadeOut, Text, UP, DOWN, LEFT, RIGHT from .cleaning import condense_ballot_ranking, remove_cand_from_ballot from .utils import ballots_by_first_cand from .elections.election_types.ranking.stv import STV @@ -153,7 +153,7 @@ def render(self, preview : bool = False) -> None: Args: preview (bool): If true, display the result in a video player immediately upon completion. """ - manimation = ElectionScene(self.candidates, self.rounds, title=self.title) + manimation = ElectionScene(deepcopy(self.candidates), deepcopy(self.rounds), title=self.title) manimation.render(preview=preview) From f2f5247278a5c7d10faabc489d6ea01d023b4ae5 Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 24 Jul 2025 16:27:18 -0400 Subject: [PATCH 09/38] Make event messages pretty. --- src/votekit/animations.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 9b872668..f132069e 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -69,7 +69,10 @@ def _make_event_list(self, election : STV) -> List[dict]: if len(elected_candidates) > 0: event_type = 'win' - message = f'Round {round_number}: {elected_candidates} Elected' + elected_candidates_str = elected_candidates[0] + for candidate_name in elected_candidates[1:]: + elected_candidates_str += ", " + candidate_name + message = f'Round {round_number}: {elected_candidates_str} Elected' support_transferred = {} if round_number == len(election): #If it's the last round, don't worry about the transferred votes support_transferred = {cand : {} for cand in elected_candidates} @@ -84,7 +87,10 @@ def _make_event_list(self, election : STV) -> List[dict]: )) elif len(eliminated_candidates) > 0: event_type = 'elimination' - message = f'Round {round_number}: {eliminated_candidates} Eliminated' + eliminated_candidates_str = eliminated_candidates[0] + for candidate_name in eliminated_candidates[1:]: + eliminated_candidates_str += ", " + candidate_name + message = f'Round {round_number}: {eliminated_candidates_str} Eliminated' support_transferred = self._get_transferred_votes(election, round_number, eliminated_candidates, 'elimination') events.append(dict( event = event_type, From 3f472d3b7adb68ce83c76391923a991460bd3a65 Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 24 Jul 2025 17:29:56 -0400 Subject: [PATCH 10/38] Formatting animations.py --- src/votekit/animations.py | 538 ++++++++++++++++++++++---------------- 1 file changed, 310 insertions(+), 228 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index f132069e..2032ecca 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,6 +1,20 @@ from copy import deepcopy import manim -from manim import Transform, Rectangle, SurroundingRectangle, Line, Create, Uncreate, FadeIn, FadeOut, Text, UP, DOWN, LEFT, RIGHT +from manim import ( + Transform, + Rectangle, + SurroundingRectangle, + Line, + Create, + Uncreate, + FadeIn, + FadeOut, + Text, + UP, + DOWN, + LEFT, + RIGHT, +) from .cleaning import condense_ballot_ranking, remove_cand_from_ballot from .utils import ballots_by_first_cand from .elections.election_types.ranking.stv import STV @@ -9,8 +23,7 @@ import logging - -class STVAnimation(): +class STVAnimation: """ A class which creates round-by-round animations of STV elections. @@ -18,27 +31,30 @@ class STVAnimation(): election (STV): An STV election to animate. title (str): Text to be displayed at the beginning of the animation as a title screen. If None, the title screen will be skipped. """ - def __init__(self, election : STV, title : (str | None) = None): + + def __init__(self, election: STV, title: str | None = None): # Extract only the salient details from the election. self.candidates = self._make_candidate_dict(election) self.rounds = self._make_event_list(election) self.title = title - def _make_candidate_dict(self, election : STV) -> dict: + def _make_candidate_dict(self, election: STV) -> dict: """ Create the dictionary of candidates and relevant facts about each one. Args: election (STV): An STV election from which to extract the candidates. - + Returns: dict: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. """ - candidates = {name : {'support' : support} for name, support in election.election_states[0].scores.items()} + candidates = { + name: {"support": support} + for name, support in election.election_states[0].scores.items() + } return candidates - - def _make_event_list(self, election : STV) -> List[dict]: + def _make_event_list(self, election: STV) -> List[dict]: """ Process an STV election into a condensed list of only the salient details from each round. @@ -49,7 +65,10 @@ def _make_event_list(self, election : STV) -> List[dict]: List[dict]: A list of dictionaries corresponding to the rounds of the election. Each dictionary records salient attributes of the corresponding round. """ events = [] - for round_number, election_round in enumerate(election.election_states[1:], start = 1): #Nothing happens in election round 0 + for round_number, election_round in enumerate( + election.election_states[1:], start=1 + ): + # Nothing happens in election round 0 remaining_candidates = [] for fset in election_round.remaining: @@ -58,51 +77,67 @@ def _make_event_list(self, election : STV) -> List[dict]: elected_candidates = [] for fset in election_round.elected: if len(fset) > 0: - name, = fset + (name,) = fset elected_candidates.append(name) eliminated_candidates = [] for fset in election_round.eliminated: if len(fset) > 0: - name, = fset + (name,) = fset eliminated_candidates.append(name) if len(elected_candidates) > 0: - event_type = 'win' + event_type = "win" elected_candidates_str = elected_candidates[0] for candidate_name in elected_candidates[1:]: elected_candidates_str += ", " + candidate_name - message = f'Round {round_number}: {elected_candidates_str} Elected' + message = f"Round {round_number}: {elected_candidates_str} Elected" support_transferred = {} - if round_number == len(election): #If it's the last round, don't worry about the transferred votes - support_transferred = {cand : {} for cand in elected_candidates} + if round_number == len(election): + # If it's the last round, don't worry about the transferred votes + support_transferred = {cand: {} for cand in elected_candidates} else: - support_transferred = self._get_transferred_votes(election, round_number, elected_candidates, 'win') - events.append(dict( - event = event_type, - candidates = elected_candidates, - support_transferred = support_transferred, - quota = election.threshold, - message = message - )) + support_transferred = self._get_transferred_votes( + election, round_number, elected_candidates, "win" + ) + events.append( + dict( + event=event_type, + candidates=elected_candidates, + support_transferred=support_transferred, + quota=election.threshold, + message=message, + ) + ) elif len(eliminated_candidates) > 0: - event_type = 'elimination' + event_type = "elimination" eliminated_candidates_str = eliminated_candidates[0] for candidate_name in eliminated_candidates[1:]: eliminated_candidates_str += ", " + candidate_name - message = f'Round {round_number}: {eliminated_candidates_str} Eliminated' - support_transferred = self._get_transferred_votes(election, round_number, eliminated_candidates, 'elimination') - events.append(dict( - event = event_type, - candidates = eliminated_candidates, - support_transferred = support_transferred, - quota = election.threshold, - message = message - )) + message = ( + f"Round {round_number}: {eliminated_candidates_str} Eliminated" + ) + support_transferred = self._get_transferred_votes( + election, round_number, eliminated_candidates, "elimination" + ) + events.append( + dict( + event=event_type, + candidates=eliminated_candidates, + support_transferred=support_transferred, + quota=election.threshold, + message=message, + ) + ) return events - - def _get_transferred_votes(self, election : STV, round_number : int, from_candidates : List[str], event_type : Literal['win', 'elimination']) -> dict[str, dict[str, float]]: + def _get_transferred_votes( + self, + election: STV, + round_number: int, + from_candidates: List[str], + event_type: Literal["win", "elimination"], + ) -> dict[str, dict[str, float]]: """ Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. @@ -115,12 +150,14 @@ def _get_transferred_votes(self, election : STV, round_number : int, from_candid Returns: dict[str, dict[str, float]]: A nested dictionary. If d is the return value, c1 was a candidate eliminated this round, and c2 is a remaining candidate, then d[c1][c2] will be the total support transferred this round from c1 to c2. """ - prev_profile, prev_state = election.get_step(round_number-1) + prev_profile, prev_state = election.get_step(round_number - 1) current_state = election.election_states[round_number] - if event_type == 'elimination': + if event_type == "elimination": if len(from_candidates) > 1: - raise ValueError(f'Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported.') + raise ValueError( + f"Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported." + ) from_candidate = from_candidates[0] result_dict = {} for to_candidate in [c for s in current_state.remaining for c in s]: @@ -128,8 +165,8 @@ def _get_transferred_votes(self, election : STV, round_number : int, from_candid current_score = int(current_state.scores[to_candidate]) result_dict[to_candidate] = current_score - prev_score return result_dict - - elif event_type == 'win': + + elif event_type == "win": ballots_by_fpv = ballots_by_first_cand(prev_profile) transfers = {} for from_candidate in from_candidates: @@ -137,7 +174,7 @@ def _get_transferred_votes(self, election : STV, round_number : int, from_candid from_candidate, prev_state.scores[from_candidate], ballots_by_fpv[from_candidate], - election.threshold + election.threshold, ) clean_ballots = [ condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) @@ -146,20 +183,22 @@ def _get_transferred_votes(self, election : STV, round_number : int, from_candid transfer_weights_from_candidate = defaultdict(float) for ballot in clean_ballots: if ballot.ranking is not None: - to_candidate, = ballot.ranking[0] + (to_candidate,) = ballot.ranking[0] transfer_weights_from_candidate[to_candidate] += ballot.weight transfers[from_candidate] = transfer_weights_from_candidate return transfers - def render(self, preview : bool = False) -> None: + def render(self, preview: bool = False) -> None: """ Renders the STV animation using Manim. Args: preview (bool): If true, display the result in a video player immediately upon completion. """ - manimation = ElectionScene(deepcopy(self.candidates), deepcopy(self.rounds), title=self.title) + manimation = ElectionScene( + deepcopy(self.candidates), deepcopy(self.rounds), title=self.title + ) manimation.render(preview=preview) @@ -172,25 +211,32 @@ class ElectionScene(manim.Scene): rounds (List[dict]): A list of dictionaries representing the rounds of the election. Each dictionary is a summary of the events in the round it represents. title (str): A string to be displayed at the beginning of the animation as a title screen. If None, the animation will skip the title screen. """ - colors = [manim.color.ManimColor(hex) for hex in [ - "#16DEBD", - "#163EDE", - "#9F34F6", - "#FF6F00", - "#8F560C", - "#E2AD00", - "#8AD412"]] + + colors = [ + manim.color.ManimColor(hex) + for hex in [ + "#16DEBD", + "#163EDE", + "#9F34F6", + "#FF6F00", + "#8F560C", + "#E2AD00", + "#8AD412", + ] + ] bar_color = manim.LIGHT_GRAY bar_opacity = 1 win_bar_color = manim.GREEN eliminated_bar_color = manim.RED ghost_opacity = 0.3 ticker_tape_height = 2 - offscreen_sentinel = '__offscreen__' + offscreen_sentinel = "__offscreen__" offscreen_candidate_color = manim.GRAY title_font_size = 48 - def __init__(self, candidates : dict[str,dict], rounds : List[dict], title : (str | None) = None): + def __init__( + self, candidates: dict[str, dict], rounds: List[dict], title: str | None = None + ): super().__init__() self.candidates = candidates self.rounds = rounds @@ -198,16 +244,15 @@ def __init__(self, candidates : dict[str,dict], rounds : List[dict], title : (st self.width = 8 self.bar_height = 3.5 / len(self.candidates) - self.font_size = 3*40 / len(self.candidates) + self.font_size = 3 * 40 / len(self.candidates) self.bar_opacity = 1 self.bar_buffer_size = 1 / len(self.candidates) self.strikethrough_thickness = self.font_size / 5 - self.max_support = 1.1 * max([round['quota'] for round in self.rounds]) + self.max_support = 1.1 * max([round["quota"] for round in self.rounds]) self.quota_line = None self.ticker_tape_line = None - def construct(self) -> None: """ Constructs the animation. @@ -220,41 +265,43 @@ def construct(self) -> None: self._draw_title(self.title) self._draw_initial_bars() self._initialize_ticker_tape() - + # Go round by round and animate the events - for round_number, round in enumerate(self.rounds): + for round_number, round in enumerate(self.rounds): # Remove the candidate from the candidate list - if round['candidates'] == self.offscreen_sentinel: #If the eliminated candidate is offscreen + if round["candidates"] == self.offscreen_sentinel: + # If the eliminated candidate is offscreen eliminated_candidates = None else: eliminated_candidates = {} - for name in round['candidates']: - eliminated_candidates[name] = (self.candidates[name]) + for name in round["candidates"]: + eliminated_candidates[name] = self.candidates[name] self.candidates.pop(name) self.wait(2) # Draw or move the quota line - self._update_quota_line(round['quota']) + self._update_quota_line(round["quota"]) self._ticker_animation_shift(round) self._ticker_animation_highlight(round) - if round['event'] == 'elimination': - if eliminated_candidates is None: #Offscreen candidate eliminated + if round["event"] == "elimination": + if eliminated_candidates is None: + # Offscreen candidate eliminated self._animate_elimination_offscreen(round) - else: #Onscreen candidate eliminated + else: + # Onscreen candidate eliminated self._animate_elimination(eliminated_candidates, round) - elif round['event'] == 'win': + elif round["event"] == "win": assert eliminated_candidates is not None self._animate_win(eliminated_candidates, round) else: raise Exception(f"Event type {round['event']} not recognized.") - - self.wait(2) + self.wait(2) - def _draw_title(self, message : str) -> None: + def _draw_title(self, message: str) -> None: """ Draw the title screen. @@ -262,8 +309,7 @@ def _draw_title(self, message : str) -> None: message (str): String that the title screen will display. """ text = manim.Tex( - r"{7cm}\centering " + message, - tex_environment="minipage" + r"{7cm}\centering " + message, tex_environment="minipage" ).scale_to_fit_width(10) self.play(Create(text)) self.wait(3) @@ -274,109 +320,130 @@ def _draw_initial_bars(self) -> None: Instantiate and draw the names and bars for each candidate. """ # Sort candidates by starting first place votes - sorted_candidates = sorted(self.candidates.keys(), key=lambda x: self.candidates[x]['support'], reverse=True) + sorted_candidates = sorted( + self.candidates.keys(), + key=lambda x: self.candidates[x]["support"], + reverse=True, + ) # Create bars for i, name in enumerate(sorted_candidates): color = self.colors[i % len(self.colors)] - self.candidates[name]['color'] = color - self.candidates[name]['bars'] = [Rectangle( - width=self._support_to_bar_width(self.candidates[name]['support']), - height=self.bar_height, - color=self.bar_color, - fill_color=color, - fill_opacity=self.bar_opacity - )] - - self.candidates[sorted_candidates[0]]['bars'][0].to_edge(UP) #First candidate goes at the top + self.candidates[name]["color"] = color + self.candidates[name]["bars"] = [ + Rectangle( + width=self._support_to_bar_width(self.candidates[name]["support"]), + height=self.bar_height, + color=self.bar_color, + fill_color=color, + fill_opacity=self.bar_opacity, + ) + ] + # First candidate goes at the top + self.candidates[sorted_candidates[0]]["bars"][0].to_edge(UP) + # The rest of the candidates go below for i, name in enumerate(sorted_candidates[1:], start=1): - self.candidates[name]['bars'][0].next_to( - self.candidates[sorted_candidates[i-1]]['bars'][0], + self.candidates[name]["bars"][0].next_to( + self.candidates[sorted_candidates[i - 1]]["bars"][0], DOWN, - buff=self.bar_buffer_size - ).align_to( - self.candidates[sorted_candidates[0]]['bars'][0], - LEFT - ) #The rest of the candidates go below + buff=self.bar_buffer_size, + ).align_to(self.candidates[sorted_candidates[0]]["bars"][0], LEFT) # Draw a large black rectangle for the background so that the ticker tape vanishes behind it frame_width = manim.config.frame_width frame_height = manim.config.frame_height - background = Rectangle(width = frame_width, - height = frame_height, - fill_color = manim.BLACK, - color = manim.BLACK, - fill_opacity = 1).shift(UP * self.ticker_tape_height).set_z_index(-1) + background = ( + Rectangle( + width=frame_width, + height=frame_height, + fill_color=manim.BLACK, + color=manim.BLACK, + fill_opacity=1, + ) + .shift(UP * self.ticker_tape_height) + .set_z_index(-1) + ) # Create and place candidate names for name in sorted_candidates: candidate = self.candidates[name] - candidate['name_text'] = Text( - name, - font_size=self.font_size, - color=candidate['color'] - ).next_to(candidate['bars'][0], LEFT, buff=0.2) + candidate["name_text"] = Text( + name, font_size=self.font_size, color=candidate["color"] + ).next_to(candidate["bars"][0], LEFT, buff=0.2) # Draw the bars and names - self.play(*[FadeIn(self.candidates[name]['name_text']) for name in sorted_candidates], FadeIn(background)) - self.play(*[Create(self.candidates[name]['bars'][0]) for name in sorted_candidates]) + self.play( + *[FadeIn(self.candidates[name]["name_text"]) for name in sorted_candidates], + FadeIn(background), + ) + self.play( + *[Create(self.candidates[name]["bars"][0]) for name in sorted_candidates] + ) def _initialize_ticker_tape(self) -> None: """ Instantiate and draw the ticker tape line and text. """ line_length = manim.config.frame_width - ticker_line = Line(start = LEFT * line_length/2, - end = RIGHT * line_length/2,) + ticker_line = Line( + start=LEFT * line_length / 2, + end=RIGHT * line_length / 2, + ) ticker_line.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height) - ticker_line.set_z_index(2) #Keep this line in front of the bars and the quota line + ticker_line.set_z_index( + 2 + ) # Keep this line in front of the bars and the quota line self.ticker_tape_line = ticker_line for i, round in enumerate(self.rounds): - new_message = Text( - round['message'], - font_size = 24, - color=manim.DARK_GRAY) - round['ticker_text'] = new_message + new_message = Text(round["message"], font_size=24, color=manim.DARK_GRAY) + round["ticker_text"] = new_message if i == 0: - new_message.to_edge(DOWN,buff=0).shift(DOWN) + new_message.to_edge(DOWN, buff=0).shift(DOWN) else: - new_message.next_to(self.rounds[i-1]['ticker_text'], DOWN) + new_message.next_to(self.rounds[i - 1]["ticker_text"], DOWN) new_message.set_z_index(-2) - + self.play(Create(ticker_line)) - self.play(*[Create(round['ticker_text']) for round in self.rounds]) + self.play(*[Create(round["ticker_text"]) for round in self.rounds]) - def _ticker_animation_shift(self, round : dict) -> None: + def _ticker_animation_shift(self, round: dict) -> None: """ Animate the shifting of the ticker tape to display the message for a given round. Args: round (dict): The round whose message will shift into view. """ - shift_to_round = round['ticker_text'].animate.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height/3) + shift_to_round = ( + round["ticker_text"] + .animate.to_edge(DOWN, buff=0) + .shift(UP * self.ticker_tape_height / 3) + ) drag_other_messages = [ - manim.MaintainPositionRelativeTo(other_round['ticker_text'], round['ticker_text']) - for other_round in self.rounds if not other_round == round + manim.MaintainPositionRelativeTo( + other_round["ticker_text"], round["ticker_text"] + ) + for other_round in self.rounds + if not other_round == round ] self.play(shift_to_round, *drag_other_messages) - def _ticker_animation_highlight(self, round : dict) -> None: + def _ticker_animation_highlight(self, round: dict) -> None: """ Play an animation graying out all ticker tape message but one. Args: round (dict): The round whose message will be highlighted. """ - highlight_message = round['ticker_text'].animate.set_color(manim.WHITE) + highlight_message = round["ticker_text"].animate.set_color(manim.WHITE) unhighlight_other_messages = [ - other_round['ticker_text'].animate.set_color(manim.DARK_GRAY) - for other_round in self.rounds if not other_round == round + other_round["ticker_text"].animate.set_color(manim.DARK_GRAY) + for other_round in self.rounds + if not other_round == round ] self.play(highlight_message, *unhighlight_other_messages) - - def _update_quota_line(self, quota : float) -> None: + def _update_quota_line(self, quota: float) -> None: """ Update the position of the quota line to reflect the given quota. If no quota line exists, create it and animate its creation. @@ -390,20 +457,22 @@ def _update_quota_line(self, quota : float) -> None: line_bottom = self.ticker_tape_line.get_top()[1] line_top = manim.config.frame_height / 2 self.quota_line = Line( - start = UP * line_top, - end = UP * line_bottom, - color=self.win_bar_color) - self.quota_line.align_to(some_candidate['bars'][0], LEFT) - self.quota_line.shift((self.width * quota/self.max_support) * RIGHT) - self.quota_line.set_z_index(1) # Keep the quota line in the front + start=UP * line_top, end=UP * line_bottom, color=self.win_bar_color + ) + self.quota_line.align_to(some_candidate["bars"][0], LEFT) + self.quota_line.shift((self.width * quota / self.max_support) * RIGHT) + self.quota_line.set_z_index(1) # Keep the quota line in the front self.play(Create(self.quota_line)) self.wait(2) else: - self.play(self.quota_line.animate.align_to(some_candidate['bars'][0], LEFT).shift((self.width * quota/self.max_support) * RIGHT)) - + self.play( + self.quota_line.animate.align_to(some_candidate["bars"][0], LEFT).shift( + (self.width * quota / self.max_support) * RIGHT + ) + ) - def _animate_win(self, from_candidates : dict[str,dict], round : dict) -> None: + def _animate_win(self, from_candidates: dict[str, dict], round: dict) -> None: """ Animate a round in which one or more candidates are elected. @@ -411,72 +480,84 @@ def _animate_win(self, from_candidates : dict[str,dict], round : dict) -> None: from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates elected this round and the values are dictionaries recording the candidate's attributes. round (dict): A dictionary recording the events of this round. """ - #Box the winners' names - winner_boxes = [SurroundingRectangle(from_candidate['name_text'], color=manim.GREEN, buff=0.1) - for from_candidate in from_candidates.values()] - + # Box the winners' names + winner_boxes = [ + SurroundingRectangle( + from_candidate["name_text"], color=manim.GREEN, buff=0.1 + ) + for from_candidate in from_candidates.values() + ] # Animate the box around the candidate name and the message text self.play(*[Create(box) for box in winner_boxes]) # Create and animate the subdivision and redistribution of winners' leftover votes for from_candidate_name, from_candidate in from_candidates.items(): - old_bars : List[Rectangle] = from_candidate['bars'] - new_bars : List[Rectangle] = [] + old_bars: List[Rectangle] = from_candidate["bars"] + new_bars: List[Rectangle] = [] transformations = [] - destinations = round['support_transferred'][from_candidate_name] - candidate_color = from_candidate['color'] + destinations = round["support_transferred"][from_candidate_name] + candidate_color = from_candidate["color"] - winner_bar = Rectangle( - width=self._support_to_bar_width(round['quota']), - height=self.bar_height, - color=self.bar_color, - fill_color=self.win_bar_color, - fill_opacity=self.bar_opacity - ).align_to(from_candidate['bars'][0], LEFT).align_to(from_candidate['bars'][0], UP) + winner_bar = ( + Rectangle( + width=self._support_to_bar_width(round["quota"]), + height=self.bar_height, + color=self.bar_color, + fill_color=self.win_bar_color, + fill_opacity=self.bar_opacity, + ) + .align_to(from_candidate["bars"][0], LEFT) + .align_to(from_candidate["bars"][0], UP) + ) - #Create a sub-bar for each destination + # Create a sub-bar for each destination for destination, votes in destinations.items(): if votes <= 0: continue sub_bar = Rectangle( - width = self._support_to_bar_width(votes), + width=self._support_to_bar_width(votes), height=self.bar_height, - color=self.bar_color, + color=self.bar_color, fill_color=candidate_color, - fill_opacity=self.bar_opacity) + fill_opacity=self.bar_opacity, + ) # The first sub-bar should start at the right end of the eliminated candidate's stack. The rest should be arranged to the left of that one if len(new_bars) == 0: - sub_bar.align_to(from_candidate['bars'][-1], RIGHT).align_to(from_candidate['bars'][-1], UP) + sub_bar.align_to(from_candidate["bars"][-1], RIGHT).align_to( + from_candidate["bars"][-1], UP + ) else: sub_bar.next_to(new_bars[-1], LEFT, buff=0) new_bars.append(sub_bar) - self.candidates[destination]['support'] += votes + self.candidates[destination]["support"] += votes + # The sub-bars will move to be next to the bars of their destination candidates transformation = sub_bar.animate.next_to( - self.candidates[destination]['bars'][-1], - RIGHT, - buff=0 - ) - transformations.append( - transformation - ) #The sub-bars will move to be next to the bars of their destination candidates + self.candidates[destination]["bars"][-1], RIGHT, buff=0 + ) + transformations.append(transformation) # Let the new sub-bar be owned by its destination candidate - self.candidates[destination]['bars'] += sub_bar - - + self.candidates[destination]["bars"] += sub_bar # Create a final short bar representing the exhausted votes - exhausted_votes = from_candidate['support'] - round['quota'] - sum(list(destinations.values())) + exhausted_votes = ( + from_candidate["support"] + - round["quota"] + - sum(list(destinations.values())) + ) exhausted_bar = Rectangle( - width = self._support_to_bar_width(exhausted_votes), + width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, - color=self.bar_color, + color=self.bar_color, fill_color=candidate_color, - fill_opacity=self.bar_opacity) + fill_opacity=self.bar_opacity, + ) assert self.quota_line is not None - exhausted_bar.align_to(self.quota_line, LEFT).align_to(from_candidate['bars'][-1], UP) + exhausted_bar.align_to(self.quota_line, LEFT).align_to( + from_candidate["bars"][-1], UP + ) exhausted_bar.set_z_index(1) # Animate the splitting of the old bar into the new sub_bars @@ -484,16 +565,16 @@ def _animate_win(self, from_candidates : dict[str,dict], round : dict) -> None: *[FadeOut(bar) for bar in old_bars], FadeIn(winner_bar), *[FadeIn(bar) for bar in new_bars], - FadeIn(exhausted_bar) + FadeIn(exhausted_bar), ) # Animate moving the sub-bars to the destination bars, and the destruction of the exhausted votes if len(transformations) > 0: self.play(*transformations, Uncreate(exhausted_bar)) - - - def _animate_elimination(self, from_candidates : dict[str,dict], round : dict) -> None: + def _animate_elimination( + self, from_candidates: dict[str, dict], round: dict + ) -> None: """ Animate a round in which a candidate was eliminated. While the interface supports multiple candidate eliminations in one round for future extensibility, this function currently only supports elimination of one candidate at a time. The from_candidates argument should have exactly one entry. @@ -502,129 +583,130 @@ def _animate_elimination(self, from_candidates : dict[str,dict], round : dict) - round (dict): A dictionary recording the events of this round. """ from_candidate = list(from_candidates.values())[0] - destinations = round['support_transferred'] + destinations = round["support_transferred"] - #Cross out the candidate name + # Cross out the candidate name cross = Line( - from_candidate['name_text'].get_left(), - from_candidate['name_text'].get_right(), + from_candidate["name_text"].get_left(), + from_candidate["name_text"].get_right(), color=manim.RED, - ) + ) cross.set_stroke(width=self.strikethrough_thickness) self.play(Create(cross)) # Create short bars that will replace the candidate's current bars - candidate_color = from_candidate['color'] - old_bars = from_candidate['bars'] - new_bars = [] #The bits to be redistributed + candidate_color = from_candidate["color"] + old_bars = from_candidate["bars"] + new_bars = [] # The bits to be redistributed transformations = [] for destination, votes in destinations.items(): if votes <= 0: continue sub_bar = Rectangle( - width = self._support_to_bar_width(votes), + width=self._support_to_bar_width(votes), height=self.bar_height, - color=self.bar_color, + color=self.bar_color, fill_color=candidate_color, - fill_opacity=self.bar_opacity) - self.candidates[destination]['support'] += votes + fill_opacity=self.bar_opacity, + ) + self.candidates[destination]["support"] += votes new_bars.append(sub_bar) transformations.append( new_bars[-1].animate.next_to( - self.candidates[destination]['bars'][-1], - RIGHT, - buff=0) - ) #The sub-bars will move to be next to the bars of their destination candidates - self.candidates[destination]['bars'].append(sub_bar) + self.candidates[destination]["bars"][-1], RIGHT, buff=0 + ) + ) # The sub-bars will move to be next to the bars of their destination candidates + self.candidates[destination]["bars"].append(sub_bar) # Create a final short bar representing the exhausted votes - exhausted_votes = from_candidate['support'] - sum(list(destinations.values())) + exhausted_votes = from_candidate["support"] - sum(list(destinations.values())) exhausted_bar = Rectangle( - width = self._support_to_bar_width(exhausted_votes), + width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, - color=self.bar_color, + color=self.bar_color, fill_color=candidate_color, - fill_opacity=self.bar_opacity) + fill_opacity=self.bar_opacity, + ) exhausted_bar.align_to(old_bars[0], LEFT).align_to(old_bars[0], UP) if len(new_bars) > 0: # The short bars should start in the same place as the old bars. Place them. rightmost_old_bar = old_bars[-1] - new_bars[0].align_to(rightmost_old_bar, RIGHT).align_to(rightmost_old_bar,UP) + new_bars[0].align_to(rightmost_old_bar, RIGHT).align_to( + rightmost_old_bar, UP + ) for i, sub_bar in enumerate(new_bars[1:], start=1): - sub_bar.next_to(new_bars[i-1], LEFT, buff=0) - + sub_bar.next_to(new_bars[i - 1], LEFT, buff=0) # Animate the splitting of the old bar into sub-bars - self.play(*[ - bar.animate.set_opacity(self.ghost_opacity) - for bar in old_bars], - *[FadeIn(bar) for bar in new_bars], - FadeIn(exhausted_bar)) + self.play( + *[bar.animate.set_opacity(self.ghost_opacity) for bar in old_bars], + *[FadeIn(bar) for bar in new_bars], + FadeIn(exhausted_bar), + ) # Animate the exhaustion of votes and moving the sub-bars to the destination bars self.play(Uncreate(exhausted_bar), *transformations) - - - def _animate_elimination_offscreen(self, round : dict) -> None: + def _animate_elimination_offscreen(self, round: dict) -> None: """ Animate a round in which offscreen candidates were eliminated. Args: round (dict): A dictionary recording the events of this round. """ - destinations = round['support_transferred'] + destinations = round["support_transferred"] # Create short bars that will replace the candidate's current bars - new_bars = [] #The bits to be redistributed + new_bars = [] # The bits to be redistributed transformations = [] for destination, votes in destinations.items(): if votes <= 0: continue sub_bar = Rectangle( - width = self._support_to_bar_width(votes), + width=self._support_to_bar_width(votes), height=self.bar_height, - color=self.bar_color, + color=self.bar_color, fill_color=self.offscreen_candidate_color, fill_opacity=self.bar_opacity, ) - self.candidates[destination]['support'] += votes + self.candidates[destination]["support"] += votes new_bars.append(sub_bar) transformations.append( new_bars[-1].animate.next_to( - self.candidates[destination]['bars'][-1], - RIGHT, - buff=0) - ) #The sub-bars will move to be next to the bars of their destination candidates - self.candidates[destination]['bars'].append(sub_bar) - + self.candidates[destination]["bars"][-1], RIGHT, buff=0 + ) + ) # The sub-bars will move to be next to the bars of their destination candidates + self.candidates[destination]["bars"].append(sub_bar) + for bar in new_bars: - bar.to_edge(DOWN).shift((self.bar_height + 2)*DOWN) + bar.to_edge(DOWN).shift((self.bar_height + 2) * DOWN) # Animate the exhaustion of votes and moving the sub-bars to the destination bars self.play(*transformations) - def _rescale_bars(self) -> None: """Re-scale the bars so they fit nicely in frame.""" - self.max_support = max([candidate['support'] for candidate in self.candidates.values()]) + self.max_support = max( + [candidate["support"] for candidate in self.candidates.values()] + ) - transformations : List[Transform] = [] + transformations: List[Transform] = [] for candidate in self.candidates.values(): - old_bar = candidate['bar'] + old_bar = candidate["bar"] new_bar = Rectangle( - width=self._support_to_bar_width(candidate['support']), + width=self._support_to_bar_width(candidate["support"]), height=self.bar_height, - color=self.bar_color, - fill_color=candidate['color'], - fill_opacity=self.bar_opacity).next_to(candidate['name_text'], RIGHT) + color=self.bar_color, + fill_color=candidate["color"], + fill_opacity=self.bar_opacity, + ).next_to(candidate["name_text"], RIGHT) bar_shortening_transformation = Transform( old_bar, new_bar, ) - transformations.append(bar_shortening_transformation) + transformations.append(bar_shortening_transformation) self.play(*transformations) - def _support_to_bar_width(self, support : float) -> float: + def _support_to_bar_width(self, support: float) -> float: """ Convert a number of votes to the width of a bar in manim coordinates representing that many votes. From e4f50687235b78df5095ff688054ce720a2e7d0e Mon Sep 17 00:00:00 2001 From: prismika Date: Fri, 25 Jul 2025 11:06:28 -0400 Subject: [PATCH 11/38] Add test for animation initialization. --- tests/test_animations.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 tests/test_animations.py diff --git a/tests/test_animations.py b/tests/test_animations.py new file mode 100644 index 00000000..4bc38250 --- /dev/null +++ b/tests/test_animations.py @@ -0,0 +1,24 @@ +from votekit.animations import STVAnimation +from votekit import Ballot, PreferenceProfile +from votekit.elections import STV + +# modified from STV wiki +# Election following the "happy path". One elimination or election per round. No ties. No exact quota matches. No funny business. +test_profile_happy = PreferenceProfile( + ballots=( + Ballot(ranking=({"Orange"}, {"Pear"}), weight=3), + Ballot(ranking=({"Pear"}, {"Strawberry"}, {"Cake"}), weight=8), + Ballot(ranking=({"Strawberry"}, {"Orange"}, {"Pear"}), weight=1), + Ballot(ranking=({"Cake"}, {"Chocolate"}), weight=3), + Ballot(ranking=({"Chocolate"}, {"Cake"}, {"Burger"}), weight=1), + Ballot(ranking=({"Burger"}, {"Chicken"}), weight=4), + Ballot(ranking=({"Chicken"}, {"Chocolate"}, {"Burger"}), weight=3), + ), + max_ranking_length=3, +) + +test_election_happy = STV(test_profile_happy, m=3) + +def test_init(): + animation = STVAnimation(test_election_happy) + From 7b287ee413334a4fcdb2d0605b671ff4fd9a6d3d Mon Sep 17 00:00:00 2001 From: prismika Date: Fri, 25 Jul 2025 11:09:16 -0400 Subject: [PATCH 12/38] Remove the unused _rescale_bars method. --- src/votekit/animations.py | 22 ---------------------- 1 file changed, 22 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 2032ecca..e4313c06 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -683,28 +683,6 @@ def _animate_elimination_offscreen(self, round: dict) -> None: # Animate the exhaustion of votes and moving the sub-bars to the destination bars self.play(*transformations) - def _rescale_bars(self) -> None: - """Re-scale the bars so they fit nicely in frame.""" - self.max_support = max( - [candidate["support"] for candidate in self.candidates.values()] - ) - - transformations: List[Transform] = [] - for candidate in self.candidates.values(): - old_bar = candidate["bar"] - new_bar = Rectangle( - width=self._support_to_bar_width(candidate["support"]), - height=self.bar_height, - color=self.bar_color, - fill_color=candidate["color"], - fill_opacity=self.bar_opacity, - ).next_to(candidate["name_text"], RIGHT) - bar_shortening_transformation = Transform( - old_bar, - new_bar, - ) - transformations.append(bar_shortening_transformation) - self.play(*transformations) def _support_to_bar_width(self, support: float) -> float: """ From 8a36597929237a600368c5bb564335885599b824 Mon Sep 17 00:00:00 2001 From: prismika Date: Fri, 25 Jul 2025 11:50:39 -0400 Subject: [PATCH 13/38] Update documentation, fix typesetting and typechecking. --- src/votekit/animations.py | 90 ++++++++++++++++++++++----------------- tests/test_animations.py | 4 +- 2 files changed, 54 insertions(+), 40 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index e4313c06..8458ad0c 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,7 +1,6 @@ from copy import deepcopy -import manim +import manim # type: ignore from manim import ( - Transform, Rectangle, SurroundingRectangle, Line, @@ -24,12 +23,11 @@ class STVAnimation: - """ - A class which creates round-by-round animations of STV elections. + """A class which creates round-by-round animations of STV elections. Args: election (STV): An STV election to animate. - title (str): Text to be displayed at the beginning of the animation as a title screen. If None, the title screen will be skipped. + title (str): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. """ def __init__(self, election: STV, title: str | None = None): @@ -39,8 +37,7 @@ def __init__(self, election: STV, title: str | None = None): self.title = title def _make_candidate_dict(self, election: STV) -> dict: - """ - Create the dictionary of candidates and relevant facts about each one. + """Create the dictionary of candidates and relevant facts about each one. Args: election (STV): An STV election from which to extract the candidates. @@ -55,8 +52,7 @@ def _make_candidate_dict(self, election: STV) -> dict: return candidates def _make_event_list(self, election: STV) -> List[dict]: - """ - Process an STV election into a condensed list of only the salient details from each round. + """Process an STV election into a condensed list of only the salient details from each round. Args: election (STV): The STV election to process. @@ -92,7 +88,7 @@ def _make_event_list(self, election: STV) -> List[dict]: for candidate_name in elected_candidates[1:]: elected_candidates_str += ", " + candidate_name message = f"Round {round_number}: {elected_candidates_str} Elected" - support_transferred = {} + support_transferred : dict[str, dict[str, float]] = {} if round_number == len(election): # If it's the last round, don't worry about the transferred votes support_transferred = {cand: {} for cand in elected_candidates} @@ -138,37 +134,41 @@ def _get_transferred_votes( from_candidates: List[str], event_type: Literal["win", "elimination"], ) -> dict[str, dict[str, float]]: - """ - Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. + """Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. Args: election (STV): The election. round_number (int): The number of the round in question. from_candidates (List[str]): A list of the names of the elected or eliminated candidates. - event_type (str): "win" if candidates were elected this round, "elimination" otherwise. + event_type (str): ``"win"`` if candidates were elected this round, ``"elimination"`` otherwise. Returns: - dict[str, dict[str, float]]: A nested dictionary. If d is the return value, c1 was a candidate eliminated this round, and c2 is a remaining candidate, then d[c1][c2] will be the total support transferred this round from c1 to c2. + dict[str, dict[str, float]]: A nested dictionary. If ``d`` is the return value, ``c1`` was a candidate eliminated this round, and ``c2`` is a remaining candidate, then ``d[c1][c2]`` will be the total support transferred this round from candidate ``c1`` to candidate ``c2``. + + + Notes: + This function supports the election, but not the elimination, of multiple candidates in one round. If ``event_type`` is ``"elimination"`` then ``from_candidates`` should have length 1. + + Raises: + ValueError: If multiple candidates are eliminated in this round. """ prev_profile, prev_state = election.get_step(round_number - 1) current_state = election.election_states[round_number] - + + transfers : dict[str, dict[str, float]] = {} if event_type == "elimination": if len(from_candidates) > 1: raise ValueError( f"Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported." ) from_candidate = from_candidates[0] - result_dict = {} + transfers = {from_candidate : {}} for to_candidate in [c for s in current_state.remaining for c in s]: prev_score = int(prev_state.scores[to_candidate]) current_score = int(current_state.scores[to_candidate]) - result_dict[to_candidate] = current_score - prev_score - return result_dict - + transfers[from_candidate][to_candidate] = current_score - prev_score elif event_type == "win": ballots_by_fpv = ballots_by_first_cand(prev_profile) - transfers = {} for from_candidate in from_candidates: new_ballots = election.transfer( from_candidate, @@ -180,21 +180,23 @@ def _get_transferred_votes( condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) for b in new_ballots ] - transfer_weights_from_candidate = defaultdict(float) + transfer_weights_from_candidate : dict[str, float] = defaultdict(float) for ballot in clean_ballots: if ballot.ranking is not None: (to_candidate,) = ballot.ranking[0] transfer_weights_from_candidate[to_candidate] += ballot.weight transfers[from_candidate] = transfer_weights_from_candidate - return transfers + + return transfers def render(self, preview: bool = False) -> None: - """ - Renders the STV animation using Manim. + """Renders the STV animation using Manim. + + The completed video will appear in the directory ``media/videos``. Args: - preview (bool): If true, display the result in a video player immediately upon completion. + preview (bool): If ``True``, display the result in a video player immediately upon completing the render. """ manimation = ElectionScene( deepcopy(self.candidates), deepcopy(self.rounds), title=self.title @@ -203,13 +205,15 @@ def render(self, preview: bool = False) -> None: class ElectionScene(manim.Scene): - """ - Class for Manim animation of an STV election. This class is instantiated by the class STVAnimation. It should not be instantiated directly. + """Class for Manim animation of an STV election. + + Notes: + This class is instantiated by the class ``STVAnimation``. It should not be instantiated directly. Args: candidates (dict[str,dict]): A dictionary mapping each candidate to a dictionary of attributes of the candidate. rounds (List[dict]): A list of dictionaries representing the rounds of the election. Each dictionary is a summary of the events in the round it represents. - title (str): A string to be displayed at the beginning of the animation as a title screen. If None, the animation will skip the title screen. + title (str): A string to be displayed at the beginning of the animation as a title screen. If ``None``, the animation will skip the title screen. """ colors = [ @@ -381,9 +385,7 @@ def _draw_initial_bars(self) -> None: ) def _initialize_ticker_tape(self) -> None: - """ - Instantiate and draw the ticker tape line and text. - """ + """Instantiate and draw the ticker tape line and text.""" line_length = manim.config.frame_width ticker_line = Line( start=LEFT * line_length / 2, @@ -575,15 +577,28 @@ def _animate_win(self, from_candidates: dict[str, dict], round: dict) -> None: def _animate_elimination( self, from_candidates: dict[str, dict], round: dict ) -> None: - """ - Animate a round in which a candidate was eliminated. While the interface supports multiple candidate eliminations in one round for future extensibility, this function currently only supports elimination of one candidate at a time. The from_candidates argument should have exactly one entry. + """Animate a round in which a candidate was eliminated. Args: from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates eliminated this round and the values are dictionaries recording the candidate's attributes. This function currently only supports one elimination at a time, so from_candidates should be a dictionary with exactly one entry. round (dict): A dictionary recording the events of this round. + + Notes: + While the interface supports multiple candidate eliminations in one round for future extensibility, this function currently only supports elimination of one candidate at a time. The from_candidates argument should have exactly one entry. + + Raises: + ValueError: If the length of ``from_candidates`` is not 1. """ + num_eliminated_candidates = len(list(from_candidates.values())) + if num_eliminated_candidates != 1: + raise ValueError( + f"Elimination round animations only support one eliminated candidate at a time. Attempted to animate {num_eliminated_candidates} eliminations in one election round." + ) + del num_eliminated_candidates + + from_candidate_name = list(from_candidates.keys())[0] from_candidate = list(from_candidates.values())[0] - destinations = round["support_transferred"] + destinations = round["support_transferred"][from_candidate_name] # Cross out the candidate name cross = Line( @@ -647,8 +662,7 @@ def _animate_elimination( self.play(Uncreate(exhausted_bar), *transformations) def _animate_elimination_offscreen(self, round: dict) -> None: - """ - Animate a round in which offscreen candidates were eliminated. + """Animate a round in which offscreen candidates were eliminated. Args: round (dict): A dictionary recording the events of this round. @@ -683,10 +697,8 @@ def _animate_elimination_offscreen(self, round: dict) -> None: # Animate the exhaustion of votes and moving the sub-bars to the destination bars self.play(*transformations) - def _support_to_bar_width(self, support: float) -> float: - """ - Convert a number of votes to the width of a bar in manim coordinates representing that many votes. + """Convert a number of votes to the width of a bar in manim coordinates representing that many votes. Args: support (float): A number of votes. diff --git a/tests/test_animations.py b/tests/test_animations.py index 4bc38250..9eb2c1bb 100644 --- a/tests/test_animations.py +++ b/tests/test_animations.py @@ -19,6 +19,8 @@ test_election_happy = STV(test_profile_happy, m=3) + def test_init(): animation = STVAnimation(test_election_happy) - + assert animation.candidates is not None + assert animation.rounds is not None \ No newline at end of file From bc2c5762b720b922cc348a4277d7fbb3f1e4841d Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 31 Jul 2025 15:00:22 -0400 Subject: [PATCH 14/38] Introduce dataclass for animation events. --- src/votekit/animations.py | 187 ++++++++++++++++++-------------------- 1 file changed, 89 insertions(+), 98 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 8458ad0c..ce8eee80 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -17,9 +17,31 @@ from .cleaning import condense_ballot_ranking, remove_cand_from_ballot from .utils import ballots_by_first_cand from .elections.election_types.ranking.stv import STV -from typing import Literal, List +from typing import Literal, List, Optional, Sequence, Mapping from collections import defaultdict import logging +from dataclasses import dataclass + + +@dataclass +class AnimationEvent: + quota : float + message : str + +@dataclass +class EliminationEvent(AnimationEvent): + candidate : str + support_transferred : Mapping[str,float] + +@dataclass +class EliminationOffscreenEvent(AnimationEvent): + support_transferred : Mapping[str,float] + +@dataclass +class WinEvent(AnimationEvent): + candidates : Sequence[str] + support_transferred : Mapping[str, Mapping[str,float]] + class STVAnimation: @@ -30,7 +52,7 @@ class STVAnimation: title (str): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. """ - def __init__(self, election: STV, title: str | None = None): + def __init__(self, election: STV, title: Optional[str] = None): # Extract only the salient details from the election. self.candidates = self._make_candidate_dict(election) self.rounds = self._make_event_list(election) @@ -51,7 +73,7 @@ def _make_candidate_dict(self, election: STV) -> dict: } return candidates - def _make_event_list(self, election: STV) -> List[dict]: + def _make_event_list(self, election: STV) -> List[AnimationEvent]: """Process an STV election into a condensed list of only the salient details from each round. Args: @@ -60,30 +82,17 @@ def _make_event_list(self, election: STV) -> List[dict]: Returns: List[dict]: A list of dictionaries corresponding to the rounds of the election. Each dictionary records salient attributes of the corresponding round. """ - events = [] + events : List[AnimationEvent] = [] for round_number, election_round in enumerate( election.election_states[1:], start=1 ): # Nothing happens in election round 0 - remaining_candidates = [] - for fset in election_round.remaining: - remaining_candidates += list(fset) - - elected_candidates = [] - for fset in election_round.elected: - if len(fset) > 0: - (name,) = fset - elected_candidates.append(name) - - eliminated_candidates = [] - for fset in election_round.eliminated: - if len(fset) > 0: - (name,) = fset - eliminated_candidates.append(name) + remaining_candidates = [c for s in election_round.remaining for c in s] + elected_candidates = [c for s in election_round.elected for c in s] + eliminated_candidates = [c for s in election_round.eliminated for c in s] - if len(elected_candidates) > 0: - event_type = "win" + if len(elected_candidates) > 0: # Win round elected_candidates_str = elected_candidates[0] for candidate_name in elected_candidates[1:]: elected_candidates_str += ", " + candidate_name @@ -97,33 +106,24 @@ def _make_event_list(self, election: STV) -> List[dict]: election, round_number, elected_candidates, "win" ) events.append( - dict( - event=event_type, - candidates=elected_candidates, - support_transferred=support_transferred, - quota=election.threshold, - message=message, - ) - ) - elif len(eliminated_candidates) > 0: - event_type = "elimination" - eliminated_candidates_str = eliminated_candidates[0] - for candidate_name in eliminated_candidates[1:]: - eliminated_candidates_str += ", " + candidate_name - message = ( - f"Round {round_number}: {eliminated_candidates_str} Eliminated" + WinEvent(quota = election.threshold, + candidates = elected_candidates, + support_transferred = support_transferred, + message=message) ) + elif len(eliminated_candidates) > 0: # Elimination round + if len(eliminated_candidates) > 1: + raise ValueError(f"Multiple-elimination rounds not supported. At most one candidate should be eliminated in each round. Candidates eliminated in round {round_number}: {eliminated_candidates}.") + eliminated_candidate = eliminated_candidates[0] + message = f"Round {round_number}: {eliminated_candidate} Eliminated" support_transferred = self._get_transferred_votes( election, round_number, eliminated_candidates, "elimination" ) events.append( - dict( - event=event_type, - candidates=eliminated_candidates, - support_transferred=support_transferred, - quota=election.threshold, - message=message, - ) + EliminationEvent(quota = election.threshold, + candidate = eliminated_candidates[0], + support_transferred = support_transferred[eliminated_candidate], + message = message) ) return events @@ -239,7 +239,7 @@ class ElectionScene(manim.Scene): title_font_size = 48 def __init__( - self, candidates: dict[str, dict], rounds: List[dict], title: str | None = None + self, candidates: dict[str, dict], rounds: List[AnimationEvent], title: str | None = None ): super().__init__() self.candidates = candidates @@ -252,10 +252,11 @@ def __init__( self.bar_opacity = 1 self.bar_buffer_size = 1 / len(self.candidates) self.strikethrough_thickness = self.font_size / 5 - self.max_support = 1.1 * max([round["quota"] for round in self.rounds]) + self.max_support = 1.1 * max([round.quota for round in self.rounds]) self.quota_line = None self.ticker_tape_line = None + self.ticker_tape : List[Text] = [] def construct(self) -> None: """ @@ -271,38 +272,28 @@ def construct(self) -> None: self._initialize_ticker_tape() # Go round by round and animate the events - for round_number, round in enumerate(self.rounds): - # Remove the candidate from the candidate list - if round["candidates"] == self.offscreen_sentinel: - # If the eliminated candidate is offscreen - eliminated_candidates = None - else: - eliminated_candidates = {} - for name in round["candidates"]: - eliminated_candidates[name] = self.candidates[name] - self.candidates.pop(name) - + for event_number, round in enumerate(self.rounds): self.wait(2) - # Draw or move the quota line - self._update_quota_line(round["quota"]) - - self._ticker_animation_shift(round) - self._ticker_animation_highlight(round) - - if round["event"] == "elimination": - if eliminated_candidates is None: - # Offscreen candidate eliminated - self._animate_elimination_offscreen(round) - else: - # Onscreen candidate eliminated - self._animate_elimination(eliminated_candidates, round) - elif round["event"] == "win": - assert eliminated_candidates is not None - self._animate_win(eliminated_candidates, round) + self._update_quota_line(round.quota) + + self._ticker_animation_shift(event_number) + self._ticker_animation_highlight(event_number) + + if isinstance(round, EliminationEvent): # Onscreen candidate eliminated + # Remove the candidate from the candidate list + eliminated_candidates = {round.candidate : self.candidates.pop(round.candidate)} + self._animate_elimination(eliminated_candidates, round) + elif isinstance(round, EliminationOffscreenEvent):# Offscreen candidate eliminated + self._animate_elimination_offscreen(round) + elif isinstance(round, WinEvent): # Election round + # Remove the candidates from the candidate list + elected_candidates = {} + for name in round.candidates: + elected_candidates[name] = self.candidates.pop(name) + self._animate_win(elected_candidates, round) else: - raise Exception(f"Event type {round['event']} not recognized.") - + raise Exception(f"Invalid type for event {round}.") self.wait(2) def _draw_title(self, message: str) -> None: @@ -314,7 +305,7 @@ def _draw_title(self, message: str) -> None: """ text = manim.Tex( r"{7cm}\centering " + message, tex_environment="minipage" - ).scale_to_fit_width(10) + ).scale_to_fit_width(10) # We do this one with a TeX minipage to get the text to wrap if it's too long. self.play(Create(text)) self.wait(3) self.play(Uncreate(text)) @@ -396,20 +387,20 @@ def _initialize_ticker_tape(self) -> None: 2 ) # Keep this line in front of the bars and the quota line self.ticker_tape_line = ticker_line - + self.ticker_tape = [] for i, round in enumerate(self.rounds): - new_message = Text(round["message"], font_size=24, color=manim.DARK_GRAY) - round["ticker_text"] = new_message + new_message = Text(round.message, font_size=24, color=manim.DARK_GRAY) if i == 0: new_message.to_edge(DOWN, buff=0).shift(DOWN) else: - new_message.next_to(self.rounds[i - 1]["ticker_text"], DOWN) + new_message.next_to(self.ticker_tape[-1], DOWN) new_message.set_z_index(-2) + self.ticker_tape.append(new_message) self.play(Create(ticker_line)) - self.play(*[Create(round["ticker_text"]) for round in self.rounds]) + self.play(*[Create(message) for message in self.ticker_tape]) - def _ticker_animation_shift(self, round: dict) -> None: + def _ticker_animation_shift(self, event_number: int) -> None: """ Animate the shifting of the ticker tape to display the message for a given round. @@ -417,31 +408,31 @@ def _ticker_animation_shift(self, round: dict) -> None: round (dict): The round whose message will shift into view. """ shift_to_round = ( - round["ticker_text"] + self.ticker_tape[event_number] .animate.to_edge(DOWN, buff=0) .shift(UP * self.ticker_tape_height / 3) ) drag_other_messages = [ manim.MaintainPositionRelativeTo( - other_round["ticker_text"], round["ticker_text"] + self.ticker_tape[i], self.ticker_tape[event_number] ) - for other_round in self.rounds - if not other_round == round + for i in range(len(self.ticker_tape)) + if i != event_number ] self.play(shift_to_round, *drag_other_messages) - def _ticker_animation_highlight(self, round: dict) -> None: + def _ticker_animation_highlight(self, event_number : int) -> None: """ Play an animation graying out all ticker tape message but one. Args: round (dict): The round whose message will be highlighted. """ - highlight_message = round["ticker_text"].animate.set_color(manim.WHITE) + highlight_message = self.ticker_tape[event_number].animate.set_color(manim.WHITE) unhighlight_other_messages = [ - other_round["ticker_text"].animate.set_color(manim.DARK_GRAY) - for other_round in self.rounds - if not other_round == round + self.ticker_tape[i].animate.set_color(manim.DARK_GRAY) + for i in range(len(self.ticker_tape)) + if i != event_number ] self.play(highlight_message, *unhighlight_other_messages) @@ -474,7 +465,7 @@ def _update_quota_line(self, quota: float) -> None: ) ) - def _animate_win(self, from_candidates: dict[str, dict], round: dict) -> None: + def _animate_win(self, from_candidates: dict[str, dict], round: WinEvent) -> None: """ Animate a round in which one or more candidates are elected. @@ -498,12 +489,12 @@ def _animate_win(self, from_candidates: dict[str, dict], round: dict) -> None: old_bars: List[Rectangle] = from_candidate["bars"] new_bars: List[Rectangle] = [] transformations = [] - destinations = round["support_transferred"][from_candidate_name] + destinations = round.support_transferred[from_candidate_name] candidate_color = from_candidate["color"] winner_bar = ( Rectangle( - width=self._support_to_bar_width(round["quota"]), + width=self._support_to_bar_width(round.quota), height=self.bar_height, color=self.bar_color, fill_color=self.win_bar_color, @@ -546,7 +537,7 @@ def _animate_win(self, from_candidates: dict[str, dict], round: dict) -> None: # Create a final short bar representing the exhausted votes exhausted_votes = ( from_candidate["support"] - - round["quota"] + - round.quota - sum(list(destinations.values())) ) exhausted_bar = Rectangle( @@ -575,7 +566,7 @@ def _animate_win(self, from_candidates: dict[str, dict], round: dict) -> None: self.play(*transformations, Uncreate(exhausted_bar)) def _animate_elimination( - self, from_candidates: dict[str, dict], round: dict + self, from_candidates: dict[str, dict], round: EliminationEvent ) -> None: """Animate a round in which a candidate was eliminated. @@ -598,7 +589,7 @@ def _animate_elimination( from_candidate_name = list(from_candidates.keys())[0] from_candidate = list(from_candidates.values())[0] - destinations = round["support_transferred"][from_candidate_name] + destinations = round.support_transferred # Cross out the candidate name cross = Line( @@ -661,13 +652,13 @@ def _animate_elimination( # Animate the exhaustion of votes and moving the sub-bars to the destination bars self.play(Uncreate(exhausted_bar), *transformations) - def _animate_elimination_offscreen(self, round: dict) -> None: + def _animate_elimination_offscreen(self, round: EliminationOffscreenEvent) -> None: """Animate a round in which offscreen candidates were eliminated. Args: round (dict): A dictionary recording the events of this round. """ - destinations = round["support_transferred"] + destinations = round.support_transferred # Create short bars that will replace the candidate's current bars new_bars = [] # The bits to be redistributed From a355b018300f0914dbceea71414a2be3a10764e5 Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 31 Jul 2025 16:11:11 -0400 Subject: [PATCH 15/38] Let users specify focused candidates. --- src/votekit/animations.py | 45 +++++++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 14 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index ce8eee80..e5096f09 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -52,8 +52,12 @@ class STVAnimation: title (str): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. """ - def __init__(self, election: STV, title: Optional[str] = None): + def __init__(self, election: STV, title: Optional[str] = None, focus : Optional[List[str]] = None): # Extract only the salient details from the election. + if focus is None: + #Focus all candidates + focus = [c for s in election.get_remaining(0) for c in s] + self.focus : List[str] = focus self.candidates = self._make_candidate_dict(election) self.rounds = self._make_event_list(election) self.title = title @@ -70,6 +74,7 @@ def _make_candidate_dict(self, election: STV) -> dict: candidates = { name: {"support": support} for name, support in election.election_states[0].scores.items() + if name in self.focus } return candidates @@ -87,8 +92,6 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: election.election_states[1:], start=1 ): # Nothing happens in election round 0 - - remaining_candidates = [c for s in election_round.remaining for c in s] elected_candidates = [c for s in election_round.elected for c in s] eliminated_candidates = [c for s in election_round.eliminated for c in s] @@ -119,12 +122,22 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: support_transferred = self._get_transferred_votes( election, round_number, eliminated_candidates, "elimination" ) - events.append( - EliminationEvent(quota = election.threshold, - candidate = eliminated_candidates[0], - support_transferred = support_transferred[eliminated_candidate], - message = message) - ) + if eliminated_candidate in self.focus: + events.append( + EliminationEvent(quota = election.threshold, + candidate = eliminated_candidate, + support_transferred = support_transferred[eliminated_candidate], + message = message) + ) + else: + events.append( + EliminationOffscreenEvent(quota = election.threshold, + support_transferred = support_transferred[eliminated_candidate], + message = message) + ) + + self.rounds = self._condense_offscreen_rounds(self.rounds) + return events def _get_transferred_votes( @@ -163,7 +176,7 @@ def _get_transferred_votes( ) from_candidate = from_candidates[0] transfers = {from_candidate : {}} - for to_candidate in [c for s in current_state.remaining for c in s]: + for to_candidate in [c for s in current_state.remaining for c in s if c in self.focus]: prev_score = int(prev_state.scores[to_candidate]) current_score = int(current_state.scores[to_candidate]) transfers[from_candidate][to_candidate] = current_score - prev_score @@ -184,11 +197,15 @@ def _get_transferred_votes( for ballot in clean_ballots: if ballot.ranking is not None: (to_candidate,) = ballot.ranking[0] - transfer_weights_from_candidate[to_candidate] += ballot.weight + if to_candidate in self.focus: + transfer_weights_from_candidate[to_candidate] += ballot.weight transfers[from_candidate] = transfer_weights_from_candidate return transfers + + def _condense_offscreen_rounds(self, rounds : List[AnimationEvent]) -> List[AnimationEvent]: + return rounds def render(self, preview: bool = False) -> None: """Renders the STV animation using Manim. @@ -239,7 +256,7 @@ class ElectionScene(manim.Scene): title_font_size = 48 def __init__( - self, candidates: dict[str, dict], rounds: List[AnimationEvent], title: str | None = None + self, candidates: dict[str, dict], rounds: List[AnimationEvent], title: Optional[str] = None ): super().__init__() self.candidates = candidates @@ -660,8 +677,8 @@ def _animate_elimination_offscreen(self, round: EliminationOffscreenEvent) -> No """ destinations = round.support_transferred - # Create short bars that will replace the candidate's current bars - new_bars = [] # The bits to be redistributed + # Create short bars that will begin offscreen + new_bars = [] transformations = [] for destination, votes in destinations.items(): if votes <= 0: From c4fb8da12f63297edf2afbfb097136abe4c5b030 Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 4 Aug 2025 12:21:20 -0400 Subject: [PATCH 16/38] Automatically condense offscreen events into one event. For instance, 'Rounds 3-7: 5 candidates eliminated.' --- src/votekit/animations.py | 69 +++++++++++++++++++++++++++++++-------- 1 file changed, 56 insertions(+), 13 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index e5096f09..f693b2db 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -21,26 +21,47 @@ from collections import defaultdict import logging from dataclasses import dataclass +from abc import ABC, abstractmethod @dataclass -class AnimationEvent: +class AnimationEvent(ABC): quota : float - message : str + @abstractmethod + def get_message(self) -> str: + pass + @dataclass class EliminationEvent(AnimationEvent): candidate : str support_transferred : Mapping[str,float] + round_number : int + def get_message(self) -> str: + return f"Round {self.round_number}: {self.candidate} eliminated." @dataclass class EliminationOffscreenEvent(AnimationEvent): support_transferred : Mapping[str,float] + round_numbers : List[int] + def get_message(self) -> str: + if len(self.round_numbers) == 1: + return f"Round {self.round_numbers[0]}: 1 candidate eliminated." + else: + message = f"Rounds {self.round_numbers[0]}-{self.round_numbers[-1]}: {len(self.round_numbers)} candidates eliminated." + return message + @dataclass class WinEvent(AnimationEvent): candidates : Sequence[str] support_transferred : Mapping[str, Mapping[str,float]] + round_number : int + def get_message(self) -> str: + candidate_string = self.candidates[0] + for candidate_name in self.candidates[1:]: + candidate_string += f", {candidate_name}" + return f"Round {self.round_number}: {candidate_string} elected." @@ -97,9 +118,6 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: if len(elected_candidates) > 0: # Win round elected_candidates_str = elected_candidates[0] - for candidate_name in elected_candidates[1:]: - elected_candidates_str += ", " + candidate_name - message = f"Round {round_number}: {elected_candidates_str} Elected" support_transferred : dict[str, dict[str, float]] = {} if round_number == len(election): # If it's the last round, don't worry about the transferred votes @@ -112,13 +130,12 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: WinEvent(quota = election.threshold, candidates = elected_candidates, support_transferred = support_transferred, - message=message) + round_number = round_number) ) elif len(eliminated_candidates) > 0: # Elimination round if len(eliminated_candidates) > 1: raise ValueError(f"Multiple-elimination rounds not supported. At most one candidate should be eliminated in each round. Candidates eliminated in round {round_number}: {eliminated_candidates}.") eliminated_candidate = eliminated_candidates[0] - message = f"Round {round_number}: {eliminated_candidate} Eliminated" support_transferred = self._get_transferred_votes( election, round_number, eliminated_candidates, "elimination" ) @@ -127,16 +144,16 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: EliminationEvent(quota = election.threshold, candidate = eliminated_candidate, support_transferred = support_transferred[eliminated_candidate], - message = message) + round_number = round_number) ) else: events.append( EliminationOffscreenEvent(quota = election.threshold, support_transferred = support_transferred[eliminated_candidate], - message = message) + round_numbers = [round_number]) ) - self.rounds = self._condense_offscreen_rounds(self.rounds) + events = self._condense_offscreen_events(events) return events @@ -204,8 +221,34 @@ def _get_transferred_votes( return transfers - def _condense_offscreen_rounds(self, rounds : List[AnimationEvent]) -> List[AnimationEvent]: - return rounds + def _condense_offscreen_events(self, rounds : List[AnimationEvent]) -> List[AnimationEvent]: + return_events : List[AnimationEvent] = [rounds[0]] + for event in rounds[1:]: + if isinstance( + return_events[-1], EliminationOffscreenEvent + ) and isinstance( + event, EliminationOffscreenEvent + ): + return_events[-1] = self._compose_offscreen_eliminations(return_events[-1], event) + else: + return_events.append(event) + return return_events + + def _compose_offscreen_eliminations(self, event1 : EliminationOffscreenEvent, event2: EliminationOffscreenEvent) -> EliminationOffscreenEvent: + support_transferred = defaultdict(float) + for key, value in event1.support_transferred.items(): + support_transferred[key] += value + for key, value in event2.support_transferred.items(): + support_transferred[key] += value + round_numbers = event1.round_numbers + event2.round_numbers + quota = event1.quota + return EliminationOffscreenEvent( + quota = quota, + support_transferred = support_transferred, + round_numbers = round_numbers + ) + + def render(self, preview: bool = False) -> None: """Renders the STV animation using Manim. @@ -406,7 +449,7 @@ def _initialize_ticker_tape(self) -> None: self.ticker_tape_line = ticker_line self.ticker_tape = [] for i, round in enumerate(self.rounds): - new_message = Text(round.message, font_size=24, color=manim.DARK_GRAY) + new_message = Text(round.get_message(), font_size=24, color=manim.DARK_GRAY) if i == 0: new_message.to_edge(DOWN, buff=0).shift(DOWN) else: From 661dc79d7d92dad4c6544216ffc3a39f2375d7ea Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 4 Aug 2025 16:12:34 -0400 Subject: [PATCH 17/38] Fix layout bug in which bars would extend off the screen on both the left and right sides. --- src/votekit/animations.py | 57 ++++++++++++++++++++++++--------------- 1 file changed, 36 insertions(+), 21 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index f693b2db..5d3faf0d 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -310,7 +310,7 @@ def __init__( self.bar_height = 3.5 / len(self.candidates) self.font_size = 3 * 40 / len(self.candidates) self.bar_opacity = 1 - self.bar_buffer_size = 1 / len(self.candidates) + self.bar_buffer_size = self.bar_height self.strikethrough_thickness = self.font_size / 5 self.max_support = 1.1 * max([round.quota for round in self.rounds]) @@ -381,28 +381,49 @@ def _draw_initial_bars(self) -> None: reverse=True, ) - # Create bars + # Assign colors for i, name in enumerate(sorted_candidates): color = self.colors[i % len(self.colors)] self.candidates[name]["color"] = color - self.candidates[name]["bars"] = [ + + + # Create candidate name text + for i, name in enumerate(sorted_candidates): + candidate = self.candidates[name] + candidate["name_text"] = Text( + name, font_size=self.font_size, color=candidate["color"] + ) + if i==0: + #First candidate goes at the top + candidate["name_text"].to_edge(UP, buff=self.bar_buffer_size) + else: + #The rest of the candidates go below, right justified + candidate["name_text"].next_to( + self.candidates[sorted_candidates[i-1]]["name_text"], + DOWN, + buff=self.bar_buffer_size + ).align_to( + self.candidates[sorted_candidates[0]]["name_text"], + RIGHT + ) + # Align candidate names to the left + group = manim.Group().add(*[ + candidate["name_text"] for candidate in self.candidates.values() + ]) + group.to_edge(LEFT) + del group + + # Create bars + for candidate in self.candidates.values(): + candidate["bars"] = [ Rectangle( - width=self._support_to_bar_width(self.candidates[name]["support"]), + width=self._support_to_bar_width(candidate["support"]), height=self.bar_height, color=self.bar_color, - fill_color=color, + fill_color=candidate["color"], fill_opacity=self.bar_opacity, - ) + ).next_to(candidate["name_text"], RIGHT, buff=0.2) ] - # First candidate goes at the top - self.candidates[sorted_candidates[0]]["bars"][0].to_edge(UP) - # The rest of the candidates go below - for i, name in enumerate(sorted_candidates[1:], start=1): - self.candidates[name]["bars"][0].next_to( - self.candidates[sorted_candidates[i - 1]]["bars"][0], - DOWN, - buff=self.bar_buffer_size, - ).align_to(self.candidates[sorted_candidates[0]]["bars"][0], LEFT) # Draw a large black rectangle for the background so that the ticker tape vanishes behind it frame_width = manim.config.frame_width @@ -419,12 +440,6 @@ def _draw_initial_bars(self) -> None: .set_z_index(-1) ) - # Create and place candidate names - for name in sorted_candidates: - candidate = self.candidates[name] - candidate["name_text"] = Text( - name, font_size=self.font_size, color=candidate["color"] - ).next_to(candidate["bars"][0], LEFT, buff=0.2) # Draw the bars and names self.play( From 1bdfb95bf829dc623030231ebbdbed20ccb3280f Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 4 Aug 2025 16:19:40 -0400 Subject: [PATCH 18/38] Automatically focus all elected candidates. --- src/votekit/animations.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 5d3faf0d..93e3f42b 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -73,12 +73,9 @@ class STVAnimation: title (str): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. """ - def __init__(self, election: STV, title: Optional[str] = None, focus : Optional[List[str]] = None): - # Extract only the salient details from the election. - if focus is None: - #Focus all candidates - focus = [c for s in election.get_remaining(0) for c in s] - self.focus : List[str] = focus + def __init__(self, election: STV, title: Optional[str] = None, focus : List[str] = []): + elected_candidates = [c for s in election.get_elected() for c in s] + self.focus : List[str] = list(set(focus).union(set(elected_candidates))) self.candidates = self._make_candidate_dict(election) self.rounds = self._make_event_list(election) self.title = title From 345e47ad226b568a8f60bfd572938173aad97ac2 Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 4 Aug 2025 17:36:46 -0400 Subject: [PATCH 19/38] Various cleaning tasks: Update type hints and documentation, change variable names to use consistent terminology, etc. --- src/votekit/animations.py | 238 +++++++++++++++++++++++--------------- 1 file changed, 146 insertions(+), 92 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 93e3f42b..0cea7179 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -26,14 +26,29 @@ @dataclass class AnimationEvent(ABC): + """ + An abstract class representing a single step of the animation, usually a single election round. + + Attributes: + quota (float): The current election threshold at the start of this event. + """ quota : float @abstractmethod def get_message(self) -> str: + """Generate a message describing the event for the viewer of the animation.""" pass @dataclass class EliminationEvent(AnimationEvent): + """ + An animation event representing a round in which a candidate was eliminated. + + Attributes: + candidate (str): The name of the eliminated candidate. + support_transferred (Mapping[str,float]): A dictionary mapping names of candidates to the amount of support they received from the elimination. + round_number (int): The round of the election process associated to this event. + """ candidate : str support_transferred : Mapping[str,float] round_number : int @@ -42,6 +57,13 @@ def get_message(self) -> str: @dataclass class EliminationOffscreenEvent(AnimationEvent): + """ + An animation event representing some number of rounds in which offscreen candidates were eliminated. + + Attributes: + support_transferred (Mapping[str,float]): A dictionary mapping names of candidates to the total amount of support they received from the eliminations. + round_numbers (List[int]): The rounds of the election process associated to this event. + """ support_transferred : Mapping[str,float] round_numbers : List[int] def get_message(self) -> str: @@ -51,9 +73,16 @@ def get_message(self) -> str: message = f"Rounds {self.round_numbers[0]}-{self.round_numbers[-1]}: {len(self.round_numbers)} candidates eliminated." return message - @dataclass class WinEvent(AnimationEvent): + """ + An animation event representing a round in which some number of candidates were elected. + + Attributes: + candidates (str): The names of the elected candidates. + support_transferred (Mapping[str, Mapping[str,float]]): A dictionary mapping pairs of candidate names to the amount of support transferred between them this round. For instance, if ``c1`` was elected this round, then ``support_transferred[c1][c2]`` will represent the amount of support that ran off from ``c1`` to candidate ``c2``. + round_number (int): The round of the election process associated to this event. + """ candidates : Sequence[str] support_transferred : Mapping[str, Mapping[str,float]] round_number : int @@ -66,44 +95,52 @@ def get_message(self) -> str: class STVAnimation: - """A class which creates round-by-round animations of STV elections. + """ + A class which creates round-by-round animations of STV elections. Args: election (STV): An STV election to animate. - title (str): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. + title (str, optional): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. Defaults to ``None``. + focus (List[str], optional): A list of names of candidates that should appear on-screen. This is useful for elections with many candidates. Note that any candidates that won the election are on-screen automatically, so passing an empty list will result in only elected candidates appearing on-screen. Defaults to an empty list. """ def __init__(self, election: STV, title: Optional[str] = None, focus : List[str] = []): + self.focus = focus elected_candidates = [c for s in election.get_elected() for c in s] - self.focus : List[str] = list(set(focus).union(set(elected_candidates))) - self.candidates = self._make_candidate_dict(election) - self.rounds = self._make_event_list(election) + focus += [name for name in elected_candidates if name not in focus] + self.candidate_dict = self._make_candidate_dict(election) + self.events = self._make_event_list(election) self.title = title - def _make_candidate_dict(self, election: STV) -> dict: - """Create the dictionary of candidates and relevant facts about each one. + def _make_candidate_dict(self, election: STV) -> dict[str, dict[str,object]]: + """ + Create the dictionary of candidates and relevant facts about each one. Args: election (STV): An STV election from which to extract the candidates. Returns: - dict: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. + dict[str, dict[str,object]]: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. """ - candidates = { + candidate_dict : dict[str, dict[str,object]] = { name: {"support": support} for name, support in election.election_states[0].scores.items() if name in self.focus } - return candidates + return candidate_dict def _make_event_list(self, election: STV) -> List[AnimationEvent]: - """Process an STV election into a condensed list of only the salient details from each round. + """ + Processes an STV election into a condensed list of animation events which roughly correspond to election rounds. Args: election (STV): The STV election to process. Returns: - List[dict]: A list of dictionaries corresponding to the rounds of the election. Each dictionary records salient attributes of the corresponding round. + List[AnimationEvent]: A list of the events of the election which are worthy of animation. + + Raises: + ValueError: If multiple candidates are eliminated in the same election round. """ events : List[AnimationEvent] = [] for round_number, election_round in enumerate( @@ -161,7 +198,8 @@ def _get_transferred_votes( from_candidates: List[str], event_type: Literal["win", "elimination"], ) -> dict[str, dict[str, float]]: - """Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. + """ + Compute the number of votes transferred from each elected or eliminated candidate to each remaining candidate. Args: election (STV): The election. @@ -172,22 +210,15 @@ def _get_transferred_votes( Returns: dict[str, dict[str, float]]: A nested dictionary. If ``d`` is the return value, ``c1`` was a candidate eliminated this round, and ``c2`` is a remaining candidate, then ``d[c1][c2]`` will be the total support transferred this round from candidate ``c1`` to candidate ``c2``. - Notes: This function supports the election, but not the elimination, of multiple candidates in one round. If ``event_type`` is ``"elimination"`` then ``from_candidates`` should have length 1. - - Raises: - ValueError: If multiple candidates are eliminated in this round. """ prev_profile, prev_state = election.get_step(round_number - 1) current_state = election.election_states[round_number] transfers : dict[str, dict[str, float]] = {} if event_type == "elimination": - if len(from_candidates) > 1: - raise ValueError( - f"Round {round_number} is eliminating multiple candidates ({len(from_candidates)}), which is not supported." - ) + assert len(from_candidates) == 1 from_candidate = from_candidates[0] transfers = {from_candidate : {}} for to_candidate in [c for s in current_state.remaining for c in s if c in self.focus]: @@ -218,9 +249,18 @@ def _get_transferred_votes( return transfers - def _condense_offscreen_events(self, rounds : List[AnimationEvent]) -> List[AnimationEvent]: - return_events : List[AnimationEvent] = [rounds[0]] - for event in rounds[1:]: + def _condense_offscreen_events(self, events : List[AnimationEvent]) -> List[AnimationEvent]: + """ + Take a list of events and condense any consecutive offscreen events into one summarizing event. For instance, if ``events`` contians three offscreen eliminations in a row, this function will condense them into one offscreen elimination of three candidates. + + Args: + events (List[AnimationEvent]): A list of animation events to be condensed. + + Returns: + List[AnimationEvent]: A condensed list of animation events. + """ + return_events : List[AnimationEvent] = [events[0]] + for event in events[1:]: if isinstance( return_events[-1], EliminationOffscreenEvent ) and isinstance( @@ -232,6 +272,16 @@ def _condense_offscreen_events(self, rounds : List[AnimationEvent]) -> List[Anim return return_events def _compose_offscreen_eliminations(self, event1 : EliminationOffscreenEvent, event2: EliminationOffscreenEvent) -> EliminationOffscreenEvent: + """ + Take two offscreen eliminations and "compose" them into a single offscreen elimination event summarizing both. + + Args: + event1 (EliminationOffscreenEvent): The first offscreen elimination event to compose. + event2 (EliminationOffscreenEvent): The second offscreen elimination event to compose. + + Returns: + EliminationOffscreenEvent: One offscreen elimination event summarizing ``event1`` and ``event2``. + """ support_transferred = defaultdict(float) for key, value in event1.support_transferred.items(): support_transferred[key] += value @@ -248,28 +298,30 @@ def _compose_offscreen_eliminations(self, event1 : EliminationOffscreenEvent, ev def render(self, preview: bool = False) -> None: - """Renders the STV animation using Manim. + """ + Renders the STV animation using Manim. The completed video will appear in the directory ``media/videos``. Args: - preview (bool): If ``True``, display the result in a video player immediately upon completing the render. + preview (bool, optional): If ``True``, display the result in a video player immediately upon completing the render. Defaults to False. """ manimation = ElectionScene( - deepcopy(self.candidates), deepcopy(self.rounds), title=self.title + deepcopy(self.candidate_dict), deepcopy(self.events), title=self.title ) manimation.render(preview=preview) class ElectionScene(manim.Scene): - """Class for Manim animation of an STV election. + """ + Class for Manim animation of an STV election. Notes: This class is instantiated by the class ``STVAnimation``. It should not be instantiated directly. Args: - candidates (dict[str,dict]): A dictionary mapping each candidate to a dictionary of attributes of the candidate. - rounds (List[dict]): A list of dictionaries representing the rounds of the election. Each dictionary is a summary of the events in the round it represents. + candidate_dict (dict[str,dict]): A dictionary mapping each candidate to a dictionary of attributes of the candidate. + events (List[AnimationEvent]): A list of animation events to be constructed and rendered. title (str): A string to be displayed at the beginning of the animation as a title screen. If ``None``, the animation will skip the title screen. """ @@ -296,20 +348,20 @@ class ElectionScene(manim.Scene): title_font_size = 48 def __init__( - self, candidates: dict[str, dict], rounds: List[AnimationEvent], title: Optional[str] = None + self, candidate_dict: dict[str, dict], events: List[AnimationEvent], title: Optional[str] = None ): super().__init__() - self.candidates = candidates - self.rounds = rounds + self.candidate_dict = candidate_dict + self.events = events self.title = title self.width = 8 - self.bar_height = 3.5 / len(self.candidates) - self.font_size = 3 * 40 / len(self.candidates) + self.bar_height = 3.5 / len(self.candidate_dict) + self.font_size = 3 * 40 / len(self.candidate_dict) self.bar_opacity = 1 self.bar_buffer_size = self.bar_height self.strikethrough_thickness = self.font_size / 5 - self.max_support = 1.1 * max([round.quota for round in self.rounds]) + self.max_support = 1.1 * max([event.quota for event in self.events]) self.quota_line = None self.ticker_tape_line = None @@ -328,29 +380,29 @@ def construct(self) -> None: self._draw_initial_bars() self._initialize_ticker_tape() - # Go round by round and animate the events - for event_number, round in enumerate(self.rounds): + # Animate each event in turn + for event_number, event in enumerate(self.events): self.wait(2) # Draw or move the quota line - self._update_quota_line(round.quota) + self._update_quota_line(event.quota) self._ticker_animation_shift(event_number) self._ticker_animation_highlight(event_number) - if isinstance(round, EliminationEvent): # Onscreen candidate eliminated + if isinstance(event, EliminationEvent): # Onscreen candidate eliminated # Remove the candidate from the candidate list - eliminated_candidates = {round.candidate : self.candidates.pop(round.candidate)} - self._animate_elimination(eliminated_candidates, round) - elif isinstance(round, EliminationOffscreenEvent):# Offscreen candidate eliminated - self._animate_elimination_offscreen(round) - elif isinstance(round, WinEvent): # Election round + eliminated_candidates = {event.candidate : self.candidate_dict.pop(event.candidate)} + self._animate_elimination(eliminated_candidates, event) + elif isinstance(event, EliminationOffscreenEvent):# Offscreen candidate eliminated + self._animate_elimination_offscreen(event) + elif isinstance(event, WinEvent): # Election round # Remove the candidates from the candidate list elected_candidates = {} - for name in round.candidates: - elected_candidates[name] = self.candidates.pop(name) - self._animate_win(elected_candidates, round) + for name in event.candidates: + elected_candidates[name] = self.candidate_dict.pop(name) + self._animate_win(elected_candidates, event) else: - raise Exception(f"Invalid type for event {round}.") + raise Exception(f"Invalid type for event {event}.") self.wait(2) def _draw_title(self, message: str) -> None: @@ -373,20 +425,20 @@ def _draw_initial_bars(self) -> None: """ # Sort candidates by starting first place votes sorted_candidates = sorted( - self.candidates.keys(), - key=lambda x: self.candidates[x]["support"], + self.candidate_dict.keys(), + key=lambda x: self.candidate_dict[x]["support"], reverse=True, ) # Assign colors for i, name in enumerate(sorted_candidates): color = self.colors[i % len(self.colors)] - self.candidates[name]["color"] = color + self.candidate_dict[name]["color"] = color # Create candidate name text for i, name in enumerate(sorted_candidates): - candidate = self.candidates[name] + candidate = self.candidate_dict[name] candidate["name_text"] = Text( name, font_size=self.font_size, color=candidate["color"] ) @@ -396,22 +448,22 @@ def _draw_initial_bars(self) -> None: else: #The rest of the candidates go below, right justified candidate["name_text"].next_to( - self.candidates[sorted_candidates[i-1]]["name_text"], + self.candidate_dict[sorted_candidates[i-1]]["name_text"], DOWN, buff=self.bar_buffer_size ).align_to( - self.candidates[sorted_candidates[0]]["name_text"], + self.candidate_dict[sorted_candidates[0]]["name_text"], RIGHT ) # Align candidate names to the left group = manim.Group().add(*[ - candidate["name_text"] for candidate in self.candidates.values() + candidate["name_text"] for candidate in self.candidate_dict.values() ]) group.to_edge(LEFT) del group # Create bars - for candidate in self.candidates.values(): + for candidate in self.candidate_dict.values(): candidate["bars"] = [ Rectangle( width=self._support_to_bar_width(candidate["support"]), @@ -440,11 +492,11 @@ def _draw_initial_bars(self) -> None: # Draw the bars and names self.play( - *[FadeIn(self.candidates[name]["name_text"]) for name in sorted_candidates], + *[FadeIn(self.candidate_dict[name]["name_text"]) for name in sorted_candidates], FadeIn(background), ) self.play( - *[Create(self.candidates[name]["bars"][0]) for name in sorted_candidates] + *[Create(self.candidate_dict[name]["bars"][0]) for name in sorted_candidates] ) def _initialize_ticker_tape(self) -> None: @@ -460,8 +512,8 @@ def _initialize_ticker_tape(self) -> None: ) # Keep this line in front of the bars and the quota line self.ticker_tape_line = ticker_line self.ticker_tape = [] - for i, round in enumerate(self.rounds): - new_message = Text(round.get_message(), font_size=24, color=manim.DARK_GRAY) + for i, event in enumerate(self.events): + new_message = Text(event.get_message(), font_size=24, color=manim.DARK_GRAY) if i == 0: new_message.to_edge(DOWN, buff=0).shift(DOWN) else: @@ -477,9 +529,9 @@ def _ticker_animation_shift(self, event_number: int) -> None: Animate the shifting of the ticker tape to display the message for a given round. Args: - round (dict): The round whose message will shift into view. + event_number (int): The index of the event whose message will shift into view. """ - shift_to_round = ( + shift_to_event = ( self.ticker_tape[event_number] .animate.to_edge(DOWN, buff=0) .shift(UP * self.ticker_tape_height / 3) @@ -491,14 +543,14 @@ def _ticker_animation_shift(self, event_number: int) -> None: for i in range(len(self.ticker_tape)) if i != event_number ] - self.play(shift_to_round, *drag_other_messages) + self.play(shift_to_event, *drag_other_messages) def _ticker_animation_highlight(self, event_number : int) -> None: """ Play an animation graying out all ticker tape message but one. Args: - round (dict): The round whose message will be highlighted. + event_number (int): The index of the event whose message will be highlighted. """ highlight_message = self.ticker_tape[event_number].animate.set_color(manim.WHITE) unhighlight_other_messages = [ @@ -515,7 +567,7 @@ def _update_quota_line(self, quota: float) -> None: Args: quota (float): The threshold number of votes necessary to be elected in the current round. """ - some_candidate = list(self.candidates.values())[0] + some_candidate = list(self.candidate_dict.values())[0] if not self.quota_line: # If the quota line doesn't exist yet, draw it. assert self.ticker_tape_line is not None @@ -537,13 +589,13 @@ def _update_quota_line(self, quota: float) -> None: ) ) - def _animate_win(self, from_candidates: dict[str, dict], round: WinEvent) -> None: + def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> None: """ Animate a round in which one or more candidates are elected. Args: from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates elected this round and the values are dictionaries recording the candidate's attributes. - round (dict): A dictionary recording the events of this round. + event (WinEvent): The event to be animated. """ # Box the winners' names winner_boxes = [ @@ -561,12 +613,12 @@ def _animate_win(self, from_candidates: dict[str, dict], round: WinEvent) -> Non old_bars: List[Rectangle] = from_candidate["bars"] new_bars: List[Rectangle] = [] transformations = [] - destinations = round.support_transferred[from_candidate_name] + destinations = event.support_transferred[from_candidate_name] candidate_color = from_candidate["color"] winner_bar = ( Rectangle( - width=self._support_to_bar_width(round.quota), + width=self._support_to_bar_width(event.quota), height=self.bar_height, color=self.bar_color, fill_color=self.win_bar_color, @@ -595,21 +647,21 @@ def _animate_win(self, from_candidates: dict[str, dict], round: WinEvent) -> Non else: sub_bar.next_to(new_bars[-1], LEFT, buff=0) new_bars.append(sub_bar) - self.candidates[destination]["support"] += votes + self.candidate_dict[destination]["support"] += votes # The sub-bars will move to be next to the bars of their destination candidates transformation = sub_bar.animate.next_to( - self.candidates[destination]["bars"][-1], RIGHT, buff=0 + self.candidate_dict[destination]["bars"][-1], RIGHT, buff=0 ) transformations.append(transformation) # Let the new sub-bar be owned by its destination candidate - self.candidates[destination]["bars"] += sub_bar + self.candidate_dict[destination]["bars"] += sub_bar # Create a final short bar representing the exhausted votes exhausted_votes = ( from_candidate["support"] - - round.quota + - event.quota - sum(list(destinations.values())) ) exhausted_bar = Rectangle( @@ -638,13 +690,14 @@ def _animate_win(self, from_candidates: dict[str, dict], round: WinEvent) -> Non self.play(*transformations, Uncreate(exhausted_bar)) def _animate_elimination( - self, from_candidates: dict[str, dict], round: EliminationEvent + self, from_candidates: dict[str, dict], event: EliminationEvent ) -> None: - """Animate a round in which a candidate was eliminated. + """ + Animate a round in which a candidate was eliminated. Args: - from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates eliminated this round and the values are dictionaries recording the candidate's attributes. This function currently only supports one elimination at a time, so from_candidates should be a dictionary with exactly one entry. - round (dict): A dictionary recording the events of this round. + from_candidates (dict[str,dict]): A dictionary in which the keys are the candidates eliminated this round and the values are dictionaries recording the candidate's attributes. + event (EliminationEvent): The event to be animated. Notes: While the interface supports multiple candidate eliminations in one round for future extensibility, this function currently only supports elimination of one candidate at a time. The from_candidates argument should have exactly one entry. @@ -661,7 +714,7 @@ def _animate_elimination( from_candidate_name = list(from_candidates.keys())[0] from_candidate = list(from_candidates.values())[0] - destinations = round.support_transferred + destinations = event.support_transferred # Cross out the candidate name cross = Line( @@ -687,14 +740,14 @@ def _animate_elimination( fill_color=candidate_color, fill_opacity=self.bar_opacity, ) - self.candidates[destination]["support"] += votes + self.candidate_dict[destination]["support"] += votes new_bars.append(sub_bar) transformations.append( new_bars[-1].animate.next_to( - self.candidates[destination]["bars"][-1], RIGHT, buff=0 + self.candidate_dict[destination]["bars"][-1], RIGHT, buff=0 ) ) # The sub-bars will move to be next to the bars of their destination candidates - self.candidates[destination]["bars"].append(sub_bar) + self.candidate_dict[destination]["bars"].append(sub_bar) # Create a final short bar representing the exhausted votes exhausted_votes = from_candidate["support"] - sum(list(destinations.values())) exhausted_bar = Rectangle( @@ -724,14 +777,14 @@ def _animate_elimination( # Animate the exhaustion of votes and moving the sub-bars to the destination bars self.play(Uncreate(exhausted_bar), *transformations) - def _animate_elimination_offscreen(self, round: EliminationOffscreenEvent) -> None: - """Animate a round in which offscreen candidates were eliminated. + def _animate_elimination_offscreen(self, event: EliminationOffscreenEvent) -> None: + """ + Animate a round in which offscreen candidates were eliminated. Args: - round (dict): A dictionary recording the events of this round. + event (EliminationOffscreenEvent) The event to be animated. """ - destinations = round.support_transferred - + destinations = event.support_transferred # Create short bars that will begin offscreen new_bars = [] transformations = [] @@ -745,14 +798,14 @@ def _animate_elimination_offscreen(self, round: EliminationOffscreenEvent) -> No fill_color=self.offscreen_candidate_color, fill_opacity=self.bar_opacity, ) - self.candidates[destination]["support"] += votes + self.candidate_dict[destination]["support"] += votes new_bars.append(sub_bar) transformations.append( new_bars[-1].animate.next_to( - self.candidates[destination]["bars"][-1], RIGHT, buff=0 + self.candidate_dict[destination]["bars"][-1], RIGHT, buff=0 ) ) # The sub-bars will move to be next to the bars of their destination candidates - self.candidates[destination]["bars"].append(sub_bar) + self.candidate_dict[destination]["bars"].append(sub_bar) for bar in new_bars: bar.to_edge(DOWN).shift((self.bar_height + 2) * DOWN) @@ -761,7 +814,8 @@ def _animate_elimination_offscreen(self, round: EliminationOffscreenEvent) -> No self.play(*transformations) def _support_to_bar_width(self, support: float) -> float: - """Convert a number of votes to the width of a bar in manim coordinates representing that many votes. + """ + Convert a number of votes to the width of a bar in manim coordinates representing that many votes. Args: support (float): A number of votes. From 3543415318dc81aaceb050a35f1d3cd4654ed6bb Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 5 Aug 2025 15:00:18 -0400 Subject: [PATCH 20/38] Various small fixes. --- src/votekit/animations.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 0cea7179..9a23512a 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -87,9 +87,7 @@ class WinEvent(AnimationEvent): support_transferred : Mapping[str, Mapping[str,float]] round_number : int def get_message(self) -> str: - candidate_string = self.candidates[0] - for candidate_name in self.candidates[1:]: - candidate_string += f", {candidate_name}" + candidate_string = ", ".join(self.candidates) return f"Round {self.round_number}: {candidate_string} elected." @@ -101,15 +99,20 @@ class STVAnimation: Args: election (STV): An STV election to animate. title (str, optional): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. Defaults to ``None``. - focus (List[str], optional): A list of names of candidates that should appear on-screen. This is useful for elections with many candidates. Note that any candidates that won the election are on-screen automatically, so passing an empty list will result in only elected candidates appearing on-screen. Defaults to an empty list. + focus (List[str], optional): A list of names of candidates that should appear on-screen. This is useful for elections with many candidates. Note that any candidates that won the election are on-screen automatically, so passing an empty list will result in only elected candidates appearing on-screen. If ``None``, focus only the elected candidates. Defaults to ``None``. """ - def __init__(self, election: STV, title: Optional[str] = None, focus : List[str] = []): + def __init__(self, election: STV, title: Optional[str] = None, focus : Optional[List[str]] = None): + if focus is None: focus = [] self.focus = focus elected_candidates = [c for s in election.get_elected() for c in s] focus += [name for name in elected_candidates if name not in focus] self.candidate_dict = self._make_candidate_dict(election) self.events = self._make_event_list(election) + if len(self.candidate_dict) == 0: + raise ValueError("Tried creating animation with no candidates.") + if len(self.events) == 0: + raise ValueError("Tried creating animation with no animation event.") self.title = title def _make_candidate_dict(self, election: STV) -> dict[str, dict[str,object]]: @@ -222,8 +225,8 @@ def _get_transferred_votes( from_candidate = from_candidates[0] transfers = {from_candidate : {}} for to_candidate in [c for s in current_state.remaining for c in s if c in self.focus]: - prev_score = int(prev_state.scores[to_candidate]) - current_score = int(current_state.scores[to_candidate]) + prev_score = prev_state.scores[to_candidate] + current_score = current_state.scores[to_candidate] transfers[from_candidate][to_candidate] = current_score - prev_score elif event_type == "win": ballots_by_fpv = ballots_by_first_cand(prev_profile) @@ -346,6 +349,8 @@ class ElectionScene(manim.Scene): offscreen_sentinel = "__offscreen__" offscreen_candidate_color = manim.GRAY title_font_size = 48 + name_bar_spacing = 0.2 + winner_box_buffer = 0.1 def __init__( self, candidate_dict: dict[str, dict], events: List[AnimationEvent], title: Optional[str] = None @@ -471,7 +476,7 @@ def _draw_initial_bars(self) -> None: color=self.bar_color, fill_color=candidate["color"], fill_opacity=self.bar_opacity, - ).next_to(candidate["name_text"], RIGHT, buff=0.2) + ).next_to(candidate["name_text"], RIGHT, buff=self.name_bar_spacing) ] # Draw a large black rectangle for the background so that the ticker tape vanishes behind it @@ -600,7 +605,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non # Box the winners' names winner_boxes = [ SurroundingRectangle( - from_candidate["name_text"], color=manim.GREEN, buff=0.1 + from_candidate["name_text"], color=manim.GREEN, buff=self.winner_box_buffer ) for from_candidate in from_candidates.values() ] @@ -656,7 +661,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non transformations.append(transformation) # Let the new sub-bar be owned by its destination candidate - self.candidate_dict[destination]["bars"] += sub_bar + self.candidate_dict[destination]["bars"].append(sub_bar) # Create a final short bar representing the exhausted votes exhausted_votes = ( From 93911817b1a182e64b82ef7e41bb832d9e8a0bfa Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 5 Aug 2025 15:36:19 -0400 Subject: [PATCH 21/38] Fixes to formatting and testing --- src/votekit/animations.py | 207 +++++++++++++++++++++++--------------- tests/test_animations.py | 4 +- 2 files changed, 129 insertions(+), 82 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 9a23512a..1419aebd 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -1,5 +1,5 @@ from copy import deepcopy -import manim # type: ignore +import manim # type: ignore from manim import ( Rectangle, SurroundingRectangle, @@ -32,7 +32,9 @@ class AnimationEvent(ABC): Attributes: quota (float): The current election threshold at the start of this event. """ - quota : float + + quota: float + @abstractmethod def get_message(self) -> str: """Generate a message describing the event for the viewer of the animation.""" @@ -49,12 +51,15 @@ class EliminationEvent(AnimationEvent): support_transferred (Mapping[str,float]): A dictionary mapping names of candidates to the amount of support they received from the elimination. round_number (int): The round of the election process associated to this event. """ - candidate : str - support_transferred : Mapping[str,float] - round_number : int + + candidate: str + support_transferred: Mapping[str, float] + round_number: int + def get_message(self) -> str: return f"Round {self.round_number}: {self.candidate} eliminated." + @dataclass class EliminationOffscreenEvent(AnimationEvent): """ @@ -64,8 +69,10 @@ class EliminationOffscreenEvent(AnimationEvent): support_transferred (Mapping[str,float]): A dictionary mapping names of candidates to the total amount of support they received from the eliminations. round_numbers (List[int]): The rounds of the election process associated to this event. """ - support_transferred : Mapping[str,float] - round_numbers : List[int] + + support_transferred: Mapping[str, float] + round_numbers: List[int] + def get_message(self) -> str: if len(self.round_numbers) == 1: return f"Round {self.round_numbers[0]}: 1 candidate eliminated." @@ -73,6 +80,7 @@ def get_message(self) -> str: message = f"Rounds {self.round_numbers[0]}-{self.round_numbers[-1]}: {len(self.round_numbers)} candidates eliminated." return message + @dataclass class WinEvent(AnimationEvent): """ @@ -83,15 +91,16 @@ class WinEvent(AnimationEvent): support_transferred (Mapping[str, Mapping[str,float]]): A dictionary mapping pairs of candidate names to the amount of support transferred between them this round. For instance, if ``c1`` was elected this round, then ``support_transferred[c1][c2]`` will represent the amount of support that ran off from ``c1`` to candidate ``c2``. round_number (int): The round of the election process associated to this event. """ - candidates : Sequence[str] - support_transferred : Mapping[str, Mapping[str,float]] - round_number : int + + candidates: Sequence[str] + support_transferred: Mapping[str, Mapping[str, float]] + round_number: int + def get_message(self) -> str: candidate_string = ", ".join(self.candidates) return f"Round {self.round_number}: {candidate_string} elected." - class STVAnimation: """ A class which creates round-by-round animations of STV elections. @@ -102,8 +111,14 @@ class STVAnimation: focus (List[str], optional): A list of names of candidates that should appear on-screen. This is useful for elections with many candidates. Note that any candidates that won the election are on-screen automatically, so passing an empty list will result in only elected candidates appearing on-screen. If ``None``, focus only the elected candidates. Defaults to ``None``. """ - def __init__(self, election: STV, title: Optional[str] = None, focus : Optional[List[str]] = None): - if focus is None: focus = [] + def __init__( + self, + election: STV, + title: Optional[str] = None, + focus: Optional[List[str]] = None, + ): + if focus is None: + focus = [] self.focus = focus elected_candidates = [c for s in election.get_elected() for c in s] focus += [name for name in elected_candidates if name not in focus] @@ -115,7 +130,7 @@ def __init__(self, election: STV, title: Optional[str] = None, focus : Optional[ raise ValueError("Tried creating animation with no animation event.") self.title = title - def _make_candidate_dict(self, election: STV) -> dict[str, dict[str,object]]: + def _make_candidate_dict(self, election: STV) -> dict[str, dict[str, object]]: """ Create the dictionary of candidates and relevant facts about each one. @@ -125,7 +140,7 @@ def _make_candidate_dict(self, election: STV) -> dict[str, dict[str,object]]: Returns: dict[str, dict[str,object]]: A dictionary whose keys are candidate names and whose values are themselves dictionaries with details about each candidate. """ - candidate_dict : dict[str, dict[str,object]] = { + candidate_dict: dict[str, dict[str, object]] = { name: {"support": support} for name, support in election.election_states[0].scores.items() if name in self.focus @@ -145,7 +160,7 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: Raises: ValueError: If multiple candidates are eliminated in the same election round. """ - events : List[AnimationEvent] = [] + events: List[AnimationEvent] = [] for round_number, election_round in enumerate( election.election_states[1:], start=1 ): @@ -153,9 +168,8 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: elected_candidates = [c for s in election_round.elected for c in s] eliminated_candidates = [c for s in election_round.eliminated for c in s] - if len(elected_candidates) > 0: # Win round - elected_candidates_str = elected_candidates[0] - support_transferred : dict[str, dict[str, float]] = {} + if len(elected_candidates) > 0: # Win round + support_transferred: dict[str, dict[str, float]] = {} if round_number == len(election): # If it's the last round, don't worry about the transferred votes support_transferred = {cand: {} for cand in elected_candidates} @@ -164,30 +178,42 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: election, round_number, elected_candidates, "win" ) events.append( - WinEvent(quota = election.threshold, - candidates = elected_candidates, - support_transferred = support_transferred, - round_number = round_number) + WinEvent( + quota=election.threshold, + candidates=elected_candidates, + support_transferred=support_transferred, + round_number=round_number, + ) ) - elif len(eliminated_candidates) > 0: # Elimination round + elif len(eliminated_candidates) > 0: # Elimination round if len(eliminated_candidates) > 1: - raise ValueError(f"Multiple-elimination rounds not supported. At most one candidate should be eliminated in each round. Candidates eliminated in round {round_number}: {eliminated_candidates}.") + raise ValueError( + f"Multiple-elimination rounds not supported. At most one candidate should be eliminated in each round. Candidates eliminated in round {round_number}: {eliminated_candidates}." + ) eliminated_candidate = eliminated_candidates[0] support_transferred = self._get_transferred_votes( election, round_number, eliminated_candidates, "elimination" ) if eliminated_candidate in self.focus: events.append( - EliminationEvent(quota = election.threshold, - candidate = eliminated_candidate, - support_transferred = support_transferred[eliminated_candidate], - round_number = round_number) + EliminationEvent( + quota=election.threshold, + candidate=eliminated_candidate, + support_transferred=support_transferred[ + eliminated_candidate + ], + round_number=round_number, + ) ) else: events.append( - EliminationOffscreenEvent(quota = election.threshold, - support_transferred = support_transferred[eliminated_candidate], - round_numbers = [round_number]) + EliminationOffscreenEvent( + quota=election.threshold, + support_transferred=support_transferred[ + eliminated_candidate + ], + round_numbers=[round_number], + ) ) events = self._condense_offscreen_events(events) @@ -218,13 +244,15 @@ def _get_transferred_votes( """ prev_profile, prev_state = election.get_step(round_number - 1) current_state = election.election_states[round_number] - - transfers : dict[str, dict[str, float]] = {} + + transfers: dict[str, dict[str, float]] = {} if event_type == "elimination": assert len(from_candidates) == 1 from_candidate = from_candidates[0] - transfers = {from_candidate : {}} - for to_candidate in [c for s in current_state.remaining for c in s if c in self.focus]: + transfers = {from_candidate: {}} + for to_candidate in [ + c for s in current_state.remaining for c in s if c in self.focus + ]: prev_score = prev_state.scores[to_candidate] current_score = current_state.scores[to_candidate] transfers[from_candidate][to_candidate] = current_score - prev_score @@ -241,40 +269,46 @@ def _get_transferred_votes( condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) for b in new_ballots ] - transfer_weights_from_candidate : dict[str, float] = defaultdict(float) + transfer_weights_from_candidate: dict[str, float] = defaultdict(float) for ballot in clean_ballots: if ballot.ranking is not None: (to_candidate,) = ballot.ranking[0] if to_candidate in self.focus: - transfer_weights_from_candidate[to_candidate] += ballot.weight + transfer_weights_from_candidate[ + to_candidate + ] += ballot.weight transfers[from_candidate] = transfer_weights_from_candidate - + return transfers - - def _condense_offscreen_events(self, events : List[AnimationEvent]) -> List[AnimationEvent]: + + def _condense_offscreen_events( + self, events: List[AnimationEvent] + ) -> List[AnimationEvent]: """ Take a list of events and condense any consecutive offscreen events into one summarizing event. For instance, if ``events`` contians three offscreen eliminations in a row, this function will condense them into one offscreen elimination of three candidates. Args: events (List[AnimationEvent]): A list of animation events to be condensed. - + Returns: List[AnimationEvent]: A condensed list of animation events. """ - return_events : List[AnimationEvent] = [events[0]] + return_events: List[AnimationEvent] = [events[0]] for event in events[1:]: - if isinstance( - return_events[-1], EliminationOffscreenEvent - ) and isinstance( + if isinstance(return_events[-1], EliminationOffscreenEvent) and isinstance( event, EliminationOffscreenEvent ): - return_events[-1] = self._compose_offscreen_eliminations(return_events[-1], event) + return_events[-1] = self._compose_offscreen_eliminations( + return_events[-1], event + ) else: return_events.append(event) return return_events - - def _compose_offscreen_eliminations(self, event1 : EliminationOffscreenEvent, event2: EliminationOffscreenEvent) -> EliminationOffscreenEvent: + + def _compose_offscreen_eliminations( + self, event1: EliminationOffscreenEvent, event2: EliminationOffscreenEvent + ) -> EliminationOffscreenEvent: """ Take two offscreen eliminations and "compose" them into a single offscreen elimination event summarizing both. @@ -285,7 +319,7 @@ def _compose_offscreen_eliminations(self, event1 : EliminationOffscreenEvent, ev Returns: EliminationOffscreenEvent: One offscreen elimination event summarizing ``event1`` and ``event2``. """ - support_transferred = defaultdict(float) + support_transferred : dict[str,float] = defaultdict(float) for key, value in event1.support_transferred.items(): support_transferred[key] += value for key, value in event2.support_transferred.items(): @@ -293,13 +327,11 @@ def _compose_offscreen_eliminations(self, event1 : EliminationOffscreenEvent, ev round_numbers = event1.round_numbers + event2.round_numbers quota = event1.quota return EliminationOffscreenEvent( - quota = quota, - support_transferred = support_transferred, - round_numbers = round_numbers + quota=quota, + support_transferred=support_transferred, + round_numbers=round_numbers, ) - - def render(self, preview: bool = False) -> None: """ Renders the STV animation using Manim. @@ -353,7 +385,10 @@ class ElectionScene(manim.Scene): winner_box_buffer = 0.1 def __init__( - self, candidate_dict: dict[str, dict], events: List[AnimationEvent], title: Optional[str] = None + self, + candidate_dict: dict[str, dict], + events: List[AnimationEvent], + title: Optional[str] = None, ): super().__init__() self.candidate_dict = candidate_dict @@ -370,7 +405,7 @@ def __init__( self.quota_line = None self.ticker_tape_line = None - self.ticker_tape : List[Text] = [] + self.ticker_tape: List[Text] = [] def construct(self) -> None: """ @@ -394,13 +429,17 @@ def construct(self) -> None: self._ticker_animation_shift(event_number) self._ticker_animation_highlight(event_number) - if isinstance(event, EliminationEvent): # Onscreen candidate eliminated + if isinstance(event, EliminationEvent): # Onscreen candidate eliminated # Remove the candidate from the candidate list - eliminated_candidates = {event.candidate : self.candidate_dict.pop(event.candidate)} + eliminated_candidates = { + event.candidate: self.candidate_dict.pop(event.candidate) + } self._animate_elimination(eliminated_candidates, event) - elif isinstance(event, EliminationOffscreenEvent):# Offscreen candidate eliminated + elif isinstance( + event, EliminationOffscreenEvent + ): # Offscreen candidate eliminated self._animate_elimination_offscreen(event) - elif isinstance(event, WinEvent): # Election round + elif isinstance(event, WinEvent): # Election round # Remove the candidates from the candidate list elected_candidates = {} for name in event.candidates: @@ -419,7 +458,9 @@ def _draw_title(self, message: str) -> None: """ text = manim.Tex( r"{7cm}\centering " + message, tex_environment="minipage" - ).scale_to_fit_width(10) # We do this one with a TeX minipage to get the text to wrap if it's too long. + ).scale_to_fit_width( + 10 + ) # We do this one with a TeX minipage to get the text to wrap if it's too long. self.play(Create(text)) self.wait(3) self.play(Uncreate(text)) @@ -440,30 +481,28 @@ def _draw_initial_bars(self) -> None: color = self.colors[i % len(self.colors)] self.candidate_dict[name]["color"] = color - # Create candidate name text for i, name in enumerate(sorted_candidates): candidate = self.candidate_dict[name] candidate["name_text"] = Text( name, font_size=self.font_size, color=candidate["color"] ) - if i==0: - #First candidate goes at the top + if i == 0: + # First candidate goes at the top candidate["name_text"].to_edge(UP, buff=self.bar_buffer_size) else: - #The rest of the candidates go below, right justified + # The rest of the candidates go below, right justified candidate["name_text"].next_to( - self.candidate_dict[sorted_candidates[i-1]]["name_text"], + self.candidate_dict[sorted_candidates[i - 1]]["name_text"], DOWN, - buff=self.bar_buffer_size + buff=self.bar_buffer_size, ).align_to( - self.candidate_dict[sorted_candidates[0]]["name_text"], - RIGHT + self.candidate_dict[sorted_candidates[0]]["name_text"], RIGHT ) # Align candidate names to the left - group = manim.Group().add(*[ - candidate["name_text"] for candidate in self.candidate_dict.values() - ]) + group = manim.Group().add( + *[candidate["name_text"] for candidate in self.candidate_dict.values()] + ) group.to_edge(LEFT) del group @@ -494,14 +533,19 @@ def _draw_initial_bars(self) -> None: .set_z_index(-1) ) - # Draw the bars and names self.play( - *[FadeIn(self.candidate_dict[name]["name_text"]) for name in sorted_candidates], + *[ + FadeIn(self.candidate_dict[name]["name_text"]) + for name in sorted_candidates + ], FadeIn(background), ) self.play( - *[Create(self.candidate_dict[name]["bars"][0]) for name in sorted_candidates] + *[ + Create(self.candidate_dict[name]["bars"][0]) + for name in sorted_candidates + ] ) def _initialize_ticker_tape(self) -> None: @@ -550,14 +594,16 @@ def _ticker_animation_shift(self, event_number: int) -> None: ] self.play(shift_to_event, *drag_other_messages) - def _ticker_animation_highlight(self, event_number : int) -> None: + def _ticker_animation_highlight(self, event_number: int) -> None: """ Play an animation graying out all ticker tape message but one. Args: event_number (int): The index of the event whose message will be highlighted. """ - highlight_message = self.ticker_tape[event_number].animate.set_color(manim.WHITE) + highlight_message = self.ticker_tape[event_number].animate.set_color( + manim.WHITE + ) unhighlight_other_messages = [ self.ticker_tape[i].animate.set_color(manim.DARK_GRAY) for i in range(len(self.ticker_tape)) @@ -605,7 +651,9 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non # Box the winners' names winner_boxes = [ SurroundingRectangle( - from_candidate["name_text"], color=manim.GREEN, buff=self.winner_box_buffer + from_candidate["name_text"], + color=manim.GREEN, + buff=self.winner_box_buffer, ) for from_candidate in from_candidates.values() ] @@ -717,7 +765,6 @@ def _animate_elimination( ) del num_eliminated_candidates - from_candidate_name = list(from_candidates.keys())[0] from_candidate = list(from_candidates.values())[0] destinations = event.support_transferred diff --git a/tests/test_animations.py b/tests/test_animations.py index 9eb2c1bb..46ca574b 100644 --- a/tests/test_animations.py +++ b/tests/test_animations.py @@ -22,5 +22,5 @@ def test_init(): animation = STVAnimation(test_election_happy) - assert animation.candidates is not None - assert animation.rounds is not None \ No newline at end of file + assert animation.candidate_dict is not None + assert animation.events is not None From 48ce6b827bef15ee7e674b0a988ac80bfa145059 Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 5 Aug 2025 17:14:25 -0400 Subject: [PATCH 22/38] Fix bug in which exhausted votes would not disappear if no votes were transferred. --- src/votekit/animations.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 1419aebd..08df4df7 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -720,7 +720,8 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non exhausted_bar = Rectangle( width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, - color=self.bar_color, + # color=self.bar_color, + color = manim.GOLD, fill_color=candidate_color, fill_opacity=self.bar_opacity, ) @@ -739,8 +740,9 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non ) # Animate moving the sub-bars to the destination bars, and the destruction of the exhausted votes + transformations.append(Uncreate(exhausted_bar)) if len(transformations) > 0: - self.play(*transformations, Uncreate(exhausted_bar)) + self.play(*transformations) def _animate_elimination( self, from_candidates: dict[str, dict], event: EliminationEvent @@ -802,7 +804,7 @@ def _animate_elimination( self.candidate_dict[destination]["bars"].append(sub_bar) # Create a final short bar representing the exhausted votes exhausted_votes = from_candidate["support"] - sum(list(destinations.values())) - exhausted_bar = Rectangle( + exhausted_bar = Rectangle(`` width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, color=self.bar_color, From 89e51114590c4d9a1103aef0f3643a14b09078d6 Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 5 Aug 2025 17:15:14 -0400 Subject: [PATCH 23/38] Typo --- src/votekit/animations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 08df4df7..e81c9b6e 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -804,7 +804,7 @@ def _animate_elimination( self.candidate_dict[destination]["bars"].append(sub_bar) # Create a final short bar representing the exhausted votes exhausted_votes = from_candidate["support"] - sum(list(destinations.values())) - exhausted_bar = Rectangle(`` + exhausted_bar = Rectangle( width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, color=self.bar_color, From 0c2e430dad780d1d7d2359b683d34bc3bb62a499 Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 5 Aug 2025 17:34:45 -0400 Subject: [PATCH 24/38] Tweak test --- tests/test_animations.py | 41 ++++++++++++++++++++++------------------ 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/tests/test_animations.py b/tests/test_animations.py index 46ca574b..2bf1c3c2 100644 --- a/tests/test_animations.py +++ b/tests/test_animations.py @@ -1,26 +1,31 @@ -from votekit.animations import STVAnimation +from votekit.animations import STVAnimation, AnimationEvent from votekit import Ballot, PreferenceProfile from votekit.elections import STV +import pytest # modified from STV wiki # Election following the "happy path". One elimination or election per round. No ties. No exact quota matches. No funny business. -test_profile_happy = PreferenceProfile( - ballots=( - Ballot(ranking=({"Orange"}, {"Pear"}), weight=3), - Ballot(ranking=({"Pear"}, {"Strawberry"}, {"Cake"}), weight=8), - Ballot(ranking=({"Strawberry"}, {"Orange"}, {"Pear"}), weight=1), - Ballot(ranking=({"Cake"}, {"Chocolate"}), weight=3), - Ballot(ranking=({"Chocolate"}, {"Cake"}, {"Burger"}), weight=1), - Ballot(ranking=({"Burger"}, {"Chicken"}), weight=4), - Ballot(ranking=({"Chicken"}, {"Chocolate"}, {"Burger"}), weight=3), - ), - max_ranking_length=3, -) -test_election_happy = STV(test_profile_happy, m=3) +@pytest.fixture +def election_happy(): + profile_happy = PreferenceProfile( + ballots=( + Ballot(ranking=({"Orange"}, {"Pear"}), weight=3), + Ballot(ranking=({"Pear"}, {"Strawberry"}, {"Cake"}), weight=8), + Ballot(ranking=({"Strawberry"}, {"Orange"}, {"Pear"}), weight=1), + Ballot(ranking=({"Cake"}, {"Chocolate"}), weight=3), + Ballot(ranking=({"Chocolate"}, {"Cake"}, {"Burger"}), weight=1), + Ballot(ranking=({"Burger"}, {"Chicken"}), weight=4), + Ballot(ranking=({"Chicken"}, {"Chocolate"}, {"Burger"}), weight=3), + ), + max_ranking_length=3, + ) + return STV(profile_happy, m=3) -def test_init(): - animation = STVAnimation(test_election_happy) - assert animation.candidate_dict is not None - assert animation.events is not None +def test_STVAnimation_init(election_happy): + animation = STVAnimation(election_happy) + assert isinstance(animation.candidate_dict, dict) + assert isinstance(animation.events, list) + assert "Pear" in animation.candidate_dict.keys() + assert animation.candidate_dict["Pear"]["support"] == 8 From 362dce7f58cb665c8433374bdbf9424104633af1 Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 5 Aug 2025 17:54:21 -0400 Subject: [PATCH 25/38] Small changes to appease poetry. --- src/votekit/animations.py | 4 ++-- tests/test_animations.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index e81c9b6e..de71bff0 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -319,7 +319,7 @@ def _compose_offscreen_eliminations( Returns: EliminationOffscreenEvent: One offscreen elimination event summarizing ``event1`` and ``event2``. """ - support_transferred : dict[str,float] = defaultdict(float) + support_transferred: dict[str, float] = defaultdict(float) for key, value in event1.support_transferred.items(): support_transferred[key] += value for key, value in event2.support_transferred.items(): @@ -721,7 +721,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, # color=self.bar_color, - color = manim.GOLD, + color=manim.GOLD, fill_color=candidate_color, fill_opacity=self.bar_opacity, ) diff --git a/tests/test_animations.py b/tests/test_animations.py index 2bf1c3c2..44a69a76 100644 --- a/tests/test_animations.py +++ b/tests/test_animations.py @@ -1,4 +1,4 @@ -from votekit.animations import STVAnimation, AnimationEvent +from votekit.animations import STVAnimation from votekit import Ballot, PreferenceProfile from votekit.elections import STV import pytest @@ -6,6 +6,7 @@ # modified from STV wiki # Election following the "happy path". One elimination or election per round. No ties. No exact quota matches. No funny business. + @pytest.fixture def election_happy(): profile_happy = PreferenceProfile( From bc7974ecb1cc46ee51fbe982a10a32cacfdbfe72 Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 1 Dec 2025 13:53:23 -0500 Subject: [PATCH 26/38] Fix exhausted bar color. --- src/votekit/animations.py | 5 ++--- tests/test_animations.py | 3 +-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index de71bff0..d37aaaeb 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -319,7 +319,7 @@ def _compose_offscreen_eliminations( Returns: EliminationOffscreenEvent: One offscreen elimination event summarizing ``event1`` and ``event2``. """ - support_transferred: dict[str, float] = defaultdict(float) + support_transferred : dict[str,float] = defaultdict(float) for key, value in event1.support_transferred.items(): support_transferred[key] += value for key, value in event2.support_transferred.items(): @@ -720,8 +720,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non exhausted_bar = Rectangle( width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, - # color=self.bar_color, - color=manim.GOLD, + color=self.bar_color, fill_color=candidate_color, fill_opacity=self.bar_opacity, ) diff --git a/tests/test_animations.py b/tests/test_animations.py index 44a69a76..2bf1c3c2 100644 --- a/tests/test_animations.py +++ b/tests/test_animations.py @@ -1,4 +1,4 @@ -from votekit.animations import STVAnimation +from votekit.animations import STVAnimation, AnimationEvent from votekit import Ballot, PreferenceProfile from votekit.elections import STV import pytest @@ -6,7 +6,6 @@ # modified from STV wiki # Election following the "happy path". One elimination or election per round. No ties. No exact quota matches. No funny business. - @pytest.fixture def election_happy(): profile_happy = PreferenceProfile( From 2758139937fc035090ece3288552bfbe860c2788 Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 1 Dec 2025 14:11:04 -0500 Subject: [PATCH 27/38] Fix formatting --- src/votekit/animations.py | 2 +- tests/test_animations.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index d37aaaeb..501ec701 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -319,7 +319,7 @@ def _compose_offscreen_eliminations( Returns: EliminationOffscreenEvent: One offscreen elimination event summarizing ``event1`` and ``event2``. """ - support_transferred : dict[str,float] = defaultdict(float) + support_transferred: dict[str, float] = defaultdict(float) for key, value in event1.support_transferred.items(): support_transferred[key] += value for key, value in event2.support_transferred.items(): diff --git a/tests/test_animations.py b/tests/test_animations.py index 2bf1c3c2..44a69a76 100644 --- a/tests/test_animations.py +++ b/tests/test_animations.py @@ -1,4 +1,4 @@ -from votekit.animations import STVAnimation, AnimationEvent +from votekit.animations import STVAnimation from votekit import Ballot, PreferenceProfile from votekit.elections import STV import pytest @@ -6,6 +6,7 @@ # modified from STV wiki # Election following the "happy path". One elimination or election per round. No ties. No exact quota matches. No funny business. + @pytest.fixture def election_happy(): profile_happy = PreferenceProfile( From 9ea2005cdb8ed0373836d16ff33020dbd18e7997 Mon Sep 17 00:00:00 2001 From: prismika Date: Mon, 1 Dec 2025 14:36:05 -0500 Subject: [PATCH 28/38] Fix bug in which an animation would fail to render if candidates eliminated offscreen transferred no votes to onscreen candidates. --- src/votekit/animations.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 501ec701..73b3689a 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -864,7 +864,8 @@ def _animate_elimination_offscreen(self, event: EliminationOffscreenEvent) -> No bar.to_edge(DOWN).shift((self.bar_height + 2) * DOWN) # Animate the exhaustion of votes and moving the sub-bars to the destination bars - self.play(*transformations) + if len(transformations) > 0: + self.play(*transformations) def _support_to_bar_width(self, support: float) -> float: """ From 23fcdfcded76cfa0da1e5516e0ce286904a9a191 Mon Sep 17 00:00:00 2001 From: prismika Date: Tue, 2 Dec 2025 17:35:56 -0500 Subject: [PATCH 29/38] Fix bug in which candidates elected while below quota were incorrectly shown to have made quota. --- src/votekit/animations.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 73b3689a..2e953358 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -669,9 +669,10 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non destinations = event.support_transferred[from_candidate_name] candidate_color = from_candidate["color"] + used_votes = min(event.quota, from_candidate["support"]) winner_bar = ( Rectangle( - width=self._support_to_bar_width(event.quota), + width=self._support_to_bar_width(used_votes), height=self.bar_height, color=self.bar_color, fill_color=self.win_bar_color, @@ -714,7 +715,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non # Create a final short bar representing the exhausted votes exhausted_votes = ( from_candidate["support"] - - event.quota + - used_votes - sum(list(destinations.values())) ) exhausted_bar = Rectangle( @@ -725,10 +726,12 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non fill_opacity=self.bar_opacity, ) assert self.quota_line is not None - exhausted_bar.align_to(self.quota_line, LEFT).align_to( - from_candidate["bars"][-1], UP - ) + if len(new_bars) > 0: + exhausted_bar.next_to(new_bars[0], RIGHT, buff=0) + else: + exhausted_bar.next_to(winner_bar, RIGHT, buff=0) exhausted_bar.set_z_index(1) + transformations.append(Uncreate(exhausted_bar)) # Animate the splitting of the old bar into the new sub_bars self.play( @@ -739,7 +742,6 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non ) # Animate moving the sub-bars to the destination bars, and the destruction of the exhausted votes - transformations.append(Uncreate(exhausted_bar)) if len(transformations) > 0: self.play(*transformations) From 5880f978cf9f30551d64974c3f2cfefe1e7d48ac Mon Sep 17 00:00:00 2001 From: jeqcho Date: Fri, 5 Dec 2025 19:07:10 -0500 Subject: [PATCH 30/38] implement schulze --- docs/social_choice_docs/scr.rst | 17 + src/votekit/elections/__init__.py | 2 + .../elections/election_types/__init__.py | 2 + .../election_types/ranking/__init__.py | 2 + .../election_types/ranking/schulze.py | 164 +++++++ .../election_types/ranking/test_schulze.py | 429 ++++++++++++++++++ 6 files changed, 616 insertions(+) create mode 100644 src/votekit/elections/election_types/ranking/schulze.py create mode 100644 tests/elections/election_types/ranking/test_schulze.py diff --git a/docs/social_choice_docs/scr.rst b/docs/social_choice_docs/scr.rst index 129ce865..5cc10a59 100644 --- a/docs/social_choice_docs/scr.rst +++ b/docs/social_choice_docs/scr.rst @@ -495,6 +495,23 @@ Just like Smith method, but user gets to choose the number of winners, :math:`m` Ties are broken with Borda scores. +Ranked Pairs +~~~~~~~~~~~~~ +A Condorcet method that ranks candidates by looking at pairwise victories. For each pair of +candidates, the "margin of victory" is computed as the difference between the number of voters +who prefer one candidate over the other. These margins are sorted from largest to smallest, +and edges are added to a directed graph in this order, skipping any edge that would create a +cycle. The final ranking is determined by the dominating tiers of this graph. + + +Schulze +~~~~~~~ +A Condorcet method based on indirect victories through "beatpaths." If Alice beats Bob +head-to-head, and Bob beats Charlie, then Alice indirectly beats Charlie. A beatpath's +strength is determined by its weakest link. Alice has a "beatpath-win" over Bob if Alice's +strongest beatpath to Bob is stronger than Bob's strongest beatpath back to Alice. The winner +is the candidate not beaten by anyone via beatpath-wins. Always elects the Condorcet winner +when one exists. Score-based diff --git a/src/votekit/elections/__init__.py b/src/votekit/elections/__init__.py index a69bca3b..4ade061f 100644 --- a/src/votekit/elections/__init__.py +++ b/src/votekit/elections/__init__.py @@ -24,6 +24,7 @@ RandomDictator, BoostedRandomDictator, RankedPairs, + Schulze, ) @@ -54,4 +55,5 @@ "RandomDictator", "BoostedRandomDictator", "RankedPairs", + "Schulze", ] diff --git a/src/votekit/elections/election_types/__init__.py b/src/votekit/elections/election_types/__init__.py index fca25c3b..751ce1bb 100644 --- a/src/votekit/elections/election_types/__init__.py +++ b/src/votekit/elections/election_types/__init__.py @@ -15,6 +15,7 @@ RandomDictator, BoostedRandomDictator, RankedPairs, + Schulze, ) @@ -44,4 +45,5 @@ "RandomDictator", "BoostedRandomDictator", "RankedPairs", + "Schulze", ] diff --git a/src/votekit/elections/election_types/ranking/__init__.py b/src/votekit/elections/election_types/ranking/__init__.py index 4b97e3c9..d412adfb 100644 --- a/src/votekit/elections/election_types/ranking/__init__.py +++ b/src/votekit/elections/election_types/ranking/__init__.py @@ -15,6 +15,7 @@ from .random_dictator import RandomDictator from .boosted_random_dictator import BoostedRandomDictator from .ranked_pairs import RankedPairs +from .schulze import Schulze __all__ = [ @@ -34,4 +35,5 @@ "RandomDictator", "BoostedRandomDictator", "RankedPairs", + "Schulze", ] diff --git a/src/votekit/elections/election_types/ranking/schulze.py b/src/votekit/elections/election_types/ranking/schulze.py new file mode 100644 index 00000000..8b6524a5 --- /dev/null +++ b/src/votekit/elections/election_types/ranking/schulze.py @@ -0,0 +1,164 @@ +import networkx as nx +import numpy as np + +from votekit.pref_profile import RankProfile +from votekit.graphs.pairwise_comparison_graph import ( + pairwise_dict, + get_dominating_tiers_digraph, +) +from votekit.utils import tiebreak_set + +from votekit.elections.election_types.ranking.abstract_ranking import RankingElection +from votekit.elections.election_state import ElectionState + + +class Schulze(RankingElection): + """ + See `Schulze Method `_ for more details. + + The Schulze method uses the widest path algorithm to determine winners. For each pair + of candidates, it computes the strength of the strongest path (where the strength of + a path is the strength of its weakest link). Candidate A is preferred to candidate B + if the strongest path from A to B is stronger than the strongest path from B to A. + + Args: + profile (RankProfile): Profile to conduct election on. + m (int, optional): Number of seats to elect. Defaults to 1. + tiebreak (str, optional): Method for breaking ties. Defaults to "lexicographic". + """ + + def __init__( + self, + profile: RankProfile, + tiebreak: str = "lexicographic", + m: int = 1, + ): + if m <= 0: + raise ValueError("m must be strictly positive") + if len(profile.candidates_cast) < m: + raise ValueError("Not enough candidates received votes to be elected.") + self.m = m + self.tiebreak = tiebreak + + def quick_tiebreak_candidates(profile: RankProfile) -> dict[str, float]: + candidate_set = frozenset(profile.candidates) + tiebroken_candidates = tiebreak_set(candidate_set, tiebreak=self.tiebreak) + + if len(tiebroken_candidates) != len(profile.candidates): + raise RuntimeError("Tiebreak did not resolve all candidates.") + + return {next(iter(c)): i for i, c in enumerate(tiebroken_candidates[::-1])} + + super().__init__( + profile, + score_function=quick_tiebreak_candidates, + sort_high_low=True, + ) + + def _is_finished(self): + """ + Check if the election is finished. + """ + # single round election + elected_cands = [c for s in self.get_elected() for c in s] + + if len(elected_cands) == self.m: + return True + return False + + def _run_step( + self, profile: RankProfile, prev_state: ElectionState, store_states=False + ) -> RankProfile: + """ + Run one step of an election from the given profile and previous state. Since this is + a single-round election, this will complete the election and return the final profile. + + The Schulze method computes the strongest paths between all pairs of candidates: + 1. Initialize p[i,j] = d[i,j] - d[j,i] (margin of victory) + 2. For each intermediate candidate k, update p[i,j] = max(p[i,j], min(p[i,k], p[k,j])) + 3. Candidate i beats j if p[i,j] > p[j,i] + + Args: + profile (RankProfile): Profile of ballots. + prev_state (ElectionState): The previous ElectionState. + store_states (bool, optional): Included for compatibility with the base class but not + used in this election type. + + Returns: + RankProfile: The profile of ballots after the round is completed. + """ + # Get pairwise comparison data: d[i,j] = number of voters who prefer i to j + pairwise = pairwise_dict(profile) + candidates = list(profile.candidates_cast) + n = len(candidates) + + # Create candidate index mapping + cand_to_idx = {cand: idx for idx, cand in enumerate(candidates)} + + # Initialize p[i,j] matrix (strongest path strengths) using NumPy + # p[i,j] represents the strength of the strongest path from i to j + p = np.zeros((n, n), dtype=np.float64) + + # Step 1: Initialize p[i,j] = d[i,j] - d[j,i] for all pairs (i != j) + # pairwise_dict returns (a, b): (weight_a, weight_b) where: + # weight_a = number of voters preferring a to b + # weight_b = number of voters preferring b to a + for (a, b), (weight_a, weight_b) in pairwise.items(): + i = cand_to_idx[a] + j = cand_to_idx[b] + # p[i,j] is the margin by which i beats j (can be negative if j beats i) + p[i, j] = weight_a - weight_b + # Also set the reverse direction + p[j, i] = weight_b - weight_a + + # Step 2: Floyd-Warshall algorithm to compute strongest paths + # For each intermediate node k, update all paths i -> j + for k in range(n): + for i in range(n): + if i != k: + for j in range(n): + if j != k and j != i: + # The strength of path i -> j through k is the minimum of + # the strengths of i -> k and k -> j + # We take the maximum of the current path and this new path + p[i, j] = max(p[i, j], min(p[i, k], p[k, j])) + + # Step 3: Build directed graph where i -> j if p[i,j] > p[j,i] + graph: nx.DiGraph = nx.DiGraph() + graph.add_nodes_from(candidates) + + for i in range(n): + for j in range(n): + if i != j and p[i, j] > p[j, i]: + graph.add_edge(candidates[i], candidates[j]) + + # Get dominating tiers from the graph + dominating_tiers = get_dominating_tiers_digraph(graph) + + tiebreak_resolutions = {} + for candidate_tier_set in dominating_tiers: + if len(candidate_tier_set) > 1: + tiebreak_resolutions[frozenset(candidate_tier_set)] = tiebreak_set( + frozenset(candidate_tier_set), tiebreak=self.tiebreak + ) + + ordered_candidates = [ + candidate + for candidate_set in dominating_tiers + for candidate in sorted(candidate_set) + ] + + elected = tuple(frozenset({c}) for c in ordered_candidates[: self.m]) + remaining = tuple(frozenset({c}) for c in ordered_candidates[self.m :]) + + if store_states: + new_state = ElectionState( + round_number=prev_state.round_number + 1, + elected=elected, + remaining=remaining, + tiebreaks=tiebreak_resolutions, + ) + + self.election_states.append(new_state) + + return profile diff --git a/tests/elections/election_types/ranking/test_schulze.py b/tests/elections/election_types/ranking/test_schulze.py new file mode 100644 index 00000000..0c68b232 --- /dev/null +++ b/tests/elections/election_types/ranking/test_schulze.py @@ -0,0 +1,429 @@ +from votekit.elections import Schulze, ElectionState +from votekit.pref_profile import ( + RankProfile, + ScoreProfile, + ProfileError, +) +from votekit.ballot import RankBallot, ScoreBallot +import pytest +import pandas as pd +import numpy as np +from time import time + + +# Wikipedia example for Schulze method +# https://en.wikipedia.org/wiki/Schulze_method +wikipedia_profile = RankProfile( + ballots=( + RankBallot( + ranking=tuple(map(frozenset, [{"A"}, {"C"}, {"B"}, {"E"}, {"D"}])), + weight=5, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"A"}, {"D"}, {"E"}, {"C"}, {"B"}])), + weight=5, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"B"}, {"E"}, {"D"}, {"A"}, {"C"}])), + weight=8, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"C"}, {"A"}, {"B"}, {"E"}, {"D"}])), + weight=3, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"C"}, {"A"}, {"E"}, {"B"}, {"D"}])), + weight=7, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"C"}, {"B"}, {"A"}, {"D"}, {"E"}])), + weight=2, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"D"}, {"C"}, {"E"}, {"B"}, {"A"}])), + weight=7, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"E"}, {"B"}, {"A"}, {"D"}, {"C"}])), + weight=8, + ), + ), + max_ranking_length=5, +) + +electowiki_profile = RankProfile( + ballots=( + RankBallot( + ranking=tuple( + map( + frozenset, + [{"Memphis"}, {"Nashville"}, {"Chattanoga"}, {"Knoxville"}], + ) + ), + weight=42, + ), + RankBallot( + ranking=tuple( + map( + frozenset, + [{"Nashville"}, {"Chattanoga"}, {"Knoxville"}, {"Memphis"}], + ) + ), + weight=26, + ), + RankBallot( + ranking=tuple( + map( + frozenset, + [{"Chattanoga"}, {"Knoxville"}, {"Nashville"}, {"Memphis"}], + ) + ), + weight=15, + ), + RankBallot( + ranking=tuple( + map( + frozenset, + [{"Knoxville"}, {"Chattanoga"}, {"Nashville"}, {"Memphis"}], + ) + ), + weight=15, + ), + ), + max_ranking_length=4, +) + +profile_with_skips = RankProfile( + ballots=( + RankBallot( + ranking=tuple(map(frozenset, [{"C"}, {"D"}, {"A"}, {}])), + weight=42, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"D"}, {"A"}, {}, {"C"}])), + weight=26, + ), + RankBallot( + ranking=tuple(map(frozenset, [{"A"}, {}, {"D"}, {"C"}])), + weight=15, + ), + RankBallot( + ranking=tuple(map(frozenset, [{}, {"A"}, {"D"}, {"C"}])), + weight=15, + ), + ), + max_ranking_length=4, +) + +test_profile_limit_case = RankProfile( + ballots=( + RankBallot( + ranking=tuple(map(frozenset, ["A", "B", "C"])), + weight=48, + ), + RankBallot( + ranking=tuple(map(frozenset, ["B", "C", "A"])), + weight=3, + ), + RankBallot( + ranking=tuple(map(frozenset, ["C", "A", "B"])), + weight=49, + ), + ), + max_ranking_length=3, +) + +borda_ambiguous_profile = RankProfile( + ballots=( + RankBallot( + ranking=tuple(map(frozenset, ["A", "B", "C"])), + weight=48, + ), + RankBallot( + ranking=tuple(map(frozenset, ["B", "C", "A"])), + weight=24, + ), + RankBallot( + ranking=tuple(map(frozenset, ["C", "A"])), + weight=28, + ), + ), + max_ranking_length=3, +) + +dominating_ambiguous_profile = RankProfile( + ballots=( + RankBallot( + ranking=tuple(map(frozenset, ["A", "B", "C", "D"])), + weight=1, + ), + RankBallot( + ranking=tuple(map(frozenset, ["A", "C", "B", "D"])), + weight=1, + ), + ), +) + +profile_tied_set = RankProfile( + ballots=( + RankBallot(ranking=tuple(map(frozenset, [{"A"}, {"B"}, {"C"}]))), + RankBallot(ranking=tuple(map(frozenset, [{"A"}, {"C"}, {"B"}]))), + RankBallot(ranking=tuple(map(frozenset, [{"B"}, {"A"}, {"C"}])), weight=2), + ), + max_ranking_length=3, +) + + +profile_cycle = RankProfile( + ballots=( + RankBallot(ranking=tuple(map(frozenset, ({"A"}, {"B"}, {"C"})))), + RankBallot(ranking=tuple(map(frozenset, ({"A"}, {"C"}, {"B"})))), + RankBallot(ranking=tuple(map(frozenset, ({"B"}, {"A"}, {"C"})))), + ), + max_ranking_length=3, +) + +profile_tied_borda = RankProfile( + ballots=( + RankBallot(ranking=tuple(map(frozenset, ({"A"}, {"B"}, {"C"})))), + RankBallot(ranking=tuple(map(frozenset, ({"A"}, {"C"}, {"B"})))), + RankBallot(ranking=tuple(map(frozenset, ({"B"}, {"A"}, {"C"})))), + RankBallot(ranking=tuple(map(frozenset, ({"B"}, {"C"}, {"A"})))), + RankBallot(ranking=tuple(map(frozenset, ({"C"}, {"A"}, {"B"})))), + RankBallot(ranking=tuple(map(frozenset, ({"C"}, {"B"}, {"A"})))), + ), + max_ranking_length=3, +) + + +def convert_to_fs_tuple(lst): + """ + Convert a list of strings to a tuple of frozensets. + """ + return tuple(map(frozenset, lst)) + + +def test_wikipedia_example(): + """Test the example from Wikipedia's Schulze method page. + + Number of voters | Order of preference + 5 | A C B E D + 5 | A D E C B + 8 | B E D A C + 3 | C A B E D + 7 | C A E B D + 2 | C B A D E + 7 | D C E B A + 8 | E B A D C + + Expected winner is E according to the Wikipedia article. + """ + e = Schulze(wikipedia_profile, m=1) + assert e.get_elected() == convert_to_fs_tuple(["E"]) + + +def test_init(): + e = Schulze(electowiki_profile, m=1) + assert e.get_elected() == (frozenset({"Nashville"}),) + + +def test_init_with_empty(): + e = Schulze(profile_with_skips, m=3) + assert e.get_elected() == convert_to_fs_tuple(["D", "A", "C"]) + + +def test_init_with_empty_errors_out_when_too_many_elected(): + with pytest.raises(ValueError): + Schulze(profile_with_skips, m=4) + + +def test_limit_case(): + e = Schulze(test_profile_limit_case, m=2) + assert e.get_elected() == convert_to_fs_tuple(["C", "A"]) + + +def test_borda_ambiguous_profile_returns_lexicographic_order(): + e = Schulze(borda_ambiguous_profile, m=1) + assert e.get_elected() == convert_to_fs_tuple(["A"]) + e = Schulze(borda_ambiguous_profile, m=2) + assert e.get_elected() == convert_to_fs_tuple(["A", "B"]) + e = Schulze(borda_ambiguous_profile, m=3) + assert e.get_elected() == convert_to_fs_tuple(["A", "B", "C"]) + + +def test_dominating_ambigous_profile_returns_lexicographic_order(): + e = Schulze(dominating_ambiguous_profile, m=1) + assert e.get_elected() == convert_to_fs_tuple(["A"]) + e = Schulze(dominating_ambiguous_profile, m=2) + assert e.get_elected() == convert_to_fs_tuple(["A", "B"]) + e = Schulze(dominating_ambiguous_profile, m=3) + assert e.get_elected() == convert_to_fs_tuple(["A", "B", "C"]) + + +def test_tied_set(): + e = Schulze(profile_tied_set, m=1) + assert e.get_elected() == convert_to_fs_tuple(["A"]) + e = Schulze(profile_tied_set, m=2) + assert e.get_elected() == convert_to_fs_tuple(["A", "B"]) + e = Schulze(profile_tied_set, m=3) + assert e.get_elected() == convert_to_fs_tuple(["A", "B", "C"]) + + +def test_profile_cycle(): + e = Schulze(profile_cycle, m=1) + assert e.get_elected() == convert_to_fs_tuple(["A"]) + e = Schulze(profile_cycle, m=2) + assert e.get_elected() == convert_to_fs_tuple(["A", "B"]) + e = Schulze(profile_cycle, m=3) + assert e.get_elected() == convert_to_fs_tuple(["A", "B", "C"]) + + +def test_tied_borda(): + e = Schulze(profile_tied_borda, m=1) + assert e.get_elected() == convert_to_fs_tuple(["A"]) + e = Schulze(profile_tied_borda, m=2) + assert e.get_elected() == convert_to_fs_tuple(["A", "B"]) + e = Schulze(profile_tied_borda, m=3) + assert e.get_elected() == convert_to_fs_tuple(["A", "B", "C"]) + + +def test_errors(): + with pytest.raises(ValueError, match="m must be strictly positive"): + Schulze(profile_tied_set, m=0) + + with pytest.raises( + ValueError, match="Not enough candidates received votes to be elected." + ): + Schulze(profile_tied_set, m=4) + + with pytest.raises(ProfileError, match="Profile must be of type RankProfile."): + Schulze(ScoreProfile(ballots=(ScoreBallot(scores={"A": 4}),))) # type: ignore + + +@pytest.mark.slow +def test_large_set_timing(): + alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + ballots_tup = tuple( + [ + RankBallot( + ranking=tuple( + map( + lambda x: frozenset({x}), + np.random.permutation(list(alphabet))[ + : np.random.randint(1, len(alphabet) + 1) + ], + ) + ), + weight=np.random.randint(1, 100), + ) + for _ in range(100_000) + ] + ) + + prof = RankProfile( + ballots=ballots_tup, + max_ranking_length=26, + ) + + start = time() + for _ in range(10): + Schulze( + prof, + m=5, + ) + end = time() + + assert ( + end - start < 120 + ), f"Schulze runtime took too long. Expected < 120 seconds, got {end - start} seconds." + + +states = [ + ElectionState( + round_number=0, + remaining=(frozenset({"A"}), frozenset({"B"}), frozenset({"C"})), + elected=(frozenset(),), + eliminated=(frozenset(),), + tiebreaks={}, + scores={"C": 0, "B": 1, "A": 2}, + ), + ElectionState( + round_number=1, + remaining=(frozenset({"B"}), frozenset({"C"})), + elected=(frozenset({"A"}),), + eliminated=(frozenset(),), + tiebreaks={frozenset({"A", "B"}): (frozenset({"A"}), frozenset({"B"}))}, + scores={}, + ), +] + + +def test_state_list(): + e = Schulze(profile_tied_set) + assert e.election_states == states + + +def test_get_profile(): + e = Schulze(profile_tied_set) + assert e.get_profile(0) == profile_tied_set + + +def test_get_step(): + e = Schulze(profile_tied_set) + profile, state = e.get_step(1) + assert profile.group_ballots(), state == (profile_tied_set, states[1]) + + +def test_get_step_does_not_extend_election_states(): + e = Schulze(profile_tied_set) + assert len(e.election_states) == 2 + + profile, state = e.get_step(1) + assert profile.group_ballots(), state == (profile_tied_set, states[1]) + assert len(e.election_states) == 2 + + profile, state = e.get_step(1) + assert len(e.election_states) == 2 + + +def test_get_elected(): + e = Schulze(profile_tied_set) + assert e.get_elected(0) == tuple() + assert e.get_elected(1) == (frozenset({"A"}),) + + +def test_get_eliminated(): + e = Schulze(profile_tied_set) + assert e.get_eliminated(0) == tuple() + assert e.get_eliminated(1) == tuple() + + +def test_get_remaining(): + e = Schulze(profile_tied_set) + assert e.get_remaining(0) == (frozenset({"A"}), frozenset({"B"}), frozenset({"C"})) + assert e.get_remaining(1) == ( + frozenset({"B"}), + frozenset({"C"}), + ) + + +def test_get_ranking(): + e = Schulze(profile_tied_set) + assert e.get_ranking(0) == (frozenset({"A"}), frozenset({"B"}), frozenset({"C"})) + assert e.get_ranking(1) == (frozenset({"A"}), frozenset({"B"}), frozenset({"C"})) + + +def test_get_status_df(): + e = Schulze(profile_tied_set) + + df_0 = pd.DataFrame( + {"Status": ["Remaining"] * 3, "Round": [0] * 3}, + index=["A", "B", "C"], + ) + df_1 = pd.DataFrame( + {"Status": ["Elected", "Remaining", "Remaining"], "Round": [1] * 3}, + index=["A", "B", "C"], + ) + + assert e.get_status_df(0).sort_index().equals(df_0) + assert e.get_status_df(1).sort_index().equals(df_1) From 8a225684bc36578f914a5277174c39c4f445c42b Mon Sep 17 00:00:00 2001 From: Jay Chooi <42904912+jeqcho@users.noreply.github.com> Date: Wed, 10 Dec 2025 21:03:44 -0500 Subject: [PATCH 31/38] Update docs/social_choice_docs/scr.rst Co-authored-by: Chris Donnay --- docs/social_choice_docs/scr.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/social_choice_docs/scr.rst b/docs/social_choice_docs/scr.rst index 5cc10a59..192d3e9d 100644 --- a/docs/social_choice_docs/scr.rst +++ b/docs/social_choice_docs/scr.rst @@ -511,7 +511,7 @@ head-to-head, and Bob beats Charlie, then Alice indirectly beats Charlie. A beat strength is determined by its weakest link. Alice has a "beatpath-win" over Bob if Alice's strongest beatpath to Bob is stronger than Bob's strongest beatpath back to Alice. The winner is the candidate not beaten by anyone via beatpath-wins. Always elects the Condorcet winner -when one exists. +when one exists. This method is capable of producing an output ranking of candidates. Score-based From cb1f6ca95ed5506da009f20413ef5cf0e4f95b95 Mon Sep 17 00:00:00 2001 From: jeqcho Date: Wed, 10 Dec 2025 21:08:11 -0500 Subject: [PATCH 32/38] use numpy broadcasting for schulze inner loop --- .../elections/election_types/ranking/schulze.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/votekit/elections/election_types/ranking/schulze.py b/src/votekit/elections/election_types/ranking/schulze.py index 8b6524a5..ccf9c078 100644 --- a/src/votekit/elections/election_types/ranking/schulze.py +++ b/src/votekit/elections/election_types/ranking/schulze.py @@ -111,17 +111,12 @@ def _run_step( # Also set the reverse direction p[j, i] = weight_b - weight_a - # Step 2: Floyd-Warshall algorithm to compute strongest paths - # For each intermediate node k, update all paths i -> j + # Step 2: Floyd-Warshall style algorithm to compute strongest (widest) paths + # Schulze requires: p[i,j] = max(p[i,j], min(p[i,k], p[k,j])) + # We use NumPy broadcasting to vectorize the inner two loops for performance. for k in range(n): - for i in range(n): - if i != k: - for j in range(n): - if j != k and j != i: - # The strength of path i -> j through k is the minimum of - # the strengths of i -> k and k -> j - # We take the maximum of the current path and this new path - p[i, j] = max(p[i, j], min(p[i, k], p[k, j])) + # p[:, k:k+1] is column k (shape n x 1), p[k:k+1, :] is row k (shape 1 x n) + p = np.maximum(p, np.minimum(p[:, k : k + 1], p[k : k + 1, :])) # Step 3: Build directed graph where i -> j if p[i,j] > p[j,i] graph: nx.DiGraph = nx.DiGraph() From b1652543f5f34c8f7cadac9bd201f3b3239a352c Mon Sep 17 00:00:00 2001 From: jeqcho Date: Wed, 10 Dec 2025 21:15:00 -0500 Subject: [PATCH 33/38] documentation changes for schulze --- docs/social_choice_docs/scr.rst | 10 ++++++---- .../elections/election_types/ranking/schulze.py | 7 ++++++- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/docs/social_choice_docs/scr.rst b/docs/social_choice_docs/scr.rst index 192d3e9d..1d369bdb 100644 --- a/docs/social_choice_docs/scr.rst +++ b/docs/social_choice_docs/scr.rst @@ -508,10 +508,12 @@ Schulze ~~~~~~~ A Condorcet method based on indirect victories through "beatpaths." If Alice beats Bob head-to-head, and Bob beats Charlie, then Alice indirectly beats Charlie. A beatpath's -strength is determined by its weakest link. Alice has a "beatpath-win" over Bob if Alice's -strongest beatpath to Bob is stronger than Bob's strongest beatpath back to Alice. The winner -is the candidate not beaten by anyone via beatpath-wins. Always elects the Condorcet winner -when one exists. This method is capable of producing an output ranking of candidates. +strength is determined by its weakest link. For example, if Alice beats Bob by 2 votes, +Bob beats Charlie by 4 votes, then the beatpath strength from Alice to Bob is 2. Alice +has a "beatpath-win" over Bob if Alice's strongest beatpath to Bob is stronger than Bob's +strongest beatpath back to Alice. The winner is the candidate not beaten by anyone via +beatpath-wins. Always elects the Condorcet winner when one exists. This method is capable +of producing an output ranking of candidates. Score-based diff --git a/src/votekit/elections/election_types/ranking/schulze.py b/src/votekit/elections/election_types/ranking/schulze.py index ccf9c078..104c1ebc 100644 --- a/src/votekit/elections/election_types/ranking/schulze.py +++ b/src/votekit/elections/election_types/ranking/schulze.py @@ -14,13 +14,18 @@ class Schulze(RankingElection): """ - See `Schulze Method `_ for more details. + See and The Schulze method uses the widest path algorithm to determine winners. For each pair of candidates, it computes the strength of the strongest path (where the strength of a path is the strength of its weakest link). Candidate A is preferred to candidate B if the strongest path from A to B is stronger than the strongest path from B to A. + The Schulze method computes the strongest paths between all pairs of candidates: + 1. Initialize p[i,j] = d[i,j] - d[j,i] (margin of victory) + 2. For each intermediate candidate k, update p[i,j] = max(p[i,j], min(p[i,k], p[k,j])) + 3. Candidate i beats j if p[i,j] > p[j,i] + Args: profile (RankProfile): Profile to conduct election on. m (int, optional): Number of seats to elect. Defaults to 1. From ebb98c86629ab157d217b104bce24fddd392cc05 Mon Sep 17 00:00:00 2001 From: jeqcho Date: Wed, 10 Dec 2025 21:21:09 -0500 Subject: [PATCH 34/38] more doc changes for schulze --- src/votekit/elections/election_types/ranking/schulze.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/votekit/elections/election_types/ranking/schulze.py b/src/votekit/elections/election_types/ranking/schulze.py index 104c1ebc..a3d664c9 100644 --- a/src/votekit/elections/election_types/ranking/schulze.py +++ b/src/votekit/elections/election_types/ranking/schulze.py @@ -18,8 +18,10 @@ class Schulze(RankingElection): The Schulze method uses the widest path algorithm to determine winners. For each pair of candidates, it computes the strength of the strongest path (where the strength of - a path is the strength of its weakest link). Candidate A is preferred to candidate B - if the strongest path from A to B is stronger than the strongest path from B to A. + a path is the strength of its weakest link). For example, if Alice beats Bob by 2 votes, + Bob beats Charlie by 4 votes, then the beatpath strength from Alice to Bob is 2. Candidate + A is preferred to candidate B if the strongest path from A to B is stronger than the + strongest path from B to A. The Schulze method computes the strongest paths between all pairs of candidates: 1. Initialize p[i,j] = d[i,j] - d[j,i] (margin of victory) From 10588c750ea127116c46850c1a385a3936a9b08a Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 11 Dec 2025 13:44:34 -0500 Subject: [PATCH 35/38] Update animation code to reflect VoteKit changes. --- src/votekit/animations.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 2e953358..4d748082 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -14,7 +14,10 @@ LEFT, RIGHT, ) -from .cleaning import condense_ballot_ranking, remove_cand_from_ballot +from .cleaning.rank_ballots_cleaning import ( + condense_rank_ballot, + remove_cand_rank_ballot, +) from .utils import ballots_by_first_cand from .elections.election_types.ranking.stv import STV from typing import Literal, List, Optional, Sequence, Mapping @@ -266,7 +269,7 @@ def _get_transferred_votes( election.threshold, ) clean_ballots = [ - condense_ballot_ranking(remove_cand_from_ballot(from_candidates, b)) + condense_rank_ballot(remove_cand_rank_ballot(from_candidates, b)) for b in new_ballots ] transfer_weights_from_candidate: dict[str, float] = defaultdict(float) From 685c578725dba25af52dc06afc8e40a260b981fd Mon Sep 17 00:00:00 2001 From: prismika Date: Thu, 11 Dec 2025 14:44:20 -0500 Subject: [PATCH 36/38] Allow users to provide candidate nicknames to appear in animations. --- src/votekit/animations.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 4d748082..23163ac5 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -51,16 +51,18 @@ class EliminationEvent(AnimationEvent): Attributes: candidate (str): The name of the eliminated candidate. + display_name (str): The candidate name to use for display purposes, such as a nickname. support_transferred (Mapping[str,float]): A dictionary mapping names of candidates to the amount of support they received from the elimination. round_number (int): The round of the election process associated to this event. """ candidate: str + display_name: str support_transferred: Mapping[str, float] round_number: int def get_message(self) -> str: - return f"Round {self.round_number}: {self.candidate} eliminated." + return f"Round {self.round_number}: {self.display_name} eliminated." @dataclass @@ -90,17 +92,19 @@ class WinEvent(AnimationEvent): An animation event representing a round in which some number of candidates were elected. Attributes: - candidates (str): The names of the elected candidates. + candidates (Sequence[str]): The names of the elected candidates. + display_names (Sequence[str]): The candidate names to use for display purposes, such as nicknames. support_transferred (Mapping[str, Mapping[str,float]]): A dictionary mapping pairs of candidate names to the amount of support transferred between them this round. For instance, if ``c1`` was elected this round, then ``support_transferred[c1][c2]`` will represent the amount of support that ran off from ``c1`` to candidate ``c2``. round_number (int): The round of the election process associated to this event. """ candidates: Sequence[str] + display_names: Sequence[str] support_transferred: Mapping[str, Mapping[str, float]] round_number: int def get_message(self) -> str: - candidate_string = ", ".join(self.candidates) + candidate_string = ", ".join(self.display_names) return f"Round {self.round_number}: {candidate_string} elected." @@ -112,6 +116,7 @@ class STVAnimation: election (STV): An STV election to animate. title (str, optional): Text to be displayed at the beginning of the animation as a title screen. If ``None``, the title screen will be skipped. Defaults to ``None``. focus (List[str], optional): A list of names of candidates that should appear on-screen. This is useful for elections with many candidates. Note that any candidates that won the election are on-screen automatically, so passing an empty list will result in only elected candidates appearing on-screen. If ``None``, focus only the elected candidates. Defaults to ``None``. + nicknames (dict[str,str]): A dictionary mapping candidate names to candidate "nicknames" to be used in the animation instead. The keys of ``nicknames`` need not contain every candidate, only the ones for which the user would like to provide a nickname. """ def __init__( @@ -119,12 +124,14 @@ def __init__( election: STV, title: Optional[str] = None, focus: Optional[List[str]] = None, + nicknames: dict[str, str] = {}, ): if focus is None: focus = [] self.focus = focus elected_candidates = [c for s in election.get_elected() for c in s] focus += [name for name in elected_candidates if name not in focus] + self.nicknames = nicknames self.candidate_dict = self._make_candidate_dict(election) self.events = self._make_event_list(election) if len(self.candidate_dict) == 0: @@ -148,6 +155,12 @@ def _make_candidate_dict(self, election: STV) -> dict[str, dict[str, object]]: for name, support in election.election_states[0].scores.items() if name in self.focus } + for name in candidate_dict.keys(): + if name in self.nicknames.keys(): + display_name = self.nicknames[name] + else: + display_name = name + candidate_dict[name]["display_name"] = display_name return candidate_dict def _make_event_list(self, election: STV) -> List[AnimationEvent]: @@ -180,10 +193,15 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: support_transferred = self._get_transferred_votes( election, round_number, elected_candidates, "win" ) + display_names = [ + str(self.candidate_dict[name]["display_name"]) + for name in elected_candidates + ] events.append( WinEvent( quota=election.threshold, candidates=elected_candidates, + display_names=display_names, support_transferred=support_transferred, round_number=round_number, ) @@ -198,10 +216,14 @@ def _make_event_list(self, election: STV) -> List[AnimationEvent]: election, round_number, eliminated_candidates, "elimination" ) if eliminated_candidate in self.focus: + display_name = str( + self.candidate_dict[eliminated_candidate]["display_name"] + ) events.append( EliminationEvent( quota=election.threshold, candidate=eliminated_candidate, + display_name=display_name, support_transferred=support_transferred[ eliminated_candidate ], @@ -488,7 +510,9 @@ def _draw_initial_bars(self) -> None: for i, name in enumerate(sorted_candidates): candidate = self.candidate_dict[name] candidate["name_text"] = Text( - name, font_size=self.font_size, color=candidate["color"] + candidate["display_name"], + font_size=self.font_size, + color=candidate["color"], ) if i == 0: # First candidate goes at the top From 46bc2ce89f7b3b1b7270740c6adcd63b9dc91295 Mon Sep 17 00:00:00 2001 From: prismika Date: Fri, 12 Dec 2025 15:45:48 -0500 Subject: [PATCH 37/38] Refactor all color decisions into a centralized color palette. --- src/votekit/animations.py | 86 +++++++++++++++++++++++---------------- 1 file changed, 51 insertions(+), 35 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 23163ac5..6025b96a 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -385,30 +385,37 @@ class ElectionScene(manim.Scene): title (str): A string to be displayed at the beginning of the animation as a title screen. If ``None``, the animation will skip the title screen. """ - colors = [ - manim.color.ManimColor(hex) - for hex in [ - "#16DEBD", - "#163EDE", - "#9F34F6", - "#FF6F00", - "#8F560C", - "#E2AD00", - "#8AD412", - ] - ] - bar_color = manim.LIGHT_GRAY + color_palette = { + "bar_fills": [ + manim.color.ManimColor(hex) + for hex in [ + "#16DEBD", + "#163EDE", + "#9F34F6", + "#FF6F00", + "#8F560C", + "#E2AD00", + "#8AD412", + ] + ], + "bar_outline": manim.LIGHT_GRAY, + "win_bar_fill": manim.GREEN, + "win_box_outline": manim.GREEN, + "offscreen_candidate_fill": manim.GRAY, + "background": manim.BLACK, + "elimination_line": manim.RED, + "ticker_tape_frosted": manim.DARK_GRAY, + "ticker_tape_highlight": manim.WHITE, + } bar_opacity = 1 - win_bar_color = manim.GREEN - eliminated_bar_color = manim.RED ghost_opacity = 0.3 ticker_tape_height = 2 - offscreen_sentinel = "__offscreen__" - offscreen_candidate_color = manim.GRAY title_font_size = 48 name_bar_spacing = 0.2 winner_box_buffer = 0.1 + offscreen_sentinel = "__offscreen__" + def __init__( self, candidate_dict: dict[str, dict], @@ -502,8 +509,9 @@ def _draw_initial_bars(self) -> None: ) # Assign colors + bar_fill_colors = self.color_palette["bar_fills"] for i, name in enumerate(sorted_candidates): - color = self.colors[i % len(self.colors)] + color = bar_fill_colors[i % len(bar_fill_colors)] self.candidate_dict[name]["color"] = color # Create candidate name text @@ -539,7 +547,7 @@ def _draw_initial_bars(self) -> None: Rectangle( width=self._support_to_bar_width(candidate["support"]), height=self.bar_height, - color=self.bar_color, + color=self.color_palette["bar_outline"], fill_color=candidate["color"], fill_opacity=self.bar_opacity, ).next_to(candidate["name_text"], RIGHT, buff=self.name_bar_spacing) @@ -552,8 +560,8 @@ def _draw_initial_bars(self) -> None: Rectangle( width=frame_width, height=frame_height, - fill_color=manim.BLACK, - color=manim.BLACK, + fill_color=self.color_palette["background"], + color=self.color_palette["background"], fill_opacity=1, ) .shift(UP * self.ticker_tape_height) @@ -589,7 +597,11 @@ def _initialize_ticker_tape(self) -> None: self.ticker_tape_line = ticker_line self.ticker_tape = [] for i, event in enumerate(self.events): - new_message = Text(event.get_message(), font_size=24, color=manim.DARK_GRAY) + new_message = Text( + event.get_message(), + font_size=24, + color=self.color_palette["ticker_tape_frosted"], + ) if i == 0: new_message.to_edge(DOWN, buff=0).shift(DOWN) else: @@ -629,10 +641,12 @@ def _ticker_animation_highlight(self, event_number: int) -> None: event_number (int): The index of the event whose message will be highlighted. """ highlight_message = self.ticker_tape[event_number].animate.set_color( - manim.WHITE + self.color_palette["ticker_tape_highlight"] ) unhighlight_other_messages = [ - self.ticker_tape[i].animate.set_color(manim.DARK_GRAY) + self.ticker_tape[i].animate.set_color( + self.color_palette["ticker_tape_frosted"] + ) for i in range(len(self.ticker_tape)) if i != event_number ] @@ -652,7 +666,9 @@ def _update_quota_line(self, quota: float) -> None: line_bottom = self.ticker_tape_line.get_top()[1] line_top = manim.config.frame_height / 2 self.quota_line = Line( - start=UP * line_top, end=UP * line_bottom, color=self.win_bar_color + start=UP * line_top, + end=UP * line_bottom, + color=self.color_palette["win_bar_fill"], ) self.quota_line.align_to(some_candidate["bars"][0], LEFT) self.quota_line.shift((self.width * quota / self.max_support) * RIGHT) @@ -679,7 +695,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non winner_boxes = [ SurroundingRectangle( from_candidate["name_text"], - color=manim.GREEN, + color=self.color_palette["win_box_outline"], buff=self.winner_box_buffer, ) for from_candidate in from_candidates.values() @@ -701,8 +717,8 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non Rectangle( width=self._support_to_bar_width(used_votes), height=self.bar_height, - color=self.bar_color, - fill_color=self.win_bar_color, + color=self.color_palette["bar_outline"], + fill_color=self.color_palette["win_bar_fill"], fill_opacity=self.bar_opacity, ) .align_to(from_candidate["bars"][0], LEFT) @@ -716,7 +732,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non sub_bar = Rectangle( width=self._support_to_bar_width(votes), height=self.bar_height, - color=self.bar_color, + color=self.color_palette["bar_outline"], fill_color=candidate_color, fill_opacity=self.bar_opacity, ) @@ -748,7 +764,7 @@ def _animate_win(self, from_candidates: dict[str, dict], event: WinEvent) -> Non exhausted_bar = Rectangle( width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, - color=self.bar_color, + color=self.color_palette["bar_outline"], fill_color=candidate_color, fill_opacity=self.bar_opacity, ) @@ -802,7 +818,7 @@ def _animate_elimination( cross = Line( from_candidate["name_text"].get_left(), from_candidate["name_text"].get_right(), - color=manim.RED, + color=self.color_palette["elimination_line"], ) cross.set_stroke(width=self.strikethrough_thickness) self.play(Create(cross)) @@ -818,7 +834,7 @@ def _animate_elimination( sub_bar = Rectangle( width=self._support_to_bar_width(votes), height=self.bar_height, - color=self.bar_color, + color=self.color_palette["bar_outline"], fill_color=candidate_color, fill_opacity=self.bar_opacity, ) @@ -835,7 +851,7 @@ def _animate_elimination( exhausted_bar = Rectangle( width=self._support_to_bar_width(exhausted_votes), height=self.bar_height, - color=self.bar_color, + color=self.color_palette["bar_outline"], fill_color=candidate_color, fill_opacity=self.bar_opacity, ) @@ -876,8 +892,8 @@ def _animate_elimination_offscreen(self, event: EliminationOffscreenEvent) -> No sub_bar = Rectangle( width=self._support_to_bar_width(votes), height=self.bar_height, - color=self.bar_color, - fill_color=self.offscreen_candidate_color, + color=self.color_palette["bar_outline"], + fill_color=self.color_palette["offscreen_candidate_fill"], fill_opacity=self.bar_opacity, ) self.candidate_dict[destination]["support"] += votes From 10ef53f7d7d99cfe4e003f6c5759240ad06c9a66 Mon Sep 17 00:00:00 2001 From: prismika Date: Sun, 28 Dec 2025 21:42:04 -0500 Subject: [PATCH 38/38] Implement light mode. --- src/votekit/animations.py | 82 ++++++++++++++++++++++++++++----------- 1 file changed, 59 insertions(+), 23 deletions(-) diff --git a/src/votekit/animations.py b/src/votekit/animations.py index 6025b96a..97043092 100644 --- a/src/votekit/animations.py +++ b/src/votekit/animations.py @@ -357,7 +357,7 @@ def _compose_offscreen_eliminations( round_numbers=round_numbers, ) - def render(self, preview: bool = False) -> None: + def render(self, preview: bool = False, color_palette: str = "dark") -> None: """ Renders the STV animation using Manim. @@ -365,9 +365,17 @@ def render(self, preview: bool = False) -> None: Args: preview (bool, optional): If ``True``, display the result in a video player immediately upon completing the render. Defaults to False. + color_palette (str, optional): A color scheme to use in the animation. Supports `'dark'` or `'light'`. Defaults to `'dark'`. """ + + manim.config.background_color = ElectionScene.color_palettes[color_palette][ + "background" + ] manimation = ElectionScene( - deepcopy(self.candidate_dict), deepcopy(self.events), title=self.title + deepcopy(self.candidate_dict), + deepcopy(self.events), + title=self.title, + color_palette=color_palette, ) manimation.render(preview=preview) @@ -383,29 +391,54 @@ class ElectionScene(manim.Scene): candidate_dict (dict[str,dict]): A dictionary mapping each candidate to a dictionary of attributes of the candidate. events (List[AnimationEvent]): A list of animation events to be constructed and rendered. title (str): A string to be displayed at the beginning of the animation as a title screen. If ``None``, the animation will skip the title screen. + color_palette (str, optional): A color scheme to use in the animation. Supports `'dark'` or `'light'`. Defaults to `'dark'`. """ - color_palette = { - "bar_fills": [ - manim.color.ManimColor(hex) - for hex in [ - "#16DEBD", - "#163EDE", - "#9F34F6", - "#FF6F00", - "#8F560C", - "#E2AD00", - "#8AD412", - ] - ], - "bar_outline": manim.LIGHT_GRAY, - "win_bar_fill": manim.GREEN, - "win_box_outline": manim.GREEN, - "offscreen_candidate_fill": manim.GRAY, - "background": manim.BLACK, - "elimination_line": manim.RED, - "ticker_tape_frosted": manim.DARK_GRAY, - "ticker_tape_highlight": manim.WHITE, + color_palettes = { + "dark": { + "bar_fills": [ + manim.color.ManimColor(hex) + for hex in [ + "#16DEBD", + "#163EDE", + "#9F34F6", + "#FF6F00", + "#8F560C", + "#E2AD00", + "#8AD412", + ] + ], + "bar_outline": manim.LIGHT_GRAY, + "win_bar_fill": manim.GREEN, + "win_box_outline": manim.GREEN, + "offscreen_candidate_fill": manim.GRAY, + "background": manim.BLACK, + "elimination_line": manim.RED, + "ticker_tape_frosted": manim.DARK_GRAY, + "ticker_tape_highlight": manim.WHITE, + }, + "light": { + "bar_fills": [ + manim.color.ManimColor(hex) + for hex in [ + "#16DEBD", + "#163EDE", + "#9F34F6", + "#FF6F00", + "#8F560C", + "#E2AD00", + "#8AD412", + ] + ], + "bar_outline": manim.BLACK, + "win_bar_fill": manim.GREEN, + "win_box_outline": manim.GREEN, + "offscreen_candidate_fill": manim.GRAY, + "background": manim.WHITE, + "elimination_line": manim.RED, + "ticker_tape_frosted": manim.LIGHT_GRAY, + "ticker_tape_highlight": manim.BLACK, + }, } bar_opacity = 1 ghost_opacity = 0.3 @@ -421,11 +454,13 @@ def __init__( candidate_dict: dict[str, dict], events: List[AnimationEvent], title: Optional[str] = None, + color_palette: str = "dark", ): super().__init__() self.candidate_dict = candidate_dict self.events = events self.title = title + self.color_palette = self.color_palettes[color_palette] self.width = 8 self.bar_height = 3.5 / len(self.candidate_dict) @@ -589,6 +624,7 @@ def _initialize_ticker_tape(self) -> None: ticker_line = Line( start=LEFT * line_length / 2, end=RIGHT * line_length / 2, + color=self.color_palette["bar_outline"], ) ticker_line.to_edge(DOWN, buff=0).shift(UP * self.ticker_tape_height) ticker_line.set_z_index(