我想编写一个函数amap
,该功能类似于构建map
,但支持Numpy广播和输出数组。它会这样工作:
>>> amap(lambda x:x**2, 1)
array(1)
>>> amap(lambda x:x**2, [1, 2])
array([1, 4])
>>> amap(lambda x,y:y**2+x**2, 1, [1, 2])
array([2, 5])
>>> amap(lambda x,y:y**2+x**2, [1, 2], [[1], [2]])
array([[2, 5], [5, 8]])
传递的功能在标量上运行。
我写了一个实现。
def amap(func, *args):
'''array version of build-in map
amap(function, scalar/sequence[, ...]) -> array
'''
args = np.broadcast(None, *args)
res = [func(*arg[1:]) for arg in args]
res = np.asarray(res).reshape(args.shape)
return res
它有效,但它在某种程度上丑陋,效率不是很高。 您有好主意吗?尤其是更numpy-stylish的东西?
直接使用numpy.vectorize
?
>>> numpy.vectorize(lambda x:x**2)(1)
array(1)
>>> numpy.vectorize(lambda x:x**2)([1,2])
array([1, 4])
>>> numpy.vectorize(lambda x, y: y**2 + x**2)(1, [1,2])
array([2, 5])
>>> numpy.vectorize(lambda x, y: y**2 + x**2)([1,2], [[1], [2]])
array([[2, 5],
[5, 8]])