Python, NumPy и Numba

Численные методы для физиков

Дмитрий В. Наумов

О чём эта лекция

Весь курс написан на Python. Сегодня — инструменты, без которых дальше не обойтись, на двух задачах с точным ответом:

  • случайное блуждание: \(\langle x_N^2\rangle = N\);
  • число \(\pi\) методом Монте-Карло: ошибка \(\propto 1/\sqrt{n}\).

Что берём с собой: рабочая среда и ноутбук, массивы NumPy, конечная точность чисел, генератор случайных чисел с seed, графики, измерение времени, Numba.

Ноутбук лекции: открыть · скачать .ipynb

Физическая задача

Случайное блуждание

Частица делает \(N\) шагов \(s_i=\pm1\) с вероятностью \(1/2\): \[ x_N=\sum_{i=1}^{N}s_i . \]

  • Модель диффузии: молекула в газе, фотон в мутной среде, цена акции.
  • Вопрос: как далеко уходит частица за \(N\) шагов?
  • Ответ проверяем моделированием \(M\) независимых блужданий.

Число \(\pi\) методом Монте-Карло

Бросаем \(n\) точек равномерно в квадрат \([0,1]^2\) и считаем долю \(k/n\) точек внутри четверти круга \(x^2+y^2<1\): \[ \frac{k}{n}\;\to\;p=\frac{\pi}{4},\qquad \hat\pi = 4\,\frac{k}{n}. \]

  • Простейший пример метода Монте-Карло: интеграл как среднее по случайным точкам.
  • Вопрос: сколько точек нужно для заданной точности?

Аналитическая точка отсчёта

Что решается точно

Блуждание. Шаги независимы, \(\langle s_i\rangle=0\), \(\langle s_i s_j\rangle=\delta_{ij}\): \[ \langle x_N\rangle=0,\qquad \langle x_N^2\rangle=\sum_{i,j}\langle s_is_j\rangle=N . \] Среднее по \(M\) блужданиям само случайно: \(\operatorname{Var}x_N^2=2N^2-2N\), поэтому относительная ошибка \(\langle x_N^2\rangle/N\) равна \(\sqrt{2/M}\).

Число \(\pi\). \(k\) — биномиальная величина с вероятностью \(p=\pi/4\): \[ \sigma_{\hat\pi}=4\sqrt{\frac{p(1-p)}{n}}\approx\frac{1.64}{\sqrt n}. \]

Зачем моделировать, если ответ известен

  • Эталон проверяет код: если моделирование не даёт \(\langle x_N^2\rangle=N\), ошибка в программе.
  • Проверенный код меняем там, где аналитики нет: блуждание с ловушками и стенками, шаги разной длины, неоднородная среда, интеграл в 10 измерениях.
  • Так устроена каждая лекция курса: сначала эталон, потом новая задача.

Рабочая среда

Python, пакеты, Jupyter

  • Python 3 и пакеты: python -m pip install numpy matplotlib numba jupyterlab.
  • Отдельное окружение на проект (venv, pyenv, conda): версии пакетов не смешиваются между проектами.
  • Jupyter: ноутбук = текст + ячейки кода + результаты. jupyter lab, Shift+Enter выполняет ячейку.
  • Без установки: Google Colab → «Файл → Загрузить блокнот»; NumPy, Matplotlib и Numba там уже есть.

Версии — в начале каждого ноутбука

import sys
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import numba

print("Python    ", sys.version.split()[0])
print("NumPy     ", np.__version__)
print("Matplotlib", matplotlib.__version__)
print("Numba     ", numba.__version__)
Python     3.12.15
NumPy      2.5.4
Matplotlib 3.11.2
Numba      0.68.0

Результат воспроизводим, только если известно, чем он получен.

NumPy

Цикл Python против массива

n = 1_000_000
values = np.arange(n, dtype=np.float64)

def sum_squares_loop(a):
    total = 0.0
    for v in a:
        total += v * v
    return total

%timeit -n 1 -r 3 sum_squares_loop(values)
%timeit -n 10 -r 3 np.sum(values * values)
140 ms ± 726 μs per loop (mean ± std. dev. of 3 runs, 1 loop each)
589 μs ± 124 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)

Цикл Python на каждом шаге проверяет типы; NumPy выполняет цикл в скомпилированном коде над непрерывным блоком памяти.

Форма, тип, срезы

