title: How to vectorize a function which contains an if statement?
body: <p>Let's say we have the following function:</p>
  
  <pre><code>def f(x, y):
      if y == 0:
          return 0
      return x/y
  </code></pre>
  
  <p>This works fine with scalar values. Unfortunately when I try to use numpy arrays for <code>x</code> and <code>y</code> the comparison <code>y == 0</code> is treated as an array operation which results in an error:</p>
  
  <pre><code>---------------------------------------------------------------------------
  ValueError                                Traceback (most recent call last)
  &lt;ipython-input-13-9884e2c3d1cd&gt; in &lt;module&gt;()
  ----&gt; 1 f(np.arange(1,10), np.arange(10,20))
  
  &lt;ipython-input-10-fbd24f17ea07&gt; in f(x, y)
        1 def f(x, y):
  ----&gt; 2     if y == 0:
        3         return 0
        4     return x/y
  
  ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
  </code></pre>
  
  <p><strike>I tried to use <code>np.vectorize</code> but it doesn't make a difference, the code still fails with the same error.</strike> <code>np.vectorize</code> is one option which gives the result I expect.</p>
  
  <p>The only solution that I can think of is to use <code>np.where</code> on the <code>y</code> array with something like:</p>
  
  <pre><code>def f(x, y):
      np.where(y == 0, 0, x/y)
  </code></pre>
  
  <p>which doesn't work for scalars.</p>
  
  <p>Is there a better way to write a function which contains an if statement? It should work with both scalars and arrays.</p>