How to use Numba to vectorize a function that takes scalars and returns a vector?

Question

I have a function that takes two scalars and returns a 1-D array. Take for example:

import numpy as np
def linear_pattern(slope,length):
    pattern = slope*np.arange(length)
    return pattern

My goal is to be able of doing many slope at a time always for the same scalar value length, i.e., I would like to run

>>> linear_pattern(np.array([3,4]),10)
array([[ 0.,  3.,  6.,  9., 12., 15., 18., 21., 24., 27.],
       [ 0.,  4.,  8., 12., 16., 20., 24., 28., 32., 36.]])

instead of

>>> linear_pattern(3,10)
array([ 0,  3,  6,  9, 12, 15, 18, 21, 24, 27])
>>> linear_pattern(4,10)
array([ 0,  4,  8, 12, 16, 20, 24, 28, 32, 36])

I am using the decorator guvectorize from Numba to do the job. The following is my implementation:

from numba import float64, guvectorize
@guvectorize( [(float64, float64, float64[:])], '(),() -> (n)' )
def linear_pattern(slope,length,pattern):
    pattern[:] = slope*np.arange(length)

However, I get an error when I try to run the following example:

>>> pattern = np.zeros(10) 
>>> linear_pattern(np.array([3,4]),10,pattern)
NameError: undefined output symbols: n

The error complains because I set the function's output to be a vector of dimension n and such dimension is not defined, i.e., I did not use n anywhere in the input dimensions. However, the inputs are all scalars which precludes the use of n.

My questions are:

1. how can I make the above code work?
2. since length is always the same scalar, I would prefer it to be a kwarg of linear_pattern instead of a arg, however, the decorator guvectorize seems not to accept functions with kwarg-type arguments. Can guvectorize accept kwarg arguments?

---

Temporarily, I got the code to work by defining a dummy variable _ with the same dimension as the output. See below:

length = 10
__ = np.zeros(length)

@guvectorize( [(float64, float64, float64[:], float64[:])], '(),(),(n) -> (n)' )
def linear_pattern(slope,length,__,pattern):
    pattern[:] = slope*np.arange(length)

>>> pattern = np.zeros(length) 
>>> linear_pattern(np.array([3,4]),length,pattern)
array([[ 0.,  3.,  6.,  9., 12., 15., 18., 21., 24., 27.],
       [ 0.,  4.,  8., 12., 16., 20., 24., 28., 32., 36.]])

---

Edit 1:

In my original question, I used a simplified version of the actual function I need to vectorize. The original one has more arguments which also need to respond to vectorization.

Take for instance the modified function below

def linear_pattern(slope, intercept, length):
pattern = slope*np.arange(length) + intercept
return pattern

I would like then to run

slopes = np.array([3, 4])
intercepts = np.array([0, 1, 2])
length = 10
linear_pattern(slopes, intercepts, length)

and that the result have dimension (len(slopes),len(intercepts),length) which, in this case, is (2,3,10). Note that length will always be a fix integer and there is no need for vectorization on that argument, in fact, the ideal solution will keep length as a kwarg-type argument.

Answer 1

If you really need to use Numba, temporary solutions live for a very long time. As a matter of fact, the use of dummy inputs is the way to go because, according to the docs, the output array is actually allocated by NumPy’s dispatch mechanism, which calls into the Numba-generated code.

Now that you need to allocate a dummy input, the length parameter is redundant. Instead, you could allocate the full output:

@nb.guvectorize([(nb.float64, nb.float64[:], nb.float64[:])], '(),(n) -> (n)')
def linear_pattern2(slope, __, pattern):
    pattern[:] = slope * np.arange(len(pattern))

>>> pattern = np.empty((2, length))
>>> linear_pattern2(np.array([3, 4]), pattern2)
array([[ 0.,  3.,  6.,  9., 12., 15., 18., 21., 24., 27.],
       [ 0.,  4.,  8., 12., 16., 20., 24., 28., 32., 36.]])

And you can prevent the unnecessary allocation of the output array using the out kwarg (automagically provided by Numpy's ufunc specification) which may be relevant for large arrays:

>>> pattern = np.empty((2, length))
>>> linear_pattern2(np.array([3, 4]), pattern, out=pattern)
>>> pattern
array([[ 0.,  3.,  6.,  9., 12., 15., 18., 21., 24., 27.],
       [ 0.,  4.,  8., 12., 16., 20., 24., 28., 32., 36.]])

The function's return and the input array are the same object:

>>> pattern = np.empty((2, length))
>>> lp = linear_pattern2(np.array([3, 4]), pattern, out=pattern)
>>> lp is pattern
True

Answer 2

If you don't mind linear_pattern returning a matrix, you could use NumPy's broadcasting (no Numba required):

import numpy as np

def linear_pattern(slope,length):
    pattern = slope * np.arange(length).reshape((1, -1))
    return pattern

print(linear_pattern(3, 10))
print(linear_pattern(np.array([[3], [4]]), 10))

Output:

[[ 0  3  6  9 12 15 18 21 24 27]]
[[ 0  3  6  9 12 15 18 21 24 27]
 [ 0  4  8 12 16 20 24 28 32 36]]

How this works

1. np.arange(10).reshape((1, -1)) is array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
2. Element-wise (!) multiplication of this by np.array([[3], [4]]) looks like this:

3                    3 * [0 1 2 3 ...]      [0 3 6  9 ...]
  * [0 1 2 3 ...] =                       = 
4                    4 * [0 1 2 3 ...]      [0 4 8 12 ...]