Elastic collision simulation

I’m trying to simulate elastic collision using the One-dimensional Newtonian equation (https://en.wikipedia.org/wiki/Elastic_collision) with pygame. The thing is that even if I transcribed the solved equation to change the velocity of the 2 bodies, it is not working (or probably there is something wrong with the code). Here it is the code:

import pygame

pygame.init()
clock = pygame.time.Clock()

screenx,screeny = 1000,800

screen = pygame.display.set_mode([screenx,screeny])
pygame.display.set_caption("Elastic collision")

# Rectangles
small= pygame.Rect(70,433,17,17)
big = pygame.Rect(220,400,50,50)
# Masses
m_small = 1
m_big = 1
# Velocity
small_vel = 0
big_vel = -1

count = 0
start = False
sumM = m_small+m_big

run = True
while run:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            run = False

    screen.fill((0,0,0))

    big.x += big_vel
    small.x+=small_vel
    # If collision between small and big...
    if big.x==small.x+17:
        start=True

    if start == True:
        # ...change velocity of small and big,
        # using the elastic collision equation at each collision
        if small.x==0 or big.bottomleft<=small.bottomright:
            small_vel = (m_small-m_big)/sumM*small_vel+2*m_big/sumM*big_vel
            big_vel = (m_big-m_small)/sumM*big_vel+2*m_small/sumM*small_vel
            count += 1
            print("Small vel:", small_vel, " Big vel:", big_vel)

    # Draw the rectangles
    pygame.draw.rect(screen, (255,255,255), small)
    pygame.draw.rect(screen, (255,255,255), big)
    pygame.draw.line(screen, (255,0,0), (0,450), (1000,450),width=1)

    pygame.display.flip()
    clock.tick(60)
    pygame.display.update()

Since the 2 masses are equal, after the collision the moving one should stop and the other one should start moving. This is not happening. I also tried to write, in the if statment where the velocity should change due to the collision, small_vel=-1 and big_vel=0 and it worked fine.

Answer

It seems to me that the problem is due to these two lines:

            small_vel = (m_small-m_big)/sumM*small_vel+2*m_big/sumM*big_vel
            big_vel = (m_big-m_small)/sumM*big_vel+2*m_small/sumM*small_vel

Here big_vel is updated according to the new value of small_vel which is not what we want. Saving the new value of small_vel in a temporary variable and restoring it after the assignment to big_vel seems to fix the problem:

            tmp_small_vel = (m_small-m_big)/sumM*small_vel+2*m_big/sumM*big_vel
            big_vel = (m_big-m_small)/sumM*big_vel+2*m_small/sumM*small_vel
            small_vel = tmp_small_vel

To be honest I haven’t looked at the math stuff and I’m not sure this is exactly what you want but the result looks at least like the video of the wikipedia page.

On a side note you should have a look at this post: i get an pygame.error: video system not initialised error everytime i run the program to fix the bug in your code occuring when you close the window (the exception raised is not the same but the source of the problem and the solution to fix it are the same).