如何在Numba中创建小长度N元组？或实现索引双向快速转换

问题描述 #

以下是一个在numpy中可用但在Numba中无法运行的简单函数：

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

如何编写一个可在Numba中运行的版本？

我尝试了以下实现，但同样在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

我在Stack Overflow上找到一个相关问题，其中提到：

通常无法在Numba函数中生成长度可变的N元组，但可为特定N生成并编译函数（当N很小，如<15时）

这似乎能解决我的问题，但该回答未说明如何为特定N生成并编译函数——难道是要编写脚本生成带jit装饰器的.py文件？考虑到维度变化不频繁，这或许可行，但不确定是否符合最佳实践，目前正准备尝试此方案，同时寻求其他解答。

请注意，我的实际问题并非仅局限于元组：

问题背景 #

我的代码中数组的维度可能偶尔变化，维度范围为1到15，但维度变化后会对该多维数组执行数万次重复操作，其中很多操作需要通过索引数组修改多维数组指定位置的值。

替代问题 #

早期版本中，我通过以下代码将多维索引转换为一维索引：

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

将多维索引的每个值与multipliers对应值相乘即可得到一维索引，这在多维转一维时运作良好。但我无法找到高效的一维转多维索引的方法：目前最快的方式是构建查找表，即一个尺寸为np.prod(array_of_sizes_in_each_dimension) × len(array_of_sizes_in_each_dimension)的two_dimensional_array，通过two_dimensional_array[one_dimensional_index]获取对应多维索引。但维度增多时，该查找表会因内存瓶颈导致代码速度骤降（如3维数组耗时8分钟，11维则需8天），因此寻求替代该查找表的高效函数。

解决方案 #

一、Numba中支持可变维度的数组索引实现 #

1. 动态生成维度专用函数

你提到的动态生成特定维度的函数是可行方案，无需手动生成.py文件，可通过Python元编程在运行时动态创建并编译对应维度的Numba函数：

import numba as nb
import numpy as np

# 缓存已编译的索引函数
index_functions = {}

def generate_index_func(dim):
    if dim in index_functions:
        return index_functions[dim]
    
    # 针对低维度直接写逻辑，高维度用exec动态生成代码
    if dim <= 3:
        if dim == 1:
            @nb.jit(nopython=True)
            def idx_func(arr, idx):
                return arr[idx[0]]
        elif dim == 2:
            @nb.jit(nopython=True)
            def idx_func(arr, idx):
                return arr[idx[0], idx[1]]
        elif dim == 3:
            @nb.jit(nopython=True)
            def idx_func(arr, idx):
                return arr[idx[0], idx[1], idx[2]]
    else:
        # 动态生成对应维度的索引代码
        code = f"""
@nb.jit(nopython=True)
def idx_func(arr, idx):
    return arr[{','.join(f'idx[{i}]' for i in range(dim))}]
"""
        local_vars = {}
        exec(code, globals(), local_vars)
        idx_func = local_vars['idx_func']
    
    index_functions[dim] = idx_func
    return idx_func

# 使用示例
arr_3d = np.random.rand(5,5,5)
idx = np.array([2,3,1])
func = generate_index_func(3)
print(func(arr_3d, idx))

这种方式利用exec动态生成对应维度的索引代码，Numba会为每个维度单独编译优化函数，既规避了可变元组的问题，又能保证运行效率。由于维度变化不频繁，首次编译的开销可忽略不计。

2. 修复循环实现

你之前的testgetvalue2函数失败是因为保留了tuple(tgvindex)的错误代码，去掉后即可正常运行：

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

这个版本支持任意维度，但性能略低于预编译的维度专用函数——循环虽能被Numba优化，但不如直接展开的索引操作高效。若维度范围固定在1-15，预编译专用函数是更优选择。

二、高效的一维索引转多维索引实现 #

无需构建内存密集的查找表，通过数学计算即可快速转换，经Numba编译后性能接近原生C水平：

@nb.jit(nopython=True)
def one_to_multi(flat_idx, shape):
    multi_idx = np.empty(len(shape), dtype=np.int64)
    remainder = flat_idx
    # 从最后一维向前计算每个维度的索引
    for i in range(len(shape)-1, -1, -1):
        multi_idx[i] = remainder % shape[i]
        remainder = remainder // shape[i]
    return multi_idx

# 使用示例
shape = (5,5,5)
flat_idx = 2*5*5 + 3*5 +1
print(one_to_multi(flat_idx, shape))  # 输出 [2,3,1]

该函数通过逐次取余和整除操作计算多维索引，时间复杂度为O(D)（D为维度数），内存开销仅为存储多维索引的数组，完全避免了查找表的内存瓶颈。

三、完整工作流 #

1. 维度变化时，预编译对应维度的数组索引/赋值函数
2. 如需批量操作，将多维索引转一维索引处理
3. 访问多维数组时，使用预编译函数或一维转多维后直接访问

此方案既解决了Numba中多维索引的兼容性问题，又规避了查找表的内存瓶颈，同时保证了数万次重复操作的执行效率。

内容的提问来源于stack exchange，提问作者Nathan Gabriel