Пакетное вычисление собственных значений и векторов 4×4 в NumPy без циклов

Векторизация вычисления собственных значений и векторов матриц 4×4 в NumPy: сборка батча, np.linalg.eigh за один вызов и сортировка без Python-циклов быстро.

Векторизация задач линейной алгебры в NumPy особенно оправдана, когда одну и ту же операцию приходится повторять для множества входных данных. Типичный случай — вычисление собственных значений и собственных векторов для структурированной матрицы 4×4, зависящей от нескольких параметров. Цель здесь — запустить весь пакет сразу, без явного цикла на Python.

Базовый вариант: разложение на собственные значения и векторы для одного вызова

Следующая функция строит матрицу 4×4 из входов (x1, y1, x2, y2), затем вычисляет и сортирует её собственные значения и собственные векторы с помощью np.linalg.eigh.

import numpy as np

def phi(px, py):
    return px + 1j * py

def phi_bar(px, py):
    return np.conj(phi(px, py))

def solve_block(u1, v1, u2, v2):
    alpha = 10
    u = u1 + u2
    v = v1 + v2
    M = np.array([
        [alpha,           phi(u, v),     phi(u, v),     phi_bar(u, v)],
        [phi_bar(u, v),   alpha,         0,             phi(u, v)],
        [phi_bar(u, v),   0,             -alpha,        phi(u, v)],
        [phi(u, v),       phi_bar(u, v), phi_bar(u, v), -alpha]
    ])

    w, V = np.linalg.eigh(M)
    order = np.argsort(w)
    return w[order], V[:, order]

Что мешает векторизации

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

Векторизованный подход

Идея в том, чтобы собрать «стек» массивов формы “batch × 4 × 4” и передать его напрямую в np.linalg.eigh. Кроме того, стоит избежать повторных вычислений одинаковых подвыражений: один раз посчитать их и переиспользовать при формировании пакета.

def solve_batched(u1_arr, v1_arr, u2_arr, v2_arr):
    alpha   = 10
    u_arr   = u1_arr + u2_arr
    v_arr   = v1_arr + v2_arr
    z       = phi(u_arr, v_arr)
    z_conj  = phi_bar(u_arr, v_arr)

    H = np.stack([
        np.stack([np.full_like(z, alpha), z,       z,       z_conj],              axis=-1),
        np.stack([z_conj,                 np.full_like(z, alpha), np.zeros_like(z), z], axis=-1),
        np.stack([z_conj,                 np.zeros_like(z),       -np.full_like(z, alpha), z], axis=-1),
        np.stack([z,                      z_conj,                 z_conj,                 -np.full_like(z, alpha)], axis=-1)
    ], axis=-2)

    evals, evecs = np.linalg.eigh(H)

    idx = np.argsort(evals, axis=-1)
    evals = np.take_along_axis(evals, idx, axis=-1)
    evecs = np.take_along_axis(evecs, idx[..., None], axis=-1)

    return evals, evecs

Почему это работает

Объединив матрицы в один массив, мы получаем разложение на собственные значения и векторы для всего пакета за единичный вызов. Предварительное вычисление phi(u, v) и его сопряжения избавляет от лишней работы при заполнении блоков. Финальный шаг повторяет базовый вариант: собственные значения и векторы последовательно сортируются вдоль последней оси.

Зачем это нужно

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

Выводы

Если матрица зависит от покомпонентных выражений над входными массивами, сначала вычислите общие промежуточные величины, затем соберите стопку 4×4-блоков по каждому входу и, наконец, один раз вызовите np.linalg.eigh для всего пакета. Постобработку оставляйте такой же, как в скалярной версии: сортируйте пары собственных значений и векторов вдоль последней оси, чтобы последующий код получал устойчивый порядок.