Как создать кортеж длины N в numba для малого N (или как быстро вычислить двунаправленные изменения между одномерными и многомерными индексами)

Вопрос:

Вот простая функция, которая работает с numpy, но не с numba:

    # @numba.jit(nopython=True, fastmath=False, parallel=False)
    def testgetvalue(tgvarray, tgvindex):
        tgvalue = tgvarray[tuple(tgvindex)]
        return tgvalue

Как мне сделать версию этой функции, которая будет работать в numba?

Я пробовал:

    @numba.jit(nopython=True, fastmath=False, parallel=False)
    def testgetvalue2(tgvarray, tgvindex):
        tgvalue = tgvarray[tuple(tgvindex)]
        currentdex = tgvindex[0]
        tgvtemp = tgvarray[currentdex]
        for idx in range(1, len(tgvindex)):
            currentdex = tgvindex[idx]
            tgvtemp =  tgvtemp[currentdex]
        return tgvalue

но это также не работает в numba

Я нашел этот вопрос, в котором ответ говорит, что это возможно:

В более общем плане, вы не можете создать N-арный кортеж, где N переменное, в функции Numba. Однако вы можете вместо этого сгенерировать и скомпилировать функцию для конкретного N. Если N очень мал, например, <15

Это решило бы мою проблему, но ответ не объясняет, как сгенерировать и затем скомпилировать функцию для конкретного N… если только предложение не заключается в том, чтобы мне написать скрипт, который создаст .py файл, определяющий функцию с декоратором jit, которую я затем смогу вызвать из .py файла. Думаю, это может сработать, учитывая нечастые изменения в измерениях. Я не уверен, можно ли это считать лучшей практикой, но я начну работать над скриптом для генерации .py файлов, пока кто-то не даст иного ответа.

Пожалуйста, продолжайте читать, прежде чем предполагать, что я дублирую что-то вроде этого вопроса, потому что моя настоящая проблема не обязательно связана с кортежами, т.е. версия вопроса в скобках.

Почему я задаю этот вопрос:
У меня есть код, в котором количество измерений объекта массива иногда может варьироваться от 1 до 15 измерений. Но как только происходит изменение, будет проведено десятки тысяч повторяющихся операций над этим многомерным массивом. Для многих из этих операций мне было бы желательно использовать массив индексов для изменения значения в месте многомерного массива.

Это ведет к альтернативному вопросу в скобках:

В предыдущей версии моего кода я преобразовал многомерные индексы в одномерные, взяв массив размеров в каждом измерении и сделав:

    multipliers = np.cumprod(array_of_sizes_in_each_dimension)
    multipliers = np.roll(multipliers, 1)
    multipliers[0] = 1

Затем я могу умножить каждое значение в многомерном индексе на соответствующее значение в multipliers, чтобы получить мой одномерный индекс. Это работает хорошо, когда я перехожу от многомерного индекса к одномерному индексу. Однако мне не приходит в голову быстрая функция для случаев, когда мне нужно перейти от одномерного индекса к многомерному индексу. В настоящее время моя самая быстрая версия – построить таблицу поиска, т.е. two_dimensional_array размером np.prod(array_of_sizes_in_each_dimension) X len(array_of_sizes_in_each_dimension), которая имеет списки всех многомерных индексов, так что two_dimensional_array[one_dimensional_index] возвращает соответствующий многомерный индекс. Это работает хорошо, когда мой многомерный массив оказывается только с несколькими измерениями и коротким в каждом из этих измерений. Однако, по мере увеличения количества измерений, размер two_dimensional_array становится таким большим, что я сталкиваюсь с проблемой памяти, и код замедляется на несколько порядков (например, 8 минут для 3-мерного массива и 8 дней для 11-мерного массива). Поэтому, если кто-то имеет идею для быстрой функции, чтобы заменить таблицу поиска two_dimensional_array, это также решит мою проблему.

На самом деле возможно создать кортеж внутри numba, если его размер известен на этапе компиляции, используя numba.np.unsafe.ndarray. Эта функция помечена как небезопасная, и я не нашел никаких следов ее в документации, поэтому не знаю, насколько она надежна в долгосрочной перспективе, но она хорошо работала для numba 0.58 до 0.60. Например, вы можете использовать ее в качестве формы для инициализации nd-мерного массива.

import numpy as np
import numba

@numba.njit
def numba_make_tuple(old_tuple):
    CONST = len(old_tuple)
    a = np.empty((CONST,), dtype=np.int64)
    for i in range(CONST):
        a[i] = i

    new_tuple = numba.np.unsafe.ndarray.to_fixed_tuple(a, CONST)
    return new_tuple


t1 = (1,)
nt1 = numba_make_tuple(t1)
print(nt1, type(t1))

t2 = (2,3,1)
nt2 = numba_make_tuple(t2)
print(nt2, type(t2))

Обратите внимание, что я использую кортеж в качестве входных данных, чтобы обеспечить знание CONST на этапе компиляции. Я не нашел ничего в официальной документации, но это обсуждается, например, на https://github.com/numba/numba/issues/8812

Попробуйте это:

import numba
import numpy as np

@numba.njit
def testgetvalue(tgvarray, tgvindex):
    tgvtemp = tgvarray
    for idx in range(len(tgvindex)):
        tgvtemp = tgvtemp[tgvindex[idx]]
    return tgvtemp