Good afternoon, I am writing a Python module for some numerical work. Since there is a lot going on there, I spent the last few days optimizing the code to improve the computation time. However, I have a question about Numba. Basically, I have a class with some fields that are numpy arrays, which I initialize as follows:
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 the Numba JIT leads to 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 instead of i. As far as I understand, if I use Numba JIT to calculate_vibr_energy, it should 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) that will only accept one integer, and also use Numba 2) Replace all uses of a function similar to this:
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 to do this?
source share