Grouping elements of a NumPy array by sum of indices

user_13104702 0
Viewed 200

I have several large numpy array of dimensions 30*30*30, on which I need to traverse the array, get the sum of each index triplet and bin these elements by this sum. For example, consider this simple 2*2 array:

test = np.array([[2,3],[0,1]])

This array has the indices [0,0],[0,1],[1,0] and [1,1]. This routine would return the list: [2,[3,0],1], because 2 in array test has index sum 0, 3 and 0 have index sum 1 and 1 has index sum 2. I know the brute force method of iterating through the NumPy array and checking the sum would work, but it is far too inefficient for my actual case with large N(=30) and several arrays. Any inputs on using NumPy routines to accomplish this grouping would be appreciated. Thank you in advance.

numpy arrays python
2 Answers

Here is one way that should be reasonably fast, but not super-fast: 30x30x30 takes 20 ms on my machine.

import numpy as np

# make example
dims = 2,3,4
a = np.arange(np.prod(dims),0,-1).reshape(dims)

# create and sort indices
idx = sum(np.ogrid[tuple(map(slice,dims))])
srt = idx.ravel().argsort(kind='stable')

# use order to arrange and split data
asrt = a.ravel()[srt]
spltpts = idx.ravel().searchsorted(np.arange(1,np.sum(dims)-len(dims)+1),sorter=srt)
out = np.split(asrt,spltpts)

# admire
out
# [array([24]), array([23, 20, 12]), array([22, 19, 16, 11,  8]), array([21, 18, 15, 10,  7,  4]), array([17, 14,  9,  6,  3]), array([13,  5,  2]), array([1])]
user_72073920

You could procedural create a list of index tuplets and use that, but may be getting into a code constant that's too large to be efficient. [(0,0),[(1,0),(0,1)],(1,1)],

So you need a function to generate these indexes on the fly for an n-demensional array.

For one dimension, a trivial count/increment 

   [(0),(1),(2),...]

The the second, use the one dimension strategy for the fist dimension, the decrement the first and increment the second to fill in. 

   [(0...)...,(1...)...,(2...)...,...] 
   [[(0,0)],[(1,0),(0,1)],[(2,0),(1,1),(0,2)],[...],...]

Notice some of these would be outside the example array, Your generator would need to include a bounds check.

Then three dimensions, give the first two demensions the treatment as above, but at the end, decrement the first dimension, increment the third, repeat until done

[[(0,0,0),...],[(1,0,0),(0,1,0),...],[(2,0,0),(1,1,0),(0,2,0),...],[...],...]
[[(0,0,0)],[(1,0,0),(0,1,0),(0,0,1)],[(2,0,0),(1,1,0),(0,2,0),(1,0,1),(0,1,1)(0,0,2)

Again need bounds checks or cleverer starting/end points to avoid trying to access outside the index, but this general algorithm is how you'd go about generating the indexes on the fly rather than having two large arrays compete for cache and i/o.

Generating the python or nympy equivalent is left as an exercise to the user.

user_23378900
Related