Функция «If» используется для массива в Python

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)

, но, очевидно, это дает ошибку в сравнении.

Можно ли это сделать без необходимости создания цикла for для X и Y?Потому что, если я хочу вызвать q (x, y) для двух двойных чисел, цикл for должен быть переопределён для этого случая ...

Редактировать: (добавляя это к вопросу какслишком долго для комментария, а не для ответа, но это может помочь другим ответить.)

Похоже, 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

Выражение r2 > 0.04 вычисляется в массив NumPy с логическими значениями, поэтому его нельзя использовать в операторе if. Вы можете автоматически интерпретировать его как числа, хотя:

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

Answer 2

Вы можете использовать map , чтобы применить вашу функцию к каждой паре элементов.