How to efficiently update a numpy ndarray given a list of indices-CodePudding

I have a 4 dimensional array called new_arr. Given a list of indices, I want to update new_arr based on an old array I have stored, old_arr. I am using a for loop to do this, but it's inefficient. My code looks something like this:

update_indices = [(2,33,1,8), (4,9,49,50), ...] #as an example 
for index in update_indices:
    i,j,k,l = index
    new_arr[i][j][k][l] = old_arr[i][j][k][l]

It's taking a very long time because update_indices is large. Is there a way I can update all of the terms at once or do this more efficiently?

CodePudding user response：

Out of curiosity I have benchmarked the various improvements posted in the comments and found that working on flat indices is fastest.

I used the following setup:

rt numpy as np

n = 57
d = 4

k = int(1e6)
dt = np.double

new_arr = np.arange(n**d, dtype=dt).reshape(d * (n,))
new_arr2 = np.arange(n**d, dtype=dt).reshape(d * (n,))
old_arr = 2*np.arange(n**d, dtype=dt).reshape(d * (n,))


update_indices = list({tuple(np.random.randint(n, size=d)) for _ in range(int(k*1.1))})[:k]

where update_indices is a list of 1e6 unique index tuples.

Using the original technique from the question

%%timeit
for index in update_indices:
    i,j,k,l = index
    new_arr[i][j][k][l] = old_arr[i][j][k][l]

takes 1.47 s ± 19.3 ms.

Direct tuple-indexing as suggested by @defladamouse

%%timeit
for index in update_indices:
    new_arr[index] = old_arr[index]

indeed gives us a speedup of 2: 778 ms ± 41.8 ms

If update_indices is not given but can be constructed as ndarray as suggested by @Jérôme Richard

update_indices_array = np.array(update_indices, dtype=np.uint32)

(the conversion itself takes 1.34 s) the path to much faster implementations is open.

In order to index numpy arrays by multidimensional list of locations we cannot use update_indices_array directly as index, but pack its columns into a tuple:

%%timeit
idx = tuple(update_indices_array.T)
new_arr2[idx] = old_arr[idx]

Giving another speedup of roughly 9: 83.5 ms ± 1.45

If we dont leave the computation of memory offsets to ndarray.__getitem__, but compute the correspondig flat indices "by hand", we can become even faster:

%%timeit
idx_weights = np.cumprod((1,)   new_arr2.shape[:0:-1])[::-1]
update_flat = update_indices_array @ idx_weights
new_arr2.ravel()[update_flat] = old_arr.ravel()[update_flat]

resulting in 41.6 ms ± 1.04 ms, another factor of 2 and a cumulative speedup factor of 35 compared with the original version. idx_weights is simply an off-by-one reverse-order cumulative product of the array dimensions.

I assume that this speedup of 2 comes from the fact that the memory offsets / flat indices are computed twice in new_arr2[idx] = old_arr[idx] and only once in update_flat = update_indices_array @ idx_weight.

CodePudding user response：

Just do:

idx = np.array([(2,33,1,8), (4,9,49,50), ...])
new_arr[idx[:,0],idx[:,1],idx[:,2],idx[:,3]] = old_arr[idx[:,0],idx[:,1],idx[:,2],idx[:,3]]

no need for a loop.