CODEWITHSUNDEEP
pythonarraysnumpy
leviathan
leviathan

Reputation: 363

NumPy nearest value along axis of multidimensional array

I'd like to create a function, that returns the nearest value in the array along a specified axis to a given value.

To get the index of the nearest value I use the following code where arr is a multidimensional array and value is the value to look for:

def nearest_index( arr, value, axis=None ):
    return ( np.abs( arr - value ) ).argmin( axis=axis )

But I struggle with using the result of this function to get the values from the array.

It is easy with 1D-arrays:

In [14]: arr_1 = np.random.randint( 10, 100, size=( 10, ) )

In [15]: arr_1
Out[15]: array([67, 49, 90, 29, 60, 80, 31, 55, 29, 10])

In [16]: nearest_index( arr_1, 50 )
Out[16]: 1

In [17]: arr_1[nearest_index( arr_1, 50 )]
Out[17]: 49

or with flattened arrays:

In [25]: arr_3 = np.random.randint( 10, 100, size=( 2, 3, 4, ) )

In [26]: arr_3
Out[26]:
array([[[85, 51, 74, 79],
        [63, 42, 27, 75],
        [89, 68, 80, 63]],

       [[85, 72, 74, 16],
        [85, 22, 47, 78],
        [44, 70, 98, 34]]])

In [27]: idx_flat = nearest_index( arr_3, 50, axis=None )

In [28]: idx_flat
Out[28]: 1

In [29]: idx = np.unravel_index( idx_flat, arr_3.shape )

In [30]: idx
Out[30]: (0, 0, 1)

In [31]: arr_3[idx]
Out[31]: 51

How can I create a function, that returns the values along the defined axis?

I tried the solution for this Question, but I only got it working for axis=-1.

Note that it is not an issue to me, if only the first occurance of the result is found if multiple elements in the array are equally near the expected value.

Upvotes: 1

Views: 1053

Answers (2)

Joonatan Samuel
Joonatan Samuel

Reputation: 651

Well it works for me.

def nearest_index(arr, value, axis=None):
    return np.argmin(np.abs( arr - value ), axis=axis)

>>> X
array([[76, 94, 56, 93, 28,  0, 44, 50, 89, 93],
       [80, 99, 29, 98, 39, 27, 55, 70, 19, 76],
       [87,  7, 28, 78, 47, 95, 34, 97, 66, 27],
       [75, 78, 82, 30, 15,  0,  2, 25, 58, 69],
       [31,  2, 34,  1, 56,  7, 87, 78, 32, 77],
       [89, 80, 76, 97, 49, 18, 62, 35, 94, 41],
       [ 2, 44, 83,  3, 64,  4, 49, 93, 46,  8],
       [51, 63, 45, 57, 77, 90, 93,  4, 26, 81],
       [43, 92, 22, 98, 93, 36, 46, 25, 35, 36],
       [30, 14, 42, 91, 86, 14, 78,  9, 37, 19]])

>>> X[nearest_index(X, 2, axis=0), np.arange(10)]
array([ 2,  2, 22,  1, 15,  0,  2,  4, 19,  8])

>>> X[np.arange(10), nearest_index(X, 2, axis=1)]
array([ 0, 19,  7,  2,  2, 18,  2,  4, 22,  9])

Upvotes: 0

Divakar
Divakar

Reputation: 221504

For a multi-dimensional array, we need to use advanced-indexing. So, for a generic n-dim array and with a specified axis, we could do something like this -

def argmin_values_along_axis(arr, value, axis):   
    argmin_idx = np.abs(arr - value).argmin(axis=axis)
    shp = arr.shape
    indx = list(np.ix_(*[np.arange(i) for i in shp]))
    indx[axis] = np.expand_dims(argmin_idx, axis=axis)
    return np.squeeze(arr[indx])

Sample runs -

In [203]: arr_3 = np.random.randint( 10, 100, size=( 2, 3, 4, ) )

In [204]: arr_3
Out[204]: 
array([[[94, 55, 26, 51],
        [82, 66, 80, 66],
        [96, 54, 93, 57]],

       [[59, 28, 95, 56],
        [47, 48, 17, 77],
        [15, 57, 57, 25]]])

In [205]: argmin_values_along_axis(arr_3, value=50, axis=0)
Out[205]: 
array([[59, 55, 26, 51],
       [47, 48, 80, 66],
       [15, 54, 57, 57]])

In [206]: argmin_values_along_axis(arr_3, value=50, axis=1)
Out[206]: 
array([[82, 54, 26, 51],
       [47, 48, 57, 56]])

In [207]: argmin_values_along_axis(arr_3, value=50, axis=2)
Out[207]: 
array([[51, 66, 54],
       [56, 48, 57]])

Upvotes: 2

Related Questions