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

我正在努力使用以下代码:

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]

我也检查了dtype.似乎vectorize函数正在从float64将类型转换为int64!

文档np.vectorize解释:

输出的数据类型vectorized是通过使用输入的第一个元素调用函数来确定的.通过指定otypes参数可以避免这种情况.

输入的第一个元素是0.0,它返回整数0,所以只要numpy知道,你需要一个整数dtype.正如您所发现的,如果您将其更改为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.        ])