I want to create a vectorized numba function to act on a numpy array. The array will be passed in as the first argument and then other important data will be passed in for the remaining arguments. How would I do this? All of the arguments seem to be iterated over like a zip using the vectorized function that I have created.
Do you have a minimal example of what you’re trying to do?
I’ll try and give one. For example, take the following code:
import numpy as np def f(a, b): return a + b example_array = np.array([1, 2, 3, 4]) result_array = f(example_array, 2)
I would like to vectorize f such that the
result_array has each value in
example_array increased by two. In this case, the constant argument is
b. In reality, the constants are a little more complicated but this should give an idea of what I am trying to achieve.
What exactly do you mean by ‘vectorize’? I’d think the numpy ‘+’ operator is already vectorized and the example above will perform well.
If you mean numba-jitted, that works as well (see below). Do you have an example of what you’d like to write that doesn’t work?
import numpy as np import numba @numba.njit def f(a, b): return a + b example_array = np.array([1, 2, 3, 4]) print(f(example_array, 2))
If you want to use the @vectorize or @guvectorize decorators, it’s documented here
These both work, the second is taken from the documentation.
import numpy as np import numba def make_vectorize_func(const_val): @numba.vectorize(nopython=True) def _(a): return a + const_val return _ example_array = np.array([1, 2, 3, 4]) f = make_vectorize_func(2) print(f(example_array)) @numba.guvectorize([(numba.int64[:], numba.int64, numba.int64[:])], '(n),()->(n)') def g(x, y, res): for i in range(x.shape): res[i] = x[i] + y out = np.empty_like(example_array) print(g(example_array, 3, out))