Как ускорить физику на Python: NumPy, Numba и векторизация

Симуляция движения 1000 частиц с гравитационным взаимодействием на чистом Python в цикле по каждому шагу — это десятки секунд на кадр. Тот же алгоритм, переписанный на векторные операции NumPy, укладывается в десятки миллисекунд, а с Numba и JIT-компиляцией счёт идёт на единицы миллисекунд. Разница в сотни раз возникает не из-за магии, а из-за того, куда именно уходит время в интерпретаторе. Ниже — три уровня ускорения физических расчётов на Python и честная цена каждого.

Почему чистый Python медленный на числах

Каждое обращение к элементу списка, каждое сложение двух float в Python — это работа с полноценными объектами. Число 3.14 — это не 8 байт в памяти, а объект PyFloatObject с заголовком, счётчиком ссылок и указателем на тип. Сложение двух таких объектов проходит через диспетчеризацию: интерпретатор проверяет типы, ищет метод __add__, создаёт новый объект под результат.

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

Правило простое: если внутренний цикл вашей симуляции написан на чистом Python и крутится миллионы раз — вы платите за интерпретатор, а не за физику.

Термин, который стоит знать

Векторизация* — приём, при котором операция применяется сразу ко всему массиву данных одной командой, а не поэлементно в Python-цикле. Тяжёлая работа уходит в скомпилированный C-код библиотеки, а интерпретатор вызывает его один раз.

Уровень 1: NumPy и векторизация

Возьмём интегрирование методом Эйлера: у нас массивы позиций и скоростей, на каждом шаге пересчитываем их. Наивная версия перебирает частицы в цикле.

import numpy as np

def step_slow(pos, vel, acc, dt):
    for i in range(len(pos)):
        vel[i] += acc[i] * dt
        pos[i] += vel[i] * dt
    return pos, vel

Векторная версия делает то же самое без цикла — операции применяются ко всем строкам массива сразу:

def step_fast(pos, vel, acc, dt):
    vel += acc * dt
    pos += vel * dt
    return pos, vel

Внутри vel += acc * dt NumPy обходит все элементы в цикле на C, без создания Python-объектов на каждое число. На массиве в 100 000 частиц вторая версия обычно быстрее первой в 50–200 раз — точное число зависит от процессора и размера данных.

Где векторизация буксует

Не всё сводится к поэлементным операциям. Расчёт попарных сил (каждая частица действует на каждую) — это матрица N×N. Её можно построить через broadcasting, но при больших N она сожрёт память: 10 000 частиц дадут матрицу на 100 миллионов элементов, около 800 МБ на float64. Здесь векторизация упирается в память, и нужен другой подход.

Уровень 2: Numba и JIT

Numba** компилирует функцию Python в машинный код при первом вызове. Достаточно навесить декоратор — и ваш «медленный» цикл начинает работать почти как C, потому что фактически им и становится.

from numba import njit

@njit
def compute_forces(pos, mass):
    n = pos.shape[0]
    acc = np.zeros_like(pos)
    for i in range(n):
        for j in range(n):
            if i == j:
                continue
            d = pos[j] - pos[i]
            r2 = d[0]**2 + d[1]**2 + 1e-9
            f = mass[j] / r2**1.5
            acc[i] += f * d
    return acc

Тот самый двойной цикл, который на чистом Python неприемлемо медленный, под @njit работает без матрицы N×N в памяти и без интерпретатора. Первый вызов включает компиляцию (это заметная задержка, доли секунды), последующие идут на полной скорости. Флаг @njit(parallel=True) вместе с numba.prange раскидывает внешний цикл по ядрам.

Уровень 3: GPU и специализированные библиотеки

Если задача — тысячи и сотни тысяч частиц, а вычисления однородны, есть смысл идти на видеокарту. CuPy повторяет API NumPy, но выполняет операции на GPU; часто достаточно заменить import numpy as np на import cupy as np и следить за пересылкой данных между CPU и GPU. Numba умеет компилировать CUDA-ядра напрямую через @cuda.jit.

GPU даёт выигрыш не всегда: пересылка массивов между оперативной памятью и памятью видеокарты стоит времени. На маленьких данных или при частом обмене с CPU карта может оказаться медленнее хорошо векторизованного NumPy.

Что выбрать под вашу задачу

Порядок действий при оптимизации

1. Сначала измерьте, а не гадайте. Профилируйте через cProfile или line_profiler и найдите функцию, где реально уходит время.
2. Переведите горячий участок на векторные операции NumPy, если он сводится к поэлементной арифметике.
3. Если алгоритм по своей природе циклический (попарные взаимодействия, зависимость шага от предыдущего) — навесьте @njit.
4. Замерьте снова. Часто уже на этом шаге ускорение достаточное и дальше идти незачем.
5. Только если данных действительно много и профиль это подтверждает — пробуйте GPU.

Главная ошибка — переписывать всё подряд на Numba или тащить GPU туда, где хватило бы одной строки NumPy. Оптимизация без замеров — это трата времени на участки, которые и так работают быстро.

* Векторизация — выполнение операции над всем массивом одной командой библиотеки вместо поэлементного цикла на Python.

** JIT (just-in-time) компиляция — перевод кода в машинные инструкции во время выполнения программы, а не заранее.