Python и Numba для векторизованных функций

Добрый день, я пишу модуль Python для некоторой числовой работы. Поскольку происходит много вещей, я потратил последние несколько дней на оптимизацию кода, чтобы улучшить время вычислений. Однако у меня есть вопрос относительно Numba. В принципе, у меня есть класс с некоторыми полями, которые являются массивами numpy, которые я инициализирую следующим образом:

def init(self):
    a = numpy.arange(0, self.max_i, 1)
    self.vibr_energy = self.calculate_vibr_energy(a)

def calculate_vibr_energy(i):
    return numpy.exp(-self.harmonic * i - self.anharmonic * (i ** 2))

Итак, код векторизован, и использование JIT Numba приводит к некоторому улучшению. Однако иногда мне нужно получить доступ к calculate_vibr_energy функция из-за пределов класса, и передать одно целое число вместо массива вместо i. Насколько я понимаю, если я использую JIT Numba на calculate_vibr_energy, он всегда должен будет принимать массив в качестве аргумента.

Итак, какой из следующих вариантов лучше: 1) Создайте новую функцию calculate_vibr_energy_single (i), которая будет принимать только одно целое число и использовать Numba на нем 2) заменить все использования функции, которые похожи на этот один:

myclass.calculate_vibr_energy(1)

С этого:

tmp = np.array([1])
myclass.calculate_vibr_energy(tmp)[0]

или есть другие, более эффективные (или, по крайней мере, более Python-ic) способы сделать это?

1 ответов


Я еще немного играл с numba, поэтому я могу ошибаться, но, насколько я понял, использование декоратора "autojit" должно давать функции, которые могут принимать аргументы любого типа.

см., например,http://numba.pydata.org/numba-doc/dev/pythonstuff.html