python - 带有 numpy 数组的条件循环

Question

我正在尝试使用 numpy 数组实现以下简单条件，但输出错误。

dt = 1.0
t = np.arange(0.0, 5.0, dt)
x = np.empty_like(t)
if np.where((t >=0) & (t < 3)):
    x = 2*t
else:
    x=4*t

我得到下面的输出

array([0., 2., 4., 6., 8.])

但我期待

array([0., 2., 4., 12., 16.])

感谢您的帮助!

Answer 1

在 docs 中查找对于 np.where:

Note: When only condition is provided, this function is a shorthand fornp.asarray(condition).nonzero(). Using nonzero directly should bepreferred, as it behaves correctly for subclasses. The rest of thisdocumentation covers only the case where all three arguments areprovided.

由于您没有提供 x 和 y 参数，因此 where 的行为类似于 nonzero。nonzero 返回 np.arrays 的 tuple，转换为 bool 时为真。所以你的代码最终评估为:

if True:
    x = 2*t

相反，您想使用:

x = np.where((t >= 0) & (t < 3), 2*t, 4*t)