Creating index array in numpy - eliminating double for loop

Question

I have some physical simulation code, written in python and using numpy/scipy. Profiling the code shows that 38% of the CPU time is spent in a single doubly nested for loop - this seems excessive, so I've been trying to cut it down.

The goal of the loop is to create an array of indices, showing which elements of a 1D array the elements of a 2D array are equal to.

indices[i,j] = where(1D_array == 2D_array[i,j])

As an example, if 1D_array = [7.2, 2.5, 3.9] and

2D_array = [[7.2, 2.5] 
            [3.9, 7.2]]

We should have

indices = [[0, 1]
           [2, 0]]

I currently have this implemented as

for i in range(ni):
    for j in range(nj):
        out[i, j] = (1D_array - 2D_array[i, j]).argmin()

The argmin is needed as I'm dealing with floating point numbers, and so the equality is not necessarily exact. I know that every number in the 1D array is unique, and that every element in the 2D array has a match, so this approach gives the correct result.

Is there any way of eliminating the double for loop?

Note:

I need the index array to perform the following operation:

f = complex_function(1D_array)
output = f[indices]

This is faster than the alternative, as the 2D array has a size of NxN compared with 1xN for the 1D array, and the 2D array has many repeated values. If anyone can suggest a different way of arriving at the same output without going through an index array, that could also be a solution

@AshwiniChaudhary, no it's not. In fact, it'll never be. I'll edit the example to remove that. — Sten, Jan 16 '15 at 13:15
For this to work, I take it that the entries in 1D_array don't repeat. Why don't you make a dictionary out of 1D_array, with the values as keys and the indexes as values? I.e. `{0:7.2, 1:2.5, 2:3.9}` Then you'd just have to apply the dict to the array. — Roberto, Jan 16 '15 at 13:32

score 2 · Answer 1 · answered Jan 16 '15 at 13:30

In pure Python you can do this using a dictionary in O(N) time, the only time penalty is going to be the Python loop involved:

>>> arr1 = np.array([7.2, 2.5, 3.9])
>>> arr2 = np.array([[7.2, 2.5], [3.9, 7.2]])
>>> indices = dict(np.hstack((arr1[:, None], np.arange(3)[:, None])))
>>> np.fromiter((indices[item] for item in arr2.ravel()), dtype=arr2.dtype).reshape(arr2.shape)
array([[ 0.,  1.],
       [ 2.,  0.]])

Alex Riley · Answer 2 · 2015-01-16T16:39:19.020

To get rid of the two Python for loops, you can do all of the equality comparisons "in one go" by adding new axes to the arrays (making them broadcastable with each other).

Bear in mind that this produces a new array containing len(arr1)*len(arr2) values. If this is a very big number, this approach could be infeasible depending on the limitations of your memory. Otherwise, it should be reasonably quick:

>>> (arr1[:,np.newaxis] == arr2[:,np.newaxis]).argmax(axis=1)
array([[0, 1],
       [2, 0]], dtype=int32)

If you need to get the index of the closest matching value in arr1 instead, use:

np.abs(arr1[:,np.newaxis] - arr2[:,np.newaxis]).argmin(axis=1)

score 1 · Accepted Answer · edited May 23 '17 at 11:54

The dictionary method that some others have suggest might work, but it requires that you know ahead of time that every element in your target array (the 2d array) has an exact match in your search array (your 1d array). Even when this should be true in principle, you still have to deal with floating point precision issues, for example try this .1 * 3 == .3.

Another approach is to use numpy's searchsorted function. searchsorted takes a sorted 1d search array and any traget array then finds the closest elements in the search array for every item in the target array. I've adapted this answer for your situation, take a look at it for a description of how the find_closest function works.

import numpy as np

def find_closest(A, target):
    order = A.argsort()
    A = A[order]

    idx = A.searchsorted(target)
    idx = np.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= target - left < right - target
    return order[idx]

array1d = np.array([7.2, 2.5, 3.9])
array2d = np.array([[7.2, 2.5],
                    [3.9, 7.2]])

indices = find_closest(array1d, array2d)
print(indices)
# [[0 1]
#  [2 0]]

Creating index array in numpy - eliminating double for loop

3 Answers3