python - Numpy vectorize 错误地将输出转换为整数

Question

我正在努力处理以下代码:

import numpy as np

e = np.linspace(0, 4, 10)

def g(x):
    if x > 1:
        return x
    else:
        return 0

vg = np.vectorize(g)

print(vg(e))

结果是这样的:

    [0 0 0 1 1 2 2 3 3 4]

我还检查了数据类型。看起来 vectorize 函数正在将类型从 float64 转换为 int64!

Answer 1

np.vectorize 的文档解释了:

The data type of the output of vectorized is determined by calling the function with the first element of the input. This can be avoided by specifying the otypes argument.

您输入的第一个元素是 0.0，它返回整数 0，因此据 numpy 所知，您需要一个整数类型。正如您所发现的，如果将其更改为 0.0，那么您就不会更改返回类型，它会正常运行。或者，您可以指定 otypes:

>>> vg = np.vectorize(g)
>>> vg(e)
array([0, 0, 0, 1, 1, 2, 2, 3, 3, 4])
>>> vg = np.vectorize(g, otypes=[np.float64])
>>> vg(e)
array([ 0.        ,  0.        ,  0.        ,  1.33333333,  1.77777778,
        2.22222222,  2.66666667,  3.11111111,  3.55555556,  4.        ])