python - 在 python 中对数组进行快速数学运算

Question

我想对数组执行一个相当简单的数学运算。让我写出这个例子:

A = numpy.ndarray((255, 255, 3), dtype=numpy.single)
# ..
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        x = simple_func1(i)
        y = simple_func2(j)
        A[i, j] = (alpha * x * y + beta * x**2 + gamma * y**2, 1, 0)

所以基本上，在 (i, j) 和该值的 3 个值之间存在映射(这是为了可视化)。我想把它卷起来并以某种方式对其进行矢量化，但我不确定如何或是否可以。谢谢。

Answer 1

这是矢量化版本:

i = arange(255)
j = arange(255)
x = simple_func1(i)
y = simple_func2(j)
y = y.reshape(-1,1)    

A = alpha * x * y + beta * x**2 + gamma * y**2 # broadcasting is your friend here

如果要用1和0填充最后一个坐标:

B = empty(A.shape+(3,))
B[:,:,0] = A
B[:,:,1] = 1 # broadcasting again
B[:,:,2] = 0