a = np.arange(12).reshape(3, 4)
print(a)
print("shape:", a.shape, " dtype:", a.dtype)
print("строка 1:   ", a[1])
print("столбец 2:  ", a[:, 2])
print("каждый 2-й: ", a[0, ::2])
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
shape: (3, 4)  dtype: int64
строка 1:    [4 5 6 7]
столбец 2:   [ 2  6 10]
каждый 2-й:  [0 2]

Срез — вид, а не копия

a = np.zeros(6)
b = a[::2]          # вид на те же данные
b[0] = 100.0
print("после b[0] = 100: a =", a)   # a тоже изменился
c = a[::2].copy()                   # независимая копия
c[1] = -1.0
print("после c[1] = -1:  a =", a)   # a не изменился
после b[0] = 100: a = [100.   0.   0.   0.   0.   0.]
после c[1] = -1:  a = [100.   0.   0.   0.   0.   0.]

Изменение среза меняет исходный массив. Нужна независимая копия — .copy().

Broadcasting

x = np.array([0.0, 1.0, 2.0])           # форма (3,)
y = np.array([[0.0], [10.0]])           # форма (2, 1)
print(x + y)                            # форма (2, 3)

pts = np.array([[0.0, 0.0], [3.0, 4.0], [6.0, 8.0]])
d = np.linalg.norm(pts[:, None, :] - pts[None, :, :], axis=-1)
print(d)                                # все попарные расстояния без цикла
[[ 0.  1.  2.]
 [10. 11. 12.]]
[[ 0.  5. 10.]
 [ 5.  0.  5.]
 [10.  5.  0.]]

Размерности сравниваются справа; размерность 1 растягивается.

Числа с плавающей точкой

Машинная точность

print(0.1 + 0.2 == 0.3, 0.1 + 0.2)
eps = np.finfo(np.float64).eps
print("eps(float64) =", eps, "  1 + eps/2 == 1:", 1.0 + eps / 2 == 1.0)
print("eps(float32) =", np.finfo(np.float32).eps)
False 0.30000000000000004
eps(float64) = 2.220446049250313e-16   1 + eps/2 == 1: True
eps(float32) = 1.1920929e-07

\(\varepsilon\) — наименьшее число, для которого \(1+\varepsilon\ne1\). Сравнивать вещественные числа на точное равенство нельзя.

Порядок суммирования важен

tenths = np.full(10_000_000, 0.1, dtype=np.float32)
print("последовательно (cumsum):", np.cumsum(tenths)[-1])
print("np.sum (попарно):        ", tenths.sum())
print("float64:                 ", tenths.astype(np.float64).sum())
последовательно (cumsum): 1.087937e+06
np.sum (попарно):         1.0000001e+06
float64:                  1000000.0149011612

Точный ответ \(10^6\). Когда сумма велика, прибавка \(0.1\) округляется; np.sum складывает попарно, и ошибка растёт медленнее.

Случайные числа

Генератор с seed

rng = np.random.default_rng(seed=2026)
print(rng.integers(0, 10, size=8))
print(np.random.default_rng(seed=2026).integers(0, 10, size=8))  # то же самое
print(rng.integers(0, 10, size=8))                                 # продолжение
[8 1 0 6 3 4 0 3]
[8 1 0 6 3 4 0 3]
[6 3 8 7 7 9 7 1]

Один генератор на расчёт, seed записан в коде: результат повторяется.

Задача 1: случайное блуждание

Моделирование

rng = np.random.default_rng(seed=1)
n_walkers, n_steps = 2000, 1000
steps = rng.choice(np.array([-1, 1], dtype=np.int8), size=(n_walkers, n_steps))
x = np.cumsum(steps, axis=1, dtype=np.int64)   # x[k, n-1] — положение k-го блуждания после n шагов
print("x.shape =", x.shape)
print("<x_N>   =", x[:, -1].mean())
print("<x_N^2> =", (x[:, -1] ** 2).mean(), " (эталон", n_steps, ")")
x.shape = (2000, 1000)
<x_N>   = -0.347
<x_N^2> = 1032.122  (эталон 1000 )

Отклонение \(\langle x_N^2\rangle\) от \(N\) сравниваем со статистической ошибкой \(\sqrt{2/M}\approx3\%\), а не с нулём.

Проверка по эталону

