在Python中用于数组的"If"函数

Question

我有一个函数：

def q(x,y):
    r2 = (x/2.)**2 + (2.0*y)**2
    if (r2 > 0.04): 
        return 0.
    else:
        return 1.5

我想用类似下面的方式来调用它：

from pylab import *
dl=0.025
X, Y = mgrid[-0.5:0.5:dl, -0.5:0.5:dl]
g(X,Y)

但显然这样会在比较时出错。

有没有办法不使用循环来处理X和Y呢？因为如果我想用两个双精度浮点数来调用q(x,y)，那就必须为这种情况重新实现循环...

编辑：（把这个加到问题里，因为评论太长了，不算答案，但可能对其他人回答有帮助。）

看起来 pylab.mgrid 和 numpy.mgrid 是一样的。

调整为 numpy 后，这段代码

import numpy
def q(x,y):
    r2 = (x/2.)**2 + (2.0*y)**2
    if (r2 > 0.04): 
        return 0.
    else:
        return 1.5

dl=0.025
X, Y = numpy.mgrid[-0.5:0.5:dl, -0.5:0.5:dl]
q(X,Y)

会出现这个错误

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

Answer 1

你可以使用 map 这个工具，把你的函数应用到每一对元素上。

Answer 2

这个表达式 r2 > 0.04 会得到一个包含布尔值的NumPy数组，也就是说它的结果是“真”或“假”的集合。所以你不能直接把它放在 if 语句里使用。不过，你可以让它自动转换成数字来使用：

def q(x,y):
    r2 = (x/2.)**2 + (2.0*y)**2
    return 1.5 * (r2 <= 0.04)