How to write conditional code compatible with both regular Python values ​​and NumPy arrays?

Question

To write "piecewise functions" in Python, I would usually use if (either in the form of a control flow or in the form of a ternary operator).

 def spam(x): return x+1 if x>=0 else 1/(1-x) 

Now, with NumPy, the mantra should avoid working with single values ​​in favor of vectorization for performance. Therefore, I believe that this would be preferable: _{As Leon points out, the following is not true.}

 def eggs(x): y = np.zeros_like(x) positive = x>=0 y[positive] = x+1 y[np.logical_not(positive)] = 1/(1-x) return y 

(Correct me if I missed something because, frankly, I find this very ugly.)

Now, of course, eggs will only work if x is actually a NumPy array, because otherwise x>=0 just prints one boolean that cannot be used for indexing (at least it doesn't do it right).

Is there a good way to write code that looks more like spam but works idiomatically on Numpy arrays, or should I just use vectorize(spam) ?

Answer 1

Use np.where . However, you will get the array as an output even for input with an open number.

 def eggs(x): y = np.asarray(x) return np.where(y>=0, y+1, 1/(1-y)) 

This works for both arrays and primes:

 >>> eggs(5) array(6.0) >>> eggs(-3) array(0.25) >>> eggs(np.arange(-3, 3)) /home/praveen/.virtualenvs/numpy3-mkl/bin/ipython3:2: RuntimeWarning: divide by zero encountered in true_divide array([ 0.25 , 0.33333333, 0.5 , 1. , 2. , 3. ]) >>> eggs(1) /home/praveen/.virtualenvs/numpy3-mkl/bin/ipython3:3: RuntimeWarning: divide by zero encountered in long_scalars # -*- coding: utf-8 -*- array(2.0) 

As ayhan notes, this raises a warning since 1/(1-x) is evaluated for the entire range. But a warning is simple: a warning. If you know what you are doing, you can ignore the warning. In this case, you select 1/(1-x) from the indices where it can never be inf , so you are safe.

Answer 2

I would use numpy.asarray (this is non-op if the argument is already a numpy array) if I want to handle both numbers and numpy arrays

 def eggs(x): x = np.asfarray(x) m = x>=0 x[m] = x[m] + 1 x[~m] = 1 / (1 - x[~m]) return x 

(here I used asfarray to force input of a floating point type, since your function requires floating point calculations).

This is less efficient than the spam feature for individual inputs, and perhaps uglier. However, this is the easiest choice.

EDIT: If you want x not to be changed (as Leon noted), you can replace np.asfarray(x) with np.array(x, dtype=np.float64) , the array constructor copies by default.