How do you vectorize a function while keeping some arguments non-vectorized?

Numba Support: How do I do ...?
1

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.

2

Do you have a minimal example of what you’re trying to do?

3

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.

4

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

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[0]):
        res[i] = x[i] + y

out = np.empty_like(example_array)
print(g(example_array, 3, out))