title: Adding condition to python function
body: <p>I'm trying to make a small example for my students plotting the electron density of an atom as a function of distance from center, r, nuclear charge, Z, and Ionization energy, I.</p>
  
  <p>The density is given by:</p>
  
  <pre><code>rho = exp(-2*Z*r) for r &lt; r_cut
  </code></pre>
  
  <p>and</p>
  
  <pre><code>rho = exp(-2*sqrt(2*I)*r (for a specified value of r_cut, I and Z).
  </code></pre>
  
  <p>I wanted to do it in a smart Pythonic way and not just loop over all r:s. I tried to do according to the code in the bottom but received an error: </p>
  
  <blockquote>
    <p>----> 5     if r &lt; r_cut:
         6         return exp(-2*z*r)
         7     else:</p>
    
    <p>ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()</p>
  </blockquote>
  
  <p>Any suggestions? 
  (I have tried to search but could not find a good answer. If there is already one please direct me.)</p>
  
  <pre><code>r = linspace(0, 3., 100)
  
  kjmol = 3.8088e10-4
  def rho_r(r, r_cut, z, I):
      if r &lt; r_cut:
          return exp(-2*z*r)
      else:
          return exp(-2*sqrt(2*I)*r)
  
  plot(r, rho_r(r, r_cut=3., z=2., I=2372*kjmol))
  </code></pre>