Рисунок 1: Отношение \(\langle x_n^2\rangle/n\) к эталону 1; полоса — статистическая ошибка \(\pm\sqrt{2/M}\).

Распределение и ловушка чётности

Рисунок 2: Распределение положения после \(N=1000\) шагов и гауссов предел.

\(x_N\) той же чётности, что \(N\): значения идут через 2, вероятность вдвое больше плотности.

Задача 2: число \(\pi\)

Векторизованная оценка

def pi_numpy(n, rng):
    xy = rng.random((n, 2))
    inside = (xy ** 2).sum(axis=1) < 1.0
    return 4.0 * inside.mean()

rng = np.random.default_rng(seed=3)
for n in [10**2, 10**4, 10**6]:
    est = pi_numpy(n, rng)
    print(f"n = {n:>8}:  pi ~ {est:.5f},  ошибка {est - np.pi:+.5f},  "
          f"ожидаемая {1.64 / np.sqrt(n):.5f}")
n =      100:  pi ~ 3.20000,  ошибка +0.05841,  ожидаемая 0.16400
n =    10000:  pi ~ 3.13760,  ошибка -0.00399,  ожидаемая 0.01640
n =  1000000:  pi ~ 3.14054,  ошибка -0.00105,  ожидаемая 0.00164

Сходимость \(1/\sqrt n\)

Рисунок 3: Ошибка оценки \(\pi\) убывает как \(1/\sqrt{n}\).

Ещё один верный знак стоит в 100 раз больше точек. При \(n=10^8\) массив координат занял бы 1.6 ГБ.

Numba

Компиляция цикла

from numba import njit

@njit
def pi_numba(n, seed):
    np.random.seed(seed)          # у Numba свой генератор; seed задаётся внутри
    hits = 0
    for _ in range(n):
        x = np.random.random()
        y = np.random.random()
        if x * x + y * y < 1.0:
            hits += 1
    return 4.0 * hits / n

Сколько это стоит

import time

def timed(f, *args):
    t0 = time.perf_counter()
    result = f(*args)
    return result, time.perf_counter() - t0

_, t_first = timed(pi_numba, 10**6, 5)
_, t_second = timed(pi_numba, 10**6, 5)
_, t_numpy = timed(pi_numpy, 10**6, np.random.default_rng(5))
est, t_big = timed(pi_numba, 10**8, 5)

print(f"Numba, первый вызов (с компиляцией), n=1e6: {t_first:8.3f} с")
print(f"Numba, второй вызов,                 n=1e6: {t_second:8.3f} с")
print(f"NumPy,                               n=1e6: {t_numpy:8.3f} с")
print(f"Numba, n=1e8: pi ~ {est:.6f}, ошибка {est - np.pi:+.6f}, время {t_big:.2f} с")
Numba, первый вызов (с компиляцией), n=1e6:    0.656 с
Numba, второй вызов,                 n=1e6:    0.020 с
NumPy,                               n=1e6:    0.030 с
Numba, n=1e8: pi ~ 3.141380, ошибка -0.000213, время 2.00 с

Когда Numba не помогает

  • Функция уже состоит из операций NumPy: они и так скомпилированы.
  • Функция вызывается один раз: компиляция дороже счёта.
  • Код использует то, чего Numba не знает: словари произвольных типов, pandas, большинство функций SciPy.

Numba нужна для циклов, которые не векторизуются (шаг зависит от предыдущего) или не помещаются в память.

Ловушки

Ловушки

  • Срез меняет исходный массив.
  • 0.1 + 0.2 != 0.3: вещественные числа сравниваются с допуском (np.isclose).
  • float32 теряет точность при накоплении; по умолчанию — float64.
  • Генератор без seed даёт неповторяемый результат; у Numba свой генератор.
  • Первый вызов функции Numba включает компиляцию: не путать с временем счёта.
  • Время выполнения зависит от машины: сравнивать только на одной.

Итоги

Что мы взяли с собой

  • Ноутбук, окружение и версии пакетов.
  • Массив NumPy: форма, тип, срезы-виды, broadcasting.
  • Конечная точность: \(\varepsilon\approx2.2\cdot10^{-16}\) для float64.
  • Генератор случайных чисел с seed.
  • Проверка по эталону: \(\langle x_N^2\rangle=N\), ошибка Монте-Карло \(\propto1/\sqrt n\).
  • Numba для циклов, которые не векторизуются.

Задачи