Частичная трассировка напрямую из вектора состояния в NumPy

Как вычислить частичную трассировку в NumPy без матрицы плотности: считаем из вектора состояния, экономим память и избегаем ошибки ArrayMemoryError.

Построение полного внешнего произведения большого комплексного вектора только ради частичной трассировки — классический способ упереться в пределы памяти. Если размерность состояния 91204 (complex128), промежуточная матрица плотности имеет размер 91204 × 91204 и мгновенно переполняет ОЗУ. Сообщение об ошибке говорит само за себя:

numpy._core._exceptions._ArrayMemoryError: Unable to allocate 124. GiB for an array with shape (91204, 91204) and data type complex128

Решение — не материализовывать полную матрицу и считать частичную трассировку напрямую из вектора. Ниже — разбор проблемы, почему она возникает и как реализовать безопасное по памяти решение.

Наивный подход, который съедает всю память

Следующий код строит полное внешнее произведение, а затем пытается взять частичную трассировку через np.einsum над преобразованным массивом. Логика понятна, но этот путь неэффективен по памяти.

import numpy as np
def keep_trace(mat, stay, axes, use_opt=False):
    stay = np.asarray(stay)
    axes = np.asarray(axes)
    n_axes = axes.size
    kept_prod = np.prod(axes[stay])
    a_idx = [k for k in range(n_axes)]
    b_idx = [n_axes + k if k in stay else k for k in range(n_axes)]
    shaped = mat.reshape(np.tile(axes, 2))
    reduced = np.einsum(shaped, a_idx + b_idx, optimize=use_opt)
    return reduced.reshape(kept_prod, kept_prod)
nB = 150
vec_len = 4*(nB+1)**2  # 91204
phi_shape = (1, vec_len)
phi = np.random.uniform(-1, 1, phi_shape) + 1.j * np.random.uniform(-1, 1, phi_shape)
phi = phi / np.linalg.norm(phi)
density = np.outer(phi, phi.transpose().conj())
trace_kept = keep_trace(density, [0], [nB+1, 2, nB+1, 2], True)

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

Матрица плотности как внешнее произведение растёт по памяти квадратично относительно размера состояния. Для 91204-мерного комплексного состояния это матрица с более чем восьмью миллиардами элементов. Её размещение требует примерно 124 ГиБ для complex128, что недостижимо в большинстве сред. При этом для вычисления частичной трассировки нужна лишь малая часть этих элементов, так что львиная доля работы и памяти тратится впустую.

Эффективный подход: считать частичную трассировку прямо из вектора

Работайте с самим вектором состояния и не создавайте явную матрицу плотности. Пусть задан nB, тогда длина вектора равна 4·(nB+1)². Идея: преобразовать этот одномерный вектор в блоки длины v = 4·(nB+1), разбить пространство на (nB+1) × (nB+1) блоков и брать след внутри каждого блока. Элемент частичной трассировки в позиции (i, j) равен сумме покомпонентных произведений между отрезками длины v исходного вектора и его комплексно-сопряжённой версии:

P[i, j] = sum(phi[i*v : (i+1)*v] * phi.conj()[j*v : (j+1)*v])

Так мы получаем те же элементы частичной трассировки, ни разу не создавая матрицу размера O(N²).

import numpy as np
nB = 150
phi_shape = (1, 4*(nB+1)**2)
phi = np.random.uniform(-1, 1, phi_shape) + 1.j * np.random.uniform(-1, 1, phi_shape)
phi = phi / np.linalg.norm(phi)
result = np.zeros((nB+1, nB+1), dtype=np.complex128)
block_len = 4*(nB+1)
for r in range(nB+1):
    for c in range(nB+1):
        result[r, c] = np.sum(
            phi[0, r*block_len:(r+1)*block_len] *
            phi.conj()[0, c*block_len:(c+1)*block_len]
        )

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

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

Итоги

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