Python和Numba用于矢量化函数

时间:2013-06-14 15:53:47

标签: python arrays numpy 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))

因此,代码是矢量化的,使用Numba的JIT会带来一些改进。但是,有时我需要从类外部访问calculate_vibr_energy函数,并传递单个整数而不是数组来代替i。 据我所知,如果我在calculate_vibr_energy上使用Numba的JIT,它必须始终以数组作为参数。

那么,以下哪个选项更好: 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 个答案:

答案 0 :(得分:1)

我只玩过numba,所以我可能会弄错,但据我所知,使用“autojit”装饰器应该提供可以接受任何类型参数的函数。

参见例如http://numba.pydata.org/numba-doc/dev/pythonstuff.html