vectorizing a non-trivial function in numpy

I have a 2d numpy array, with shape (100,3)

a = np.random((100,3))


and I want to call a function fun:

b = fun(a)


This function is currently defined:

@np.vectorize
def fun(r):
if r <= 0.0: return 0.0
else:        return min(2.0/(1.0 + r), 2.0*r/(1.0 + r))


and this works fun. But I would like to express this in a more numpythonic (and therefore more efficient) way. I tried:

def fun(r):
ans = zeros_like(r)
ans[r > 0.0] = min(2.0/(1.0 + r), 2.0*r/(1.0 + r))
return ans


but that don't work as the min is not doing what I want. I also tried:

def fun(r):
return np.piecewise(r
, [r <= 0.0, r > 0.0]
, [0.0, lambda x: min(2.0/(1.0 + x), 2.0*x/(1.0 + x))]
)


but that doesn't work for a similar reason. So how should I do it?

-
There's a difference between min, np.min and np.minimum. In this case, it looks like you want np.minimum. – Christian Clason Feb 7 '13 at 14:33
Thanks for the suggestion, but np.min would find the min of the whole array, which is the incorrect behaviour. np.minimum might be the correct one to use, but it says ValueError: array is not broadcastable to correct shape. – js947 Feb 7 '13 at 18:25
np.minimum(2./(1.+r),2.*r/(1.+r)) works for me; the problem is with the assignment: ans[r>0] returns all entries where r is nonzero as an 1d array. How about (r>0)*np.minimum(...)? – Christian Clason Feb 7 '13 at 22:40

    def f(r):
return 2 * np.clip(r, 0, 1) / (r + 1)

-
clipper, Welcome to SciComp. Usually when answering it is useful to explain your answer. Even if the code is correct, if the asker doesn't understand it, it may not help them. – Godric Seer Feb 8 '13 at 3:07
Seems pretty self-explanatory. – David Ketcheson Feb 8 '13 at 9:57

Thanks to both @clipper and @christianclason, I figured out something that works:

def fun(r):
np.clip(r, 0.0, r, out=r)
return 2.0 * np.minimum(1.0, r) / (1.0 + r)


and my program gets about a 20% reduction in total run time! Unfortunately, the code now looks 20% more rubbish.

-
Define "rubbish". – Christian Clason Feb 8 '13 at 15:47
Your call to np.clip modifies the array r, which is passed by reference; hence, if you call fun(a), the array a will remain clipped afterward in the calling function (and not just inside fun). – Christian Clason Feb 8 '13 at 15:49
more rubbish: less like mathematics, less like fortran – js947 Feb 11 '13 at 13:38