How to write conditional code that's compatible with both plain Python values and NumPy arrays?
<p>For writing âpiecewise functionsâ in Python, I'd normally use <code>if</code> (in either the control-flow or ternary-operator form).</p>

<pre><code>def spam(x):
    return x+1 if x&gt;=0 else 1/(1-x)
</code></pre>

<p>Now, with NumPy, the mantra is to avoid working on single values in favour of vectorisation, for performance. So I reckon something like this would be preferred:<sub><em>As Leon remarks, the following is <strong>wrong</strong></em></sub></p>

<pre><code>def eggs(x):
    y = np.zeros_like(x)
    positive = x&gt;=0
    y[positive] = x+1
    y[np.logical_not(positive)] = 1/(1-x)
    return y
</code></pre>

<p>(Correct me if I've missed something here, because frankly I find this very ugly.)</p>

<p>Now, of course <code>eggs</code> will <em>only</em> work if <code>x</code> is actually a NumPy array, because otherwise <code>x&gt;=0</code> simply yields a single boolean, which can't be used for indexing (at least doesn't do the right thing).</p>

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

<p>I would use <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.asarray.html" rel="nofollow">numpy.asarray</a> (which is a no-op if the argument is already an numpy array) if I want to handle both numbers and numpy arrays</p>

<pre><code>def eggs(x):
    x = np.asfarray(x)
    m = x&gt;=0
    x[m] = x[m] + 1
    x[~m] = 1 / (1 - x[~m])
    return x
</code></pre>

<p>(here I used <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.asfarray.html" rel="nofollow">asfarray</a> to enforce a floating-point type, since your function requires floating-point computations).</p>

<p>This is less efficient than your spam function for single inputs, and arguably uglier. However it seems to be the easiest choice.</p>

<p>EDIT: If you want to ensure that <code>x</code> is not modified (as pointed out by Leon) you can replace <code>np.asfarray(x)</code> by <code>np.array(x, dtype=np.float64)</code>, the array constructor copies by default.</p>

python arrays numpy vectorization