Python and Numba for vectorized functions

Question

Good day, I'm writing a Python module for some numeric work. Since there's a lot of stuff going on, I've been spending the last few days optimizing code to improve calculations times. However, I have a question concerning Numba. Basically, I have a class with some fields which are numpy arrays, which I initialize in the following way:

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

So, the code is vectorized, and using Numba's JIT results in some improvement. However, sometimes I need to access the calculate_vibr_energy function from outside the class, and pass a single integer instead of an array in place of i. As far as I understand, if I use Numba's JIT on the calculate_vibr_energy, it will have to always take an array as an argument.

So, which of the following options is better: 1) Create a new function calculate_vibr_energy_single(i), which will only take a single integer number, and use Numba on it too 2) Replace all usages of the function that are similar to this one:

myclass.calculate_vibr_energy(1)

with this:

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

Or are there other, more efficient (or at least, more Python-ic) ways of doing that?

Answer 1

I have only played a little with numba yet so I may be mistaken, but as far as I've understood it, using the "autojit" decorator should give functions that can take arguments of any type.

See e.g. http://numba.pydata.org/numba-doc/dev/pythonstuff.html