Replace for loops that exchange attributes of objects with numpy function

Talleros 0
Viewed 32

I made a bouncing ball simulation in Matplotlib, the update function exchange the move vectors (attributes) of the balls when they collide. Is it possible to replace the for loops in the update function with a Numpy function or other way to increase the FPS?

def update(self):
    for i in range(len(balls)):
        ball = balls[i]
        for j in range(len(balls)):
            other_ball = balls[j]
            if ball != other_ball:
                distance = ((other_ball.xy[0] - ball.xy[0]) ** 2 + (other_ball.xy[1] - ball.xy[1]) ** 2) ** 0.5
                if distance < 15:
                    # exchange move vectors
                    temp_vx = ball.v[0]
                    temp_vy = ball.v[1]
                    ball.v[0] = other_ball.v[0]
                    ball.v[1] = other_ball.v[1]
                    other_ball.v[0] = temp_vx
                    other_ball.v[1] = temp_vy
nested-loops matplotlib attributes numpy python
1 Answers

You can store all balls' positions in a single np.array of shape len(balls)x2, and similarly holds their velocities in another such array.
You update will then look like:

from scipy.spatial import distance

def update(balls_xy : np.array, balls_v : np.array):
  dist = distance.cdist(balls_xy, balls_xy)
  i, j = np.argwhere(dist < 15)
  # only one representative per colliding pair
  sel = i < j
  i = i[sel]
  j = j[sel]
  # switch the directions
  balls_v[i, :], balls_v[j, :] = balls_v[j, :], balls_v[i, :]
  return balls_v

see scipy.spatial.distance.cdist for more details.

usero1c70
Related