from constants import * from scene.scene import Scene from animation.animation import Animation from animation.transform import Transform, MoveToTarget from animation.simple_animations import UpdateFromFunc from mobject.mobject import Mobject from mobject.vectorized_mobject import VGroup, VMobject, VectorizedPoint from mobject.svg_mobject import SVGMobject from mobject.tex_mobject import TextMobject, TexMobject, Brace from topics.geometry import Circle, Line, Rectangle, Square, Arc, Polygon from utils.bezier import interpolate from utils.color import color_gradient, average_color from utils.config_ops import digest_config from utils.iterables import tuplify from utils.space_ops import center_of_mass EPSILON = 0.0001 class SampleSpaceScene(Scene): def get_sample_space(self, **config): self.sample_space = SampleSpace(**config) return self.sample_space def add_sample_space(self, **config): self.add(self.get_sample_space(**config)) def get_division_change_animations( self, sample_space, parts, p_list, dimension = 1, new_label_kwargs = None, **kwargs ): if new_label_kwargs is None: new_label_kwargs = {} anims = [] p_list = sample_space.complete_p_list(p_list) space_copy = sample_space.copy() vect = DOWN if dimension == 1 else RIGHT parts.generate_target() for part, p in zip(parts.target, p_list): part.replace(space_copy, stretch = True) part.stretch(p, dimension) parts.target.arrange_submobjects(vect, buff = 0) parts.target.move_to(space_copy) anims.append(MoveToTarget(parts)) if hasattr(parts, "labels") and parts.labels is not None: label_kwargs = parts.label_kwargs label_kwargs.update(new_label_kwargs) new_braces, new_labels = sample_space.get_subdivision_braces_and_labels( parts.target, **label_kwargs ) anims += [ Transform(parts.braces, new_braces), Transform(parts.labels, new_labels), ] return anims def get_horizontal_division_change_animations(self, p_list, **kwargs): assert(hasattr(self.sample_space, "horizontal_parts")) return self.get_division_change_animations( self.sample_space, self.sample_space.horizontal_parts, p_list, dimension = 1, **kwargs ) def get_vertical_division_change_animations(self, p_list, **kwargs): assert(hasattr(self.sample_space, "vertical_parts")) return self.get_division_change_animations( self.sample_space, self.sample_space.vertical_parts, p_list, dimension = 0, **kwargs ) def get_conditional_change_anims( self, sub_sample_space_index, value, post_rects = None, **kwargs ): parts = self.sample_space.horizontal_parts sub_sample_space = parts[sub_sample_space_index] anims = self.get_division_change_animations( sub_sample_space, sub_sample_space.vertical_parts, value, dimension = 0, **kwargs ) if post_rects is not None: anims += self.get_posterior_rectangle_change_anims(post_rects) return anims def get_top_conditional_change_anims(self, *args, **kwargs): return self.get_conditional_change_anims(0, *args, **kwargs) def get_bottom_conditional_change_anims(self, *args, **kwargs): return self.get_conditional_change_anims(1, *args, **kwargs) def get_prior_rectangles(self): return VGroup(*[ self.sample_space.horizontal_parts[i].vertical_parts[0] for i in range(2) ]) def get_posterior_rectangles(self, buff = MED_LARGE_BUFF): prior_rects = self.get_prior_rectangles() areas = [ rect.get_width()*rect.get_height() for rect in prior_rects ] total_area = sum(areas) total_height = prior_rects.get_height() post_rects = prior_rects.copy() for rect, area in zip(post_rects, areas): rect.stretch_to_fit_height(total_height * area/total_area) rect.stretch_to_fit_width( area/rect.get_height() ) post_rects.arrange_submobjects(DOWN, buff = 0) post_rects.next_to( self.sample_space, RIGHT, buff ) return post_rects def get_posterior_rectangle_braces_and_labels( self, post_rects, labels, direction = RIGHT, **kwargs ): return self.sample_space.get_subdivision_braces_and_labels( post_rects, labels, direction, **kwargs ) def update_posterior_braces(self, post_rects): braces = post_rects.braces labels = post_rects.labels for rect, brace, label in zip(post_rects, braces, labels): brace.stretch_to_fit_height(rect.get_height()) brace.next_to(rect, RIGHT, SMALL_BUFF) label.next_to(brace, RIGHT, SMALL_BUFF) def get_posterior_rectangle_change_anims(self, post_rects): def update_rects(rects): new_rects = self.get_posterior_rectangles() Transform(rects, new_rects).update(1) if hasattr(rects, "braces"): self.update_posterior_braces(rects) return rects anims = [UpdateFromFunc(post_rects, update_rects)] if hasattr(post_rects, "braces"): anims += map(Animation, [ post_rects.labels, post_rects.braces ]) return anims class SampleSpace(Rectangle): CONFIG = { "height" : 3, "width" : 3, "fill_color" : DARK_GREY, "fill_opacity" : 1, "stroke_width" : 0.5, "stroke_color" : LIGHT_GREY, ## "default_label_scale_val" : 1, } def add_title(self, title = "Sample space", buff = MED_SMALL_BUFF): ##TODO, should this really exist in SampleSpaceScene title_mob = TextMobject(title) if title_mob.get_width() > self.get_width(): title_mob.scale_to_fit_width(self.get_width()) title_mob.next_to(self, UP, buff = buff) self.title = title_mob self.add(title_mob) def add_label(self, label): self.label = label def complete_p_list(self, p_list): new_p_list = list(tuplify(p_list)) remainder = 1.0 - sum(new_p_list) if abs(remainder) > EPSILON: new_p_list.append(remainder) return new_p_list def get_division_along_dimension(self, p_list, dim, colors, vect): p_list = self.complete_p_list(p_list) colors = color_gradient(colors, len(p_list)) last_point = self.get_edge_center(-vect) parts = VGroup() for factor, color in zip(p_list, colors): part = SampleSpace() part.set_fill(color, 1) part.replace(self, stretch = True) part.stretch(factor, dim) part.move_to(last_point, -vect) last_point = part.get_edge_center(vect) parts.add(part) return parts def get_horizontal_division( self, p_list, colors = [GREEN_E, BLUE_E], vect = DOWN ): return self.get_division_along_dimension(p_list, 1, colors, vect) def get_vertical_division( self, p_list, colors = [MAROON_B, YELLOW], vect = RIGHT ): return self.get_division_along_dimension(p_list, 0, colors, vect) def divide_horizontally(self, *args, **kwargs): self.horizontal_parts = self.get_horizontal_division(*args, **kwargs) self.add(self.horizontal_parts) def divide_vertically(self, *args, **kwargs): self.vertical_parts = self.get_vertical_division(*args, **kwargs) self.add(self.vertical_parts) def get_subdivision_braces_and_labels( self, parts, labels, direction, buff = SMALL_BUFF, min_num_quads = 1 ): label_mobs = VGroup() braces = VGroup() for label, part in zip(labels, parts): brace = Brace( part, direction, min_num_quads = min_num_quads, buff = buff ) if isinstance(label, Mobject): label_mob = label else: label_mob = TexMobject(label) label_mob.scale(self.default_label_scale_val) label_mob.next_to(brace, direction, buff) braces.add(brace) label_mobs.add(label_mob) parts.braces = braces parts.labels = label_mobs parts.label_kwargs = { "labels" : label_mobs.copy(), "direction" : direction, "buff" : buff, } return VGroup(parts.braces, parts.labels) def get_side_braces_and_labels(self, labels, direction = LEFT, **kwargs): assert(hasattr(self, "horizontal_parts")) parts = self.horizontal_parts return self.get_subdivision_braces_and_labels(parts, labels, direction, **kwargs) def get_top_braces_and_labels(self, labels, **kwargs): assert(hasattr(self, "vertical_parts")) parts = self.vertical_parts return self.get_subdivision_braces_and_labels(parts, labels, UP, **kwargs) def get_bottom_braces_and_labels(self, labels, **kwargs): assert(hasattr(self, "vertical_parts")) parts = self.vertical_parts return self.get_subdivision_braces_and_labels(parts, labels, DOWN, **kwargs) def add_braces_and_labels(self): for attr in "horizontal_parts", "vertical_parts": if not hasattr(self, attr): continue parts = getattr(self, attr) for subattr in "braces", "labels": if hasattr(parts, subattr): self.add(getattr(parts, subattr)) def __getitem__(self, index): if hasattr(self, "horizontal_parts"): return self.horizontal_parts[index] elif hasattr(self, "vertical_parts"): return self.vertical_parts[index] return self.split()[index] class BarChart(VGroup): CONFIG = { "height" : 4, "width" : 6, "n_ticks" : 4, "tick_width" : 0.2, "label_y_axis" : True, "y_axis_label_height" : 0.25, "max_value" : 1, "bar_colors" : [BLUE, YELLOW], "bar_fill_opacity" : 0.8, "bar_stroke_width" : 3, "bar_names" : [], "bar_label_scale_val" : 0.75, } def __init__(self, values, **kwargs): VGroup.__init__(self, **kwargs) if self.max_value is None: self.max_value = max(values) self.add_axes() self.add_bars(values) self.center() def add_axes(self): x_axis = Line(self.tick_width*LEFT/2, self.width*RIGHT) y_axis = Line(MED_LARGE_BUFF*DOWN, self.height*UP) ticks = VGroup() heights = np.linspace(0, self.height, self.n_ticks+1) values = np.linspace(0, self.max_value, self.n_ticks+1) for y, value in zip(heights, values): tick = Line(LEFT, RIGHT) tick.scale_to_fit_width(self.tick_width) tick.move_to(y*UP) ticks.add(tick) y_axis.add(ticks) self.add(x_axis, y_axis) self.x_axis, self.y_axis = x_axis, y_axis if self.label_y_axis: labels = VGroup() for tick, value in zip(ticks, values): label = TexMobject(str(np.round(value, 2))) label.scale_to_fit_height(self.y_axis_label_height) label.next_to(tick, LEFT, SMALL_BUFF) labels.add(label) self.y_axis_labels = labels self.add(labels) def add_bars(self, values): buff = float(self.width) / (2*len(values) + 1) bars = VGroup() for i, value in enumerate(values): bar = Rectangle( height = (value/self.max_value)*self.height, width = buff, stroke_width = self.bar_stroke_width, fill_opacity = self.bar_fill_opacity, ) bar.move_to((2*i+1)*buff*RIGHT, DOWN+LEFT) bars.add(bar) bars.set_color_by_gradient(*self.bar_colors) bar_labels = VGroup() for bar, name in zip(bars, self.bar_names): label = TexMobject(str(name)) label.scale(self.bar_label_scale_val) label.next_to(bar, DOWN, SMALL_BUFF) bar_labels.add(label) self.add(bars, bar_labels) self.bars = bars self.bar_labels = bar_labels def change_bar_values(self, values): for bar, value in zip(self.bars, values): bar_bottom = bar.get_bottom() bar.stretch_to_fit_height( (value/self.max_value)*self.height ) bar.move_to(bar_bottom, DOWN) def copy(self): return self.deepcopy() ### Cards ### class DeckOfCards(VGroup): def __init__(self, **kwargs): possible_values = map(str, range(1, 11)) + ["J", "Q", "K"] possible_suits = ["hearts", "diamonds", "spades", "clubs"] VGroup.__init__(self, *[ PlayingCard(value = value, suit = suit, **kwargs) for value in possible_values for suit in possible_suits ]) class PlayingCard(VGroup): CONFIG = { "value" : None, "suit" : None, "key" : None, ##String like "8H" or "KS" "height" : 2, "height_to_width" : 3.5/2.5, "card_height_to_symbol_height" : 7, "card_width_to_corner_num_width" : 10, "card_height_to_corner_num_height" : 10, "color" : LIGHT_GREY, "turned_over" : False, "possible_suits" : ["hearts", "diamonds", "spades", "clubs"], "possible_values" : map(str, range(2, 11)) + ["J", "Q", "K", "A"], } def __init__(self, key = None, **kwargs): VGroup.__init__(self, key = key, **kwargs) def generate_points(self): self.add(Rectangle( height = self.height, width = self.height/self.height_to_width, stroke_color = WHITE, stroke_width = 2, fill_color = self.color, fill_opacity = 1, )) if self.turned_over: self.set_fill(DARK_GREY) self.set_stroke(LIGHT_GREY) contents = VectorizedPoint(self.get_center()) else: value = self.get_value() symbol = self.get_symbol() design = self.get_design(value, symbol) corner_numbers = self.get_corner_numbers(value, symbol) contents = VGroup(design, corner_numbers) self.design = design self.corner_numbers = corner_numbers self.add(contents) def get_value(self): value = self.value if value is None: if self.key is not None: value = self.key[:-1] else: value = random.choice(self.possible_values) value = string.upper(str(value)) if value == "1": value = "A" if value not in self.possible_values: raise Exception("Invalid card value") face_card_to_value = { "J" : 11, "Q" : 12, "K" : 13, "A" : 14, } try: self.numerical_value = int(value) except: self.numerical_value = face_card_to_value[value] return value def get_symbol(self): suit = self.suit if suit is None: if self.key is not None: suit = dict([ (string.upper(s[0]), s) for s in self.possible_suits ])[string.upper(self.key[-1])] else: suit = random.choice(self.possible_suits) if suit not in self.possible_suits: raise Exception("Invalud suit value") self.suit = suit symbol_height = float(self.height) / self.card_height_to_symbol_height symbol = SuitSymbol(suit, height = symbol_height) return symbol def get_design(self, value, symbol): if value == "A": return self.get_ace_design(symbol) if value in map(str, range(2, 11)): return self.get_number_design(value, symbol) else: return self.get_face_card_design(value, symbol) def get_ace_design(self, symbol): design = symbol.copy().scale(1.5) design.move_to(self) return design def get_number_design(self, value, symbol): num = int(value) n_rows = { 2 : 2, 3 : 3, 4 : 2, 5 : 2, 6 : 3, 7 : 3, 8 : 3, 9 : 4, 10 : 4, }[num] n_cols = 1 if num in [2, 3] else 2 insertion_indices = { 5 : [0], 7 : [0], 8 : [0, 1], 9 : [1], 10 : [0, 2], }.get(num, []) top = self.get_top() + symbol.get_height()*DOWN bottom = self.get_bottom() + symbol.get_height()*UP column_points = [ interpolate(top, bottom, alpha) for alpha in np.linspace(0, 1, n_rows) ] design = VGroup(*[ symbol.copy().move_to(point) for point in column_points ]) if n_cols == 2: space = 0.2*self.get_width() column_copy = design.copy().shift(space*RIGHT) design.shift(space*LEFT) design.add(*column_copy) design.add(*[ symbol.copy().move_to( center_of_mass(column_points[i:i+2]) ) for i in insertion_indices ]) for symbol in design: if symbol.get_center()[1] < self.get_center()[1]: symbol.rotate_in_place(np.pi) return design def get_face_card_design(self, value, symbol): from topics.characters import PiCreature sub_rect = Rectangle( stroke_color = BLACK, fill_opacity = 0, height = 0.9*self.get_height(), width = 0.6*self.get_width(), ) sub_rect.move_to(self) # pi_color = average_color(symbol.get_color(), GREY) pi_color = symbol.get_color() pi_mode = { "J" : "plain", "Q" : "thinking", "K" : "hooray" }[value] pi_creature = PiCreature( mode = pi_mode, color = pi_color, ) pi_creature.scale_to_fit_width(0.8*sub_rect.get_width()) if value in ["Q", "K"]: prefix = "king" if value == "K" else "queen" crown = SVGMobject(file_name = prefix + "_crown") crown.set_stroke(width = 0) crown.set_fill(YELLOW, 1) crown.stretch_to_fit_width(0.5*sub_rect.get_width()) crown.stretch_to_fit_height(0.17*sub_rect.get_height()) crown.move_to(pi_creature.eyes.get_center(), DOWN) pi_creature.add_to_back(crown) to_top_buff = 0 else: to_top_buff = SMALL_BUFF*sub_rect.get_height() pi_creature.next_to(sub_rect.get_top(), DOWN, to_top_buff) # pi_creature.shift(0.05*sub_rect.get_width()*RIGHT) pi_copy = pi_creature.copy() pi_copy.rotate(np.pi, about_point = sub_rect.get_center()) return VGroup(sub_rect, pi_creature, pi_copy) def get_corner_numbers(self, value, symbol): value_mob = TextMobject(value) width = self.get_width()/self.card_width_to_corner_num_width height = self.get_height()/self.card_height_to_corner_num_height value_mob.scale_to_fit_width(width) value_mob.stretch_to_fit_height(height) value_mob.next_to( self.get_corner(UP+LEFT), DOWN+RIGHT, buff = MED_LARGE_BUFF*width ) value_mob.set_color(symbol.get_color()) corner_symbol = symbol.copy() corner_symbol.scale_to_fit_width(width) corner_symbol.next_to( value_mob, DOWN, buff = MED_SMALL_BUFF*width ) corner_group = VGroup(value_mob, corner_symbol) opposite_corner_group = corner_group.copy() opposite_corner_group.rotate( np.pi, about_point = self.get_center() ) return VGroup(corner_group, opposite_corner_group) class SuitSymbol(SVGMobject): CONFIG = { "height" : 0.5, "fill_opacity" : 1, "stroke_width" : 0, "red" : "#D02028", "black" : BLACK, } def __init__(self, suit_name, **kwargs): digest_config(self, kwargs) suits_to_colors = { "hearts" : self.red, "diamonds" : self.red, "spades" : self.black, "clubs" : self.black, } if suit_name not in suits_to_colors: raise Exception("Invalid suit name") SVGMobject.__init__(self, file_name = suit_name, **kwargs) color = suits_to_colors[suit_name] self.set_stroke(width = 0) self.set_fill(color, 1) self.scale_to_fit_height(self.height)