CODEWITHSUNDEEP
pythonpython-3.xnumpy
Igor Rivin
Igor Rivin

Reputation: 4864

numpy Broadcasting for user functions

In numpy, if a is an ndarray, then, something like np.sin(a) takes sin of all the entries of ndarray. What if I need to define my own function (for a stupid example, f(x) = sin(x) if x<1 else cos(x)) with broadcasting behavior?

Upvotes: 1

Views: 268

Answers (2)

sacuL
sacuL

Reputation: 51395

Use np.where:

np.where(a<1,np.cos(a), np.sin(a))

Example:

a = [-1,1,2,-2]

>>> np.where(a<1,np.cos(a), np.sin(a))
array([-0.84147098,  0.84147098,  0.90929743, -0.90929743])

If you have more than one conditions, use np.select

Upvotes: 2

Hans Musgrave
Hans Musgrave

Reputation: 7131

You could define your own function f = lambda x: sin(x) if x<1 else cos(x) and then use numpy's builtin vectorizer f_broadcasting = np.vectorize(f).

This doesn't offer any speed improvements (and the additional overhead can slow down small problems), but it gives you the desired broadcasting behavior.

Upvotes: 2

Related Questions