CODEWITHSUNDEEP
pythonarraysnumpy
numberCruncher
numberCruncher

Reputation: 645

Making for loop with index arrays faster

I have the following problem: I have index arrays with repeating indices and would like to add values to an array like this:

grid_array[xidx[:],yidx[:],zidx[:]] += data[:]

However, as I have repeated indices this does not work as it should because numpy will create a temporary array which results in the data for the repeated indices being assigned several times instead of being added to each other (see http://docs.scipy.org/doc/numpy/user/basics.indexing.html).

A for loop like

for i in range(0,n):
    grid_array[xidx[i],yidx[i],zidx[i]] += data[i]

will be way to slow. Is there a way I can still use the vectorization of numpy? Or is there another way to make this assignment faster?

Thanks for your help

Upvotes: 3

Views: 182

Answers (4)

Bi Rico
Bi Rico

Reputation: 25813

How about using bincount?

import numpy as np

flat_index = np.ravel_multi_index([xidx, yidx, zidx], grid_array.shape)
datasum = np.bincount(flat_index, data, minlength=grid_array.size)
grid_array += datasum.reshape(grid_array.shape)

Upvotes: 3

hpaulj
hpaulj

Reputation: 231335

This is a buffering issue. The .at provides unbuffered action http://docs.scipy.org/doc/numpy/reference/generated/numpy.ufunc.at.html#numpy.ufunc.at

np.add.at(grid_array, (xidx,yidx,zidx),data)

Upvotes: 2

numberCruncher
numberCruncher

Reputation: 645

I think I found a possible solution:

def assign(xidx,yidx,zidx,data):
    grid_array[xidx,yidx,zidx] += data
    return

map(assign,xidx,yidx,zidx,sn.part0['mass'])

Upvotes: -1

Kasravnd
Kasravnd

Reputation: 107287

For add an array to elements of a nested array you just can do grid_array[::]+=data :

>>> grid_array=np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> data=np.array([3,3,3])
>>> grid_array[::]+=data
>>> grid_array
array([[ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

Upvotes: 0

Related Questions