CODEWITHSUNDEEP
pythonnumpydynamic-arraysnumba
nivniv
nivniv

Reputation: 3722

dynamically growing array in numba jitted functions

It seems that numpy.resize is not supported in numba.

What is the best way to use dynamically growing arrays with numba.jit in nopython mode?

So far the best I could do is define and resize the arrays outside the jitted function, is there a better (and neater) option?

Upvotes: 3

Views: 3429

Answers (3)

duburcqa
duburcqa

Reputation: 1131

To dynamically increase the size of an existing array (and therefore do it in-place), numpy.ndarray.resize must be used instead of numpy.resize. This method is NOT implemented in Python, and is not available in Numba, so it just cannot be done.

Upvotes: 1

MSeifert
MSeifert

Reputation: 152667

numpy.resize is a 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)

For 1D arrays this would be straight-forward to implement yourself. Unfortunatly it's a lot more complicated for ND arrays because some operations aren't supported in nopython numba functions: isinstance, reshape and the tuple-multiplication. Here is the 1D equivalent:

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

when you don't want that "repeat input" behaviour and only use it for increasing the size it's even easier:

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

These functions are decorated with numba.njit and can therefore be called in any numba functions in nopython mode.

A word of caution though: Generally you don't want to resize - or if you do then make sure you choose an approach that has amoritzed O(1) cost (Wikipedia link). If you can estimate the maximum length then it's better to pre-allocate a correctly sized (or slightly overallocated) array immediatly.

Upvotes: 5

JoshAdel
JoshAdel

Reputation: 68682

Typically the strategy I employ is to just allocate more than enough array storage to accommodate the calculation and then keep track of the final index/indices used, and then slice the array down to the actual size before returning. This assumes that you know beforehand what the maximum size you could possibly grow the array to is. The thought is that in most of my own applications, memory is cheap but resizing and switching between python and jitted functions a lot is expensive.

Upvotes: 1

Related Questions