vectorizing a non-trivial function in numpy

Question

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?

Answer 1

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

Answer 2

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.