from constants import * from scene.scene import Scene from animation.animation import Animation from animation.transform import MoveToTarget from animation.transform import Transform from animation.update import UpdateFromFunc from mobject.mobject import Mobject from mobject.svg_mobject import SVGMobject from mobject.tex_mobject import Brace from mobject.tex_mobject import TexMobject from mobject.tex_mobject import TextMobject from mobject.vectorized_mobject import VGroup from mobject.vectorized_mobject import VMobject from mobject.vectorized_mobject import VectorizedPoint from topics.geometry import Arc from topics.geometry import Circle from topics.geometry import Line from topics.geometry import Polygon from topics.geometry import Rectangle from topics.geometry import Square from utils.bezier import interpolate from utils.color import average_color from utils.color import color_gradient 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)