Быстрое префиксное заполнение строк в NumPy: цикл + Numba

Нужно быстро заполнять строки NumPy с префиксными срезами разной длины? Покажем, как скомпилировать цикл с Numba и получить 10x ускорение без смены алгоритма.

Эффективно заполнять большие массивы NumPy становится непросто, когда в каждой строке требуется обновить разное количество элементов. Если попытаться векторизовать такой шаблон, быстро упираешься в ограничения индексации или неожиданные нюансы broadcasting. Простой цикл на Python даёт верный результат, но при очень большом числе итераций (значительно больше десяти миллионов) работает мучительно медленно. Ниже — аккуратный способ ускорить это без смены алгоритма.

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

Имеется массив m×n и три массива 1×m, которые определяют заполнение. Для каждого индекса из заранее отсортированного списка вы вычисляете значение и записываете его в префиксный срез строки, выбранной по бин-индексу угла. По сути это выглядит так: «для заданной угловой строки записать одно и то же значение в столбцы [0:elev)». Попытка выполнить такое вещание без цикла приводит к проблемам — например, к нецелочисленным индексам или несовместимым формам. Прямой цикл работает, но при десятках миллионов итераций превращается в узкое место.

Базовый код: правильно, но медленно

Ниже показан цикл, который даёт корректный результат, но на очень больших входных данных занимает слишком много времени.

import numpy as np
# grid: массив m x n
# heights, radii, ang_bins: массивы 1 x m
# order_idx: индексный массив (отсортирован заранее)
def fill_naive(grid, order_idx, heights, radii, ang_bins):
    for j in order_idx:
        stop = heights[j]
        val = 1 - (radii[j] / 10813)
        grid[int(ang_bins[j]), 0:stop] = val

Пытаясь убрать цикл, можно попытаться использовать продвинутую индексацию по «рваным» срезам и наткнуться на хорошо известное ограничение NumPy.

IndexError: arrays used as indices must be of integer (or boolean) type

Это происходит потому, что для каждой строки требуется свой диапазон по столбцам, а стандартная индексация NumPy не допускает выражение с переменной длиной срезов на строку в одном присваивании.

Почему здесь сложно векторизовать

Суть операции — «рваная» запись: каждую выбранную строку нужно заполнить лишь до границы, заданной высотой (elevation). Это не прямоугольный блок, который можно присвоить разом, и одновременно требуются целочисленные индексы строк. Вместе эти два условия делают очевидные векторизованные подходы либо некорректными, либо громоздкими — и без существенной переработки они не избавляют от цикла.

Быстро и просто: скомпилируйте цикл с Numba

Отказ от цикла не обязателен. Скомпилируйте его с помощью @njit из Numba, чтобы убрать накладные расходы интерпретатора, сохранив ту же логику. Код останется понятным, а скорость заметно вырастет.

from numba import njit
@njit
def paint_spans(canvas, idx_sorted, elev_vec, rad_vec, theta_vec):
    for j in idx_sorted:
        end = elev_vec[j]
        value = 1 - (rad_vec[j] / 10813)
        canvas[int(theta_vec[j]), 0:end] = value

Вызовите эту функцию с вашими массивами и заранее отсортированным списком индексов. Алгоритм не меняется; цикл выполняется на нативной скорости. На практике это уже даёт примерно десятикратное ускорение.

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

Когда список индексов насчитывает многие миллионы элементов, основную задержку создаёт интерпретатор, даже если итерации простые. Шаблон доступа к данным здесь плохо укладывается в «чистую» векторизацию NumPy, потому что у каждой строки свой диапазон заполнения. JIT-компиляция цикла с Numba обходит обе проблемы и при этом держит код лаконичным и поддерживаемым.

Выводы

Если нужно присваивать срезы переменной длины по строкам, не пытайтесь насильно векторизовать то, что к этому не располагает. Сохраните понятный цикл и скомпилируйте его. Убедитесь, что индексы строк — целые, срез по строке задан как [0:elev), а значение вычисляется как 1 - (r / 10813). С @njit вы сохраните корректность и получите нужную производительность без лишних усложнений.