Numba векторизация для функции без ввода - PullRequest
Новая
       36

Numba векторизация для функции без ввода

0 голосов
/ 02 июля 2018

Я хочу распараллелить функцию, используя numba.vectorize, но моя функция не требует ввода. В настоящее время я использую фиктивный массив и фиктивный ввод для своей функции, которая никогда не используется.

Есть ли более элегантный / быстрый способ (возможно, без использования numba.vectorize)?

Пример кода (не мой настоящий код, только для демонстрации того, как я отбрасываю ввод): 

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)

1 Ответ

0 голосов
/ 14 июля 2018

Если вы просто будете иметь дело с одномерными массивами, то вы можете использовать следующее, где массив должен создаваться вне функции. Кажется, нет никакой причины использовать vectorize здесь, вы можете достичь цели просто с помощью jit, хотя вы должны явно записать цикл над элементами массива, используя это. Если ваш массив всегда будет 1D, то вы можете использовать: 

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)

Аналогичным образом вы можете работать с произвольными массивами, используя атрибут flat (и обязательно используйте .size для получения общего количества элементов в массиве): 

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

и, наконец, вы можете создать свой массив внутри своей функции, если вам нужен новый массив каждый раз, когда вы запускаете функцию (используйте вышеупомянутое вместо, если вы будете вызывать функцию повторно и хотите перезаписать один и тот же массив, сохраняя, таким образом, время перераспределять один и тот же массив снова и снова). 

@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)
...