CODEWITHSUNDEEP
pythonnumpy
safsoufa
safsoufa

Reputation: 17

Apply function to result of numpy.where

Assume I have an array:

a = [1,2,3,4,0,-1,-2,3,4]

Using np.where(a < 0) it returns a list of indices which elements a are < 0

How to apply a function for those elements of a?

Upvotes: 0

Views: 12722

Answers (4)

MSeifert
MSeifert

Reputation: 152725

In general I would recommend following @ev-br approach using boolean masks. But it's also possible with np.where if you use all three arguments. The second argument specifies the value chosen at the indices where the condition is True and the third where it's False:

>>> import numpy as np
>>> a = np.array([1,2,3,4,0,-1,-2,3,4])
>>> np.where(a < 0, 1000, a)   # replace values below 0 with 1000
array([   1,    2,    3,    4,    0, 1000, 1000,    3,    4])

If you want to apply a numpy-ufunc (for example np.sin) just replace the 1000:

>>> np.where(a < 0, np.sin(a), a)
array([ 1.,  2.,  3.,  4.,  0., -0.84147098, -0.90929743,  3.,  4.])

Alternativly (this requires that the array already has the correct dtype to store the result of the function) you could use the indices returned by np.where to apply the result on:

>>> a = np.array([1,2,3,4,0,-1,-2,3,4], dtype=float) # must be floating point now
>>> idx = np.where(a < 0)
>>> a[idx] = np.sin(a[idx])
>>> a
array([ 1.,  2.,  3.,  4.,  0., -0.84147098, -0.90929743,  3.,  4.])

Upvotes: 2

gold_cy
gold_cy

Reputation: 14226

I don't work with numpy a lot but did you mean something along these lines?

my_arr = np.array(a)

def my_func(my_array):
    for elem in np.where(my_array < 0):
       my_array[elem] =  my_array[elem] + 1 * 3
       return my_array

np.apply_along_axis(my_func, 0, my_arr)
rray([1, 2, 3, 4, 0, 2, 1, 3, 4])

Upvotes: -1

boot-scootin
boot-scootin

Reputation: 12515

Something like this ought to work:

square = lambda x: x**2
applied_func_array = [square(x) for x in a if x < 0]

Or with numpy.vectorize:

vec_square = np.vectorize(square)
vec_square(less_than_zero)

Which yields:

Out[220]: 
array([[1],
       [4]])

Upvotes: -1

ev-br
ev-br

Reputation: 26070

If you convert your list into a numpy array, it gets easier: you can index arrays with boolean arrays:

In [2]: a = np.asarray([1,2,3,4,0,-1,-2,3,4])

In [3]: a[a < 0]
Out[3]: array([-1, -2])

In [4]: np.sin(a[a < 0])
Out[4]: array([-0.84147098, -0.90929743])

In [5]: a[a < 0]**2
Out[5]: array([1, 4])

The key here is that a < 0 is an array itself:

In [6]: a < 0
Out[6]: array([False, False, False, False, False,  True,  True, False, False], dtype=bool)

Upvotes: 2

Related Questions