Добавление условия к функции python

Question

Я пытаюсь сделать небольшой пример для своих учеников, показывающих электронную плотность атома как функцию расстояния от центра, r, заряда ядра, Z и энергии ионизации, I.

плотность задается следующим образом:

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

и

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

Я хотел сделать это умным способом Pythoni c, а не просто l oop по всем r: s. Я пытался сделать в соответствии с кодом внизу, но получил ошибку:

----> 5, если r

ValueError: Значение истинности массива с более чем одним элементом неоднозначно. Используйте a.any () или a.all ()

Есть предложения? (Я пытался найти, но не смог найти хороший ответ. Если он уже есть, напишите мне.)

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

Я сделал это ...

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

Интересно, странно ли это

Answer 2

Полагаю, это поможет вам:

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