有没有Pythonic的方法将标量和0维数组转换为1维数组？

Question

比如，我想把'a'改成0，如果'a'小于5。

def foo(a):
    return 0 if a < 5 else a

为了让它适用于numpy数组，我把它改成：

def foo2(a):
    a[a < 5] = 0
    return a

问题是我希望这个函数能同时适用于单个数值和数组。

这个iscale()函数可以用来测试'a'是不是单个数值，但对于0维数组，比如array(12)，它会返回假。

有没有什么比较“python风”的方法，可以把单个数值和0维数组都转成1维数组，而不改变其他的多维数组呢？

Answer 1

如果你不介意这个函数即使接收到一个单一的数值也会返回一个数组，我会倾向于这样做：

import numpy as np

def foo(a,k=5):
    b = np.array(a)
    if not b.ndim:
        b = np.array([a])
    b[b < k] = 0
    return b

print(foo(3))
print(foo(6))
print(foo([1,2,3,4,5,6,7,8,9]))
print(foo(np.array([1,2,3,4,5,6,7,8,9])))

... 这样会产生以下结果：

[0]
[6]
[0 0 0 0 5 6 7 8 9]
[0 0 0 0 5 6 7 8 9]

从测试示例中可以看到，如果这个函数接收到的是一个普通的Python列表，而不是numpy数组或单一数值，它就能按预期工作。

在这个过程中创建两个数组可能看起来有点浪费，但首先，创建b可以防止这个函数产生不必要的副作用。想想看：

def foobad(a,k=5):
    a[a < k] = 0
    return a

x = np.array([1,2,3,4,5,6,7,8,9])
foobad(x)
print (x)

... 输出：

[0 0 0 0 5 6 7 8 9]

... 这可能不是函数使用者想要的结果。其次，如果第二个数组的创建是因为函数接收到的是一个单一的数值，那么它只会从一个包含1个元素的列表中创建数组，这个过程应该非常快。

Answer 2

你可以使用 numpy.vectorize，配合原来的标量实现。

@np.vectorize
def foo(a):
   return 0 if a < 5 else a

foo(3)
=> array(0)
foo(7)
=> array(7)
foo([[3,7], [8,-1]])
=> array([[0, 7],
          [8, 0]])

需要注意的是，当你使用 vectorize 时，你会牺牲一些速度（你的计算不再是在numpy/C层面上进行向量化），换来的是简单性（你可以用简单的形式来写你的函数，适用于标量）。

Answer 3

好吧，我有一个看起来有效的解决办法。

def foo3(a):
    return a * (a >= 5)

foo3(4)
=> 0

foo3(6)
=> 6

foo3(np.array(3))
=> 0

foo3(np.array(6))
=> 6

foo3(np.array([1, 5]))
=> array([0, 5])

这个方法运行得很好，但我不确定这样做是否安全。

Answer 4

使用 np.atleast_1d 这个函数。

这个方法对任何输入都有效，无论是单个数字（标量）还是数组：

def foo(a):
    a = np.atleast_1d(a)
    a[a < 5] = 0
    return a

不过要注意，如果你输入的是单个数字，它会返回一个一维的数组。