Numpy Arrays comparison and indexing

Question

I have 2 arrays of unequal size:

>>> np.size(array1)
4004001
>>> np.size(array2)
1000 

Now, each element in array2 needs to be compared to all the elements in array1, to find the element which has the nearest value to that of this element in array2. Upon finding this value, I need to store it in a different array of size 1000 - one of a size corresponding to array2.

The tedious and crude way of doing it could be using a for loop and taking each element from Array 2, subtracting its absolute value from array 1 elements and then taking the minimum value- this is going to make my code really slow.

I'd like to use numpy vectorized operations to do this but i've kind of hit a wall.

Answer 1

To make full use of the numpy parallelism we need vectorized functions. Further all values are found in the same array (array1) using the same criterium (nearest). Therefore, it is possible to make a special function for searching in array1 specifically.

However, to make the solution more reusable it is better to make a more general solution and then transform it into a more specific one. Thus, as a general approach to find the closest value, we start with this find nearest solution. Then we turn that into a more specific and vectorize it, to allow it to work on multiple element at once:

import math
import numpy as np
from functools import partial

def find_nearest_sorted(array,value):
    idx = np.searchsorted(array, value, side="left")
    if idx > 0 and (idx == len(array) or math.fabs(value - array[idx-1]) < math.fabs(value - array[idx])):
        return array[idx-1]
    else:
        return array[idx]

array1 = np.random.rand(4004001)
array2 = np.random.rand(1000)

array1_sorted = np.sort(array1)

# Partially apply array1 to find function, to turn the general function
# into a specific, working with array1 only.
find_nearest_in_array1 = partial(find_nearest_sorted, array1_sorted)

# Vectorize specific function to allow us to apply it to all elements of
# array2, the numpy way.
vectorized_find = np.vectorize(find_nearest_in_array1)

output = vectorized_find(array2)

Hopefully this is what you wanted, a new vector, mapping the data in array2 to the nearest values in array1.

Answer 2

import numpy as np
a = np.random.random(size=4004001).astype(np.float16)
b = np.random.random(size=1000).astype(np.float16)
#use numpy broadcasting to compare pairwise difference and then find the min arg in a for each element in b. Finally extract elements from a using the argmin array as indexes. 
output = a[np.argmin(np.abs(b[:,None] -a),axis=1)]

This solution while simple can be very memory intensive. It may need a bit further optimisation if using it on large arrays.

Answer 3

The most "numpythonic" way is is to use broadcasting. This is a quick and easy way to calculate a distance matrix, for which you can then take the argmin of the absolute value.

array1 = np.random.rand(4004001)
array2 = np.random.rand(1000)

# Calculate distance matrix (on truncated array1 for memory reasons)
dmat = array1[:400400] - array2[:,None]

# Take the abs of the distance matrix and work out the argmin along the  last axis
ix = np.abs(dmat).argmin(axis=1)

shape of dmat:

(1000, 400400)

shape of ix and contents:

(1000,)    
array([237473, 166831,  72369,  11663,  22998,  85179, 231702, 322752, ...])

However, it's memory hungry if you do this operation in one go, and actually doesn't work on my 8GB machine for the size of arrays that you specify, which is why I reduced the size of array1.

To make it work within memory constraints, simply slice one of the arrays into chunks and apply broadcasting on each chunk in turn (or parallelise). In this case, I've sliced array2 into 10 chunks:

# Define number of chunks and calculate chunk size
n_chunks = 10
chunk_len = array2.size // n_chunks

# Preallocate output array
out = np.zeros(1000)

for i in range(n_chunks):
    s = slice(i*chunk_len, (i+1)*chunk_len)
    out[s] = np.abs(array1 - array2[s, None]).argmin(axis=1)