接受 numpy ndarray 或序列作为参数的 Python 函数

Question

我见过一些 python 函数，它们通常接收 (n,2) 形状的 numpy ndarray 作为参数，但也可以“自动”接收 (2,n) 甚至 len(2) 序列（元组或列表）。

它是如何通过python实现的？是否有统一的良好实践来检查和处理这些情况（例如，numpy 和 scipy 模块中的函数），或者每个开发人员实现他认为最好的？

我只是想避免 ifs/elifs 的链（可能是嵌套的）链，以防有众所周知的更好的方法。

感谢您的帮助。

Answer 1

您可以使用numpy.asarray 函数将任何类似序列的输入转换为数组：

>>> import numpy
>>> numpy.asarray([1,2,3])
array([1, 2, 3])
>>> numpy.asarray(numpy.array([2,3]))
array([2, 3])
>>> numpy.asarray(1)
array(1)
>>> numpy.asarray((2,3))
array([2, 3])
>>> numpy.asarray({1:3,2:4})
array({1: 3, 2: 4}, dtype=object)

请务必注意，正如文档所述，如果输入已经是 ndarray，则不会执行复制。这真的很好，因为你可以传入一个现有的数组，它只返回相同的数组。

将其转换为 numpy 数组后，如果需要，只需检查长度即可。比如：

>>> def f(x):
...    x = numpy.asarray(x)
...    if len(x) != 2:
...       raise Exception("invalid argument")
... 
>>> f([1,2])
>>> f([1,2,3])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 4, in f
Exception: invalid argument

更新： 既然你问了，这里有一个“魔法”函数，除了 *args 作为数组之外：

>>> def f(*args):
...    args = numpy.asarray(args[0]) if len(args) == 1 else numpy.asarray(args)
...    return args
... 
>>> f(7,3,5)
array([7, 3, 5])
>>> f([1,2,3])
array([1, 2, 3])
>>> f((2,3,4))
array([2, 3, 4])
>>> f(numpy.array([1,2,3]))
array([1, 2, 3])