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?
-
You may be introducing an XY problem with your "stupid" example. Perhaps if you say what you're actually trying to do we could vectorze better?Daniel F– Daniel F2018-09-21 07:07:13 +00:00Commented Sep 21, 2018 at 7:07
-
@DanielF What is an "XY problem"?Igor Rivin– Igor Rivin2018-09-21 13:00:40 +00:00Commented Sep 21, 2018 at 13:00
Add a comment
|
2 Answers
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.
1 Comment
Igor Rivin
This answers my actual question (though the use I had in mind was something like the other answer addresses...), thanks! I was afraid of the performance problem...
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
5 Comments
Igor Rivin
I am a bit puzzled by the behavior of this. If you set
a = np.array([[1, 0], [3, 4]]) and then np.where(a==0, 0, 1/a), you get a division by zero warning (the final result is correct), which makes one think that it actually does 2x the work needed. Is this correct?
sacuL
Interesting to note! I guess it does perform the operation on the whole array, but only returns it where the condition is met (or not met)... I don't know the inner workings, but the function is often efficient enough for me.
Eric
@Igor: Yes, that's correct. If it matters, you can use
cond = a > 1, ret = np.cos(a, where=cond), ret = np.sin(a, where=~cond, out=ret), but you might find it's slowerIgor Rivin
@Eric, thanks! I guess this is the price we pay for the "Python model" (where Python is the glue for underlying monolythic C++ engine). As for speed, since in what I am doing, at some point one computes the eigenvalues and eigenvectors of the array, the slowdown is irrelevant (though amusing).
Eric
Note that CPU-level vectorization can sometimes mean it's faster to compute two contiguous blocks of memory and use
np.where to pick between them than it is to compute half of each block.