关于 Numpy 数组作为函数参数的一个小问题

注1：本文同步也发表在我的独立博客中。

注2：本文的所写的只是我自己想出来的一个办法，很可能不是最优或者最标准的做法，如果有高手看到，欢迎指正！

在计算过程中遇到这样的一个问题，需要对一个 nx*ny 的网格上每一个节点调用一次函数，然后返回一个 nx*ny 的数组。

我之前就写过类似的程序，numpy 可以直接调用普通的计算程序，比如：

def x_square(x):
    return x*x
a = np.arange(nx*ny).reshape(nx,ny)
print x_square(a)

直接把 ndarray 对象作为参数传入函数，就可以得到一个正确的结果。

于是，想当然地，也针对上面的需求，写了这样一个程序：

ff = np.linspace(1.0,2.0,nx*ny).reshape(nx,ny)
def test(x, y):
    x=np.int_(x); y=np.int_(y) # to be used as index, x & y should be integer
    return ff[x,y]
print test(np.arange(nx), np.arange(ny))

运行上面的程序的时候，就出现了 "shape mismatch: objects cannot be broadcast to a single shape" 的错误，百思不得其解。在 stackoverflow上提问，总算弄明白了这个问题。原来 ndarray 对象作为参数传入函数的时候，如果同时传入多个参数，这些数组可以被当成普通的数来计算，但前提是这些数组的维数都相同，或者是0维的数（也其实相当于任意维），如果维数不同是无法处理的。

至于上面这个问题，我想到的一个有点丑陋的做法是用 fromfunction 函数，再新建一个中间函数：

def tt(x0, x1, y0, y1):
    dx = x1-x0; dy = y1-y0
    return np.fromfunction(lambda a,b: test(a+x0,b+y0), (dx, dy))

这样，调用函数的格式有所不同，但毕竟目的是达到了。

不知是否还有更简单直接的办法，现在的做法总觉得有些 dirty。