CODEWITHSUNDEEP
pythonnumpyvectorization
abcd
abcd

Reputation: 10751

Vectorizing a function that finds the nearest value in an array

I've written the following function:

import numpy as np


def _find_nearest(array, value):
    """Find the index in array whose element is nearest to value.

    Parameters
    ----------
    array : np.array
      The array.

    value : number
      The value.

    Returns
    -------
    integer
      The index in array whose element is nearest to value.

    """
    if array.argmax() == array.size - 1 and value > array.max():
        return array.size
    return (np.abs(array - value)).argmin()

I'd like to vectorize this function, so that I can pass several values at once. That is, I'd like to have value be an array, and have _find_nearest return, rather than a single index, the indices for each of the values in the submitted value_array.

Can anyone see a way to do this?

Upvotes: 1

Views: 108

Answers (1)

Kartik
Kartik

Reputation: 8683

Inside the parent function, where both the value and the array are visible, you can use a lambda to enable the vectorization. I shall call the parent function main

def main():
    value = np.random.rand(10, 1)
    array = np.random.rand(100, 100)
    vec_nearest = lambda x: _find_nearest(array, x)
    np.vectorize(vec_nearest)(value)

This will work on one array, and multiple values of vector. The return will be an array.

Upvotes: 1

Related Questions