使用Numba的解决方案
使用Cython,CodeSurgeon已经给出了很好的答案。在这个答案中,我不想展示使用Numba的另一种方法。
我创建了三个版本。在
naive_numba我只添加了一个功能装饰。在
improved_Numba我手动组合的循环中(每个矢量化命令实际上都是一个循环)。在
improved_Numba_p我已经并行化了功能。请注意,使用并行加速器时,显然存在一个错误,不允许定义常量值。还应注意,并行化版本仅对较大的输入阵列有利。但是,您还可以添加一个小的包装器,该包装器根据输入数组的大小调用单线程或并行化版本。
代码dtype = float64
import numba as nbimport numpy as npimport time@nb.njit(fastmath=True)def naive_Numba(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile): DustJ, DustF, DustG, DustH, DustI = 250.0, 633.0, 2.513, -2.2e-3, -2.8e-6 IceI, IceC, IceD, IceE, IceF, IceG, IceH = 273.16, 1.843e5, 1.6357e8, 3.5519e9, 1.6670e2, 6.4650e4, 1.6935e6 delta = u-DustJ result_dust = DustF+DustG*delta+DustH*delta**2+DustI*(delta**3); x= u/IceI; result_ice = (x**3)*(IceC+IceD*(x**2)+IceE*(x**6))/(1+IceF*(x**2)+IceG*(x**4)+IceH*(x**8)) return (DensityIceProfile*result_ice+DensityDustProfile*result_dust)/DensityProfile#error_model='numpy' sets divison by 0 to NaN instead of throwing a exception, this allows vectorization@nb.njit(fastmath=True,error_model='numpy')def improved_Numba(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile): DustJ, DustF, DustG, DustH, DustI = 250.0, 633.0, 2.513, -2.2e-3, -2.8e-6 IceI, IceC, IceD, IceE, IceF, IceG, IceH = 273.16, 1.843e5, 1.6357e8, 3.5519e9, 1.6670e2, 6.4650e4, 1.6935e6 res=np.empty(u.shape[0],dtype=u.dtype) for i in range(u.shape[0]): delta = u[i]-DustJ result_dust = DustF+DustG*delta+DustH*delta**2+DustI*(delta**3); x= u[i]/IceI result_ice = (x**3)*(IceC+IceD*(x**2)+IceE*(x**6))/(1+IceF*(x**2)+IceG*(x**4)+IceH*(x**8)) res[i]=(DensityIceProfile[i]*result_ice+DensityDustProfile[i]*result_dust)/DensityProfile[i] return res#there is obviously a bug in Numba (declaring const values in the function)@nb.njit(fastmath=True,parallel=True,error_model='numpy')def improved_Numba_p(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile,DustJ, DustF, DustG, DustH, DustI,IceI, IceC, IceD, IceE, IceF, IceG, IceH): res=np.empty((u.shape[0]),dtype=u.dtype) for i in nb.prange(u.shape[0]): delta = u[i]-DustJ result_dust = DustF+DustG*delta+DustH*delta**2+DustI*(delta**3); x= u[i]/IceI result_ice = (x**3)*(IceC+IceD*(x**2)+IceE*(x**6))/(1+IceF*(x**2)+IceG*(x**4)+IceH*(x**8)) res[i]=(DensityIceProfile[i]*result_ice+DensityDustProfile[i]*result_dust)/DensityProfile[i] return resu=np.array(np.random.rand(1000000),dtype=np.float32)PorosityProfile=np.array(np.random.rand(1000000),dtype=np.float32)DensityIceProfile=np.array(np.random.rand(1000000),dtype=np.float32)DensityDustProfile=np.array(np.random.rand(1000000),dtype=np.float32)DensityProfile=np.array(np.random.rand(1000000),dtype=np.float32)DustJ, DustF, DustG, DustH, DustI = 250.0, 633.0, 2.513, -2.2e-3, -2.8e-6IceI, IceC, IceD, IceE, IceF, IceG, IceH = 273.16, 1.843e5, 1.6357e8, 3.5519e9, 1.6670e2, 6.4650e4, 1.6935e6#don't measure compilation overhead on first callres=improved_Numba_p(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile,DustJ, DustF, DustG, DustH, DustI,IceI, IceC, IceD, IceE, IceF, IceG, IceH) for i in range(1000): res=improved_Numba_p(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile,DustJ, DustF, DustG, DustH, DustI,IceI, IceC, IceD, IceE, IceF, IceG, IceH)print(time.time()-t1)print(time.time()-t1)
性能
Arraysize np.random.rand(100)Numpy 46.8µsnaive Numba 3.1µsimproved Numba: 1.62µsimproved_Numba_p: 17.45µs#Arraysize np.random.rand(1000000)Numpy 255.8msnaive Numba 18.6msimproved Numba: 6.13msimproved_Numba_p: 3.54ms
代码dtype = np.float32
如果np.float32足够,则必须在函数中显式声明所有常量值给float32。否则,Numba将使用float64。
@nb.njit(fastmath=True,error_model='numpy')def improved_Numba(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile): DustJ, DustF, DustG, DustH, DustI = nb.float32(250.0), nb.float32(633.0), nb.float32(2.513), nb.float32(-2.2e-3), nb.float32(-2.8e-6) IceI, IceC, IceD, IceE, IceF, IceG, IceH = nb.float32(273.16), nb.float32(1.843e5), nb.float32(1.6357e8), nb.float32(3.5519e9), nb.float32(1.6670e2), nb.float32(6.4650e4), nb.float32(1.6935e6) res=np.empty(u.shape[0],dtype=u.dtype) for i in range(u.shape[0]): delta = u[i]-DustJ result_dust = DustF+DustG*delta+DustH*delta**2+DustI*(delta**3) x= u[i]/IceI result_ice = (x**3)*(IceC+IceD*(x**2)+IceE*(x**6))/(nb.float32(1)+IceF*(x**2)+IceG*(x**4)+IceH*(x**8)) res[i]=(DensityIceProfile[i]*result_ice+DensityDustProfile[i]*result_dust)/DensityProfile[i] return res@nb.njit(fastmath=True,parallel=True,error_model='numpy')def improved_Numba_p(u, PorosityProfile, DensityIceProfile, DensityDustProfile, DensityProfile): res=np.empty((u.shape[0]),dtype=u.dtype) DustJ, DustF, DustG, DustH, DustI = nb.float32(250.0), nb.float32(633.0), nb.float32(2.513), nb.float32(-2.2e-3), nb.float32(-2.8e-6) IceI, IceC, IceD, IceE, IceF, IceG, IceH = nb.float32(273.16), nb.float32(1.843e5), nb.float32(1.6357e8), nb.float32(3.5519e9), nb.float32(1.6670e2), nb.float32(6.4650e4), nb.float32(1.6935e6) for i in nb.prange(u.shape[0]): delta = u[i]-DustJ result_dust = DustF+DustG*delta+DustH*delta**2+DustI*(delta**3) x= u[i]/IceI result_ice = (x**3)*(IceC+IceD*(x**2)+IceE*(x**6))/(nb.float32(1)+IceF*(x**2)+IceG*(x**4)+IceH*(x**8)) res[i]=(DensityIceProfile[i]*result_ice+DensityDustProfile[i]*result_dust)/DensityProfile[i] return res
性能
Arraysize np.random.rand(100).astype(np.float32)Numpy 29.3µsimproved Numba: 1.33µsimproved_Numba_p: 18µsArraysize np.random.rand(1000000).astype(np.float32)Numpy 117msimproved Numba: 2.46msimproved_Numba_p: 1.56ms
与@CodeSurgeon提供的Cython版本的比较并不十分公平,因为他没有使用启用的AVX2和FMA3指令编译该功能。Numba默认使用-march =
native进行编译,这会在我的Core i7-4xxx上启用AVX2和FMA3指令。
但是,如果您不希望分发已编译的Cython版本的代码,就会产生这种感觉,因为如果启用了该优化功能,默认情况下,它将不会在Haswell之前的处理器(或所有Pentium和Celerons)上运行。应该可以编译多个代码路径,但这是编译器的依赖和更多的工作。



