Мне в универе выдали проект один простой, но почему-то очень страдает оптимизация. суть проста: Отобразить на экране N шариков разного размера, которые бы двигались с разной скоростью в разных направлениях без особой просадки ФПС. Вроде бы задачка несложная, но при отображении больше 200 шариков, фпс начинает заметно проседать. Подскажите, пожалуйста, чем и как это можно исправить? Я слышал про намбу, и нампай, но до этого не работал с ними, так что не знаю, как будет правильнее оптимизировать свой код под эти модули. Может, сможете мне помочь?

У меня есть две функции, одна отвечает за создание этих самых шариков так, чтобы они не накладывались друг на друга, и вторая, собственно, за рендер и обнаружение коллизий. для граф. движка использую пайгем, но думаю, что его уже навряд ли как смогу оптимизировать, ибо основная проблема в том, что эти две функции слишком тяжёлые.

import numpy as np
import numba as nb
import random as rd


res = (1152, 648)

colors = [
    (0, 0, 0),
    (29, 43, 83),
    (126, 37, 83),
    (0, 135, 81),
    (171, 82, 54),
    (95, 87, 79),
    (194, 195, 199),
    (255, 241, 232),
    (255, 0, 77),
    (255, 163, 0),
    (255, 236, 39),
    (0, 228, 54),
    (41, 173, 255),
    (131, 118, 156),
    (255, 119, 168),
    (255, 204, 170),
]

Create all the balls

def init(balls):

bals = list()
while balls > 0:
    size = rd.random() * 10 + 5
    ball = [
        (rd.random() * 0.9 + 0.05) * res[0],
        (rd.random() * 0.9 + 0.05) * res[1],
        rd.random() * 2 * np.pi,
        size,
        rd.choice(colors[1:]),
        20 - size
    ]

    cols = False
    for bal in bals:
        if (ball[0] - bal[0]) ** 2 + (ball[1] - bal[1]) ** 2 <= (ball[3] + bal[3]) ** 2:
            cols = True

    if not cols:
        bals.append(ball)
        balls -= 1

return bals

Main render func

def render(balls):

for num in range(len(balls)):
    ball = balls[num]

    # Detect collision with walls
    if ball[0] + ball[3] >= res[0] or ball[0] - ball[3] <= 0:
        ball[2] = - ball[2]
    if ball[1] + ball[3] >= res[1] or ball[1] - ball[3] <= 0:
        ball[2] = np.pi - ball[2]

    # Detect collision between balls
    for bal in range(num):
        bal = balls[bal]
        if (ball[0] - bal[0]) ** 2 + (ball[1] - bal[1]) ** 2 <= (ball[3] + bal[3]) ** 2:

            if (ball[0] - bal[0]) * (ball[1] - bal[1]) > 0:
                a = np.arccos(np.abs(ball[0] - bal[0]) / (ball[3] + bal[3]))
                ball[2] = 2 * a - ball[2]
                bal[2] = 2 * a - bal[2]

            elif (ball[0] - bal[0]) * (ball[1] - bal[1]) < 0:
                a = np.pi - np.arccos(np.abs(ball[0] - bal[0]) / (ball[3] + bal[3]))
                ball[2] = 2 * a - ball[2]
                bal[2] = 2 * a - bal[2]

    # Move ball
    ball[0] += np.sin(ball[2]) * ball[5]
    ball[1] += np.cos(ball[2]) * ball[5]

P.S. я знаю, что модель поведения шариков после столкновения друг с другом не доработана до конца и может глючить (если кто сможет и захочет всё это дело протестировать, в комментариях оставлю ссылку на репозиторий). Но у меня нет идей, как по-другому должны вести себя шарики при столкновении, так как максимум математического из своего мозга я уже выжал.