Try this:
from math import sqrt, sin, cos
import tkinter as tk
import time
# For more info about this read: https://stackoverflow.com/a/17985217/11106801
def _create_circle(self, x, y, r, **kwargs):
return self.create_oval(x-r, y-r, x+r, y+r, **kwargs)
tk.Canvas.create_circle = _create_circle
WINDOW_WIDTH = 500
WINDOW_HEIGHT = 500
# The coefficient of restitution
# Set to 1 for perfectly elastic collitions
e = 1
class Ball:
def __init__(self, mass:float, r:float, x:float, y:float,
vx:float=0, vy:float=0, **kwargs):
"""
This is a class that defines what a ball is and how it interacts
with other balls.
Arguments:
mass:float # The mass of the ball
------------------------------------------------------
r:float # The radius of the ball, must be >0
------------------------------------------------------
x:float # The x position of the ball
# must be >0 and <WINDOW_WIDTH
y:float # The y position of the ball
# must be >0 and <WINDOW_HEIGHT
------------------------------------------------------
vx:float # The x velocity of the ball
vy:float # The y velocity of the ball
------------------------------------------------------
**kwargs # All of the args to be passed in to `create_circle`
"""
self.m = mass
self.r = r
self.x = x
self.y = y
self.vx = vx
self.vy = vy
self.kwargs = kwargs
def display(self, canvas:tk.Canvas) -> int:
"""
Displays the ball on the screen and returns the canvas_id (which
is a normal python int).
"""
canvas_id = canvas.create_circle(self.x, self.y, self.r, **self.kwargs)
return canvas_id
def move(self, all_balls:list, dt:float) -> (float, float):
"""
This moves the ball according to `self.vx` and `self.vy`.
It also checks for collisions with other balls.
Arguments:
all_balls:list # A list of all balls that are in the
# simulation. It can include this ball.
---------------------------------------------------
dt:float # delta time - used to move the balls
Returns:
dx:float # How much the ball has moved in the x direction
dy:float # How much the ball has moved in the y direction
Note: This function isn't optimised in any way. If you optimise
it, it should run faster.
"""
# Check if there are any collitions:
for other, _ in all_balls:
# Skip is `ball` is the same as this ball
if id(other) == id(self):
continue
# Check if there is a collision:
distance_squared = (other.x - self.x)**2 + (other.y - self.y)**2
if distance_squared <= (other.r + self.r)**2:
# Now the fun part - calulating the resultant velocity.
# I am assuming you already know all of the reasoning
# behind the math (if you don't ask physics.stackexchange.com)
# First I will find the normal vector of the balls' radii.
# That is just the unit vector in the direction of the
# balls' radii
ball_radii_vector_x = other.x - self.x
ball_radii_vector_y = other.y - self.y
abs_ball_radii_vector = sqrt(ball_radii_vector_x**2 +\
ball_radii_vector_y**2)
nx = ball_radii_vector_x / abs_ball_radii_vector **2*2
ny = ball_radii_vector_y / abs_ball_radii_vector **2*2
# Now I will calculate the tangent
tx = -ny
ty = nx
""" Only for debug
print("n =", (nx, ny), "\t t =", (tx, ty))
print("u1 =", (self.vx, self.vy),
"\t u2 =", (other.vx, other.vy))
#"""
# Now I will split the balls' velocity vectors to the sum
# of 2 vectors parallel to n and t
# self_velocity = λ*n + μ*t
# other_velocity = a*n + b*t
λ = (self.vx*ty - self.vy*tx) / (nx*ty - ny*tx)
# Sometimes `tx` can be 0 if so we are going to use `ty` instead
try:
μ = (self.vx - λ*nx) / tx
except ZeroDivisionError:
μ = (self.vy - λ*ny) / ty
""" Only for debug
print("λ =", λ, "\t μ =", μ)
#"""
a = (other.vx*ty - other.vy*tx) / (nx*ty - ny*tx)
# Sometimes `tx` can be 0 if so we are going to use `ty` instead
try:
b = (other.vx - a*nx) / tx
except ZeroDivisionError:
b = (other.vy - a*ny) / ty
""" Only for debug
print("a =", a, "\t b =", b)
#"""
self_u = λ
other_u = a
sum_mass_u = self.m*self_u + other.m*other_u
sum_masses = self.m + other.m
# Taken from en.wikipedia.org/wiki/Inelastic_collision
self_v = (e*other.m*(other_u-self_u) + sum_mass_u)/sum_masses
other_v = (e*self.m*(self_u-other_u) + sum_mass_u)/sum_masses
self.vx = self_v*nx + μ*tx
self.vy = self_v*ny + μ*ty
other.vx = other_v*nx + b*tx
other.vy = other_v*ny + b*ty
print("v1 =", (self.vx, self.vy),
"\t v2 =", (other.vx, other.vy))
# Move the ball
dx = self.vx * dt
dy = self.vy * dt
self.x += dx
self.y += dy
return dx, dy
class Simulator:
def __init__(self):
self.balls = [] # Contains tuples of (<Ball>, <canvas_id>)
self.root = tk.Tk()
self.root.resizable(False, False)
self.canvas = tk.Canvas(self.root, width=WINDOW_WIDTH,
height=WINDOW_HEIGHT)
self.canvas.pack()
def step(self, dt:float) -> None:
"""
Steps the simulation as id `dt` seconds have passed
"""
for ball, canvas_id in self.balls:
dx, dy = ball.move(self.balls, dt)
self.canvas.move(canvas_id, dx, dy)
def run(self, dt:float, total_time:float) -> None:
"""
This function keeps steping `dt` seconds until `total_time` has
elapsed.
Arguments:
dt:float # The time resolution in seconds
total_time:float # The number of seconds to simulate
Note: This function is just for proof of concept. It is badly written.
"""
for i in range(int(total_time//dt)):
self.step(dt)
self.root.update()
time.sleep(dt)
def add_ball(self, *args, **kwargs) -> None:
"""
Adds a ball by passing all of the args and keyword args to `Ball`.
It also displays the ball and appends it to the list of balls.
"""
ball = Ball(*args, **kwargs)
canvas_id = ball.display(self.canvas)
self.balls.append((ball, canvas_id))
app = Simulator()
app.add_ball(mass=1, r=10, x=20, y=250, vx=100, vy=0, fill="red")
app.add_ball(mass=1, r=10, x=240, y=250, vx=0, vy=0, fill="blue")
app.add_ball(mass=1, r=10, x=480, y=250, vx=0, vy=0, fill="yellow")
app.run(0.01, 10)
That code has a lot of comments describing how it works. If you still have any questions, ask me. Also the move method isn't optimised. The run method isn't great and can throw an error if it's still running and the user closes the window. I will try to find a better approch for that method. If you find any bugs, please let me know. I will try to fix them.
