Search code examples
pythonnumpyvectorization

Numpy cannot `vectorize` a function


I want to apply this simple function to numpy arrays fast:

def f (x):
    return max(0,1-abs(x))

Just for clarity's sake, here is the plot:

import numpy as np
import matplotlib.pyplot as plt
x = np.linspace(start=-4,stop=4, num=100)
plt.plot(x, list(map(f,x)))

function plot

However, np.vectorize(f)(x) returns a vector of 0

Yes, I know that

The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.

but it should at least return a correct value, no?

PS. I managed to vectorize f by hand:

def f (x):
    z = np.zeros(x.shape)
    return np.fmax(z,1-np.abs(x), out=z)

Solution

  • This is because np.vectorize is inferring the return type as an integer dtype. Note this part from the docs:

    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.

    In your case, the output of max(0, 1 - abs(-4.)) is an int object, 0. You could either change your function to use a float, or specify the outypes parameter explicitly:

    np.vectorize(f, otypes=[np.float64])(x)