python-3.xnumba

How to use numba.vectorize function with an argument being a sequence of floats


I am trying to optimize some code and looking for ways to make make a numpy universal function with numba.

The original function has the following signature

import numpy as np
from numba import njit, vectorize
from typing import Sequence

@njit
def f(x: float, y: Sequence[float], z:float = 1e-14) -> float:
    """Computations are only exemplarily"""

    a = np.sum(y)
    if np.abs(x - a) < z:
        return 0.
    else:
        return np.abs(x - a)

which works as expected.

I would like use numba.vectorize to create an universal function s.t. f can be called with

x: np.ndarray | float, y: Sequence[np.ndarray | float], z: float

Also, f is created in some other part of the code, i.e. the length of y varies, but is known at the time of definition of f.

One obvious solution is to make another function which takes (possibly) vectorized input and do some operations inside using numba.prange to fil up the result vector.

But I am looking for a more elegant solution using numba.vectorize and I don't want to deal with broadcasting, since x can be a float, or an element of y can be a float.

Best regards,

I tried

 f_v = vectorize(
     [
         numba.float64(
             numba.float64,
             numba.float64[:],
             numba.float64,
         )
     ],
     nopython=True,
)(f)

But it failed to compile, no matching signature found.

Edit: f performs simple arithmetics on scalars. I inserted some dummy computations since the originals are rather lengthy. The vectorized version should perform those computations component-wise and return an array.


Solution

  • You can use guvectorize instead of vectorize.

    @numba.guvectorize("(float64, float64[:], float64, float64[:])", "(),(n),()->()", target="parallel")
    def f_vectorized(x, y, z, out):
        out[0] = f(x, y, z)
    
    
    result1 = f_vectorized(1.0, np.array([1.0, 2.0]), 1e-14)
    result2 = f_vectorized(np.array([1.0]), np.array([[1.0, 2.0]]), 1e-14)
    result3 = f_vectorized(np.array([1.0, 2.0, 3.0]), np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 1e-14)
    result4 = f_vectorized(np.array([1.0, 2.0, 3.0]), [np.array([1.0, 2.0]), np.array([3.0, 4.0]), np.array([5.0, 6.0])], 1e-14)
    

    However, there are two issues to consider.

    The first is obvious: z cannot have a default value. To get around this, you have to create another function.

    def f_vectorized_with_default_value(x, y, z=1e-14):
        return f_vectorized(x, y, z)
    

    The second is more critical: it does not work well if y is a python list of numpy arrays. That is, this case:

    result4 = f_vectorized(np.array([1.0, 2.0, 3.0]), [np.array([1.0, 2.0]), np.array([3.0, 4.0]), np.array([5.0, 6.0])], 1e-14)
    

    Here is the benchmark code.

    import time
    import timeit
    import warnings
    
    import numpy as np
    from typing import Sequence
    import numba
    
    
    @numba.njit(cache=True)
    def f(x: float, y: Sequence[float], z: float = 1e-14) -> float:
        """Computations are only exemplarily"""
    
        a = np.sum(y)
        if np.abs(x - a) < z:
            return 0.0
        else:
            return np.abs(x - a)
    
    
    @numba.njit(parallel=True, cache=True)
    def f_prange(x, y, z):
        """Implementation with prange as a baseline."""
        n = len(x)
        out = np.empty(n, dtype=x.dtype)
        for i in numba.prange(n):
            out[i] = f(x[i], y[i], z)
        return out
    
    
    def f_loop(x, y, z):
        """Another baseline."""
        n = len(x)
        out = np.empty(n, dtype=x.dtype)
        for i in range(n):
            out[i] = f(x[i], y[i], z)
        return out
    
    
    @numba.guvectorize("(float64, float64[:], float64, float64[:])", "(),(n),()->()", target="parallel", cache=True)
    def f_vectorized(x, y, z, out):
        out[0] = f(x, y, z)
    
    
    def f_vectorized_with_default_value(x, y, z=1e-14):
        return f_vectorized(x, y, z)
    
    
    def test(title, func, x, y, z, expected, number=10):
        assert np.array_equal(func(x, y, z), expected)
        elapsed = timeit.timeit(lambda: func(x, y, z), number=number) / number
        print(f"{title}: {elapsed}")
    
    
    def main():
        rng = np.random.default_rng(0)
        n = 1000000
        m = 1000
    
        x = rng.random(size=(n,), dtype=np.float64)
        y_as_numpy = rng.random(size=(n, m), dtype=np.float64)
        y_as_list = [y_as_numpy[i] for i in range(n)]
        z = 1e-14
    
        started = time.perf_counter()
        _ = np.array(y_as_list)
        print("merge into a np array:", time.perf_counter() - started)
    
        started = time.perf_counter()
        y_as_typed_list = numba.typed.List(y_as_list)
        print("convert to typed list:", time.perf_counter() - started)
    
        expected = f(x[0], y_as_numpy[0], z)
        test("f(single)", f, x[0], y_as_numpy[0], z, expected)
        test("f_vectorized(single)", f_vectorized, x[0], y_as_numpy[0], z, expected)
    
        expected = f_prange(x, y_as_numpy, z)
        test("f_prange(numpy)", f_prange, x, y_as_numpy, z, expected)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", numba.NumbaPendingDeprecationWarning)
            test("f_prange(python_list)", f_prange, x, y_as_list, z, expected)
        test("f_prange(typed_list)", f_prange, x, y_as_typed_list, z, expected)
        test("f_loop(python_list)", f_loop, x, y_as_list, z, expected)
        test("f_vectorized(numpy)", f_vectorized, x, y_as_numpy, z, expected)
        test("f_vectorized(python_list)", f_vectorized, x, y_as_list, z, expected)
        test("f_vectorized(typed_list)", f_vectorized, x, y_as_typed_list, z, expected)
    
    
    if __name__ == "__main__":
        main()
    
    

    Result:

    merge into a np array: 1.3248698
    convert to typed list: 1.6472616999999996
    f(single): 1.430000000013365e-06
    f_vectorized(single): 1.6320000000025204e-05
    f_prange(numpy): 0.18999052999999985
    f_prange(python_list): 7.21668874
    f_prange(typed_list): 0.20048094999999932 + (convert to typed list)
    f_loop(python_list): 1.3346305299999996
    f_vectorized(numpy): 0.18892754000000025
    f_vectorized(python_list): 1.8655723899999999
    f_vectorized(typed_list): 3.323696179999999 + (convert to typed list)
    

    As you can see from the last two lines of results, it is very slow. It works, but slow. In fact, it is faster to call the f function sequentially in a python loop (not numba.prange, pure python loop).

    If y is a 2D numpy array in the first place, this will not bother you. f_vectorized should work fine. But if it is a python list, you might have to find another trick.