Numpy vectorize wrongly converts the output to be integer

问题

I am struggling with the following code:

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))

the result looks like this:

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

I also checked the dtype. It seems that the vectorize function is conveting the type to int64 from float64!

回答1:

The documentation for np.vectorize explains:

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.

The first element of your input is 0.0, which returns the integer 0, so as far as numpy knows, you want an integer dtype. As you discovered, if you change this to 0.0 so you're not changing the return type, it'll behave. Alternatively you can specify 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.        ])

回答2:

@amiragha

Python is looking at your input variables as ints as well, so if you wanted to keep everything float64 you would need to specify all necessary "numbers" as floats.

来源：https://stackoverflow.com/questions/60876903/trivial-changes-in-numpy-vectorized-functions-domain-produce-huge-discrepancies