传入numpy数组作为lambda参数时如何实现内部条件判断

报错原因 #

Python原生三元表达式a if 条件 else b仅支持单值布尔判断：当传入Numpy数组x时，x < np.pi会生成和x长度一致的布尔数组，Python无法直接将整个布尔数组转换为单个布尔值用于逻辑判断，因此抛出歧义真值错误。

解决方案1：使用np.where（性能最优，推荐） #

Numpy内置的np.where是专门针对数组的条件判断函数，语法为np.where(条件, 满足条件时的返回值, 不满足条件时的返回值)，可以完美替代原有三元逻辑，修改后的代码如下：

import numpy as np
import matplotlib.pyplot as plt    

def my_func():
    return lambda x: np.where(x < np.pi, np.sin(x), np.cos(x))
    
x = np.linspace(0, 2*np.pi, 1000)
y = my_func()
plt.plot(x, y(x))
plt.show()

解决方案2：使用np.vectorize做函数向量化（写法更贴近原有逻辑） #

如果你希望保留原生三元表达式的写法，可以用np.vectorize将lambda函数转换为支持数组输入的向量化函数，不过该方法本质是循环遍历数组元素，性能比np.where低，仅适合小数据量场景：

import numpy as np
import matplotlib.pyplot as plt    

def my_func():
    f = lambda x: np.sin(x) if x < np.pi else np.cos(x)
    return np.vectorize(f)
    
x = np.linspace(0, 2*np.pi, 1000)
y = my_func()
plt.plot(x, y(x))
plt.show()

内容的提问来源于stack exchange，提问作者nsk