3b1b-manim/topics/probability.py

653 lines
21 KiB
Python

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)