Adding condition to python function

Question

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.

The density is given by:

rho = exp(-2*Z*r) for r < r_cut

and

rho = exp(-2*sqrt(2*I)*r (for a specified value of r_cut, I and Z).

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:

----> 5 if r < r_cut: 6 return exp(-2*z*r) 7 else:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Any suggestions? (I have tried to search but could not find a good answer. If there is already one please direct me.)

r = linspace(0, 3., 100)

kjmol = 3.8088e10-4
def rho_r(r, r_cut, z, I):
    if r < 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))

Answer 1

I suppose this should do the trick for you:

import numpy as np
r = np.linspace(0, 3., 100)

kjmol = 3.8088e10-4
def rho_r(r, r_cut, z, I):
    res=np.exp(-2*np.sqrt(2*I)*r)
    res[r < r_cut]=np.exp(-2*z*r[r<r_cut])
    return res

Answer 2

I did this...

def rho_r(r, r_cut, z, I):
    rho = where((r < r_cut), exp(-2*z*r), exp(-2*sqrt(2*I)*r) )
    return rho

I wonder if this is to weird