Как векторизованно решить миллионы независимых систем 2×2 в NumPy

Как заменить медленный цикл и SymPy на пакетное численное решение: собираем системы 2×2 в батчи, используем NumPy np.linalg.solve и ускоряем расчёты в разы.

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

Постановка задачи

Нагрузка состоит из независимых линейных систем с двумя переменными, заданных векторами коэффициентов A, B, C, A', B', C'. Для каждого индекса i система имеет вид:

A[i] * x + B[i] * y = -C[i]
A'[i] * x + B'[i] * y = -C'[i]

Последовательный подход, который не масштабируется

Простейшая реализация — цикл на Python и символический решатель для каждой системы. Это удобно, но накладные расходы становятся запредельными уже на 10M+ строк.

from sympy import symbols, Eq, solve

u_sym, v_sym = symbols('u_sym v_sym')
res_x = []
res_y = []
for idx in range(len(A)):
    sol_pair = solve([
        Eq(A[idx] * u_sym + B[idx] * v_sym, -C[idx]),
        Eq(A_prime[idx] * u_sym + B_prime[idx] * v_sym, -C_prime[idx])
    ], [u_sym, v_sym])
    res_x.append(sol_pair[u_sym])
    res_y.append(sol_pair[v_sym])

Такой шаблон делает ровно то, что описано: решает по одной системе в интерпретаторе Python, каждый раз вызывая символический движок. При миллионах строк это перестаёт быть жизнеспособным.

Откуда реально берётся торможение

Затраты идут из двух источников. Во‑первых, многократные вызовы символического решателя на крошечных плотных системах на порядки тяжелее, чем численные процедуры. Во‑вторых, сам по себе цикл Python на каждой итерации добавляет накладные расходы, которые накапливаются. Когда системы независимы и одинаковой структуры, естественное решение — собрать их в пакеты и применить векторизованный численный метод.

Векторизованное решение с пакетным численным решателем

Переход на пакетный численный решатель избавляет и от построчных накладных расходов Python, и от символического стека. Идея проста: собрать все левые части 2×2 в один массив, правые части — в другой, затем передать их пакетному решателю. Компактный пример:

import numpy as np

# Две независимые системы 2×2 в одном пакете
# Система 1: [[1, 0], [0, 2]] * [x, y]^T = [2, 4]^T
# Система 2: [[2, 0], [0, 1]] * [x, y]^T = [6, 1]^T
batch_lhs = [
    [[1, 0], [0, 2]],
    [[2, 0], [0, 1]],
]

batch_rhs = [
    [[2], [4]],
    [[6], [1]],
]

sol = np.linalg.solve(batch_lhs, batch_rhs)
print(sol)

Результат:

array([[[2.],
        [2.]],

       [[3.],
        [1.]]])

В ещё более компактном виде:

import numpy as np

lhs_pack = [[[1, 0], [0, 2]], [[2, 0], [0, 1]]]
rhs_pack = [[[2], [4]], [[6], [1]]]
print(np.linalg.solve(lhs_pack, rhs_pack).tolist())

что выводит:

[[[2.0], [2.0]], [[3.0], [1.0]]]

В одном из запусков пакетный np.linalg.solve отработал за 0,23 секунды.

От векторов коэффициентов к пакетным массивам

Имея исходные векторы A, B, C, A', B', C', можно напрямую собрать блоки левой и правой части для всех индексов. Каждый блок 2×2 строится из соответствующих коэффициентов, а правая часть — это столбец констант:

import numpy as np

# Предположим, что A, B, C, A_prime, B_prime, C_prime — одномерные массивы/списки одинаковой длины
n = len(A)

# Построить [ [A[i], B[i]], [A_prime[i], B_prime[i]] ] для каждого i
lhs_batched = np.array([
    [
        [A[i],       B[i]      ],
        [A_prime[i], B_prime[i]]
    ]
    for i in range(n)
], dtype=float)

# Построить [ [-C[i]], [-C_prime[i]] ] для каждого i
rhs_batched = np.array([
    [
        [-C[i]      ],
        [-C_prime[i]]
    ]
    for i in range(n)
], dtype=float)

# Решить все системы за один вызов
solutions = np.linalg.solve(lhs_batched, rhs_batched)

# Извлечь x и y как одномерные массивы
x_vals = solutions[:, 0, 0]
y_vals = solutions[:, 1, 0]

Так сохраняется точная алгебра исходной записи и применяется ко всему набору данных одним вызовом.

Ещё более простой вариант для систем 2×2

Для системы 2×2 можно просто выписать аналитическое решение — без какого‑либо линейного решателя. Это полезно, когда хочется ещё сильнее снизить зависимости или получить максимальный контроль над вычислениями.

Почему это важно

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

Итоги

Если вам нужно решать огромное количество небольших независимых систем, избегайте построчной символики. Представьте их в пакетной числовой форме и используйте векторизованный решатель. Для задач 2×2 подойдёт и замкнутая формула. В итоге код становится проще, поведение — предсказуемее, а производительность — соразмерной объёму данных.