在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个值之间存在映射关系（这是为了可视化）。 我想将其整合并进行向量化处理，但我不确定如何处理或是否可以。谢谢。

- Mike Axiak

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

- Olivier Verdier

Answer 2

你需要修改simple_funcN函数，使其接受数组作为输入，并创建数组作为输出。之后，你可以查看numpy.meshgrid()或这里的cartesian()函数来构建坐标数组。然后，你应该能够使用坐标数组用一行代码填充A。

- janneb