Смежность из pandas без плотной матрицы: список рёбер для PyTorch

Как из большого pandas DataFrame получить смежность без MemoryError: векторизация через stack/crosstab и прямое формирование списка рёбер для PyTorch — быстро

Построение плотной матрицы смежности из большого pandas DataFrame кажется простой задачей — пока всё не упирается в память. Когда индекс охватывает range(132000), материализовать квадратную матрицу 132k × 132k становится нереалистично, особенно если конечная цель — список рёбер для PyTorch. Есть более аккуратный и эффективный путь, который полностью избегает плотного промежуточного представления.

Воспроизведение сценария

Рассмотрим DataFrame, где в каждой строке хранится небольшой набор целочисленных значений (с некоторыми NaN). Задача — для каждой строки отметить, какие индексы встречались среди её значений, в итоге получив данные смежности.

import pandas as pd
import numpy as np
from numpy.random import default_rng
# пример фрейма
df_demo = pd.DataFrame(index=[i for i in range(0, 10)], columns=list('abcd'))
for ridx in df_demo.index:
    df_demo.loc[ridx] = default_rng().choice(10, size=4, replace=False)
# добавим NaN
df_demo.loc[1, 'b'] = np.nan
df_demo.loc[3, 'd'] = np.nan

Наивный подход — пройтись по строкам и заполнить квадратную 0/1‑матрицу, у которой строки и столбцы соответствуют одному и тому же множеству индексов.

# наивная квадратная матрица смежности
adj_square = pd.DataFrame(index=df_demo.index, columns=df_demo.index)
for ridx in df_demo.index:
    adj_square.loc[ridx, df_demo.loc[ridx].dropna().values] = 1
adj_square = adj_square.replace(np.nan, 0)

Почему тут возникают проблемы

Для крошечных данных это сработает, но масштаб не выдерживает. Плотная матрица 132k × 132k — это массивный объект, который на типичном оборудовании вызовет MemoryError задолго до того, как получится сделать что-то полезное. Корень проблемы не в самом цикле, а в попытке материализовать полный квадрат смежности, который в большинстве случаев крайне разрежен.

Скажите — чему равно 132_000**2?

Если в итоге вам нужен тензор со списком рёбер формы (2, number_of_edges), заполнять плотный DataFrame и расходовать память попусту не нужно.

Векторизованный путь в pandas (когда матрица действительно нужна)

Если всё-таки требуется квадратное представление смежности в pandas, используйте векторизованный метод без Python-циклов. Сплющите фрейм с помощью stack и поручите crosstab напрямую построить индикаторную матрицу.

# векторизованная матрица смежности через crosstab
series_flat = df_demo.stack()
adj_compact = (
    pd.crosstab(series_flat.index.get_level_values(0), series_flat.values)
      .rename_axis(index=None, columns=None)
)

Это даёт 0/1‑матрицу смежности: строки — исходные индексы, столбцы — значения, встречавшиеся в соответствующих строках. По сравнению с ручными циклами решение компактнее и быстрее. Но если область индексов — 132k, плотный объект 132k × 132k всё равно может выйти за пределы памяти.

Эффективный путь: сразу собрать список рёбер для PyTorch

Пропустите квадратную матрицу и сформируйте ровно то, что требуется PyTorch. В стеке Series уже есть индексы строк и соответствующие им значения — эти пары и есть рёбра.

import torch
row_ids = series_flat.index.get_level_values(0)
coords = torch.tensor([row_ids, series_flat.values], dtype=torch.int32)

В результате coords — тензор формы (2, number_of_edges). Если позже понадобится разреженный квадратный тензор, его можно построить напрямую, без плотного промежуточного шага.

sparse_mat = torch.sparse_coo_tensor(coords, torch.ones(len(series_flat)))

Если во входных данных встречаются дубли координат, дополнительно выполните coalesce.

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

Работа с крупными данными смежности требует учитывать разреженность. Построение полной плотной матрицы для 132k индексов не просто медленно — в условиях ограниченной памяти это практически нереализуемо. Прямое представление в виде списка рёбер соответствует целевому формату и укладывается в разумные ресурсы памяти и вычислений. Даже если нужен квадратный вид, правильная абстракция — разреженный тензор.

Выводы

При большой области индексов отдавайте предпочтение работе с «длинными» данными (stack) и формируйте минимальное представление, которое ожидает ваша библиотека. Используйте crosstab только когда матрица действительно необходима, а для больших масштабов выбирайте разреженные тензоры вместо плотных DataFrame. Эта небольшая смена подхода превращает вылет по памяти в компактный, дружелюбный к GPU конвейер.