如何创建一个接受numpy数组,可迭代数组或标量的numpy函数？

Question

假设我有这个:

def incrementElements(x):
   return x+1

但我想修改它,以便它可以采用numpy数组,可迭代或标量,并将参数提升为numpy数组并为每个元素添加1.

我怎么能这样做？我想我可以测试参数类,但这似乎是一个坏主意.如果我这样做:

def incrementElements(x):
   return numpy.array(x)+1

它适用于数组或迭代,但不适用于标量.这里的问题是numpy.array(x)标量x产生一些奇怪的对象,它由一个numpy数组包含但不是一个"真正的"数组; 如果我添加一个标量,结果将降级为标量.

Answer 1

你可以试试

def incrementElements(x):
    x = np.asarray(x)
    return x+1

np.asarray(x)是等价的np.array(x, copy=False),意味着标量或可迭代将转换为a ndarray,但如果x已经是a ndarray,则不会复制其数据.

如果你传递一个标量并想要一个ndarray输出(不是标量),你可以使用:

def incrementElements(x):
    x = np.array(x, copy=False, ndmin=1)
    return x

该ndmin=1参数将强制数组至少具有一个维度.使用ndmin=2至少2个维度,等等.您也可以使用它的等价物np.atleast_1d(或np.atleast_2d2D版本......)