CODEWITHSUNDEEP
pythonnumpyvectorizationnumba
Slizzered
Slizzered

Reputation: 899

Numba vectorize for function with no input

I want to parallelize a function using numba.vectorize, but my function doesn't take any input. Currently, I use a dummy array and dummy input for my function that is never used.

Is there a more elegant/fast way (possibly without using numba.vectorize)?

Code example (not my actual code, only for demonstration how I discard input):

import numpy as np
from numba import vectorize

@vectorize(["int32(int32)"], nopython=True)
def particle_path(discard_me):
    x = 0
    for _ in range(10):
        x += np.random.uniform(0, 1)
    return np.int32(x)

arr = particle_path(np.empty(1024, dtype=np.int32))
print(arr)

Upvotes: 2

Views: 373

Answers (1)

Justin Lanfranchi
Justin Lanfranchi

Reputation: 370

If you'll simply be dealing with 1D arrays, then you can use the following, where the array must be instantiated outside the function. There doesn't seem to be any reason to use vectorize here, you can achieve the goal simply with jit although you do have to explicitly write the loop over the array elements using this. If your array will always be 1D, then you can use:

import numpy as np
from numba import jit

@jit(nopython=True)
def particle_path(out):
    for i in range(len(out)):
        x = 0
        for _ in range(10):
            x += np.random.uniform(0, 1)
        out[i] = x

arr = np.empty(1024, dtype=np.int32)
particle_path(arr)

You can similarly deal with any-dimensional arrays using the flat attribute (and make sure to use .size to get total number of elements in the array):

@jit(nopython=True)
def particle_path(out):
    for i in range(out.size):
        x = 0
        for _ in range(10):
            x += np.random.uniform(0, 1)
        out.flat[i] = x

arr = np.empty(1024, dtype=np.int32)
particle_path(arr)

and finally you can create your array inside your function if you need a new array each time you run the function (use the above instead if you'll be calling the function repeatedly and want to overwrite the same array, hence saving the time to re-allocate the same array over and over again).

@jit(nopython=True)
def particle_path(num):
    out = np.empty(shape=num, dtype=np.int32)
    for i in range(num):
        x = 0
        for _ in range(10):
            x += np.random.uniform(0, 1)
        out[i] = x
    return out

arr = particle_path(1024)

Upvotes: 2

Related Questions