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) 这对数字和那个值的三个数值之间有一种对应关系（这是为了方便理解）。我想把这些数据整合起来，并且想办法让它变得更高效，但我不太确定该怎么做或者是否能做到。谢谢。

Answer 1

你需要把simple_funcN这个函数改成可以接收数组作为输入，并且返回数组作为输出。之后，你可以看看numpy.meshgrid()这个函数，或者这里的cartesian()函数，用来生成坐标数组。接下来，你就可以用这些坐标数组，用一行代码来填充A了。

Answer 2

这是一个向量化的版本：

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