假设我有一个适用于一维数组的功能:
def opfunc(a, b):
if len(a.shape) != 1 and len(b.shape) != 1:
raise ValueError('opfunc: both inputs must be 1-d arrays.')
return 2 * b + 2 * (a[0] - a[-1])
我希望在循环或嵌套循环中重复使用它,比如
nx1, nx2, nx3 = a.shape
c = np.zeros(a.shape)
for ix1 in range(nx1):
for ix2 in range(nx2):
avec = a[ix1, ix2, :]
bvec = b[ix1, ix2, :]
cvec = opfunc(avec, bvec)
c[ix1, ix2, :] = cvec
我可以通过将这些语句放在一个函数中并使用opfunc
numba.jit
进行装饰来加快速度:
@numba.jit(nopython=True)
def opfunc(a, b):
if len(a.shape) != 1 and len(b.shape) != 1:
raise ValueError('opfunc: both inputs must be 1-d arrays.')
return 2 * b + 2 * (a[0] - a[-1])
@numba.jit(nopython=True)
def opfunc1(a, b):
nx1, nx2, nx3 = a.shape
c = np.zeros(a.shape)
for ix1 in range(nx1):
for ix2 in range(nx2):
avec = a[ix1, ix2, :]
bvec = b[ix1, ix2, :]
cvec = opfunc(avec, bvec)
c[ix1, ix2, :] = cvec
return c
这里确定加速。
如果在Fortran中编写opfunc
$ !cat fortran_test.f90
module test_fortran
implicit none
contains
subroutine opfunc(a, b, c)
implicit none
real (kind=8), dimension(6), intent(in) :: a, b
real (kind=8), dimension(6), intent(out) :: c
c = 2 * b + 2 !* (a(1) - a(ubound(a)))
end subroutine opfunc
end module test_fortran
并已使用f2py导入Python?
import fortran_test
opfuncx = fortran_test.test_fortran.opfunc
@numba.jit(nopython=True)
def opfunc1x(a, b):
nx1, nx2, nx3 = a.shape
c = np.zeros(a.shape)
for ix1 in range(nx1):
for ix2 in range(nx2):
avec = a[ix1, ix2, :]
bvec = b[ix1, ix2, :]
cvec = opfuncx(avec, bvec)
c[ix1, ix2, :] = cvec
return c
实际上,我正在
TypingError: Failed at nopython (nopython frontend)
Untyped global name 'opfuncx'
File "<ipython-input-48-c51d5b4fdad5>", line 13
运行该功能时。我猜这是因为opfuncx
没有用numba.jit
装饰。有可能这样做吗?