python - 如何制作一个接受 numpy 数组、可迭代对象或标量的 numpy 函数？

Question

假设我有这个:

def incrementElements(x):
   return x+1

但我想修改它，以便它可以采用 numpy 数组、可迭代对象或标量，并将参数提升为 numpy 数组并将每个元素加 1。

我该怎么做？我想我可以测试参数类，但这似乎是个坏主意。如果我这样做:

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

它适用于数组或可迭代对象，但不适用于标量。这里的问题是标量 x 的 numpy.array(x) 产生一些奇怪的对象，它包含在 numpy 数组中，但不是“真正的”数组；如果我向它添加一个标量，结果将降级为一个标量。

Answer 1

你可以试试

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

np.asarray(x) 等同于 np.array(x, copy=False)，这意味着标量或可迭代对象将被转换为ndarray，但是如果x已经是一个ndarray，它的数据将不会被复制。

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

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

ndmin=1 参数将强制数组至少有一个维度。至少对 2 个维度使用 ndmin=2，依此类推。您还可以使用其等效的 np.atleast_1d(或 np.atleast_2d 用于 2D 版本...)