From 0fdb8aefdf52c7c049ec0925d1663b20ba49d0a2 Mon Sep 17 00:00:00 2001
From: Grant Sanderson <grant@3blue1brown.com>
Date: Sat, 16 Mar 2019 22:13:55 -0700
Subject: [PATCH] Start to implement a pendulum

---
 active_projects/ode/all_part1_scenes.py |   4 +-
 active_projects/ode/part1/pendulum.py   | 129 +++++++++++++++++++++---
 2 files changed, 117 insertions(+), 16 deletions(-)

diff --git a/active_projects/ode/all_part1_scenes.py b/active_projects/ode/all_part1_scenes.py
index 5e117116..d4b6ae50 100644
--- a/active_projects/ode/all_part1_scenes.py
+++ b/active_projects/ode/all_part1_scenes.py
@@ -1,6 +1,6 @@
-# from active_projects.ode.part1 import *
+from active_projects.ode.part1.pendulum import *
 
 OUTPUT_DIRECTORY = "ode/part1"
 ALL_SCENE_CLASSES = [
-
+    PendulumTest
 ]
diff --git a/active_projects/ode/part1/pendulum.py b/active_projects/ode/part1/pendulum.py
index 76af75a2..5005f73e 100644
--- a/active_projects/ode/part1/pendulum.py
+++ b/active_projects/ode/part1/pendulum.py
@@ -6,45 +6,146 @@ class Pendulum(VGroup):
         "length": 3,
         "gravity": 9.8,
         "weight_diameter": 0.5,
+        "initial_theta": 0.3,
+        "omega": 0,
+        "damping": 0.1,
+        "top_point": 2 * UP,
         "rod_style": {
-            "stroke_width": 2,
+            "stroke_width": 3,
             "stroke_color": LIGHT_GREY,
+            "sheen_direction": UP,
+            "sheen_factor": 1,
         },
         "weight_style": {
             "stroke_width": 0,
             "fill_opacity": 1,
-            "fill_color": LIGHT_GREY,
+            "fill_color": DARK_GREY,
             "sheen_direction": UL,
             "sheen_factor": 0.5,
-        }
+        },
+        "dashed_line_config": {
+            "num_dashes": 25,
+            "stroke_color": WHITE,
+            "stroke_width": 2,
+        },
+        "angle_arc_config": {
+            "radius": 1,
+            "stroke_color": WHITE,
+            "stroke_width": 2,
+        },
+        "velocity_vector_config": {
+            "color": RED,
+        },
+        "n_steps_per_frame": 10,
     }
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
+        self.create_fixed_point()
         self.create_rod()
+        self.create_weight()
+        self.rotating_group = VGroup(self.rod, self.weight)
+        self.create_dashed_line()
+        self.create_angle_arc()
+
+        self.set_theta(self.initial_theta)
+        self.update()
+
+    def create_fixed_point(self):
+        self.fixed_point_tracker = VectorizedPoint(self.top_point)
 
     def create_rod(self):
-        pass
+        rod = self.rod = Line(UP, DOWN)
+        rod.set_height(self.length)
+        rod.set_style(**self.rod_style)
+        rod.move_to(self.get_fixed_point(), UP)
+        self.add(rod)
 
     def create_weight(self):
-        pass
+        weight = self.weight = Circle()
+        weight.set_width(self.weight_diameter)
+        weight.set_style(**self.weight_style)
+        weight.move_to(self.rod.get_end())
+        self.add(weight)
 
+    def create_dashed_line(self):
+        line = self.dashed_line = DashedLine(
+            self.get_fixed_point(),
+            self.get_fixed_point() + self.length * DOWN,
+            **self.dashed_line_config
+        )
+        self.add_to_back(line)
+
+    def create_angle_arc(self):
+        self.angle_arc = always_redraw(lambda: Arc(
+            arc_center=self.get_fixed_point(),
+            start_angle=-90 * DEGREES,
+            angle=self.get_theta(),
+            **self.angle_arc_config,
+        ))
+        self.add(self.angle_arc)
+
+    def add_velocity_vector(self):
+        def make_vector():
+            vector = Vector(0.5 * self.get_omega() * RIGHT)
+        self.velocity_vector = always_redraw(make_vector)
+        self.add(self.velocity_vector)
+        return self
+
+    #
     def get_theta(self):
-        pass
+        theta = self.rod.get_angle() - self.dashed_line.get_angle()
+        theta = (theta + PI) % TAU - PI
+        return theta
 
-    def set_theta(self):
-        pass
+    def set_theta(self, theta):
+        self.rotating_group.rotate(
+            theta - self.get_theta()
+        )
+        self.rotating_group.shift(
+            self.get_fixed_point() - self.rod.get_start(),
+        )
+        return self
 
-    def get_angular_velocity(self):
-        pass
+    def get_omega(self):
+        return self.omega
 
-    def set_angular_velocity(self):
-        pass
+    def set_omega(self, omega):
+        self.omega = omega
+        return self
 
     def get_fixed_point(self):
-        pass
+        return self.fixed_point_tracker.get_location()
+
+    #
+    def start_swinging(self):
+        self.add_updater(Pendulum.update_by_gravity)
+
+    def end_swinging(self):
+        self.remove_updater(Pendulum.update_by_gravity)
+
+    def update_by_gravity(self, dt):
+        theta = self.get_theta()
+        omega = self.get_omega()
+        nspf = self.n_steps_per_frame
+        for x in range(nspf):
+            d_theta = omega * dt / nspf
+            d_omega = op.add(
+                -self.damping * omega,
+                -(self.gravity / self.length) * np.sin(theta),
+            ) * dt / nspf
+            theta += d_theta
+            omega += d_omega
+        self.set_theta(theta)
+        self.set_omega(omega)
+        return self
 
 
 class PendulumTest(Scene):
     def construct(self):
-        pass
+        pendulum = Pendulum(
+            initial_theta=150 * DEGREES,
+        )
+        self.add(pendulum)
+        pendulum.start_swinging()
+        self.wait(10)