在numba jitted函数中动态增长阵列

Question

似乎在numba中不受支持。

在nopython模式下使用numba.jit动态增长数组的最佳方式是什么？

到目前为止，我能做的最好的事情是定义并调整jitted函数之外的数组的大小，是否有更好的（和更整齐的）选项？

Answer 1

numpy.resize是pure python function：

import numpy as np 

def resize(a, new_shape): 
    """I did some minor changes so it all works with just `import numpy as np`.""" 
    if isinstance(new_shape, (int, np.core.numerictypes.integer)): 
     new_shape = (new_shape,) 
    a = np.ravel(a) 
    Na = len(a) 
    if not Na: 
     return np.zeros(new_shape, a.dtype) 
    total_size = np.multiply.reduce(new_shape) 
    n_copies = int(total_size/Na) 
    extra = total_size % Na 

    if total_size == 0: 
     return a[:0] 

    if extra != 0: 
     n_copies = n_copies+1 
     extra = Na-extra 

    a = np.concatenate((a,)*n_copies) 
    if extra > 0: 
     a = a[:-extra] 

    return np.reshape(a, new_shape) 

对于一维数组，这将是直着自己实现。不幸的是，ND阵列要复杂得多，因为一些操作在nopython numba函数中不受支持：isinstance,reshape和元组乘法。这里是1D相当于：

import numpy as np 
import numba as nb 

@nb.njit 
def resize(a, new_size): 
    new = np.zeros(new_size, a.dtype) 
    idx = 0 
    while True: 
     newidx = idx + a.size 
     if newidx > new_size: 
      new[idx:] = a[:new_size-newidx] 
      break 
     new[idx:newidx] = a 
     idx = newidx 
    return new 

，当你不希望这样“重复输入”行为，只会用它来增加大小那就更简单了：

@nb.njit 
def resize(a, new_size): 
    new = np.zeros(new_size, a.dtype) 
    new[:a.size] = a 
    return new 

这些功能都装饰与numba.njit，因此可以在nopython模式中的任何numba函数中调用。

---

注意的一点是，虽然：一般来说，你不想来调整 - 或者如果你然后确保你选择有amoritzed O(1) cost (Wikipedia link)的方法。如果你可以估计最大长度，那么最好立即预先分配一个正确大小（或略微超额分配）的数组。

Answer 2

通常，我使用的策略是只分配足够多的数组存储以适应计算，然后跟踪最终使用的索引/索引，然后在返回之前将数组切片降至实际大小。这假定您事先知道您可能将阵列增长到的最大大小。我的想法是，在我的大多数应用程序中，内存很便宜，但是调整大小并在python和jitted函数之间切换很贵。