Initial commit

simulation/__init__.py

@@ -0,0 +1,60 @@
+from itertools import combinations
+from superturtle.movement import fly, no_delay
+from superturtle.animation import animate
+from simulation.gravity import get_gravity_force
+from simulation.collision import collide2d
+from simulation.drawing import draw_boundary
+
+class Simulation:
+    def __init__(self, objects, g=None, e=None, boundary=None):
+        self.objects = objects
+        self.g = g
+        self.e = e
+        self.boundary = boundary
+        self.collisions = set()
+
+    def run(self, *args):
+        for frame in animate(*args):
+            for a, b in combinations(self.objects, 2):
+                if self.g is not None:
+                    force = get_gravity_force(a, b, self.g)
+                    a.acceleration += force / a.mass
+                    b.acceleration -= force / b.mass
+                if self.e is not None:
+                    if a.collided_with(b):
+                        self.collide(a, b)
+                    else:
+                        self.clear_collision(a, b)
+            for obj in self.objects:
+                if self.boundary:
+                    draw_boundary(self.boundary)
+                    self.bounce_object_if_outside_boundary(obj)
+                obj.update()
+                fly(obj.position.x, obj.position.y)
+                obj.draw()
+
+    def collide(self, a, b):
+        if not self.currently_colliding(a, b):
+            self.collisions.add((a, b))
+            av, bv = collide2d(a, b, self.e)
+            a.velocity = av
+            b.velocity = bv
+
+    def currently_colliding(self, a, b):
+        return (a, b) in self.collisions or (b, a) in self.collisions
+
+    def clear_collision(self, a, b):
+        if (a, b) in self.collisions:
+            self.collisions.remove((a, b))
+        if (b, a) in self.collisions:
+            self.collisions.remove((b, a))
+
+    def bounce_object_if_outside_boundary(self, obj):
+        if obj.position.x + obj.radius > self.boundary:
+            obj.velocity.x = -abs(obj.velocity.x)
+        if obj.position.x - obj.radius < -self.boundary:
+            obj.velocity.x = abs(obj.velocity.x)
+        if obj.position.y + obj.radius > self.boundary:
+            obj.velocity.y = -abs(obj.velocity.y)
+        if obj.position.y - obj.radius < -self.boundary:
+            obj.velocity.y = abs(obj.velocity.y)

simulation/ball.py

@@ -0,0 +1,37 @@
+from geometry.point import Point
+from geometry.vector import Vector
+from geometry.validation import expect_type
+from simulation.drawing import centered_circle
+from turtle import begin_fill, end_fill, fillcolor
+
+class Ball:
+    def __init__(self, position, velocity=None, mass=1, radius=20, color="black"):
+        expect_type(position, Point)
+        self.position = position
+        self.velocity = velocity or Vector()
+        expect_type(self.velocity, Vector)
+        self.acceleration = Vector()
+        self.radius = radius
+        self.mass = mass
+        self.color = color
+
+    def update(self):
+        """Update the ball's position and velocity at this moment. 
+        
+        """
+        self.velocity += self.acceleration
+        self.position += self.velocity
+        self.acceleration = Vector()
+
+    def collided_with(self, other_ball):
+        """Checks whether this ball collided with another ball.
+        """
+        return (self.position - other_ball.position).mag() <= self.radius + other_ball.radius
+
+    def draw(self):
+        "Draws a colored ball."
+        fillcolor(self.color)
+        begin_fill()
+        centered_circle(self.radius)
+        end_fill()
+

simulation/collision.py

@@ -0,0 +1,39 @@
+def collide2d(a, b, e):
+    """Returns the velocities of balls a and b after collision. 
+
+    e is the *coefficient of restitution*, a parameter between 0 and 1 
+    which determines how much the objects bounce off each other. 
+    When e is 0, the two objects completely stick together. 
+    When e is 1, the two objects bounce off each other completely. 
+
+    This is a 2-dimensional collision because a and b each have a 2d 
+    velocity vector, like two balls moving on a plane. We can separate
+    each velocity vector into two components: a normal vector (pointed 
+    straight at the other ball) and a tangent vector (perpendicular to
+    the normal). The tangent vectors don't participate in the collision
+    at all, so we just need to calculate how the normal vectors change
+    after the collision. This is easier, since it's a 1-dimensional 
+    collision.
+
+    Afterwards we combine the updated normal vector with the tangent 
+    vector to get the resulting velocity vector for each ball.
+    """
+    a_to_b = b.position - a.position
+    av_normal = a.velocity.project(a_to_b)
+    av_tangent = a.velocity - av_normal
+    bv_normal = b.velocity.project(-a_to_b)
+    bv_tangent = b.velocity - bv_normal
+    av_normal_after, bv_normal_after = collide1d(av_normal, a.mass, bv_normal, b.mass, e)
+    av_after = av_normal_after + av_tangent
+    bv_after = bv_normal_after + bv_tangent
+    return av_after, bv_after
+
+def collide1d(av, am, bv, bm, e):
+    """Returns the velicoties of a and b after a 1-dimensional collision, 
+    where a and b are headed directly at each other. This formula is from Wikipedia: 
+    https://en.wikipedia.org/wiki/Elastic_collision#One-dimensional_Newtonian
+    """
+    velocity_center_of_mass = (av * am + bv * bm) / (am + bm)
+    av_after = velocity_center_of_mass * (1 + e) - av * e
+    bv_after = velocity_center_of_mass * (1 + e) - bv * e
+    return av_after, bv_after

simulation/drawing.py

@@ -0,0 +1,31 @@
+from turtle import *
+
+def centered_circle(radius):
+    """Draws a circle, centered on the turtle's current position.
+    The built-in circle function annoyingly draws a circle whose center
+    is one radius left of the turtle.
+    """
+    right(90)
+    penup()
+    forward(radius)
+    pendown()
+    left(90)
+    circle(radius)
+    left(90)
+    penup()
+    forward(radius)
+    pendown()
+    right(90)
+
+def draw_boundary(distance):
+    """Draws the boundary, a square `distance` from the origin on each side.
+    """
+    penup()
+    goto(-distance, -distance)
+    pendown()
+    goto(distance, -distance)
+    goto(distance, distance)
+    goto(-distance, distance)
+    goto(-distance, -distance)
+
+

simulation/gravity.py

@@ -0,0 +1,13 @@
+
+def get_gravity_force(a, b, g):
+    """Gravity is a force of attraction between objects. 
+    The attraction between two objects is equal to their masses
+    multipled together, divided by the square of the distance between them, 
+    times a constant g. This is the real-life definition of gravity. 
+    """
+    a_to_b = b.position - a.position
+    distance = a_to_b.mag()
+    force = a_to_b.scale(g * a.mass * b.mass / (distance * distance))
+    return force
+    a.acceleration += force
+    b.acceleration -= force