Обобщающий код Cython для n-мерных массивов

Я пишу много кода на Cython, который нужно обобщить на n измерений (в большинстве случаев это будет n = 1, 2 или 3).

Возьмите следующий код в качестве примера, это простая сумма элементов.

cimport cython
from cython cimport floating

cpdef floating nd_sum(floating [:, :, :] arr):
    cdef:
        Py_ssize_t [3] N = [arr.shape[0], arr.shape[1], arr.shape[2]]
        Py_ssize_t [3] i
        floating total = 0

    for i[0] in range(N[0]):
        for i[1] in range(N[1]):
            for i[2] in range(N[2]):
                total += arr[i[0], i[1], i[2]]

    return total

Концептуально обобщение на n измерений вполне очевидно. Тем не менее, я не могу реализовать это в коде. Я не хотел бы, чтобы в основном повторять этот код для каждого случая n= 1,2,3,...

В общем, мне также нужно иметь возможность доступа к соседям каждого элемента массива в каждом измерении (например, n- мерных сверток ядра и т. Д.), Что делает невозможным уплощение массива.

Мои основные проблемы:

  • Как мне сказать Cython ожидать массива с произвольными измерениями в качестве входных данных? Я полагаю, мне придется прибегнуть к np.ndarray Вот...

  • Как я могу получить элемент из n- мерного массива с массивом из n индексов?

  • Как я могу обобщить конструкцию из n вложенных циклов? Здесь рекурсия кажется неизбежной, но как мне получить доступ ко всем n индексам таким образом?

0 ответов

Другие вопросы по тегам