Numpy：编写适用于一维和二维数组的函数的好方法是什么

Question

我想编写一个函数 f(x)，它执行以下操作：返回 g(x)，如果 x 是一维 numpy 数组如果 x 是二维 numpy 数组，则返回 np.array([g(xi) for xi in x])是否有特定于 numpy 的函数可以在一行中执行此操作，而无需编写 if 语句？

Answer 1

从一维向量生成标量的函数 g(x) 可以扩展到任意更高的维度，如下所示：

import numpy as np

def myfunc(x):

    return sum(x)

def f( g, x ):

    if len(x.shape) == 1:

        return g(x)

    if len(x.shape) > 1:

        return np.array( [f(g,v) for v in x] )

# Test with one dimensional input

res = f( myfunc, np.array( [0.,1.,2.] ) )

print( res )

# Test with two dimensional input

res = f( myfunc, np.array( [[0.,1.,2.],[3.,4.,5.]] ) )

print( res )

# And, still more dimensions

res = f( myfunc, np.ones( (3,2,2) ) )

print( res )

产生，

3.0

[  3.  12.]

[[ 2.  2.]

 [ 2.  2.]

 [ 2.  2.]]