向python函数添加条件答案

【问题标题】：Adding condition to python function向python函数添加条件
【发布时间】：2020-06-07 07:31:26
【问题描述】：

我正在尝试为我的学生制作一个小示例，将原子的电子密度绘制为距中心距离 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).

我想以一种聪明的 Pythonic 方式来完成它，而不仅仅是循环遍历所有 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))

【问题讨论】：

您在寻找numpy.where吗？
这能回答你的问题吗？ Conditional operations on numpy arrays
有很多答案可以解释这种歧义值错误。 x<3 用于数组产生多个真/假值，这样的数组。但是if 仅适用于一个值。 if 分支中没有隐含循环。

标签： python numpy

【解决方案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

我不知道这是不是很奇怪

【讨论】：

【解决方案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

【讨论】：